{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "active-ipynb",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "try:\n",
    "    from openmdao.utils.notebook_utils import notebook_mode  # noqa: F401\n",
    "except ImportError:\n",
    "    !python -m pip install openmdao[notebooks]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# SubmodelComp\n",
    "\n",
    "`SubmodelComp` provides a way to evaluate an OpenMDAO system within an OpenMDAO system.\n",
    "\n",
    "# SubmodelComp Options"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import openmdao.api as om\n",
    "om.show_options_table(om.SubmodelComp(problem=om.Problem()))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# SubmodelComp Constructor\n",
    "\n",
    "The call signature for the `SubmodelComp` constructor is\n",
    "\n",
    "```{eval-rst}\n",
    "    .. automethod:: openmdao.components.submodel_comp.SubmodelComp.__init__\n",
    "        :noindex:\n",
    "```\n",
    "\n",
    "# Using the SubmodelComp\n",
    "\n",
    "`SubmodelComp` allows you to add a component that contains a complete `Problem` in itself. The only required argument is `problem`. This takes an instantiated problem that already has a model. That way, the user can declare any desired options for the problem and have it used for the underlying subproblem.\n",
    "\n",
    "Other arguments include `inputs`, `outputs`, and `reports`. Input and output argument formats are discussed in the next section, but if the inputs and outputs are not specified in the `SubmodelComp`'s instantiation, they will have to be added later using the `add_input` and `add_output` methods. The `reports` option controls report generation for the underlying subproblem. `SubmodelComp` performs more efficiently if it doesn't have to write reports, but if reports are desired, then the `reports` argument can be set to `True`.\n",
    "\n",
    "`SubmodelComp` provides a way to evaluate an OpenMDAO model within an OpenMDAO system. Unlike a typical subsystem, using a subproblem exposes only a limited number of inputs and outputs of the underlying model to the parent system. This may be beneficial from a performance standpoint when there are many inputs and outputs of the internal model that aren't needed in the top-level model.\n",
    "\n",
    "Note that this means that nonlinear and linear solvers are needed on the subproblem's model if implicit behavior is present, since the internal variables are invisible to solvers in the parent system.\n",
    "\n",
    "At this time, OpenMDAO does not evaluate gradients through the optimization process, so a gradient-based driver at the outer level cannot access derivatives across the optimization on the inner level.\n",
    "\n",
    "## Inputs and Outputs\n",
    "\n",
    "Inputs and outputs can be passed into `SubmodelComp` as a list of `str`, `tuple`, or both. If the list element is `str`, then it must specify the promoted name of a variable in the model. All inputs and outputs you wish to use must be promoted to the model level. If the list element is `tuple`, then the first element of the `tuple` must be the group level promoted name or the absolute name of the variable, and the second element is an alias that `SubmodelComp` will use to refer to the variable. For example:\n",
    "\n",
    "    (path.to.var, desired_var_name)\n",
    "    or (var, desired_var_name) if the variable is already promoted at the group level\n",
    "\n",
    "Input and output names can also include wildcards if the type is `str`.\n",
    "\n",
    "If not specified at instantiation, inputs and outputs will have to be declared using the `add_input` and `add_output` methods.\n",
    "\n",
    "# Example: Using SubmodelComp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import openmdao.api as om\n",
    "from numpy import pi\n",
    "\n",
    "p = om.Problem()\n",
    "\n",
    "# create a model that computes x = r*cos(theta)\n",
    "submodel1 = om.Group()\n",
    "submodel1.add_subsystem('sub1_ivc_r', om.IndepVarComp('r', 1.),\n",
    "                        promotes_outputs=['r'])\n",
    "submodel1.add_subsystem('sub1_ivc_theta', om.IndepVarComp('theta', pi),\n",
    "                        promotes_outputs=['theta'])\n",
    "submodel1.add_subsystem('subComp1', om.ExecComp('x = r*cos(theta)'),\n",
    "                        promotes_inputs=['r', 'theta'],\n",
    "                        promotes_outputs=['x'])\n",
    "\n",
    "# create a model that computes y = r*sin(theta)\n",
    "submodel2 = om.Group()\n",
    "submodel2.add_subsystem('sub2_ivc_r', om.IndepVarComp('r', 2),\n",
    "                        promotes_outputs=['r'])\n",
    "submodel2.add_subsystem('sub2_ivc_theta', om.IndepVarComp('theta', pi/2),\n",
    "                        promotes_outputs=['theta'])\n",
    "submodel2.add_subsystem('subComp2', om.ExecComp('y = r*sin(theta)'),\n",
    "                        promotes_inputs=['r', 'theta'],\n",
    "                        promotes_outputs=['y'])\n",
    "\n",
    "# wrap the x and y models as sub-problems\n",
    "subprob1 = om.Problem(model=submodel1)\n",
    "subprob2 = om.Problem(model=submodel2)\n",
    "\n",
    "subcomp1 = om.SubmodelComp(problem=subprob1,\n",
    "                            inputs=['r', 'theta'], outputs=['x'])\n",
    "subcomp2 = om.SubmodelComp(problem=subprob2,\n",
    "                            inputs=['r', 'theta'], outputs=['y'])\n",
    "\n",
    "# add both sub-problems and a Component that uses their outputs\n",
    "p.model.add_subsystem('sub1', subcomp1, \n",
    "                      promotes_inputs=['r','theta'],\n",
    "                      promotes_outputs=['x'])\n",
    "p.model.add_subsystem('sub2', subcomp2, \n",
    "                      promotes_inputs=['r','theta'],\n",
    "                      promotes_outputs=['y'])\n",
    "p.model.add_subsystem('supComp', om.ExecComp('z = x**2 + y'),\n",
    "                      promotes_inputs=['x', 'y'],\n",
    "                      promotes_outputs=['z'])\n",
    "\n",
    "p.model.set_input_defaults('r', 1)\n",
    "p.model.set_input_defaults('theta', pi)\n",
    "\n",
    "p.setup(force_alloc_complex=True)\n",
    "\n",
    "p.run_model()\n",
    "\n",
    "cpd = p.check_partials(method='cs', out_stream=None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "allow-assert",
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "assert(sorted(cpd['sub1'].keys()) == [('x', 'r'), ('x', 'theta')])\n",
    "assert(sorted(cpd['sub2'].keys()) == [('y', 'r'), ('y', 'theta')])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "source": [
    "# Example: Variable Aliases"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import openmdao.api as om\n",
    "\n",
    "# create a Problem that will be nested in another Problem\n",
    "subprob = om.Problem()\n",
    "submodel = subprob.model\n",
    "submodel.add_subsystem('subsys', om.ExecComp('z = x**2 + y**2'))\n",
    "\n",
    "# wrap that Problem in a SubmodelComp, renaming the inputs and output\n",
    "subcomp = om.SubmodelComp(problem=subprob,\n",
    "                          inputs=[('subsys.x', 'a'), ('subsys.y', 'b')],\n",
    "                          outputs=[('subsys.z', 'c')])\n",
    "\n",
    "# create a top-level Problem that contains the SubmodelComp\n",
    "p = om.Problem()\n",
    "p.model.add_subsystem('subcomp', subcomp, promotes_inputs=['a', 'b'], promotes_outputs=['c'])\n",
    "p.setup()\n",
    "\n",
    "p.set_val('a', 1)\n",
    "p.set_val('b', 2)\n",
    "\n",
    "p.run_model()\n",
    "\n",
    "inputs = p.model.list_inputs()\n",
    "outputs = p.model.list_outputs()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "from openmdao.utils.assert_utils import assert_near_equal\n",
    "\n",
    "inputs = dict(inputs)\n",
    "assert_near_equal(inputs['subcomp.a']['val'], 1.)\n",
    "assert_near_equal(inputs['subcomp.b']['val'], 2.)\n",
    "\n",
    "outputs = dict(outputs)\n",
    "assert_near_equal(outputs['subcomp.c']['val'], 5.)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Example: Using Variables with the Same Name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import openmdao.api as om\n",
    "\n",
    "p = om.Problem()\n",
    "\n",
    "model = om.Group()\n",
    "\n",
    "model.add_subsystem('x1Comp', om.ExecComp('x1 = x*3'))\n",
    "model.add_subsystem('x2Comp', om.ExecComp('x2 = x**3'))\n",
    "model.add_subsystem('model', om.ExecComp('z = x1**2 + x2**2'))\n",
    "\n",
    "model.connect('x1Comp.x1', 'model.x1')\n",
    "model.connect('x2Comp.x2', 'model.x2')\n",
    "\n",
    "subprob = om.Problem()\n",
    "subprob.model.add_subsystem('submodel', model)\n",
    "\n",
    "comp = om.SubmodelComp(problem=subprob, \n",
    "                       inputs=[('submodel.x1Comp.x', 'x'), ('submodel.x2Comp.x', 'y')],\n",
    "                       outputs=[('submodel.model.z', 'z')])\n",
    "\n",
    "p.model.add_subsystem('comp', comp)\n",
    "\n",
    "p.model.set_input_defaults('comp.x', 1)\n",
    "p.model.set_input_defaults('comp.y', 2)\n",
    "\n",
    "p.setup()\n",
    "\n",
    "p.run_model()\n",
    "\n",
    "inputs = p.model.list_inputs()\n",
    "outputs = p.model.list_outputs()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "from openmdao.utils.assert_utils import assert_near_equal\n",
    "\n",
    "inputs = dict(inputs)\n",
    "assert_near_equal(inputs['comp.x']['val'], 1.)\n",
    "assert_near_equal(inputs['comp.y']['val'], 2.)\n",
    "\n",
    "outputs = dict(outputs)\n",
    "assert_near_equal(outputs['comp.z']['val'], 73.)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Example: Adding Inputs and Outputs Using Wildcards"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import openmdao.api as om\n",
    "\n",
    "p = om.Problem()\n",
    "\n",
    "submodel = om.Group()\n",
    "submodel.add_subsystem('sub', om.ExecComp('z = x1**2 + x2**2 + x3**2'), promotes=['*'])\n",
    "\n",
    "subprob = om.Problem()\n",
    "subprob.model.add_subsystem('submodel', submodel, promotes=['*'])\n",
    "\n",
    "comp = om.SubmodelComp(problem=subprob, inputs=['x*'], outputs=['*'])\n",
    "\n",
    "p.model.add_subsystem('comp', comp, promotes_inputs=['*'], promotes_outputs=['*'])\n",
    "p.setup()\n",
    "\n",
    "p.set_val('x1', 1)\n",
    "p.set_val('x2', 2)\n",
    "p.set_val('x3', 3)\n",
    "\n",
    "p.run_model()\n",
    "\n",
    "inputs = p.model.list_inputs()\n",
    "outputs = p.model.list_outputs()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "from openmdao.utils.assert_utils import assert_near_equal\n",
    "\n",
    "inputs = dict(inputs)\n",
    "assert_near_equal(inputs['comp.x1']['val'], 1.)\n",
    "assert_near_equal(inputs['comp.x2']['val'], 2.)\n",
    "assert_near_equal(inputs['comp.x3']['val'], 3.)\n",
    "\n",
    "outputs = dict(outputs)\n",
    "assert_near_equal(outputs['comp.z']['val'], 14.)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Example: Adding Inputs and Outputs After Instantiation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import openmdao.api as om\n",
    "\n",
    "p = om.Problem()\n",
    "\n",
    "# create a model that computes x = r*cos(theta)\n",
    "submodel1 = om.Group()\n",
    "submodel1.add_subsystem('sub1_ivc_r', om.IndepVarComp('r', 1.),\n",
    "                        promotes_outputs=['r'])\n",
    "submodel1.add_subsystem('sub1_ivc_theta', om.IndepVarComp('theta', pi),\n",
    "                        promotes_outputs=['theta'])\n",
    "submodel1.add_subsystem('subComp1', om.ExecComp('x = r*cos(theta)'))\n",
    "\n",
    "# create a model that computes y = r*sin(theta)\n",
    "submodel2 = om.Group()\n",
    "submodel2.add_subsystem('sub2_ivc_r', om.IndepVarComp('r', 2),\n",
    "                        promotes_outputs=['r'])\n",
    "submodel2.add_subsystem('sub2_ivc_theta', om.IndepVarComp('theta', pi/2),\n",
    "                        promotes_outputs=['theta'])\n",
    "submodel2.add_subsystem('subComp2', om.ExecComp('y = r*sin(theta)'))\n",
    "\n",
    "# wrap the x and y models as sub-problems\n",
    "subprob1 = om.Problem()\n",
    "subprob1.model.add_subsystem('submodel1', submodel1, promotes=['*'])\n",
    "subprob2 = om.Problem()\n",
    "subprob2.model.add_subsystem('submodel2', submodel2, promotes=['*'])\n",
    "\n",
    "comp1 = om.SubmodelComp(problem=subprob1)\n",
    "comp2 = om.SubmodelComp(problem=subprob2)\n",
    "\n",
    "# add inputs and outputs for the sub-problems\n",
    "comp1.add_input('subComp1.r', name='r')\n",
    "comp1.add_input('subComp1.theta', name='theta')\n",
    "comp2.add_input('subComp2.r', name='r')\n",
    "comp2.add_input('subComp2.theta', name='theta')\n",
    "\n",
    "comp1.add_output('subComp1.x', name='x')\n",
    "comp2.add_output('subComp2.y', name='y')\n",
    "\n",
    "# add both sub-problems and a Component that uses their outputs\n",
    "p.model.add_subsystem('sub1', comp1, \n",
    "                      promotes_inputs=['r','theta'],\n",
    "                      promotes_outputs=['x'])\n",
    "p.model.add_subsystem('sub2', comp2, \n",
    "                      promotes_inputs=['r','theta'],\n",
    "                      promotes_outputs=['y'])\n",
    "p.model.add_subsystem('supComp', om.ExecComp('z = x**2 + y'),\n",
    "                      promotes_inputs=['x', 'y'],\n",
    "                      promotes_outputs=['z'])\n",
    "\n",
    "p.model.set_input_defaults('r', 1)\n",
    "p.model.set_input_defaults('theta', pi)\n",
    "\n",
    "p.setup(force_alloc_complex=True)\n",
    "\n",
    "p.run_model()\n",
    "\n",
    "cpd = p.check_partials(method='cs', out_stream=None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output",
     "allow-assert"
    ]
   },
   "outputs": [],
   "source": [
    "assert(sorted(cpd['sub1'].keys()) == [('x', 'r'), ('x', 'theta')])\n",
    "assert(sorted(cpd['sub2'].keys()) == [('y', 'r'), ('y', 'theta')])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Example: Adding Inputs and Outputs in Configure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import openmdao.api as om\n",
    "\n",
    "class Subsys1(om.Group):\n",
    "    def setup(self):\n",
    "        model = om.Group()\n",
    "        comp = om.ExecComp('x = r*cos(theta)')\n",
    "        model.add_subsystem('comp', comp, promotes_inputs=['r', 'theta'],\n",
    "                            promotes_outputs=['x'])\n",
    "        subprob = om.Problem()\n",
    "        subprob.model.add_subsystem('model', model)\n",
    "        subprob.model.promotes('model', any=['*'])\n",
    "        self.add_subsystem('submodel1', om.SubmodelComp(problem=subprob))\n",
    "\n",
    "    def configure(self):\n",
    "        self._get_subsystem('submodel1').add_input('r')\n",
    "        self._get_subsystem('submodel1').add_input('theta')\n",
    "        self._get_subsystem('submodel1').add_output('x')\n",
    "\n",
    "        self.promotes('submodel1', ['r', 'theta', 'x'])\n",
    "\n",
    "class Subsys2(om.Group):\n",
    "    def setup(self):\n",
    "        model = om.Group()\n",
    "        comp = om.ExecComp('y = r*sin(theta)')\n",
    "        model.add_subsystem('comp', comp, promotes_inputs=['r', 'theta'],\n",
    "                            promotes_outputs=['y'])\n",
    "        subprob = om.Problem()\n",
    "        subprob.model.add_subsystem('model', model)\n",
    "        subprob.model.promotes('model', any=['*'])\n",
    "        self.add_subsystem('submodel2', om.SubmodelComp(problem=subprob))\n",
    "\n",
    "    def configure(self):\n",
    "        self._get_subsystem('submodel2').add_input('r')\n",
    "        self._get_subsystem('submodel2').add_input('theta')\n",
    "        self._get_subsystem('submodel2').add_output('y')\n",
    "\n",
    "        self.promotes('submodel2', ['r', 'theta', 'y'])\n",
    "\n",
    "p = om.Problem()\n",
    "p.model.add_subsystem('sub1', Subsys1(), \n",
    "                      promotes_inputs=['r', 'theta'],\n",
    "                      promotes_outputs=['x'])\n",
    "p.model.add_subsystem('sub2', Subsys2(), \n",
    "                      promotes_inputs=['r', 'theta'],\n",
    "                      promotes_outputs=['y'])\n",
    "p.model.add_subsystem('supComp', om.ExecComp('z = x**2 + y'),\n",
    "                      promotes_inputs=['x', 'y'],\n",
    "                      promotes_outputs=['z'])\n",
    "\n",
    "p.setup(force_alloc_complex=True)\n",
    "\n",
    "p.set_val('r', 1)\n",
    "p.set_val('theta', pi)\n",
    "\n",
    "p.run_model()\n",
    "cpd = p.check_partials(method='cs', out_stream=None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "assert(sorted(cpd['sub1.submodel1'].keys()) == [('x', 'r'), ('x', 'theta')])\n",
    "assert(sorted(cpd['sub2.submodel2'].keys()) == [('y', 'r'), ('y', 'theta')])"
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Tags",
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.11.3"
  },
  "orphan": true
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
