{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "active-ipynb",
     "remove-output"
    ]
   },
   "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",
   "metadata": {},
   "source": [
    "# LinearUserDefined\n",
    "\n",
    "LinearUserDefined is a solver that lets you define a custom method for performing a linear solve on a component. The default method is named \"solve_linear\", but you can give it any name by passing in the function or method handle to the \"solve_function\" attribute.\n",
    "\n",
    "The function needs to have the following signature:\n",
    "\n",
    "```python\n",
    "    def my_solve_function(d_outputs, d_residuals, mode):\n",
    "        r\"\"\"\n",
    "        Apply inverse jac product. The model is assumed to be in an unscaled state.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        d_outputs: Vector\n",
    "            unscaled, dimensional quantities read via d_outputs[key]\n",
    "        d_residuals: Vector\n",
    "            unscaled, dimensional quantities read via d_residuals[key]\n",
    "        mode: str\n",
    "            either 'fwd' or 'rev'\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        None or bool or (bool, float, float)\n",
    "            The bool is the failure flag; and the two floats are absolute and relative error.\n",
    "        \"\"\"\n",
    "```\n",
    "\n",
    "Here is a rather contrived example where an identity preconditioner is used by giving the component's \"mysolve\"\n",
    "method to a LinearUserDefined solver."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import openmdao.api as om\n",
    "\n",
    "from openmdao.utils.array_utils import evenly_distrib_idxs\n",
    "\n",
    "\n",
    "class CustomSolveImplicit(om.ImplicitComponent):\n",
    "\n",
    "    def setup(self):\n",
    "\n",
    "        self.add_input('a', val=10., units='m')\n",
    "\n",
    "        rank = self.comm.rank\n",
    "        GLOBAL_SIZE = 15\n",
    "        sizes, offsets = evenly_distrib_idxs(self.comm.size, GLOBAL_SIZE)\n",
    "\n",
    "        self.add_output('states', shape=int(sizes[rank]))\n",
    "\n",
    "        self.add_output('out_var', shape=1)\n",
    "        self.local_size = sizes[rank]\n",
    "\n",
    "        self.linear_solver = om.PETScKrylov()\n",
    "        self.linear_solver.precon = om.LinearUserDefined(solve_function=self.mysolve)\n",
    "\n",
    "    def solve_nonlinear(self, i, o):\n",
    "        o['states'] = i['a']\n",
    "\n",
    "        local_sum = np.zeros(1)\n",
    "        local_sum[0] = np.sum(o['states'])\n",
    "        tmp = np.zeros(1)\n",
    "\n",
    "        o['out_var'] = tmp[0]\n",
    "\n",
    "    def apply_nonlinear(self, i, o, r):\n",
    "        r['states'] = o['states'] - i['a']\n",
    "\n",
    "        local_sum = np.zeros(1)\n",
    "        local_sum[0] = np.sum(o['states'])\n",
    "        global_sum = np.zeros(1)\n",
    "\n",
    "        r['out_var'] = o['out_var'] - tmp[0]\n",
    "\n",
    "    def apply_linear(self, i, o, d_i, d_o, d_r, mode):\n",
    "        if mode == 'fwd':\n",
    "            if 'states' in d_o:\n",
    "                d_r['states'] += d_o['states']\n",
    "\n",
    "                local_sum = np.array([np.sum(d_o['states'])])\n",
    "                global_sum = np.zeros(1)\n",
    "                self.comm.Allreduce(local_sum, global_sum, op=MPI.SUM)\n",
    "                d_r['out_var'] -= global_sum\n",
    "\n",
    "            if 'out_var' in d_o:\n",
    "                    d_r['out_var'] += d_o['out_var']\n",
    "\n",
    "            if 'a' in d_i:\n",
    "                    d_r['states'] -= d_i['a']\n",
    "\n",
    "        elif mode == 'rev':\n",
    "            if 'states' in d_o:\n",
    "                d_o['states'] += d_r['states']\n",
    "\n",
    "                tmp = np.zeros(1)\n",
    "                if self.comm.rank == 0:\n",
    "                    tmp[0] = d_r['out_var'].copy()\n",
    "                self.comm.Bcast(tmp, root=0)\n",
    "\n",
    "                d_o['states'] -= tmp\n",
    "\n",
    "            if 'out_var' in d_o:\n",
    "                d_o['out_var'] += d_r['out_var']\n",
    "\n",
    "            if 'a' in d_i:\n",
    "                    d_i['a'] -= np.sum(d_r['states'])\n",
    "\n",
    "    def mysolve(self, d_outputs, d_residuals, mode):\n",
    "        r\"\"\"\n",
    "        Apply inverse jac product. The model is assumed to be in an unscaled state.\n",
    "\n",
    "        If mode is:\n",
    "            'fwd': d_residuals \\|-> d_outputs\n",
    "\n",
    "            'rev': d_outputs \\|-> d_residuals\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        d_outputs : Vector\n",
    "            unscaled, dimensional quantities read via d_outputs[key]\n",
    "        d_residuals : Vector\n",
    "            unscaled, dimensional quantities read via d_residuals[key]\n",
    "        mode: str\n",
    "            either 'fwd' or 'rev'\n",
    "        \"\"\"\n",
    "        # Note: we are just preconditioning with Identity as a proof of concept.\n",
    "        if mode == 'fwd':\n",
    "            d_outputs.set_vec(d_residuals)\n",
    "        elif mode == 'rev':\n",
    "            d_residuals.set_vec(d_outputs)\n",
    "\n",
    "prob = om.Problem()\n",
    "\n",
    "\n",
    "prob.model.add_subsystem('icomp', CustomSolveImplicit(), promotes=['*'])\n",
    "prob.model.set_input_defaults('a', 10., units='m')\n",
    "\n",
    "model = prob.model\n",
    "\n",
    "model.linear_solver = om.PETScKrylov()\n",
    "model.linear_solver.precon = om.LinearRunOnce()\n",
    "\n",
    "prob.setup(mode='rev', check=False)\n",
    "prob.run_model()\n",
    "jac = prob.compute_totals(of=['out_var'], wrt=['a'], return_format='dict')\n",
    "\n",
    "print(jac['out_var']['a'][0][0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "from openmdao.utils.assert_utils import assert_near_equal\n",
    "\n",
    "assert_near_equal(15.0, jac['out_var']['a'][0][0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## LinearUserDefined Options"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input"
    ]
   },
   "outputs": [],
   "source": [
    "om.show_options_table(\"openmdao.solvers.linear.user_defined.LinearUserDefined\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## LinearUserDefined Constructor\n",
    "\n",
    "The call signature for the `LinearUserDefined` constructor is:\n",
    "\n",
    "```{eval-rst}\n",
    "    .. automethod:: openmdao.solvers.linear.user_defined.LinearUserDefined.__init__\n",
    "        :noindex:\n",
    "```"
   ]
  }
 ],
 "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": 4
}
