{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "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]"
   ],
   "id": "451c28fb"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# In-Memory Assembly of Jacobians\n",
    "\n",
    "When you have groups, or entire models, that are small enough to fit the entire Jacobian into memory, you can have OpenMDAO actually assemble the partial-derivative Jacobian in memory. In many cases this can yield a substantial speed up over the default, [matrix-free](theory-assembled-vs-matrix-free) implementation in OpenMDAO."
   ],
   "id": "94c877bd"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```{Note}\n",
    "Assembled Jacobians are especially effective when you have a deeply-nested hierarchy with a large number of components and/or variables. See the [matrix-free](theory-assembled-vs-matrix-free) for more details on how to best select which type of Jacobian to use.\n",
    "\n",
    "```"
   ],
   "id": "707ef830"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To use an assembled Jacobian, you set the `assemble_jac` option of the linear solver that will use it to True. The type of the assembled Jacobian will be determined by the value of `options['assembled_jac_type']` in the solver’s containing system. There are two options of `assembled_jac_type` to choose from, `dense` and `csc`.\n",
    "\n"
   ],
   "id": "a251f7d6"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```{Note}\n",
    "`csc` is an abbreviation for compressed sparse column. `csc` is one of many sparse storage schemes that allocate contiguous storage in memory for the nonzero elements of the matrix, and perhaps a limited number of zeros. For more information, see [Compressed sparse column](https://en.wikipedia.org/wiki/Sparse_matrix).\n",
    "```"
   ],
   "id": "14a54f63"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For example:\n",
    "\n",
    "    model.options['assembled_jac_type'] = 'dense'\n",
    "    model.linear_solver = DirectSolver(assemble_jac=True)"
   ],
   "id": "6595c006"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "‘csc’ is the default, and you should try that first if you’re not sure of which one to use. Most problems, even if they have dense sub-Jacobians from each component, are fairly sparse at the model level and the [DirectSolver](../../building_blocks/solvers/direct_solver) will usually be much faster with a sparse factorization."
   ],
   "id": "1ff5c3d6"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```{Note}\n",
    "You are allowed to use multiple assembled Jacobians at multiple different levels of your model hierarchy. This may be useful if you have nested nonlinear solvers to converge very difficult problems.\n",
    "```"
   ],
   "id": "1a832b7c"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Usage Example\n",
    "\n",
    "In the following example, borrowed from the [newton solver tutorial](../../../advanced_user_guide/models_implicit_components/models_with_solvers_implicit), we assemble the Jacobian at the same level of the model hierarchy as the [NewtonSolver](../../building_blocks/solvers/newton)  and [DirectSolver](../../building_blocks/solvers/direct_solver). In general, you will always do the assembly at the same level as the linear solver that will make use of the Jacobian matrix."
   ],
   "id": "20a3da7b"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "from openmdao.utils.notebook_utils import get_code\n",
    "from myst_nb import glue\n",
    "glue(\"code_src56\", get_code(\"openmdao.test_suite.scripts.circuit_analysis.Resistor\"), display=False)"
   ],
   "id": "f7a374b8"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{Admonition} `Resistor` class definition \n",
    ":class: dropdown\n",
    "\n",
    "{glue:}`code_src56`\n",
    ":::"
   ],
   "id": "d906a09f"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "from openmdao.utils.notebook_utils import get_code\n",
    "from myst_nb import glue\n",
    "glue(\"code_src57\", get_code(\"openmdao.test_suite.scripts.circuit_analysis.Diode\"), display=False)"
   ],
   "id": "4d1ec254"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{Admonition} `Diode` class definition \n",
    ":class: dropdown\n",
    "\n",
    "{glue:}`code_src57`\n",
    ":::"
   ],
   "id": "e56dbbe7"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "from openmdao.utils.notebook_utils import get_code\n",
    "from myst_nb import glue\n",
    "glue(\"code_src58\", get_code(\"openmdao.test_suite.scripts.circuit_analysis.Node\"), display=False)"
   ],
   "id": "d7ad3e5d"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{Admonition} `Node` class definition \n",
    ":class: dropdown\n",
    "\n",
    "{glue:}`code_src58`\n",
    ":::"
   ],
   "id": "22e18deb"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import openmdao.api as om\n",
    "from openmdao.test_suite.scripts.circuit_analysis import Resistor, Diode, Node\n",
    "\n",
    "class Circuit(om.Group):\n",
    "\n",
    "    def setup(self):\n",
    "        self.add_subsystem('n1', Node(n_in=1, n_out=2), promotes_inputs=[('I_in:0', 'I_in')])\n",
    "        self.add_subsystem('n2', Node())  # leaving defaults\n",
    "\n",
    "        self.add_subsystem('R1', Resistor(R=100.), promotes_inputs=[('V_out', 'Vg')])\n",
    "        self.add_subsystem('R2', Resistor(R=10000.))\n",
    "        self.add_subsystem('D1', Diode(), promotes_inputs=[('V_out', 'Vg')])\n",
    "\n",
    "        self.connect('n1.V', ['R1.V_in', 'R2.V_in'])\n",
    "        self.connect('R1.I', 'n1.I_out:0')\n",
    "        self.connect('R2.I', 'n1.I_out:1')\n",
    "\n",
    "        self.connect('n2.V', ['R2.V_out', 'D1.V_in'])\n",
    "        self.connect('R2.I', 'n2.I_in:0')\n",
    "        self.connect('D1.I', 'n2.I_out:0')\n",
    "\n",
    "        self.nonlinear_solver = om.NewtonSolver(solve_subsystems=False)\n",
    "        self.nonlinear_solver.options['iprint'] = 2\n",
    "        self.nonlinear_solver.options['maxiter'] = 20\n",
    "        ##################################################################\n",
    "        # Assemble at the group level. Default assembled jac type is 'csc'\n",
    "        ##################################################################\n",
    "        self.options['assembled_jac_type'] = 'csc'\n",
    "        self.linear_solver = om.DirectSolver(assemble_jac=True)\n",
    "\n",
    "p = om.Problem()\n",
    "model = p.model\n",
    "\n",
    "model.add_subsystem('circuit', Circuit())\n",
    "\n",
    "p.setup()\n",
    "\n",
    "p.set_val('circuit.I_in', 0.1)\n",
    "p.set_val('circuit.Vg', 0.)\n",
    "\n",
    "# set some initial guesses\n",
    "p.set_val('circuit.n1.V', 10.)\n",
    "p.set_val('circuit.n2.V', 1.)\n",
    "\n",
    "p.run_model()"
   ],
   "id": "cca6a3ad"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(p['circuit.n1.V'])\n",
    "print(p['circuit.n2.V'])\n",
    "print(p['circuit.R1.I'])\n",
    "print(p['circuit.R2.I'])\n",
    "print(p['circuit.D1.I'])\n",
    "\n",
    "# sanity check: should sum to .1 Amps\n",
    "print(p['circuit.R1.I'] + p['circuit.D1.I'])"
   ],
   "id": "5fb0e8f8"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    },
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "from openmdao.utils.assert_utils import assert_near_equal\n",
    "\n",
    "assert_near_equal(p['circuit.n1.V'], 9.90804735, tolerance=1e-5)\n",
    "assert_near_equal(p['circuit.n2.V'], 0.71278185, tolerance=1e-5)\n",
    "assert_near_equal(p['circuit.R1.I'], 0.09908047, tolerance=1e-5)\n",
    "assert_near_equal(p['circuit.R2.I'], 0.00091953, tolerance=1e-5)\n",
    "assert_near_equal(p['circuit.D1.I'], 0.00091953, tolerance=1e-5)"
   ],
   "id": "c45866b6"
  }
 ],
 "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.1"
  },
  "orphan": true
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
