{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    },
    "pycharm": {
     "name": "#%%\n"
    },
    "tags": [
     "remove-input",
     "remove-output",
     "active-ipynb"
    ]
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "from ipyparallel import Client, error  # noqa: F401\n",
    "cluster=Client(profile=\"mpi\")\n",
    "view=cluster[:]\n",
    "view.block=True\n",
    "\n",
    "try:\n",
    "    from openmdao.utils.notebook_utils import notebook_mode  # noqa: F401\n",
    "except ImportError:\n",
    "    !python -m pip install openmdao[notebooks]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Adding Constraints\n",
    "\n",
    "To add a constraint to an optimization, use the `add_constraint` method on System."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```{eval-rst}\n",
    "    .. automethod:: openmdao.core.system.System.add_constraint\n",
    "        :noindex:\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Specifying units\n",
    "\n",
    "You can specify units when adding a constraint. When this is done, the constraint value is converted from the target output’s units to the desired unit before giving it to the optimizer. If you also specify scaling, that scaling is applied after the unit conversion. Moreover, the upper and lower limits in the constraint definition should be specified using these units.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    },
    "pycharm": {
     "name": "#%%\n"
    },
    "tags": [
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "import openmdao.api as om\n",
    "\n",
    "prob = om.Problem()\n",
    "model = prob.model\n",
    "\n",
    "model.add_subsystem('comp1', om.ExecComp('y1 = 2.0*x',\n",
    "                                         x={'val': 2.0, 'units': 'degF'},\n",
    "                                         y1={'val': 2.0, 'units': 'degF'}),\n",
    "                    promotes=['x', 'y1'])\n",
    "\n",
    "model.add_subsystem('comp2', om.ExecComp('y2 = 3.0*x',\n",
    "                                         x={'val': 2.0, 'units': 'degF'},\n",
    "                                         y2={'val': 2.0, 'units': 'degF'}),\n",
    "                    promotes=['x', 'y2'])\n",
    "\n",
    "model.set_input_defaults('x', 35.0, units='degF')\n",
    "\n",
    "model.add_design_var('x', units='degC', lower=0.0, upper=100.0)\n",
    "model.add_constraint('y1', units='degC', lower=0.0, upper=100.0)\n",
    "model.add_objective('y2', units='degC')\n",
    "\n",
    "prob.setup()\n",
    "prob.run_driver()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    },
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "print('Model variables')\n",
    "print(prob.get_val('x', indices=[0]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    },
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "print(prob.get_val('comp2.y2', indices=[0]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    },
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "print(prob.get_val('comp1.y1', indices=[0]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    },
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "print('Driver variables')\n",
    "dv = prob.driver.get_design_var_values()\n",
    "print(dv['x'][0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    },
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "obj = prob.driver.get_objective_values(driver_scaling=True)\n",
    "print(obj['y2'][0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    },
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "con = prob.driver.get_constraint_values(driver_scaling=True)\n",
    "print(con['y1'][0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    },
    "pycharm": {
     "name": "#%%\n"
    },
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from openmdao.utils.assert_utils import assert_near_equal\n",
    "\n",
    "assert_near_equal(prob.get_val('x', indices=[0]), 35.)\n",
    "assert_near_equal(prob.get_val('y2', indices=[0]), 105.)\n",
    "assert_near_equal(prob.get_val('y1', indices=[0]), 70.)\n",
    "assert_near_equal(dv['x'][0], 1.6666666666666983)\n",
    "assert_near_equal(obj['y2'][0], 40.555555555555586)\n",
    "assert_near_equal(con['y1'][0], 21.111111111111143)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using the output of a distributed component as a constraint\n",
    "\n",
    "You can use an output of a distributed component as a constraint or an objective. OpenMDAO automatically collects the values from all processors and provides them to the driver.\n",
    "\n",
    "Here is an example where we perform optimization on a model that contains a `DistParabFeature` component that is distributed. The output is declared as a inequality constraint."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```{note}\n",
    "This feature requires MPI, and may not be able to be run on Colab or Binder.\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "hide-output"
    ]
   },
   "outputs": [],
   "source": [
    "import openmdao.api as om\n",
    "om.display_source(\"openmdao.test_suite.components.paraboloid_distributed.DistParabFeature\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%px\n",
    "\n",
    "import numpy as np\n",
    "import openmdao.api as om\n",
    "\n",
    "from openmdao.test_suite.components.paraboloid_distributed import DistParabFeature\n",
    "\n",
    "size = 7\n",
    "\n",
    "prob = om.Problem()\n",
    "model = prob.model\n",
    "\n",
    "ivc = om.IndepVarComp()\n",
    "ivc.add_output('x', np.ones(size))\n",
    "ivc.add_output('y', -1.42 * np.ones(size))\n",
    "\n",
    "model.add_subsystem('p', ivc, promotes=['*'])\n",
    "model.add_subsystem(\"parab\", DistParabFeature(arr_size=size), promotes=['*'])\n",
    "\n",
    "model.add_design_var('x', lower=-50.0, upper=50.0)\n",
    "model.add_constraint('f_xy', lower=0.0)\n",
    "model.add_objective('f_sum', index=-1)\n",
    "\n",
    "prob.driver = om.pyOptSparseDriver(optimizer='SLSQP')\n",
    "prob.setup()\n",
    "        \n",
    "prob.run_driver()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%px\n",
    "\n",
    "desvar = prob.get_val('p.x', get_remote=True)\n",
    "obj = prob.get_val('f_sum', get_remote=True)\n",
    "\n",
    "print(desvar)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%px\n",
    "\n",
    "print(obj)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    },
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "%%px\n",
    "\n",
    "from openmdao.utils.assert_utils import assert_near_equal\n",
    "\n",
    "assert_near_equal(prob.get_val('p.x'), np.array([2.65752672, 2.60433212, 2.51005989, 1.91021257, 1.3100763,  0.70992863, 0.10978096]), 1e-6 )\n",
    "assert_near_equal(prob.get_val('f_sum'),11.50150011, 1e-6 )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Adding multiple constraints on an array variable.\n",
    "\n",
    "Sometimes you have a variable and you would like to constrain difference parts of it in different ways. OpenMDAO maintains the constraint data in a dictionary that keys on the variable pathname, so we need to specify an additional \"alias\" argument when we declare any additional constraints.\n",
    "\n",
    "In the following example, the variable \"exec.z\" has an equality constraint on the first point and a lower bound on the end point. We constrain them by giving the second one an alias to differentiate it from the first one."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "p = om.Problem()\n",
    "\n",
    "exec = om.ExecComp(['y = x**2',\n",
    "                    'z = a + x**2'],\n",
    "                a={'shape': (1,)},\n",
    "                y={'shape': (101,)},\n",
    "                x={'shape': (101,)},\n",
    "                z={'shape': (101,)})\n",
    "\n",
    "p.model.add_subsystem('exec', exec)\n",
    "\n",
    "p.model.add_design_var('exec.a', lower=-1000, upper=1000)\n",
    "p.model.add_objective('exec.y', index=50)\n",
    "p.model.add_constraint('exec.z', indices=[0], equals=25)\n",
    "p.model.add_constraint('exec.z', indices=[-1], lower=20, alias=\"End_Constraint\")\n",
    "\n",
    "p.driver = om.pyOptSparseDriver()\n",
    "p.driver.options['optimizer'] = 'SLSQP'\n",
    "\n",
    "p.setup()\n",
    "\n",
    "p.set_val('exec.x', np.linspace(-10, 10, 101))\n",
    "\n",
    "p.run_driver()\n",
    "\n",
    "print(p.get_val('exec.z')[0], 25)\n",
    "print(p.get_val('exec.z')[50], -75)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "assert_near_equal(p.get_val('exec.z')[0], 25)\n",
    "assert_near_equal(p.get_val('exec.z')[50], -75)"
   ]
  }
 ],
 "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.10"
  },
  "orphan": true
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
