Solve the Sellar MDO problem with JAX.
Note
Click here to download the full example code
Solve the Sellar MDO problem with JAX.
from __future__ import annotations
from gemseo import configure_logger
from gemseo import create_scenario
from gemseo.core.mdo_functions.mdo_function import MDOFunction
from gemseo.problems.mdo.sellar.sellar_design_space import SellarDesignSpace
from gemseo_jax.jax_chain import JAXChain
from gemseo_jax.problems.sellar.sellar_1 import JAXSellar1
from gemseo_jax.problems.sellar.sellar_2 import JAXSellar2
from gemseo_jax.problems.sellar.sellar_system import JAXSellarSystem
configure_logger()
Out:
<RootLogger root (INFO)>
Create the disciplines:
sellar_1 = JAXSellar1()
sellar_2 = JAXSellar2()
sellar_system = JAXSellarSystem()
Make a JAXChain to assemble the 3 without reconverting to NumPy:
disciplines = [sellar_1, sellar_2, sellar_system]
jax_chain = JAXChain(disciplines, name="SellarChain")
Add the differentiated outputs to reduce the computation graph of the Jacobian:
jax_chain.add_differentiated_outputs(["obj", "c_1", "c_2"])
Compile functions, this takes an extra compilation time, but lowers the cost of re-evaluation:
jax_chain.compile_jit()
Out:
INFO - 08:28:54: Compilation of the output function SellarChain: 0:00:00.039976 seconds.
INFO - 08:28:54: Compilation of the Jacobian function SellarChain: 0:00:00.058221 seconds.
Create the MDO scenario with an MDF formulation:
design_space = SellarDesignSpace()
scenario = create_scenario(
jax_chain,
"obj",
design_space,
formulation_name="MDF",
main_mda_settings={"inner_mda_name": "MDAGaussSeidel"},
)
scenario.add_constraint(["c_1", "c_2"], MDOFunction.ConstraintType.INEQ)
Out:
INFO - 08:28:54: Variable y_1 was removed from the Design Space, it is not an input of any discipline.
Execute the scenario and post-process the results:
scenario.execute(algo_name="SLSQP", max_iter=10)
scenario.post_process(post_name="OptHistoryView", save=False, show=True)
Out:
INFO - 08:28:54: *** Start MDOScenario execution ***
INFO - 08:28:54: MDOScenario
INFO - 08:28:54: Disciplines: SellarChain
INFO - 08:28:54: MDO formulation: MDF
INFO - 08:28:54: Optimization problem:
INFO - 08:28:54: minimize obj(x_1, x_2, x_shared)
INFO - 08:28:54: with respect to x_1, x_2, x_shared
INFO - 08:28:54: under the inequality constraints
INFO - 08:28:54: c_1_c_2(x_1, x_2, x_shared) <= 0
INFO - 08:28:54: over the design space:
INFO - 08:28:54: +-------------+-------------+-------+-------------+-------+
INFO - 08:28:54: | Name | Lower bound | Value | Upper bound | Type |
INFO - 08:28:54: +-------------+-------------+-------+-------------+-------+
INFO - 08:28:54: | x_1 | 0 | 1 | 10 | float |
INFO - 08:28:54: | x_2 | 0 | 1 | 10 | float |
INFO - 08:28:54: | x_shared[0] | -10 | 4 | 10 | float |
INFO - 08:28:54: | x_shared[1] | 0 | 3 | 10 | float |
INFO - 08:28:54: +-------------+-------------+-------+-------------+-------+
INFO - 08:28:54: Solving optimization problem with algorithm SLSQP:
INFO - 08:28:55: 10%|█ | 1/10 [00:00<00:02, 3.31 it/sec, feas=True, obj=23]
INFO - 08:28:55: 20%|██ | 2/10 [00:00<00:01, 6.58 it/sec, feas=True, obj=56.9]
INFO - 08:28:55: 30%|███ | 3/10 [00:00<00:00, 9.81 it/sec, feas=True, obj=27.1]
INFO - 08:28:55: 40%|████ | 4/10 [00:00<00:00, 12.98 it/sec, feas=True, obj=23.9]
INFO - 08:28:55: 50%|█████ | 5/10 [00:00<00:00, 16.12 it/sec, feas=True, obj=23.2]
INFO - 08:28:55: 60%|██████ | 6/10 [00:00<00:00, 19.21 it/sec, feas=True, obj=23]
INFO - 08:28:55: 70%|███████ | 7/10 [00:00<00:00, 22.27 it/sec, feas=True, obj=23]
INFO - 08:28:55: 80%|████████ | 8/10 [00:00<00:00, 25.27 it/sec, feas=True, obj=23]
INFO - 08:28:55: 90%|█████████ | 9/10 [00:00<00:00, 28.24 it/sec, feas=True, obj=23]
INFO - 08:28:55: 100%|██████████| 10/10 [00:00<00:00, 31.17 it/sec, feas=True, obj=23]
INFO - 08:28:55: Optimization result:
INFO - 08:28:55: Optimizer info:
INFO - 08:28:55: Status: None
INFO - 08:28:55: Message: Maximum number of iterations reached. GEMSEO stopped the driver.
INFO - 08:28:55: Solution:
INFO - 08:28:55: The solution is feasible.
INFO - 08:28:55: Objective: 22.952625867476453
INFO - 08:28:55: Standardized constraints:
INFO - 08:28:55: c_1_c_2 = [-14.79259005 -13.76295031]
INFO - 08:28:55: Design space:
INFO - 08:28:55: +-------------+-------------+-------------------+-------------+-------+
INFO - 08:28:55: | Name | Lower bound | Value | Upper bound | Type |
INFO - 08:28:55: +-------------+-------------+-------------------+-------------+-------+
INFO - 08:28:55: | x_1 | 0 | 1 | 10 | float |
INFO - 08:28:55: | x_2 | 0 | 1 | 10 | float |
INFO - 08:28:55: | x_shared[0] | -10 | 4.000000000000002 | 10 | float |
INFO - 08:28:55: | x_shared[1] | 0 | 3 | 10 | float |
INFO - 08:28:55: +-------------+-------------+-------------------+-------------+-------+
INFO - 08:28:55: *** End MDOScenario execution ***
<gemseo.post.opt_history_view.OptHistoryView object at 0x7f75eb4b86e0>
Total running time of the script: ( 0 minutes 1.030 seconds)
Download Python source code: plot_sellar.py