{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc7e3027",
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output",
     "active-ipynb"
    ]
   },
   "outputs": [],
   "source": [
    "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",
   "id": "25b7f5f2",
   "metadata": {},
   "source": [
    "# Approximating Partial Derivatives\n",
    "\n",
    "OpenMDAO allows you to specify analytic derivatives for your models, but it is not a requirement. If certain partial derivatives are not available, you can ask the framework to approximate the derivatives by using the `declare_partials` method inside `setup_partials`, and give it a method that is either ‘fd’ for finite difference or ‘cs’ for complex step.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98b87456",
   "metadata": {},
   "source": [
    "```{eval-rst}\n",
    "    .. automethod:: openmdao.core.component.Component.declare_partials\n",
    "        :noindex:\n",
    "```\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "404246bb",
   "metadata": {},
   "source": [
    "## Usage\n",
    "\n",
    "You may use glob patterns as arguments to `to` and `wrt`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e807dfd",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "import openmdao.api as om\n",
    "\n",
    "class FDPartialComp(om.ExplicitComponent):\n",
    "    def setup(self):\n",
    "        self.add_input('x', shape=(4,))\n",
    "        self.add_input('y', shape=(2,))\n",
    "        self.add_input('y2', shape=(2,))\n",
    "        self.add_output('f', shape=(2,))\n",
    "\n",
    "    def setup_partials(self):\n",
    "        self.declare_partials('f', 'y*', method='fd')\n",
    "        self.declare_partials('f', 'x', method='fd')\n",
    "\n",
    "    def compute(self, inputs, outputs):\n",
    "        f = outputs['f']\n",
    "\n",
    "        x = inputs['x']\n",
    "        y = inputs['y']\n",
    "\n",
    "        f[0] = x[0] + y[0]\n",
    "        f[1] = np.dot([0, 2, 3, 4], x) + y[1]\n",
    "\n",
    "model = om.Group()\n",
    "model.add_subsystem('example', FDPartialComp())\n",
    "\n",
    "problem = om.Problem(model=model)\n",
    "problem.setup()\n",
    "problem.run_model()\n",
    "totals = problem.compute_totals(['example.f'], ['example.x', 'example.y'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5865e28c",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(totals['example.f', 'example.x'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50c26c47",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(totals['example.f', 'example.y'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a88edd3",
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "from openmdao.utils.assert_utils import assert_near_equal\n",
    "import numpy as np\n",
    "\n",
    "assert_near_equal(totals['example.f', 'example.x'],\n",
    "                 np.array([[1., -0., -0., -0.], [-0.,  2.,  3.,  4.]]), tolerance=1e-9)\n",
    "assert_near_equal(totals['example.f', 'example.y'],\n",
    "                 np.array([[1., -0.], [-0.,  1.]]), tolerance=1e-9)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "ce28a62a",
   "metadata": {},
   "source": [
    "For finite difference approximations (method='fd'), we have three (optional) parameters: the form, step size, and the step_calc. The form should be one of the following:\n",
    "\n",
    "* form='forward' (default): Approximates the derivative as $\\frac{\\partial f}{\\partial x} \\approx \\frac{f(x+\\delta, y) - f(x,y)}{||\\delta||}$. Error scales like $||\\delta||$.\n",
    "\n",
    "* form='backward': Approximates the derivative as $\\frac{\\partial f}{\\partial x} \\approx \\frac{f(x,y) - f(x-\\delta, y) }{||\\delta||}$. Error scales like $||\\delta||$.\n",
    "\n",
    "* form='central': Approximates the derivative as $\\frac{\\partial f}{\\partial x} \\approx \\frac{f(x+\\delta, y) - f(x-\\delta,y)}{2||\\delta||}$. Error scales like $||\\delta||^2$, but requires an extra function evaluation.\n",
    "\n",
    "The step size can be any nonzero number, but should be positive (one can change the form to perform backwards finite difference formulas), small enough to reduce truncation error, but large enough to avoid round-off error. Choosing a step size can be highly problem dependent, but for double precision floating point numbers and reasonably bounded derivatives, 10−6 can be a good place to start. \n",
    "\n",
    "The step_calc can be either ‘abs’ for absolute or one of ('rel', 'rel_avg', 'rel_legacy', 'rel_element') for different forms of relative stepping. This determines whether the stepsize ie absolute or a percentage of the input value. The following table details how the relative step is calculated:\n",
    "\n",
    "| step_calc       | Step size is scaled by |\n",
    "| :---            |    :----:                                       |\n",
    "| \"rel_avg\"       | Average absolute value of the vector.           |\n",
    "| \"rel_element\"   | Absolute value of each vector element.          |\n",
    "| \"rel_legacy\"    | Norm of the vector.                             |\n",
    "| \"rel\"           | Same as \"rel_avg\".                              |"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e451769e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "import openmdao.api as om\n",
    "\n",
    "class FDPartialComp(om.ExplicitComponent):\n",
    "\n",
    "    def setup(self):\n",
    "        self.add_input('x', shape=(4,))\n",
    "        self.add_input('y', shape=(2,))\n",
    "        self.add_input('y2', shape=(2,))\n",
    "        self.add_output('f', shape=(2,))\n",
    "\n",
    "    def setup_partials(self):\n",
    "        self.declare_partials('f', 'y*', method='fd', form='backward', step=1e-6)\n",
    "        self.declare_partials('f', 'x', method='fd', form='central', step=1e-4)\n",
    "\n",
    "    def compute(self, inputs, outputs):\n",
    "        f = outputs['f']\n",
    "\n",
    "        x = inputs['x']\n",
    "        y = inputs['y']\n",
    "\n",
    "        f[0] = x[0] + y[0]\n",
    "        f[1] = np.dot([0, 2, 3, 4], x) + y[1]\n",
    "\n",
    "model = om.Group()\n",
    "model.add_subsystem('example', FDPartialComp())\n",
    "\n",
    "problem = om.Problem(model=model)\n",
    "problem.setup()\n",
    "problem.run_model()\n",
    "totals = problem.compute_totals(['example.f'], ['example.x', 'example.y'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "00e074e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(totals['example.f', 'example.x'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b06a12c",
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "assert_near_equal(totals['example.f', 'example.x'],\n",
    "                 np.array([[1., -0., -0., -0.], [-0.,  2.,  3.,  4.]]), tolerance=1e-9)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "900dca07",
   "metadata": {},
   "source": [
    "## Complex Step\n",
    "If you have a pure python component (or an external code that can support complex inputs and outputs), then you can also choose to use complex step to calculate the Jacobian of that component. This will give more accurate derivatives that are less sensitive to the step size. Like finite difference, complex step runs your component using the `apply_nonlinear` or `solve_nonlinear` functions, but it applies a step in the complex direction. You can activate it using the `declare_partials` method inside `setup_partials` and giving it a method of ‘cs’. In many cases, this will require no other changes to your code, as long as all of the calculation in your `solve_nonlinear` and `apply_nonlinear` support complex numbers. During a complex step, the incoming inputs vector will return a complex number when a variable is being stepped. Likewise, the outputs and residuals vectors will accept complex values. If you are allocating temporary numpy arrays, remember to conditionally set their dtype based on the dtype in the outputs vector.\n",
    "\n",
    "Here is how to turn on complex step for all input/output pairs in the Sellar problem:\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b3b6c31",
   "metadata": {},
   "outputs": [],
   "source": [
    "class SellarDis1CS(om.ExplicitComponent):\n",
    "    \"\"\"\n",
    "    Component containing Discipline 1 -- no derivatives version.\n",
    "    Uses Complex Step\n",
    "    \"\"\"\n",
    "\n",
    "    def setup(self):\n",
    "\n",
    "        # Global Design Variable\n",
    "        self.add_input('z', val=np.zeros(2))\n",
    "\n",
    "        # Local Design Variable\n",
    "        self.add_input('x', val=0.)\n",
    "\n",
    "        # Coupling parameter\n",
    "        self.add_input('y2', val=1.0)\n",
    "\n",
    "        # Coupling output\n",
    "        self.add_output('y1', val=1.0)\n",
    "\n",
    "    def setup_partials(self):\n",
    "        # Finite difference all partials.\n",
    "        self.declare_partials('*', '*', method='cs')\n",
    "\n",
    "    def compute(self, inputs, outputs):\n",
    "        \"\"\"\n",
    "        Evaluates the equation\n",
    "        y1 = z1**2 + z2 + x1 - 0.2*y2\n",
    "        \"\"\"\n",
    "        z1 = inputs['z'][0]\n",
    "        z2 = inputs['z'][1]\n",
    "        x1 = inputs['x']\n",
    "        y2 = inputs['y2']\n",
    "\n",
    "        outputs['y1'] = z1**2 + z2 + x1 - 0.2*y2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f5fa1f7d",
   "metadata": {},
   "outputs": [],
   "source": [
    "class SellarDis2CS(om.ExplicitComponent):\n",
    "    \"\"\"\n",
    "    Component containing Discipline 2 -- no derivatives version.\n",
    "    Uses Complex Step\n",
    "    \"\"\"\n",
    "\n",
    "    def setup(self):\n",
    "        # Global Design Variable\n",
    "        self.add_input('z', val=np.zeros(2))\n",
    "\n",
    "        # Coupling parameter\n",
    "        self.add_input('y1', val=1.0)\n",
    "\n",
    "        # Coupling output\n",
    "        self.add_output('y2', val=1.0)\n",
    "\n",
    "    def setup_partials(self):\n",
    "        # Finite difference all partials.\n",
    "        self.declare_partials('*', '*', method='cs')\n",
    "\n",
    "    def compute(self, inputs, outputs):\n",
    "        \"\"\"\n",
    "        Evaluates the equation\n",
    "        y2 = y1**(.5) + z1 + z2\n",
    "        \"\"\"\n",
    "\n",
    "        z1 = inputs['z'][0]\n",
    "        z2 = inputs['z'][1]\n",
    "        y1 = inputs['y1']\n",
    "\n",
    "        # Note: this may cause some issues. However, y1 is constrained to be\n",
    "        # above 3.16, so lets just let it converge, and the optimizer will\n",
    "        # throw it out\n",
    "        if y1.real < 0.0:\n",
    "            y1 *= -1\n",
    "\n",
    "        outputs['y2'] = y1**.5 + z1 + z2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f16a1ee",
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "from openmdao.utils.assert_utils import assert_near_equal\n",
    "\n",
    "prob = om.Problem()\n",
    "model = prob.model\n",
    "\n",
    "model.add_subsystem('px', om.IndepVarComp('x', 1.0), promotes=['x'])\n",
    "model.add_subsystem('pz', om.IndepVarComp('z', np.array([5.0, 2.0])), promotes=['z'])\n",
    "\n",
    "model.add_subsystem('d1', SellarDis1CS(), promotes=['x', 'z', 'y1', 'y2'])\n",
    "model.add_subsystem('d2', SellarDis2CS(), promotes=['z', 'y1', 'y2'])\n",
    "\n",
    "model.add_subsystem('obj_cmp', om.ExecComp('obj = x**2 + z[1] + y1 + exp(-y2)',\n",
    "                                           z=np.array([0.0, 0.0]), x=0.0),\n",
    "                    promotes=['obj', 'x', 'z', 'y1', 'y2'])\n",
    "\n",
    "model.add_subsystem('con_cmp1', om.ExecComp('con1 = 3.16 - y1'), promotes=['con1', 'y1'])\n",
    "model.add_subsystem('con_cmp2', om.ExecComp('con2 = y2 - 24.0'), promotes=['con2', 'y2'])\n",
    "\n",
    "prob.model.nonlinear_solver = om.NewtonSolver(solve_subsystems=False)\n",
    "prob.model.linear_solver = om.DirectSolver(assemble_jac=False)\n",
    "\n",
    "prob.model.approx_totals(method='cs')\n",
    "prob.setup()\n",
    "prob.set_solver_print(level=0)\n",
    "prob.run_model()\n",
    "\n",
    "assert_near_equal(prob['y1'], 25.58830273, tolerance=1e-6)\n",
    "assert_near_equal(prob['y2'], 12.05848819, tolerance=1e-6)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "511b2d4e",
   "metadata": {},
   "source": [
    "Sometimes you need to know when you are under a complex step so that your component can correctly handle complex inputs (e.g, in case you need to allocate intermediate arrays as complex.) All `Components` and `Groups` provide the attribute `under_complex_step` that you can use to tell if you are under a complex step. In the following example, we print out the incoming complex value when the “compute” method is called while computing this component’s derivatives under complex step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3604265a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import openmdao.api as om\n",
    "\n",
    "class SimpleComp(om.ExplicitComponent):\n",
    "\n",
    "    def setup(self):\n",
    "        self.add_input('x', val=1.0)\n",
    "        self.add_output('y', val=1.0)\n",
    "\n",
    "        self.declare_partials(of='y', wrt='x', method='cs')\n",
    "\n",
    "    def compute(self, inputs, outputs):\n",
    "        outputs['y'] = 3.0*inputs['x']\n",
    "\n",
    "        if self.under_complex_step:\n",
    "            print(\"Under complex step\")\n",
    "            print(\"x\", inputs['x'])\n",
    "            print(\"y\", outputs['y'])\n",
    "\n",
    "prob = om.Problem()\n",
    "prob.model.add_subsystem('comp', SimpleComp())\n",
    "\n",
    "prob.model.add_design_var('comp.x', lower=-100, upper=100)\n",
    "prob.model.add_objective('comp.y')\n",
    "\n",
    "prob.setup(force_alloc_complex=True)\n",
    "\n",
    "prob.run_model()\n",
    "\n",
    "prob.compute_totals(of=['comp.y'], wrt=['comp.x'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bfe12886",
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "from openmdao.utils.assert_utils import assert_near_equal\n",
    "import numpy as np\n",
    "\n",
    "J = prob.compute_totals(of=['comp.y'], wrt=['comp.x'])\n",
    "assert_near_equal(J[('comp.y', 'comp.x')], np.array([[3.]]), tolerance=1e-9)"
   ]
  }
 ],
 "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": 5
}
