{
 "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": [
    "# Discrete Variables\n",
    "\n",
    "There may be times when it’s necessary to pass variables that are not floats or float arrays between components. These variables can be declared as discrete variables. A discrete variable can be any picklable python object.\n",
    "\n",
    "In explicit and implicit components, the user must call `add_discrete_input` and `add_discrete_output` to declare discrete variables in the `setup` method.\n",
    "\n",
    "## Methods for Adding Discrete Variables\n",
    "\n",
    "Here are the methods used to add discrete variables to components.\n",
    "\n",
    "```{eval-rst}\n",
    "    .. automethod:: openmdao.core.component.Component.add_discrete_input\n",
    "        :noindex:\n",
    "```\n",
    "```{eval-rst}\n",
    "    .. automethod:: openmdao.core.component.Component.add_discrete_output\n",
    "        :noindex:\n",
    "```\n",
    "\n",
    "## Discrete Variable Considerations\n",
    "\n",
    "Discrete variables, like continuous ones, can be connected to each other using the `connect` function or by promoting an input and an output to the same name. The type of the output must be a valid subclass of the type of the input or the connection will raise an exception.\n",
    "\n",
    "```{warning}\n",
    "If a model computes derivatives and any of those derivatives depend on the value of a discrete output variable, an exception will be raised.\n",
    "```\n",
    "\n",
    "If a component or group contains discrete variables, then the discrete inputs and/or outputs will be passed to the relevant API functions. In general, if nonlinear inputs are passed to a function, then a discrete inputs argument will be added. If nonlinear outputs are passed, then a discrete outputs argument will be added. The signatures of the affected functions are shown below:\n",
    "\n",
    "```{eval-rst}\n",
    "    .. automethod:: openmdao.core.explicitcomponent.ExplicitComponent.compute\n",
    "        :noindex:\n",
    "```\n",
    "```{eval-rst}\n",
    "    .. automethod:: openmdao.core.explicitcomponent.ExplicitComponent.compute_jacvec_product\n",
    "        :noindex:\n",
    "```\n",
    "```{eval-rst}\n",
    "    .. automethod:: openmdao.core.explicitcomponent.ExplicitComponent.compute_partials\n",
    "        :noindex:\n",
    "```\n",
    "```{eval-rst}\n",
    "    .. automethod:: openmdao.core.implicitcomponent.ImplicitComponent.apply_nonlinear\n",
    "        :noindex:\n",
    "```\n",
    "```{eval-rst}\n",
    "    .. automethod:: openmdao.core.implicitcomponent.ImplicitComponent.guess_nonlinear\n",
    "        :noindex:\n",
    "```\n",
    "```{eval-rst}\n",
    "    .. automethod:: openmdao.core.implicitcomponent.ImplicitComponent.linearize\n",
    "        :noindex:\n",
    "```\n",
    "```{eval-rst}\n",
    "    .. automethod:: openmdao.core.group.Group.guess_nonlinear\n",
    "        :noindex:\n",
    "```\n",
    "\n",
    "## Discrete Variable Examples\n",
    "\n",
    "An example is given below that shows an explicit component that has a discrete input along with continuous inputs and outputs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import openmdao.api as om\n",
    "\n",
    "\n",
    "class BladeSolidity(om.ExplicitComponent):\n",
    "    def setup(self):\n",
    "\n",
    "        # Continuous Inputs\n",
    "        self.add_input('r_m', 1.0, units=\"ft\", desc=\"Mean radius\")\n",
    "        self.add_input('chord', 1.0, units=\"ft\", desc=\"Chord length\")\n",
    "\n",
    "        # Discrete Inputs\n",
    "        self.add_discrete_input('num_blades', 2, desc=\"Number of blades\")\n",
    "\n",
    "        # Continuous Outputs\n",
    "        self.add_output('blade_solidity', 0.0, desc=\"Blade solidity\")\n",
    "\n",
    "    def compute(self, inputs, outputs, discrete_inputs, discrete_outputs):\n",
    "\n",
    "        num_blades = discrete_inputs['num_blades']\n",
    "        chord = inputs['chord']\n",
    "        r_m = inputs['r_m']\n",
    "\n",
    "        outputs['blade_solidity'] = chord / (2.0 * np.pi * r_m / num_blades)\n",
    "\n",
    "# build the model\n",
    "prob = om.Problem()\n",
    "\n",
    "prob.model.add_subsystem('SolidityComp', BladeSolidity(),\n",
    "                         promotes_inputs=['r_m', 'chord', 'num_blades'])\n",
    "\n",
    "prob.setup()\n",
    "\n",
    "prob.set_val('num_blades', 2)\n",
    "prob.set_val('r_m', 3.2)\n",
    "prob.set_val('chord', .3)\n",
    "\n",
    "prob.run_model()\n",
    "\n",
    "# minimum value\n",
    "print(prob['SolidityComp.blade_solidity'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "from openmdao.utils.assert_utils import assert_near_equal\n",
    "assert_near_equal(prob['SolidityComp.blade_solidity'], 0.02984155, 1e-4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Similarly, discrete variables can be added to implicit components."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import openmdao.api as om\n",
    "\n",
    "class ImpWithInitial(om.ImplicitComponent):\n",
    "    \"\"\"\n",
    "    An implicit component to solve the quadratic equation: x^2 - 4x + 3\n",
    "    (solutions at x=1 and x=3)\n",
    "    \"\"\"\n",
    "    def setup(self):\n",
    "        self.add_input('a', val=1.)\n",
    "        self.add_input('b', val=-4.)\n",
    "        self.add_discrete_input('c', val=3)\n",
    "        self.add_output('x', val=5.)\n",
    "\n",
    "        self.declare_partials(of='*', wrt='*')\n",
    "\n",
    "    def apply_nonlinear(self, inputs, outputs, residuals, discrete_inputs, discrete_outputs):\n",
    "        a = inputs['a']\n",
    "        b = inputs['b']\n",
    "        c = discrete_inputs['c']\n",
    "        x = outputs['x']\n",
    "        residuals['x'] = a * x ** 2 + b * x + c\n",
    "\n",
    "    def linearize(self, inputs, outputs, partials, discrete_inputs, discrete_outputs):\n",
    "        a = inputs['a']\n",
    "        b = inputs['b']\n",
    "        x = outputs['x']\n",
    "\n",
    "        partials['x', 'a'] = x ** 2\n",
    "        partials['x', 'b'] = x\n",
    "        partials['x', 'x'] = 2 * a * x + b\n",
    "\n",
    "    def guess_nonlinear(self, inputs, outputs, resids, discrete_inputs, discrete_outputs):\n",
    "        # Default initial state of zero for x takes us to x=1 solution.\n",
    "        # Here we set it to a value that will take us to the x=3 solution.\n",
    "        outputs['x'] = 5\n",
    "\n",
    "prob = om.Problem()\n",
    "model = prob.model\n",
    "\n",
    "model.add_subsystem('comp', ImpWithInitial())\n",
    "\n",
    "model.comp.nonlinear_solver = om.NewtonSolver(solve_subsystems=False)\n",
    "model.comp.linear_solver = om.ScipyKrylov()\n",
    "\n",
    "prob.setup()\n",
    "prob.run_model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(prob.get_val('comp.x'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "assert_near_equal(prob.get_val('comp.x'), 3., 1e-4)"
   ]
  }
 ],
 "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
}
