{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Solve the Sellar MDO problem with JAX.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from __future__ import annotations\n\nfrom gemseo import configure_logger\nfrom gemseo import create_scenario\nfrom gemseo.core.mdo_functions.mdo_function import MDOFunction\nfrom gemseo.problems.mdo.sellar.sellar_design_space import SellarDesignSpace\n\nfrom gemseo_jax.jax_chain import JAXChain\nfrom gemseo_jax.problems.sellar.sellar_1 import JAXSellar1\nfrom gemseo_jax.problems.sellar.sellar_2 import JAXSellar2\nfrom gemseo_jax.problems.sellar.sellar_system import JAXSellarSystem\n\nconfigure_logger()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Create the disciplines:\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "sellar_1 = JAXSellar1()\nsellar_2 = JAXSellar2()\nsellar_system = JAXSellarSystem()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Make a `JAXChain` to assemble the 3 without reconverting to NumPy:\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "disciplines = [sellar_1, sellar_2, sellar_system]\njax_chain = JAXChain(disciplines, name=\"SellarChain\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Add the differentiated outputs to reduce the computation graph of the Jacobian:\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "jax_chain.add_differentiated_outputs([\"obj\", \"c_1\", \"c_2\"])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Compile functions, this takes an extra compilation time, but lowers the cost of\nre-evaluation:\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "jax_chain.compile_jit()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Create the MDO scenario with an MDF formulation:\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "design_space = SellarDesignSpace()\nscenario = create_scenario(\n    jax_chain,\n    \"obj\",\n    design_space,\n    formulation_name=\"MDF\",\n    main_mda_settings={\"inner_mda_name\": \"MDAGaussSeidel\"},\n)\nscenario.add_constraint([\"c_1\", \"c_2\"], MDOFunction.ConstraintType.INEQ)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Execute the scenario and post-process the results:\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "scenario.execute(algo_name=\"SLSQP\", max_iter=10)\nscenario.post_process(post_name=\"OptHistoryView\", save=False, show=True)"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.9.19"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}