{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2199e974",
   "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": "a0eae820",
   "metadata": {},
   "source": [
    "# Matrix Free Total Derivatives\n",
    "\n",
    "The `compute_jacvec_product` method of Problem can be used to compute a matrix free total Jacobian vector product. It’s analogous to the way that the `compute_jacvec_product` method of `System` can be used to compute partial Jacobian vector products."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b4213239",
   "metadata": {},
   "source": [
    "```{eval-rst}\n",
    "    .. automethod:: openmdao.core.problem.Problem.compute_jacvec_product\n",
    "        :noindex:\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "223dfcbc",
   "metadata": {},
   "source": [
    "Below is an example of a component that embeds a sub-problem and uses `compute_jacvec_product` on that sub-problem to compute its Jacobian. The `SubProbComp` component computes derivatives in both ‘fwd’ and ‘rev’ directions, but in a realistic scenario, it would only compute them in a single direction.\n",
    "\n",
    "The model that `SubProbComp` is intended to emulate is one that looks like this:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e2508bb1",
   "metadata": {},
   "source": [
    "![Model using 3 ExecComps](cjvp_xdsm.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aad7aeb9",
   "metadata": {},
   "source": [
    "Instead of using 3 identical `ExecComps` as shown above and having OpenMDAO automatically compute the total derivatives for us, `SubProbComp` will use just a single `ExecComp` and will compute its derivatives internally. The model contained in the sub-problem looks like this:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b4b1daee",
   "metadata": {},
   "source": [
    "![Model using 3 ExecComps](cjvp_small_xdsm.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0a140fea",
   "metadata": {},
   "source": [
    "The code for `SubProbComp` is shown below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "568633f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import openmdao.api as om\n",
    "\n",
    "\n",
    "class SubProbComp(om.ExplicitComponent):\n",
    "    \"\"\"\n",
    "    This component contains a sub-Problem with a component that will be solved over num_nodes\n",
    "    points instead of creating num_nodes instances of that same component and connecting them\n",
    "    together.\n",
    "    \"\"\"\n",
    "    def __init__(self, input_size, num_nodes, mode, **kwargs):\n",
    "        super().__init__(**kwargs)\n",
    "        self.prob = None\n",
    "        self.size = input_size\n",
    "        self.num_nodes = num_nodes\n",
    "        self.mode = mode\n",
    "\n",
    "    def _setup_subprob(self):\n",
    "        self.prob = p = om.Problem(comm=self.comm)\n",
    "        model = self.prob.model\n",
    "\n",
    "        model.add_subsystem('comp', get_comp(self.size))\n",
    "\n",
    "        p.setup()\n",
    "        p.final_setup()\n",
    "\n",
    "    def setup(self):\n",
    "        self._setup_subprob()\n",
    "\n",
    "        self.add_input('x', np.zeros(self.size - 1))\n",
    "        self.add_input('inp', val=0.0)\n",
    "        self.add_output('out', val=0.0)\n",
    "        self.declare_partials('*', '*')\n",
    "\n",
    "    def compute(self, inputs, outputs):\n",
    "        p = self.prob\n",
    "        p['comp.x'] = inputs['x']\n",
    "        p['comp.inp'] = inputs['inp']\n",
    "        inp = inputs['inp']\n",
    "        for i in range(self.num_nodes):\n",
    "            p['comp.inp'] = inp\n",
    "            p.run_model()\n",
    "            inp = p['comp.out']\n",
    "\n",
    "        outputs['out'] = p['comp.out']\n",
    "\n",
    "    def _compute_partials_fwd(self, inputs, partials):\n",
    "        p = self.prob\n",
    "        x = inputs['x']\n",
    "        p['comp.x'] = x\n",
    "        p['comp.inp'] = inputs['inp']\n",
    "\n",
    "        seed = {'comp.x':np.zeros(x.size), 'comp.inp': np.zeros(1)}\n",
    "        p.run_model()\n",
    "        p.model._linearize(None)\n",
    "        for rhsname in seed:\n",
    "            for rhs_i in range(seed[rhsname].size):\n",
    "                seed['comp.x'][:] = 0.0\n",
    "                seed['comp.inp'][:] = 0.0\n",
    "                seed[rhsname][rhs_i] = 1.0\n",
    "                for i in range(self.num_nodes):\n",
    "                    p.model._vectors['output']['linear'].set_val(0.0)\n",
    "                    p.model._vectors['residual']['linear'].set_val(0.0)\n",
    "                    jvp = p.compute_jacvec_product(of=['comp.out'], wrt=['comp.x','comp.inp'], mode='fwd', seed=seed)\n",
    "                    seed['comp.inp'][:] = jvp['comp.out']\n",
    "\n",
    "                if rhsname == 'comp.x':\n",
    "                    partials[self.pathname + '.out', self.pathname +'.x'][0, rhs_i] = jvp[self.pathname + '.out']\n",
    "                else:\n",
    "                    partials[self.pathname + '.out', self.pathname + '.inp'][0, 0] = jvp[self.pathname + '.out']\n",
    "\n",
    "    def _compute_partials_rev(self, inputs, partials):\n",
    "        p = self.prob\n",
    "        p['comp.x'] = inputs['x']\n",
    "        p['comp.inp'] = inputs['inp']\n",
    "        seed = {'comp.out': np.ones(1)}\n",
    "\n",
    "        stack = []\n",
    "        comp = p.model.comp\n",
    "        comp._inputs['inp'] = inputs['inp']\n",
    "        # store the inputs to each comp (the comp at each node point) by doing nonlinear solves\n",
    "        # and storing what the inputs are for each node point.  We'll set these inputs back\n",
    "        # later when we linearize about each node point.\n",
    "        for i in range(self.num_nodes):\n",
    "            stack.append(comp._inputs['inp'][0])\n",
    "            comp._inputs['x'] = inputs['x']\n",
    "            comp._solve_nonlinear()\n",
    "            comp._inputs['inp'] = comp._outputs['out']\n",
    "\n",
    "        for i in range(self.num_nodes):\n",
    "            p.model._vectors['output']['linear'].set_val(0.0)\n",
    "            p.model._vectors['residual']['linear'].set_val(0.0)\n",
    "            comp._inputs['inp'] = stack.pop()\n",
    "            comp._inputs['x'] = inputs['x']\n",
    "            p.model._linearize(None)\n",
    "            jvp = p.compute_jacvec_product(of=['comp.out'], wrt=['comp.x','comp.inp'], mode='rev', seed=seed)\n",
    "            seed['comp.out'][:] = jvp['comp.inp']\n",
    "\n",
    "            # all of the comp.x's are connected to the same indepvarcomp, so we have\n",
    "            # to accumulate their contributions together\n",
    "            partials[self.pathname + '.out', self.pathname + '.x'] += jvp['comp.x']\n",
    "\n",
    "            # this one doesn't get accumulated because each comp.inp contributes to the\n",
    "            # previous comp's .out (or to comp.inp in the case of the first comp) only.\n",
    "            # Note that we have to handle this explicitly here because normally in OpenMDAO\n",
    "            # we accumulate derivatives when we do reverse transfers.  We can't do that\n",
    "            # here because we only have one instance of our component, so instead of\n",
    "            # accumulating into separate 'comp.out' variables for each comp instance,\n",
    "            # we would be accumulating into a single comp.out variable, which would make\n",
    "            # our derivative too big.\n",
    "            partials[self.pathname + '.out', self.pathname + '.inp'] = jvp['comp.inp']\n",
    "\n",
    "    def compute_partials(self, inputs, partials):\n",
    "        # note that typically you would only have to define partials for one direction,\n",
    "        # either fwd OR rev, not both.\n",
    "        if self.mode == 'fwd':\n",
    "            self._compute_partials_fwd(inputs, partials)\n",
    "        else:\n",
    "            self._compute_partials_rev(inputs, partials)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "563199d9",
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from openmdao.utils.assert_utils import assert_near_equal\n",
    "\n",
    "def get_comp(size):\n",
    "    return om.ExecComp('out = sum(3*x**2) + inp', x=np.ones(size - 1), inp=0.0, out=0.0)\n",
    "\n",
    "size = 5\n",
    "p = om.Problem()\n",
    "model = p.model\n",
    "comp = model.add_subsystem('comp', om.IndepVarComp('x', val=np.zeros(size - 1)))\n",
    "comp.add_output('inp', val=0.0)\n",
    "\n",
    "C1 = model.add_subsystem('C1', get_comp(size))\n",
    "C2 = model.add_subsystem('C2', get_comp(size))\n",
    "C3 = model.add_subsystem('C3', get_comp(size))\n",
    "\n",
    "model.connect('comp.x', ['C1.x', 'C2.x', 'C3.x'])\n",
    "model.connect('comp.inp', 'C1.inp')\n",
    "model.connect('C1.out', 'C2.inp')\n",
    "model.connect('C2.out', 'C3.inp')\n",
    "p.setup(mode='fwd')\n",
    "\n",
    "p['comp.x'] = np.arange(size-1, dtype=float) + 1. #np.random.random(size - 1)\n",
    "p['comp.inp'] = np.array([7.])  #np.random.random(1)[0]\n",
    "p.final_setup()\n",
    "\n",
    "p2 = om.Problem()\n",
    "\n",
    "comp = p2.model.add_subsystem('comp', SubProbComp(input_size=size, num_nodes=3, mode='fwd'))\n",
    "\n",
    "p2.setup(mode='fwd')\n",
    "\n",
    "p2['comp.x'] = p['comp.x']\n",
    "p2['comp.inp'] = p['comp.inp']\n",
    "\n",
    "p2.run_model()\n",
    "J2 = p2.compute_totals(of=['comp.out'], wrt=['comp.x', 'comp.inp'], return_format='array')\n",
    "\n",
    "p.run_model()\n",
    "J = p.compute_totals(of=['C3.out'], wrt=['comp.x', 'comp.inp'], return_format='array')\n",
    "\n",
    "assert_near_equal(p['C3.out'], p2['comp.out'])\n",
    "np.testing.assert_allclose(J2, J)"
   ]
  }
 ],
 "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.5"
  },
  "orphan": true
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
