{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Compare JAX and NumPy for the resolution of the Sobieski's SSBJ problem.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from __future__ import annotations\n\nfrom gemseo import configure\nfrom gemseo import configure_logger\nfrom gemseo.problems.mdo.sobieski.core.design_space import SobieskiDesignSpace\nfrom gemseo.problems.mdo.sobieski.disciplines import SobieskiAerodynamics\nfrom gemseo.problems.mdo.sobieski.disciplines import SobieskiMission\nfrom gemseo.problems.mdo.sobieski.disciplines import SobieskiPropulsion\nfrom gemseo.problems.mdo.sobieski.disciplines import SobieskiStructure\nfrom gemseo.scenarios.doe_scenario import DOEScenario\nfrom gemseo.scenarios.mdo_scenario import MDOScenario\n\nfrom gemseo_jax.problems.sobieski.chain import JAXSobieskiChain\n\n# Deactivate some checkers to speed up calculations in presence of cheap disciplines.\nconfigure(False, False, True, False, False, False, False)\n\nconfigure_logger()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## DOE\n\n### MDF with MDAJacobi\n\n#### Solve the Sobieski's SSBJ problem with JAX\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "chain = JAXSobieskiChain()\ndoe_scenario = DOEScenario(\n    [chain],\n    \"y_4\",\n    SobieskiDesignSpace(),\n    maximize_objective=True,\n    formulation_name=\"MDF\",\n)\ndoe_scenario.add_constraint(\"g_1\", \"ineq\")\ndoe_scenario.add_constraint(\"g_2\", \"ineq\")\ndoe_scenario.add_constraint(\"g_3\", \"ineq\")\ndoe_scenario.execute(algo_name=\"OT_OPT_LHS\", n_samples=100)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "#### Solve the Sobieski's SSBJ problem with NumPy\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "doe_scenario = DOEScenario(\n    [\n        SobieskiStructure(),\n        SobieskiAerodynamics(),\n        SobieskiPropulsion(),\n        SobieskiMission(),\n    ],\n    \"y_4\",\n    SobieskiDesignSpace(),\n    maximize_objective=True,\n    formulation_name=\"MDF\",\n)\ndoe_scenario.add_constraint(\"g_1\", \"ineq\")\ndoe_scenario.add_constraint(\"g_2\", \"ineq\")\ndoe_scenario.add_constraint(\"g_3\", \"ineq\")\ndoe_scenario.execute(algo_name=\"OT_OPT_LHS\", n_samples=100)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "#### Conclusion\nJAX is about 3 times faster than NumPy\nin the case of a sampling loop with a Jacobi MDA.\n\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### MDF with MDAGaussSeidel\n\n#### Solve the Sobieski's SSBJ problem with JAX\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "chain = JAXSobieskiChain()\ndoe_scenario = DOEScenario(\n    [chain],\n    \"y_4\",\n    SobieskiDesignSpace(),\n    maximize_objective=True,\n    formulation_name=\"MDF\",\n    main_mda_settings={\"inner_mda_name\": \"MDAGaussSeidel\"},\n)\ndoe_scenario.add_constraint(\"g_1\", \"ineq\")\ndoe_scenario.add_constraint(\"g_2\", \"ineq\")\ndoe_scenario.add_constraint(\"g_3\", \"ineq\")\ndoe_scenario.execute(algo_name=\"OT_OPT_LHS\", n_samples=100)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "#### Solve the Sobieski's SSBJ problem with NumPy\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "doe_scenario = DOEScenario(\n    [\n        SobieskiStructure(),\n        SobieskiAerodynamics(),\n        SobieskiPropulsion(),\n        SobieskiMission(),\n    ],\n    \"y_4\",\n    SobieskiDesignSpace(),\n    maximize_objective=True,\n    formulation_name=\"MDF\",\n    main_mda_settings={\"inner_mda_name\": \"MDAGaussSeidel\"},\n)\ndoe_scenario.add_constraint(\"g_1\", \"ineq\")\ndoe_scenario.add_constraint(\"g_2\", \"ineq\")\ndoe_scenario.add_constraint(\"g_3\", \"ineq\")\ndoe_scenario.execute(algo_name=\"OT_OPT_LHS\", n_samples=100)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "#### Conclusion\nJAX is about 3 times faster than NumPy\nin the case of a sampling loop with a Gauss-Seidel MDA.\n\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Optimization\n\n### MDF with MDAJacobi\n\n#### Solve the Sobieski's SSBJ problem with JAX\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "chain = JAXSobieskiChain()\nmod_scenario = MDOScenario(\n    [chain],\n    \"y_4\",\n    SobieskiDesignSpace(),\n    maximize_objective=True,\n    formulation_name=\"MDF\",\n)\nmod_scenario.add_constraint(\"g_1\", \"ineq\")\nmod_scenario.add_constraint(\"g_2\", \"ineq\")\nmod_scenario.add_constraint(\"g_3\", \"ineq\")\nmod_scenario.execute(algo_name=\"NLOPT_COBYLA\", max_iter=100)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "#### Solve the Sobieski's SSBJ problem with NumPy\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "mod_scenario = MDOScenario(\n    [\n        SobieskiStructure(),\n        SobieskiAerodynamics(),\n        SobieskiPropulsion(),\n        SobieskiMission(),\n    ],\n    \"y_4\",\n    SobieskiDesignSpace(),\n    maximize_objective=True,\n    formulation_name=\"MDF\",\n)\nmod_scenario.add_constraint(\"g_1\", \"ineq\")\nmod_scenario.add_constraint(\"g_2\", \"ineq\")\nmod_scenario.add_constraint(\"g_3\", \"ineq\")\nmod_scenario.execute(algo_name=\"NLOPT_COBYLA\", max_iter=100)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "#### Conclusion\nJAX is about 4 times faster than NumPy\nin the case of an optimization loop with a Jacobi MDA.\n\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### MDF with MDAGaussSeidel\n\n#### Solve the Sobieski's SSBJ problem with JAX\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "chain = JAXSobieskiChain()\nmod_scenario = MDOScenario(\n    [chain],\n    \"y_4\",\n    SobieskiDesignSpace(),\n    maximize_objective=True,\n    formulation_name=\"MDF\",\n    main_mda_settings={\"inner_mda_name\": \"MDAGaussSeidel\"},\n)\nmod_scenario.add_constraint(\"g_1\", \"ineq\")\nmod_scenario.add_constraint(\"g_2\", \"ineq\")\nmod_scenario.add_constraint(\"g_3\", \"ineq\")\nmod_scenario.execute(algo_name=\"NLOPT_COBYLA\", max_iter=100)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "#### Solve the Sobieski's SSBJ problem with NumPy\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "mod_scenario = MDOScenario(\n    [\n        SobieskiStructure(),\n        SobieskiAerodynamics(),\n        SobieskiPropulsion(),\n        SobieskiMission(),\n    ],\n    \"y_4\",\n    SobieskiDesignSpace(),\n    maximize_objective=True,\n    formulation_name=\"MDF\",\n    main_mda_settings={\"inner_mda_name\": \"MDAGaussSeidel\"},\n)\nmod_scenario.add_constraint(\"g_1\", \"ineq\")\nmod_scenario.add_constraint(\"g_2\", \"ineq\")\nmod_scenario.add_constraint(\"g_3\", \"ineq\")\nmod_scenario.execute(algo_name=\"NLOPT_COBYLA\", max_iter=100)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "#### Conclusion\nJAX is about 3 times faster than NumPy\nin the case of an optimization loop with a Gauss-Seidel MDA.\n\n"
      ]
    }
  ],
  "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
}