{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "sufficient-exclusive",
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output",
     "active-ipynb"
    ]
   },
   "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",
   "id": "spread-asbestos",
   "metadata": {},
   "source": [
    "# Approximating Semi-Total Derivatives\n",
    "There are times where it makes sense to approximate the derivatives for an entire group in one shot. You can turn on the approximation by calling `approx_totals` on any `Group`.\n",
    "\n",
    "```{eval-rst}\n",
    "    .. automethod:: openmdao.core.group.Group.approx_totals\n",
    "        :noindex:\n",
    "```\n",
    "\n",
    "The default method for approximating semi-total derivatives is the finite difference method. When you call the `approx_totals` method on a group, OpenMDAO will generate an approximate Jacobian for the entire group during the linearization step before derivatives are calculated. OpenMDAO automatically figures out which inputs and output pairs are needed in this Jacobian. When `solve_linear` is called from any system that contains this system, the approximated Jacobian is used for the derivatives in this system.\n",
    "\n",
    "The derivatives approximated in this manner are total derivatives of outputs of the group with respect to inputs. If any components in the group contain implicit states, then you must have an appropriate solver (such as `NewtonSolver`) inside the group to solve the implicit relationships.\n",
    "\n",
    "Here is a classic example of where you might use an approximation like finite difference. In this example, we could just approximate the partials on components CompOne and CompTwo separately. However, CompTwo has a vector input that is 25 wide, so it would require 25 separate executions under finite difference. If we instead approximate the total derivatives on the whole group, we only have one input, so just one extra execution.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "emerging-fantasy",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "import openmdao.api as om\n",
    "\n",
    "class CompOne(om.ExplicitComponent):\n",
    "\n",
    "    def setup(self):\n",
    "        self.add_input('x', val=0.0)\n",
    "        self.add_output('y', val=np.zeros(25))\n",
    "        self._exec_count = 0\n",
    "\n",
    "    def compute(self, inputs, outputs):\n",
    "        x = inputs['x']\n",
    "        outputs['y'] = np.arange(25) * x\n",
    "        self._exec_count += 1\n",
    "\n",
    "class CompTwo(om.ExplicitComponent):\n",
    "\n",
    "    def setup(self):\n",
    "        self.add_input('y', val=np.zeros(25))\n",
    "        self.add_output('z', val=0.0)\n",
    "        self._exec_count = 0\n",
    "\n",
    "    def compute(self, inputs, outputs):\n",
    "        y = inputs['y']\n",
    "        outputs['z'] = np.sum(y)\n",
    "        self._exec_count += 1\n",
    "\n",
    "prob = om.Problem()\n",
    "model = prob.model\n",
    "\n",
    "model.set_input_defaults('x', 0.0)\n",
    "\n",
    "model.add_subsystem('comp1', CompOne(), promotes=['x', 'y'])\n",
    "comp2 = model.add_subsystem('comp2', CompTwo(), promotes=['y', 'z'])\n",
    "\n",
    "model.linear_solver = om.ScipyKrylov()\n",
    "model.approx_totals()\n",
    "\n",
    "prob.setup()\n",
    "prob.run_model()\n",
    "\n",
    "of = ['z']\n",
    "wrt = ['x']\n",
    "derivs = prob.compute_totals(of=of, wrt=wrt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "mathematical-authorization",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(derivs['z', 'x'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "worth-italic",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(comp2._exec_count)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "prospective-boards",
   "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(derivs['z', 'x'], np.array([[300.]]), tolerance=1e-9)\n",
    "assert_near_equal(comp2._exec_count, 2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "played-venue",
   "metadata": {},
   "source": [
    "The same arguments are used for both partial and total derivative approximation specifications. Here we set the finite difference step size, the form to central differences, and the step_calc to relative instead of absolute."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "assigned-dinner",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "import openmdao.api as om\n",
    "\n",
    "class CompOne(om.ExplicitComponent):\n",
    "\n",
    "    def setup(self):\n",
    "        self.add_input('x', val=1.0)\n",
    "        self.add_output('y', val=np.zeros(25))\n",
    "        self._exec_count = 0\n",
    "\n",
    "    def compute(self, inputs, outputs):\n",
    "        x = inputs['x']\n",
    "        outputs['y'] = np.arange(25) * x\n",
    "        self._exec_count += 1\n",
    "\n",
    "class CompTwo(om.ExplicitComponent):\n",
    "\n",
    "    def setup(self):\n",
    "        self.add_input('y', val=np.zeros(25))\n",
    "        self.add_output('z', val=0.0)\n",
    "        self._exec_count = 0\n",
    "\n",
    "    def compute(self, inputs, outputs):\n",
    "        y = inputs['y']\n",
    "        outputs['z'] = np.sum(y)\n",
    "        self._exec_count += 1\n",
    "\n",
    "prob = om.Problem()\n",
    "model = prob.model\n",
    "model.add_subsystem('comp1', CompOne(), promotes=['x', 'y'])\n",
    "model.add_subsystem('comp2', CompTwo(), promotes=['y', 'z'])\n",
    "\n",
    "model.linear_solver = om.ScipyKrylov()\n",
    "model.approx_totals(method='fd', step=1e-7, form='central', step_calc='rel')\n",
    "\n",
    "prob.setup()\n",
    "prob.run_model()\n",
    "\n",
    "of = ['z']\n",
    "wrt = ['x']\n",
    "derivs = prob.compute_totals(of=of, wrt=wrt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "linear-measurement",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(derivs['z', 'x'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "respected-assignment",
   "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(derivs['z', 'x'], np.array([[300.]]), tolerance=1e-8)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "photographic-costs",
   "metadata": {},
   "source": [
    "(complex-step-guidelines)=\n",
    "## Complex Step\n",
    "\n",
    "You can also complex step your model or group, though there are some important restrictions.\n",
    "\n",
    "**All components must support complex calculations in solve_nonlinear:**\n",
    "\n",
    "Under complex step, a component’s inputs are complex, all stages of the calculation will operate on complex inputs to produce complex outputs, and the final value placed into outputs is complex. Most Python functions already support complex numbers, so pure Python components will generally satisfy this requirement. Take care with functions like abs, which effectively squelches the complex part of the argument.\n",
    "\n",
    "**If you complex step around a solver that requires gradients, the solver must not get its gradients from complex step:**\n",
    "\n",
    "When you complex step around a nonlinear solver that requires gradients (like Newton), the nonlinear solver must solve a complex linear system rather than a real one. Most of OpenMDAO’s linear solvers (with the exception of PETScKrylov) support the solution of such a system. However, when you linearize the submodel to determine derivatives around a complex point, the application of complex step loses some of its robust properties when compared to real-valued finite difference (in particular, you get subtractive cancelation which causes increased inaccuracy for smaller step sizes.) When OpenMDAO encounters this situation, it will warn the user, and the inner approximation will automatically switch over to using finite difference with default settings.\n",
    "\n",
    "**Care must be taken with iterative solver tolerances; you may need to adjust the stepsize for complex step:**\n",
    "\n",
    "If you are using an iterative nonlinear solver, and you don’t converge it tightly, then the complex stepped linear system may have trouble converging as well. You may need to tighten the convergence of your solvers and increase the step size used for complex step. To prevent the nonlinear solvers from ignoring a tiny complex step, a tiny offset is added to the states to nudge it off the solution, allowing it to reconverge with the complex step. You can also turn this behavior off by setting the “cs_reconverge” to False.\n",
    "\n",
    "Similarly, an iterative linear solver may also require adjusting the step size, particularly if you are using the ScipyKrylov solver."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "enclosed-official",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "import openmdao.api as om\n",
    "\n",
    "class CompOne(om.ExplicitComponent):\n",
    "\n",
    "    def setup(self):\n",
    "        self.add_input('x', val=0.0)\n",
    "        self.add_output('y', val=np.zeros(25))\n",
    "        self._exec_count = 0\n",
    "\n",
    "    def compute(self, inputs, outputs):\n",
    "        x = inputs['x']\n",
    "        outputs['y'] = np.arange(25) * x\n",
    "        self._exec_count += 1\n",
    "\n",
    "class CompTwo(om.ExplicitComponent):\n",
    "\n",
    "    def setup(self):\n",
    "        self.add_input('y', val=np.zeros(25))\n",
    "        self.add_output('z', val=0.0)\n",
    "        self._exec_count = 0\n",
    "\n",
    "    def compute(self, inputs, outputs):\n",
    "        y = inputs['y']\n",
    "        outputs['z'] = np.sum(y)\n",
    "        self._exec_count += 1\n",
    "\n",
    "prob = om.Problem()\n",
    "model = prob.model\n",
    "model.set_input_defaults('x', 0.0)\n",
    "\n",
    "model.add_subsystem('comp1', CompOne(), promotes=['x', 'y'])\n",
    "model.add_subsystem('comp2', CompTwo(), promotes=['y', 'z'])\n",
    "\n",
    "model.linear_solver = om.ScipyKrylov()\n",
    "model.approx_totals(method='cs')\n",
    "\n",
    "prob.setup()\n",
    "prob.run_model()\n",
    "\n",
    "of = ['z']\n",
    "wrt = ['x']\n",
    "derivs = prob.compute_totals(of=of, wrt=wrt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "afraid-evolution",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(derivs['z', 'x'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "forty-experiment",
   "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(derivs['z', 'x'], np.array([[300.]]), tolerance=1e-8)"
   ]
  }
 ],
 "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
}
