{
 "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\n",
    "except ImportError:\n",
    "    !python -m pip install openmdao[notebooks]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# EQConstraintComp\n",
    "\n",
    "`EQConstraintComp` is a specialized component that provides a simple way to implement\n",
    "equality constraints.\n",
    "\n",
    "You can add one or more outputs to an `EQConstraintComp` that compute the difference\n",
    "between a pair of input values for the purposes of driving the two inputs to equality. It\n",
    "computes the output value as:\n",
    "\n",
    "$$\n",
    "  name_{output} = \\frac{ name_{mult} \\times name_{lhs} - name_{rhs} }{f_{norm}(name_{rhs})}\n",
    "$$\n",
    "\n",
    "The normalization function $f_{norm}(name_{rhs})$ takes one of the following forms:\n",
    "\n",
    "$$\n",
    "  f_{norm}(name_{rhs}) =\n",
    "  \\begin{cases}\n",
    "    \\left| name_{rhs} \\right|,      & \\text{if normalize and } \\left| name_{rhs} \\right| \\geq 2 \\\\\n",
    "    0.25 name_{rhs}^2 + 1,      & \\text{if normalize and } \\left| name_{rhs} \\right| < 2 \\\\\n",
    "    1,      & \\text{if not normalize}\n",
    "  \\end{cases}\n",
    "$$ \n",
    "\n",
    "The following inputs and outputs are associated with each output variable.\n",
    "\n",
    "```{eval-rst}\n",
    "    =========== ======= ====================================================\n",
    "    Name        I/O     Description\n",
    "    =========== ======= ====================================================\n",
    "    {name}      output  output variable\n",
    "    lhs:{name}  input   left-hand side of difference equation\n",
    "    rhs:{name}  input   right-hand side of difference equation\n",
    "    mult:{name} input   left-hand side multiplier of difference equation\n",
    "    =========== ======= ====================================================\n",
    "```\n",
    "\n",
    "The default value for the `rhs:{name}` input can be set to via the\n",
    "`rhs_val` argument (see arguments below). If the rhs value is fixed (e.g. 0),\n",
    "then the input can be left unconnected. The `lhs:{name}` input must always have\n",
    "something connected to it.\n",
    "\n",
    "The multiplier is optional and will default to 1.0 if not connected.\n",
    "\n",
    "`EQConstraintComp` supports vectorized outputs. Simply provide a default\n",
    "value or shape when adding the output that reflects the correct shape.\n",
    "\n",
    "You can provide the arguments to create an output variable when instantiating an\n",
    "`EQConstraintComp` or you can use the ``add_eq_output`` method to create one\n",
    "or more outputs after instantiation.  The constructor accepts all the same arguments\n",
    "as the ``add_eq_output`` method:\n",
    "\n",
    "```{eval-rst}\n",
    "    .. automethod:: openmdao.components.eq_constraint_comp.EQConstraintComp.add_eq_output\n",
    "       :noindex:\n",
    "```\n",
    "\n",
    "Note that the `kwargs` arguments can include any of the keyword arguments normally available\n",
    "when creating an output variable with the\n",
    "[add_output](../../../_srcdocs/packages/core/component.html#openmdao.core.component.Component.add_output) method of a `Component`.\n",
    "\n",
    "\n",
    "## Example: Sellar IDF\n",
    "\n",
    "The following example shows an Individual Design Feasible (IDF) architecture for the\n",
    "[Sellar](../../../basic_user_guide/multidisciplinary_optimization/sellar.ipynb) problem that demonstrates the use of an `EQConstraintComp`.\n",
    "\n",
    "In IDF, the direct coupling between the disciplines is removed and the coupling variables\n",
    "are added to the optimizer’s design variables. The algorithm calls for two new equality\n",
    "constraints that enforce the coupling between the disciplines. This ensures that the final\n",
    "optimized solution is feasible, though it is achieved through the optimizer instead of\n",
    "using a solver.  The two new equality constraints are implemented in this example with\n",
    "an `EQConstraintComp`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import openmdao.api as om\n",
    "\n",
    "\n",
    "class SellarIDF(om.Group):\n",
    "    \"\"\"\n",
    "    Individual Design Feasible (IDF) architecture for the Sellar problem.\n",
    "    \"\"\"\n",
    "    def setup(self):\n",
    "        # construct the Sellar model with `y1` and `y2` as independent variables\n",
    "\n",
    "        self.set_input_defaults('x', 5.)\n",
    "        self.set_input_defaults('y1', 5.)\n",
    "        self.set_input_defaults('y2', 5.)\n",
    "        self.set_input_defaults('z', np.array([2., 0.]))\n",
    "\n",
    "        self.add_subsystem('d1', SellarDis1withDerivatives(), promotes_inputs=['x', 'z', 'y2'])\n",
    "        self.add_subsystem('d2', SellarDis2withDerivatives(), promotes_inputs=['y1', 'z'])\n",
    "\n",
    "        self.add_subsystem('obj_cmp', om.ExecComp('obj = x**2 + z[1] + y1 + exp(-y2)',\n",
    "                           x=0., z=np.array([0., 0.])), promotes_inputs=['x', 'z', 'y1', 'y2'])\n",
    "\n",
    "        self.add_subsystem('con_cmp1', om.ExecComp('con1 = 3.16 - y1'), promotes_inputs=['y1'])\n",
    "        self.add_subsystem('con_cmp2', om.ExecComp('con2 = y2 - 24.0'), promotes_inputs=['y2'])\n",
    "\n",
    "        # rather than create a cycle by connecting d1.y1 to d2.y1 and d2.y2 to d1.y2\n",
    "        # we will constrain y1 and y2 to be equal for the two disciplines\n",
    "\n",
    "        equal = om.EQConstraintComp()\n",
    "        self.add_subsystem('equal', equal, promotes_inputs=[('lhs:y1', 'y1'), ('lhs:y2', 'y2')])\n",
    "\n",
    "        equal.add_eq_output('y1', add_constraint=True)\n",
    "        equal.add_eq_output('y2', add_constraint=True)\n",
    "\n",
    "        self.connect('d1.y1', 'equal.rhs:y1')\n",
    "        self.connect('d2.y2', 'equal.rhs:y2')\n",
    "\n",
    "        # the driver will effectively solve the cycle\n",
    "        # by satisfying the equality constraints\n",
    "\n",
    "        self.add_design_var('x', lower=0., upper=5.)\n",
    "        self.add_design_var('y1', lower=0., upper=5.)\n",
    "        self.add_design_var('y2', lower=0., upper=5.)\n",
    "        self.add_design_var('z', lower=np.array([-5., 0.]), upper=np.array([5., 5.]))\n",
    "        self.add_objective('obj_cmp.obj')\n",
    "        self.add_constraint('con_cmp1.con1', upper=0.)\n",
    "        self.add_constraint('con_cmp2.con2', upper=0.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "from openmdao.utils.notebook_utils import get_code\n",
    "from myst_nb import glue\n",
    "glue(\"code_src14\", get_code(\"openmdao.test_suite.components.sellar_feature.SellarIDF\"), display=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{Admonition} `SellarIDF` class definition \n",
    ":class: dropdown\n",
    "\n",
    "{glue:}`code_src14`\n",
    ":::"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from openmdao.test_suite.components.sellar_feature import SellarIDF\n",
    "\n",
    "prob = om.Problem(model=SellarIDF())\n",
    "prob.driver = om.ScipyOptimizeDriver(optimizer='SLSQP', disp=True)\n",
    "prob.setup()\n",
    "prob.run_driver();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(prob.get_val('x'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print([prob.get_val('y1'), prob.get_val('d1.y1')])\n",
    "print([prob.get_val('y2'), prob.get_val('y2')])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(prob.get_val('z'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(prob.get_val('obj_cmp.obj'))"
   ]
  },
  {
   "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",
    "assert_near_equal(prob.get_val('x'), 0., 1e-5)\n",
    "\n",
    "assert_near_equal([prob.get_val('y1'), prob.get_val('d1.y1')], [[3.16], [3.16]], 1e-5)\n",
    "assert_near_equal([prob.get_val('y2'), prob.get_val('y2')], [[3.7552778], [3.7552778]], 1e-5)\n",
    "\n",
    "assert_near_equal(prob.get_val('z'), [1.977639, 0.], 1e-5)\n",
    "\n",
    "assert_near_equal(prob.get_val('obj_cmp.obj'), 3.18339395045, 1e-5)"
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Tags",
  "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.8.1"
  },
  "orphan": true
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
