{
 "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  # noqa: F401\n",
    "except ImportError:\n",
    "    !python -m pip install openmdao[notebooks]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Computing Partial Derivatives of Components Using JAX\n",
    "\n",
    "To truly take advantage of OpenMDAO, the user needs to compute partial derivatives for any `Component`\n",
    "that they write.  This can be done using finite difference, but that can have issues with accuracy\n",
    "and performance.  Using complex step is another option which has good accuracy but isn't always \n",
    "possible because it requires the component's computations to be compatible with complex numbers. \n",
    "In some cases, the user can provide analytic partial derivatives, which likely has good performance \n",
    "but can be difficult to determine depending on the complexity of the component.\n",
    "\n",
    "This notebook describes another method, which is to use the optional third-party \n",
    "[JAX](https://jax.readthedocs.io/en/latest/index.html) library, to \n",
    "automatically differentiate native Python and NumPy functions.  To simplify jax usage within OpenMDAO, \n",
    "we've created two component classes, [JaxExplicitComponent](jax_explicitcomp_api.ipynb) and \n",
    "[JaxImplicitComponent](jax_implicitcomp_api.ipynb).  These components require only the definition of \n",
    "a `compute_primal` method that replaces the `compute` method for `JaxExplicitComponent` and the\n",
    "`apply_nonlinear` method for `JaxImplicitComponent`.\n",
    "\n",
    "This notebook will describe in more detail how to create and use a JaxExplicitComponent or \n",
    "JaxImplicitComponent and will give examples.\n",
    "\n",
    "Before going further, it's a good idea to aquaint yourself with some of jax's 'sharp edges' \n",
    "[here](https://docs.jax.dev/en/latest/notebooks/Common_Gotchas_in_JAX.html). This will hopefully \n",
    "make the process of creating a `JaxExplicitComponent` or `JaxImplicitComponent` a less frustrating one."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The use of JAX is optional for OpenMDAO so if not already installed, the user needs to install it by \n",
    "issuing *one* of the following commands at your operating system command prompt:\n",
    "```\n",
    "pip install jax jaxlib\n",
    "pip install openmdao[jax]\n",
    "pip install openmdao[all]\n",
    "```\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install jax"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The JAX library includes a NumPy-like API, `jax.numpy`, which implements the NumPy API using the primitives in JAX. Almost anything that can be done with NumPy can be done with `jax.numpy`. JAX arrays are similar to NumPy arrays, but they are designed to work with accelerators such as GPUs and TPUs. `jax.numpy` is typically imported as `jnp`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The default for JAX is to do single precision computations. OpenMDAO uses double precision, so this line of code is needed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "jax.config.update(\"jax_enable_x64\", True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Automatic Determination of Derivative Direction\n",
    "\n",
    "`JaxExplicitComponent` and `JaxImplicitComponent` automatically determine the direction they will\n",
    "use to compute their partial jacobians based on their jacobian's shape.  If there are more rows\n",
    "than columns in the jacobian, they'll use forward mode. Otherwise they'll use reverse mode.\n",
    "The number of columns in the `JaxExplicitComponent`'s jacobian is equal to the size of its inputs\n",
    "vector, and the number of columns in the `JaxImplicitComponent`'s jacobian is equal to combined size\n",
    "of its inputs **and** outputs vectors.\n",
    "Note that this automatic determination of derivative direction only occurs if the `matrix_free`\n",
    "attribute is False.\n",
    "\n",
    "\n",
    "## Self Statics\n",
    "\n",
    "When jax compiles a function, it assumes that the only variables that can change are those that are\n",
    "passed into the function as arguments and any internal variables that depend on those arguments.  \n",
    "All other variables are treated as static. But what if \n",
    "our jax component has an option or attribute that contributes to the output of our `compute_primal`\n",
    "function?  Since that option or attribute doesn't get passed into the function as an argument, jax\n",
    "doesn't know about it. In that case, we must be able to detect when those 'static' options or \n",
    "attributes change so that we can tell jax to recompile the function.  Otherwise the outputs of the \n",
    "function won't reflect the current values of the static options and attributes.\n",
    "\n",
    "In `JaxExplicitComponent` and `JaxImplicitComponent`, we add a method called `get_self_statics` to\n",
    "handle this situation.  `get_self_statics` is a simple method that returns a tuple containing\n",
    "any option or attribute in your component that will affect the output of your `compute_primal` method.\n",
    "If your component doesn't have any of these 'self static' variables then you don't have to define\n",
    "`get_self_statics`.\n",
    "\n",
    "Here's a simple example.  Suppose my component has an option called 'mult1' and an attribute called\n",
    "'mult2', and they're used in `compute_primal` as follows:\n",
    "\n",
    "\n",
    "```python\n",
    "def compute_primal(self, x):\n",
    "    return x * self.options['mult1'] * self.mult2\n",
    "```\n",
    "\n",
    "In this case, we would be required to define the `get_self_statics` method shown below:\n",
    "\n",
    "```python\n",
    "def get_self_statics(self):\n",
    "    return (self.options['mult1'], self.mult2)\n",
    "```\n",
    "\n",
    "Doing this will allow the `compute_primal` to be recompiled whenever `self.options['mult1']` or\n",
    "`self.mult2` change.\n",
    "\n",
    "\n",
    "Note that **not** all of a component's options and/or attributes need to be returned from `get_self_statics`.\n",
    "Only those that are referenced inside of `compute_primal` **and** affect its outputs should be returned.\n",
    "\n",
    "Note also that you **do not** return any variable that you've added to your component via `add_input`,\n",
    "`add_output`, `add_discrete_input`, or `add_discrete_output`, even if you think that they won't\n",
    "change during a run for some reason.  Jax already knows about all of them and can handle changes to\n",
    "them properly.\n",
    "\n",
    "\n",
    "## Configuration Options\n",
    "\n",
    "[JaxExplicitComponent](jax_explicitcomp_api.ipynb) and [JaxImplicitComponent](jax_implicitcomp_api.ipynb) \n",
    "both have the following options:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "import openmdao.api as om\n",
    "class SimpleJaxComp(om.JaxExplicitComponent):\n",
    "    def compute_primal(self, x):\n",
    "        y = 2.*x\n",
    "        return y\n",
    "\n",
    "comp = SimpleJaxComp()\n",
    "om.show_options_table(comp)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Debugging\n",
    "\n",
    "While normally you want the `use_jit` option to be True for performance reasons, if you want to debug\n",
    "your `compute_primal` method it often helps to set `use_jit` to False. This will allow you to put\n",
    "print statements in your `compute_primal` or to set breakpoints inside it with a python debugger.\n",
    "\n",
    "\n",
    "## Examples\n",
    "\n",
    "- [JaxExplicitComponent Example](jax_explicitcomp_api)\n",
    "- [JaxImplicitComponent Example](jax_implicitcomp_api)\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.11.10"
  },
  "orphan": true
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
