{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50f80b17",
   "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": "7aa77ba4",
   "metadata": {},
   "source": [
    "# Sparse Partial Derivatives\n",
    "\n",
    "When a partial derivative is sparse (few nonzero entries compared to the total size of the matrix), it may be advantageous to utilize a format that stores only the nonzero entries. To use sparse partial derivatives, they must first be declared with the sparsity pattern in setup_partials using the declare_partials method.\n",
    "\n",
    "## Usage\n",
    "\n",
    "To specify the sparsity pattern in the AIJ format (alternatively known as COO format), use the `rows` and `cols` arguments to `declare_partials`. For example, to declare a sparsity pattern of nonzero entries in the (0, 0), (1, 1), (1, 2), and (1,3) positions, one would use rows=[0, 1, 1, 1], cols=[0, 1, 2, 3]. When using `compute_partials`, you do not need to pass the sparsity pattern again. Instead, you simply give the values for the entries in the same order as given in `declare_partials`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a123ebf9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "import openmdao.api as om\n",
    "\n",
    "class SparsePartialComp(om.ExplicitComponent):\n",
    "    def setup(self):\n",
    "        self.add_input('x', shape=(4,))\n",
    "        self.add_output('f', shape=(2,))\n",
    "\n",
    "    def setup_partials(self):\n",
    "        self.declare_partials(of='f', wrt='x',\n",
    "                              rows=[0, 1, 1, 1],\n",
    "                              cols=[0, 1, 2, 3])\n",
    "\n",
    "    def compute_partials(self, inputs, partials):\n",
    "        # Corresponds to the [(0,0), (1,1), (1,2), (1,3)] entries.\n",
    "        partials['f', 'x'] = [1., 2., 3., 4.]\n",
    "\n",
    "model = om.Group()\n",
    "model.add_subsystem('example', SparsePartialComp())\n",
    "\n",
    "problem = om.Problem(model=model)\n",
    "problem.setup()\n",
    "problem.run_model()\n",
    "totals = problem.compute_totals(['example.f'], ['example.x'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e747bff5",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(totals['example.f', 'example.x'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b232e007",
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "from openmdao.utils.assert_utils import assert_near_equal\n",
    "assert_near_equal(totals['example.f', 'example.x'],\n",
    "                  np.array([[1., -0., -0., -0.], [-0.,  2.,  3.,  4.]]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f1743be",
   "metadata": {},
   "source": [
    "If only some of your Jacobian entries change across iterations, or if you wish to avoid creating intermediate arrays, you may update the entries in-place."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f175558",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "import openmdao.api as om\n",
    "\n",
    "class SparsePartialComp(om.ExplicitComponent):\n",
    "    def setup(self):\n",
    "        self.add_input('x', shape=(4,))\n",
    "        self.add_output('f', shape=(2,))\n",
    "\n",
    "    def setup_partials(self):\n",
    "        self.declare_partials(of='f', wrt='x',\n",
    "                              rows=[0, 1, 1, 1],\n",
    "                              cols=[0, 1, 2, 3])\n",
    "\n",
    "    def compute_partials(self, inputs, partials):\n",
    "        pd = partials['f', 'x']\n",
    "\n",
    "        # Corresponds to the (0, 0) entry\n",
    "        pd[0] = 1.\n",
    "\n",
    "        # (1,1) entry\n",
    "        pd[1] = 2.\n",
    "\n",
    "        # (1, 2) entry\n",
    "        pd[2] = 3.\n",
    "\n",
    "        # (1, 3) entry\n",
    "        pd[3] = 4\n",
    "\n",
    "\n",
    "model = om.Group()\n",
    "model.add_subsystem('example', SparsePartialComp())\n",
    "\n",
    "problem = om.Problem(model=model)\n",
    "problem.setup()\n",
    "problem.run_model()\n",
    "totals = problem.compute_totals(['example.f'], ['example.x'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77b7b28f",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(totals['example.f', 'example.x'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1501370",
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "assert_near_equal(totals['example.f', 'example.x'],\n",
    "                  np.array([[1., -0., -0., -0.], [-0.,  2.,  3.,  4.]]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b4df1c2",
   "metadata": {},
   "source": [
    "If your partial derivative is constant and sparse, or if you simply wish to provide an initial value for the derivative, you can pass in the values using the `val` argument. If you are using the AIJ format, `val` should receive the nonzero entries in the same order as given for `rows` and `cols`. Alternatively, you may provide a Scipy sparse matrix, from which the sparsity pattern is deduced."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e88835c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import scipy as sp\n",
    "\n",
    "import openmdao.api as om\n",
    "\n",
    "class SparsePartialComp(om.ExplicitComponent):\n",
    "    def setup(self):\n",
    "        self.add_input('x', shape=(4,))\n",
    "        self.add_input('y', shape=(2,))\n",
    "        self.add_output('f', shape=(2,))\n",
    "\n",
    "    def setup_partials(self):\n",
    "        self.declare_partials(of='f', wrt='x',\n",
    "                              rows=[0, 1, 1, 1],\n",
    "                              cols=[0, 1, 2, 3],\n",
    "                              val=[1., 2., 3., 4.])\n",
    "        self.declare_partials(of='f', wrt='y', val=sp.sparse.eye(2, format='csc'))\n",
    "\n",
    "    def compute_partials(self, inputs, partials):\n",
    "        pass\n",
    "\n",
    "model = om.Group()\n",
    "model.add_subsystem('example', SparsePartialComp())\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": "cfb0f18a",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(totals['example.f', 'example.x'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c071c1e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(totals['example.f', 'example.y'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66a7e36b",
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "assert_near_equal(totals['example.f', 'example.x'],\n",
    "                  np.array([[1., -0., -0., -0.], [-0.,  2.,  3.,  4.]]))\n",
    "assert_near_equal(totals['example.f', 'example.y'],\n",
    "                  np.array([[1., -0.], [-0., 1.]]))"
   ]
  }
 ],
 "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.8"
  },
  "orphan": true
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
