{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "f2f921e7",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-02T14:44:55.170723Z",
     "iopub.status.busy": "2026-10-02T14:44:55.170449Z",
     "iopub.status.idle": "2026-10-02T14:44:55.176137Z",
     "shell.execute_reply": "2026-10-02T14:44:55.175254Z"
    },
    "papermill": {
     "duration": 0.008733,
     "end_time": "2026-10-02T14:44:55.176811+00:00",
     "exception": false,
     "start_time": "2026-10-02T14:44:55.168078+00:00",
     "status": "completed"
    },
    "tags": [
     "remove-input",
     "active-ipynb",
     "remove-output"
    ]
   },
   "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": "code",
   "execution_count": 2,
   "id": "0c911e64",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-02T14:44:55.228341Z",
     "iopub.status.busy": "2026-10-02T14:44:55.228131Z",
     "iopub.status.idle": "2026-10-02T14:44:55.791983Z",
     "shell.execute_reply": "2026-10-02T14:44:55.791166Z"
    },
    "papermill": {
     "duration": 0.566656,
     "end_time": "2026-10-02T14:44:55.792617+00:00",
     "exception": false,
     "start_time": "2026-10-02T14:44:55.225961+00:00",
     "status": "completed"
    },
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: jax in /home/runner/work/OpenMDAO/OpenMDAO/.pixi/envs/dev/lib/python3.13/site-packages (0.11.0)\r\n",
      "Requirement already satisfied: jaxlib<=0.11.0,>=0.11.0 in /home/runner/work/OpenMDAO/OpenMDAO/.pixi/envs/dev/lib/python3.13/site-packages (from jax) (0.11.0)\r\n",
      "Requirement already satisfied: ml_dtypes>=0.5.0 in /home/runner/work/OpenMDAO/OpenMDAO/.pixi/envs/dev/lib/python3.13/site-packages (from jax) (0.5.4)\r\n",
      "Requirement already satisfied: numpy>=2.1 in /home/runner/work/OpenMDAO/OpenMDAO/.pixi/envs/dev/lib/python3.13/site-packages (from jax) (2.4.6)\r\n",
      "Requirement already satisfied: opt_einsum in /home/runner/work/OpenMDAO/OpenMDAO/.pixi/envs/dev/lib/python3.13/site-packages (from jax) (3.4.0)\r\n",
      "Requirement already satisfied: scipy>=1.15 in /home/runner/work/OpenMDAO/OpenMDAO/.pixi/envs/dev/lib/python3.13/site-packages (from jax) (1.17.0)\r\n"
     ]
    }
   ],
   "source": [
    "!pip install jax"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6d239b7",
   "metadata": {
    "papermill": {
     "duration": 0.000876,
     "end_time": "2026-10-02T14:44:55.794686+00:00",
     "exception": false,
     "start_time": "2026-10-02T14:44:55.793810+00:00",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Computing Partial Derivatives using JaxImplicitComponent\n",
    "\n",
    "This notebook gives an example of using jax to do automatic differentiation (AD) for a linear system component.\n",
    "The example contains a `JaxImplicitComponent` called `LinSysComp`.  A static option called 'adder' has been added to `LinSysComp` in order to demonstrate how to handle what we call\n",
    "'self statics' in a jax component. Comments interspersed in the code provide some explanations and guidance.\n",
    "\n",
    "Here is an overview of the steps that need to be taken to make use of jax for your `JaxImplicitComponent`. \n",
    "\n",
    "1. Inherit your component from `JaxImplicitComponent`.\n",
    "\n",
    "2. Write a method named `compute_primal` to compute the residuals from the inputs and outputs. This \n",
    "method is the same as what you would normally write for the `apply_nonlinear` method of an \n",
    "`ImplicitComponent`, but it takes as its arguments the actual individual input and output variables \n",
    "rather than a dictionary of the inputs and outputs, and returns the residuals as a tuple. This allows \n",
    "us to use JAX's AD capabilities on this method. Ordering of the inputs and outputs is critical.\n",
    "The order of the inputs and outputs passed into the method and the residuals returned from the method \n",
    "**must** match the order that they are declared in the component.  Note that all of the inputs are\n",
    "passed first, followed by the outputs. Also, discrete inputs, if any, are \n",
    "passed individually as arguments after the output variables.\n",
    "\n",
    "3. By default your component will jit the `compute_primal` method. If for some reason you don't want\n",
    "this, then you can set `self.options['use_jit']` to False.  This can be useful when debugging as it allows you to put print statements inside of your `compute_primal` method.\n",
    "\n",
    "4. For a typical component, that's it.  You can skip step 5.\n",
    "\n",
    "5. If your `compute_primal` depends on variables that are 'static' according to jax, i.e., they affect the output of your compute_primal but are not passed in as arguments, you'll need to define a `get_self_statics` method on your component that \n",
    "returns a tuple containing all such variables.  The returned tuple must be hashable.  If these static values \n",
    "ever change, jax will recompile the `compute_primal` function.  In `LinSysComp` below, there is one \n",
    "static option variable, `self.options['adder']`.\n",
    "\n",
    "\n",
    "## Linear System Component Example\n",
    "\n",
    "The following component is a JaxImplicitComponent representing a linear system `A(x+adder) = b`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "3eb59362",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-02T14:44:55.806202Z",
     "iopub.status.busy": "2026-10-02T14:44:55.806027Z",
     "iopub.status.idle": "2026-10-02T14:44:57.011493Z",
     "shell.execute_reply": "2026-10-02T14:44:57.010956Z"
    },
    "papermill": {
     "duration": 1.216292,
     "end_time": "2026-10-02T14:44:57.012164+00:00",
     "exception": false,
     "start_time": "2026-10-02T14:44:55.795872+00:00",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import openmdao.api as om\n",
    "import numpy as np\n",
    "import jax.numpy as jnp  # noqa\n",
    "\n",
    "\n",
    "class LinSysComp(om.JaxImplicitComponent):\n",
    "    def initialize(self):\n",
    "        self.options.declare('size', default=1, types=int)\n",
    "\n",
    "        # adder is a 'static' value that is constant during derivative computation, but it\n",
    "        # can be changed by the user between runs, so any jitted jax functions need to be re-jitted\n",
    "        # if it changes\n",
    "        self.options.declare('adder', default=0.0, types=float)\n",
    "\n",
    "    def setup(self):\n",
    "        size = self.options['size']\n",
    "\n",
    "        shape = (size, )\n",
    "\n",
    "        self.add_input(\"A\", val=np.eye(size))\n",
    "        self.add_input(\"b\", val=np.ones(shape))\n",
    "        self.add_output(\"x\", shape=shape, val=.1)\n",
    "\n",
    "    def setup_partials(self):\n",
    "        # Because this is an ImplicitComponent, we have to define a linear and nonlinear solver\n",
    "        if self.matrix_free:\n",
    "            # if we're runnning in matrix_free mode, don't use a direct solver\n",
    "            self.linear_solver = om.ScipyKrylov()\n",
    "        else:\n",
    "            self.linear_solver = om.DirectSolver()\n",
    "        self.nonlinear_solver = om.NewtonSolver(solve_subsystems=False)\n",
    "\n",
    "    # we have a static variable, self.options['adder'] that we use in compute_primal, so we need to\n",
    "    # define the get_self_statics method to tell OpenMDAO about it.  Also note that the return value\n",
    "    # must be a tuple, so we follow the value with a comma since we have only one value.\n",
    "    def get_self_statics(self):\n",
    "        return (self.options['adder'], )\n",
    "\n",
    "    # compute_primal replaces the apply_nonlinear method\n",
    "    def compute_primal(self, A, b, x):\n",
    "        return A.dot(x + self.options['adder']) - b\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c3bbc20d",
   "metadata": {
    "papermill": {
     "duration": 0.000794,
     "end_time": "2026-10-02T14:44:57.014004+00:00",
     "exception": false,
     "start_time": "2026-10-02T14:44:57.013210+00:00",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "The rest of this code is standard OpenMDAO code. The code can be run as normal."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "546fcd8a",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-02T14:44:57.016374Z",
     "iopub.status.busy": "2026-10-02T14:44:57.016097Z",
     "iopub.status.idle": "2026-10-02T14:44:58.585080Z",
     "shell.execute_reply": "2026-10-02T14:44:58.584299Z"
    },
    "papermill": {
     "duration": 1.571169,
     "end_time": "2026-10-02T14:44:58.585944+00:00",
     "exception": false,
     "start_time": "2026-10-02T14:44:57.014775+00:00",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1790952298.195743] [runnervm8df0l:9151 :0]        ib_iface.c:1269 UCX  ERROR mana_0: iface 0x559282b77970 failed to create UD QP TX wr:256 sge:6 inl:64 resp:0 RX wr:4096 sge:1 resp:0 failed: Operation not supported\n",
      "[1790952298.196002] [runnervm8df0l:9151 :0]      ucp_worker.c:1412 UCX  ERROR uct_iface_open(ud_verbs/mana_0:1) failed: Input/output error\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[runnervm8df0l:09151] pml_ucx.c:313  Error: Failed to create UCP worker\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[-4.  9. -4.]\n",
      "[-5.  8. -5.]\n"
     ]
    }
   ],
   "source": [
    "A = np.array([[1., 1., 1.], [1., 2., 3.], [0., 1., 3.]])\n",
    "b = np.array([1, 2, -3])\n",
    "\n",
    "prob = om.Problem()\n",
    "\n",
    "ivc = prob.model.add_subsystem('ivc', om.IndepVarComp())\n",
    "ivc.add_output('A', A)\n",
    "ivc.add_output('b', b)\n",
    "\n",
    "lingrp = prob.model.add_subsystem('lingrp', om.Group())\n",
    "lin = lingrp.add_subsystem('lin', LinSysComp(size=b.size))\n",
    "\n",
    "prob.model.connect('ivc.A', 'lingrp.lin.A')\n",
    "prob.model.connect('ivc.b', 'lingrp.lin.b')\n",
    "\n",
    "prob.setup()\n",
    "prob.set_solver_print(level=0)\n",
    "\n",
    "prob.set_val('ivc.A', A)\n",
    "prob.set_val('ivc.b', b)\n",
    "\n",
    "prob.run_model()\n",
    "\n",
    "print(prob['lingrp.lin.x'])\n",
    "\n",
    "# changing the adder value should change the resulting x value\n",
    "lin.options['adder'] = 1.0\n",
    "\n",
    "prob.run_model()\n",
    "\n",
    "print(prob['lingrp.lin.x'])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "c9077a8b",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-02T14:44:58.589659Z",
     "iopub.status.busy": "2026-10-02T14:44:58.589501Z",
     "iopub.status.idle": "2026-10-02T14:44:58.677885Z",
     "shell.execute_reply": "2026-10-02T14:44:58.677253Z"
    },
    "papermill": {
     "duration": 0.091228,
     "end_time": "2026-10-02T14:44:58.678690+00:00",
     "exception": false,
     "start_time": "2026-10-02T14:44:58.587462+00:00",
     "status": "completed"
    },
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "\n",
    "from openmdao.utils.assert_utils import assert_near_equal, assert_check_totals, assert_check_partials\n",
    "\n",
    "assert_near_equal(prob['lingrp.lin.x'], np.array([-5, 8, -5]))\n",
    "\n",
    "assert_check_totals(prob.check_totals(of=['lingrp.lin.x'], wrt=['ivc.b', 'ivc.A'],\n",
    "                                      abs_err_tol=1e-4, out_stream=None), atol=1e-4,)\n",
    "assert_check_partials(prob.check_partials(out_stream=None))\n",
    "\n",
    "lin.options['adder'] = 0.0\n",
    "\n",
    "prob.run_model()\n",
    "\n",
    "assert_near_equal(prob['lingrp.lin.x'], np.array([-4, 9, -4]))\n",
    "\n",
    "assert_check_totals(prob.check_totals(of=['lingrp.lin.x'], wrt=['ivc.b', 'ivc.A'],\n",
    "                                     abs_err_tol=1e-4, out_stream=None), atol=1e-4)\n",
    "assert_check_partials(prob.check_partials(out_stream=None))\n"
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Tags",
  "kernelspec": {
   "display_name": "py311forge",
   "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.13.14"
  },
  "orphan": true,
  "papermill": {
   "default_parameters": {},
   "duration": 4.686379,
   "end_time": "2026-10-02T14:44:59.195101+00:00",
   "environment_variables": {},
   "exception": null,
   "input_path": "/home/runner/work/OpenMDAO/OpenMDAO/openmdao/docs/openmdao_book/features/experimental/jax_implicitcomp_api.ipynb",
   "output_path": "/home/runner/work/OpenMDAO/OpenMDAO/openmdao/docs/_executed_book/features/experimental/jax_implicitcomp_api.ipynb",
   "parameters": {},
   "start_time": "2026-10-02T14:44:54.508722+00:00",
   "version": "2.7.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}