{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e569d931",
   "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",
   "id": "111e8af7",
   "metadata": {},
   "source": [
    "(sec:openmdao_jax_subpackage)=\n",
    "# Composable functions via `jax` (`openmdao.jax`)\n",
    "\n",
    "Certain functions are useful in a gradient-based optimization context, such as smooth activation \n",
    "functions or differentiable maximum/minimum functions.\n",
    "\n",
    "Rather than provide a component that forces a user to structure their system in a certain way and \n",
    "add more components than necessary, the `openmdao.jax` package is intended to provide a universal \n",
    "source for _composable_ functions that users can use within their own components.\n",
    "\n",
    "Functions in `openmdao.jax` are built using the [jax](https://github.com/google/jax) Python package.\n",
    "This allows users to develop components that use these functions, along with other code written with \n",
    "`jax`, and leverage capabilities of `jax` like automatic differentiation, vectorization, and \n",
    "just-in-time compilation.  For most users, these functions will be used within the `compute_primal` \n",
    "method of a [JaxExplicitComponent](../../experimental/jax_explicitcomp_api.ipynb) or \n",
    "[JaxImplicitComponent](../../experimental/jax_implicitcomp_api.ipynb), but users who are proficient in jax \n",
    "can also write their own custom components using these functions if necessary.\n",
    "\n",
    "Many of these functions are focused on providing differentiable forms of strictly non-differentiable \n",
    "functions, such as step responses, absolute value, and minimums or maximums. Near regions where the \n",
    "nominal functions would have invalid derivatives, these functions are smooth but will not perfectly \n",
    "match their non-smooth counterparts."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "86efbe58",
   "metadata": {},
   "source": [
    "## Available Functions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a201b2ba",
   "metadata": {},
   "source": [
    "```{eval-rst}\n",
    "    .. autofunction:: openmdao.jax.act_tanh\n",
    "        :noindex:\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "613fa1f1",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import openmdao.jax_funcs as omj\n",
    "\n",
    "fig, ax = plt.subplots(2, 2, figsize=(8, 8))\n",
    "fig.suptitle('Impact of different parameters on act_tanh')\n",
    "x = np.linspace(0, 1, 1000)\n",
    "\n",
    "mup001 = omj.act_tanh(x, mu=0.001, z=0.5, a=0, b=1)\n",
    "mup01 = omj.act_tanh(x, mu=0.01, z=0.5, a=0, b=1)\n",
    "mup1 = omj.act_tanh(x, mu=0.1, z=0.5, a=0, b=1)\n",
    "\n",
    "ax[0, 0].plot(x, mup001, label=r'$\\mu$ = 0.001')\n",
    "ax[0, 0].plot(x, mup01, label=r'$\\mu$ = 0.01')\n",
    "ax[0, 0].plot(x, mup1, label=r'$\\mu$ = 0.1')\n",
    "ax[0, 0].legend()\n",
    "ax[0, 0].grid()\n",
    "\n",
    "zp5 = omj.act_tanh(x, mu=0.01, z=0.5, a=0, b=1)\n",
    "zp4 = omj.act_tanh(x, mu=0.01, z=0.4, a=0, b=1)\n",
    "zp6 = omj.act_tanh(x, mu=0.01, z=0.6, a=0, b=1)\n",
    "\n",
    "ax[0, 1].plot(x, zp4, label=r'$z$ = 0.4')\n",
    "ax[0, 1].plot(x, zp5, label=r'$z$ = 0.5')\n",
    "ax[0, 1].plot(x, zp6, label=r'$z$ = 0.6')\n",
    "ax[0, 1].legend()\n",
    "ax[0, 1].grid()\n",
    "\n",
    "a0 = omj.act_tanh(x, mu=0.01, z=0.5, a=0, b=1)\n",
    "ap2 = omj.act_tanh(x, mu=0.01, z=0.5, a=0.2, b=1)\n",
    "ap4 = omj.act_tanh(x, mu=0.01, z=0.5, a=0.4, b=1)\n",
    "\n",
    "ax[1, 0].plot(x, a0, label=r'$a$ = 0.0')\n",
    "ax[1, 0].plot(x, ap2, label=r'$a$ = 0.2')\n",
    "ax[1, 0].plot(x, ap4, label=r'$a$ = 0.4')\n",
    "ax[1, 0].legend()\n",
    "ax[1, 0].grid()\n",
    "\n",
    "bp6 = omj.act_tanh(x, mu=0.01, z=0.5, a=0, b=.6)\n",
    "bp8 = omj.act_tanh(x, mu=0.01, z=0.5, a=0, b=.8)\n",
    "b1 = omj.act_tanh(x, mu=0.01, z=0.5, a=0, b=1)\n",
    "\n",
    "ax[1, 1].plot(x, bp6, label=r'$b$ = 0.6')\n",
    "ax[1, 1].plot(x, bp8, label=r'$b$ = 0.8')\n",
    "ax[1, 1].plot(x, b1, label=r'$b$ = 1.0')\n",
    "ax[1, 1].legend()\n",
    "ax[1, 1].grid()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "716b50cc",
   "metadata": {},
   "source": [
    "```{eval-rst}\n",
    "    .. autofunction:: openmdao.jax.smooth_abs\n",
    "        :noindex:\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "308f5047",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1, figsize=(4, 4))\n",
    "fig.suptitle('Impact of different parameters on smooth_abs')\n",
    "x = np.linspace(-0.2, 0.2, 1000)\n",
    "\n",
    "mup001 = omj.smooth_abs(x, mu=0.001)\n",
    "mup01 = omj.smooth_abs(x, mu=0.01)\n",
    "mup1 = omj.smooth_abs(x, mu=0.1)\n",
    "\n",
    "ax.plot(x, mup001, label=r'$\\mu$ = 0.001')\n",
    "ax.plot(x, mup01, label=r'$\\mu$ = 0.01')\n",
    "ax.plot(x, mup1, label=r'$\\mu$ = 0.1')\n",
    "ax.legend()\n",
    "ax.grid()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "268fc9f7",
   "metadata": {},
   "source": [
    "```{eval-rst}\n",
    "    .. autofunction:: openmdao.jax.smooth_max\n",
    "        :noindex:\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "810cb8c2",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1, figsize=(4, 4))\n",
    "fig.suptitle('Impact of different parameters on smooth_max of sin and cos')\n",
    "x = np.linspace(0.5, 1, 1000)\n",
    "\n",
    "sin = np.sin(x)\n",
    "cos = np.cos(x)\n",
    "\n",
    "mup001 = omj.smooth_max(sin, cos, mu=0.001)\n",
    "mup01 = omj.smooth_max(sin, cos, mu=0.01)\n",
    "mup1 = omj.smooth_max(sin, cos, mu=0.1)\n",
    "\n",
    "ax.plot(x, sin, '--', label=r'$\\sin{x}$')\n",
    "ax.plot(x, cos, '--', label=r'$\\cos{x}$')\n",
    "ax.plot(x, mup01, label=r'$\\mu$ = 0.01')\n",
    "ax.plot(x, mup1, label=r'$\\mu$ = 0.1')\n",
    "ax.legend()\n",
    "ax.grid()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77b3f93a",
   "metadata": {},
   "source": [
    "```{eval-rst}\n",
    "    .. autofunction:: openmdao.jax.smooth_min\n",
    "        :noindex:\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7322fd31",
   "metadata": {
    "scrolled": true,
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1, figsize=(4, 4))\n",
    "fig.suptitle('Impact of different parameters on smooth_min of sin and cos')\n",
    "x = np.linspace(0.5, 1, 1000)\n",
    "\n",
    "sin = np.sin(x)\n",
    "cos = np.cos(x)\n",
    "\n",
    "mup001 = omj.smooth_min(sin, cos, mu=0.001)\n",
    "mup01 = omj.smooth_min(sin, cos, mu=0.01)\n",
    "mup1 = omj.smooth_min(sin, cos, mu=0.1)\n",
    "\n",
    "ax.plot(x, sin, '--', label=r'$\\sin{x}$')\n",
    "ax.plot(x, cos, '--', label=r'$\\cos{x}$')\n",
    "ax.plot(x, mup01, label=r'$\\mu$ = 0.01')\n",
    "ax.plot(x, mup1, label=r'$\\mu$ = 0.1')\n",
    "ax.legend(ncol=2)\n",
    "ax.grid()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5087f654",
   "metadata": {},
   "source": [
    "```{eval-rst}\n",
    "    .. autofunction:: openmdao.jax.smooth_round\n",
    "        :noindex:\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8a697c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1, figsize=(8, 4))\n",
    "fig.suptitle('Impact of different parameters on smooth_round of a step function')\n",
    "x = np.linspace(2, -2, 1000)\n",
    "x_round = np.round(x)\n",
    "\n",
    "mup01 = omj.smooth_round(x, mu=0.05)\n",
    "mup1 = omj.smooth_round(x, mu=0.1)\n",
    "\n",
    "ax.plot(x, x_round, '--', label=r'$np.round(x)$')\n",
    "\n",
    "ax.plot(x, mup01, label=r'$\\mu$ = 0.05')\n",
    "ax.plot(x, mup1, label=r'$\\mu$ = 0.1')\n",
    "ax.legend(ncol=2)\n",
    "ax.grid()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5828da93",
   "metadata": {},
   "source": [
    "```{eval-rst}\n",
    "    .. autofunction:: openmdao.jax.ks_max\n",
    "        :noindex:\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf24ddab",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "from openmdao.jax_funcs import ks_max\n",
    "\n",
    "fig, ax = plt.subplots(1, 1, figsize=(4, 4))\n",
    "fig.suptitle('Impact of different parameters on ks_max')\n",
    "y = np.random.random(100)\n",
    "x = np.linspace(0, 1, 100)\n",
    "\n",
    "rho1 = ks_max(y, rho=10.)\n",
    "rho10 = ks_max(y, rho=100.)\n",
    "rho100 = ks_max(y, rho=1000.)\n",
    "\n",
    "ax.plot(x, y, '.', label='y')\n",
    "ax.plot(x, rho1 * np.ones_like(x), label='ks_max(y, rho=10)')\n",
    "ax.plot(x, rho10 * np.ones_like(x), label='ks_max(y, rho=100)')\n",
    "ax.legend(ncol=1)\n",
    "ax.grid()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13ea9f69",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "from openmdao.jax_funcs import ks_min\n",
    "\n",
    "fig, ax = plt.subplots(1, 1, figsize=(4, 4))\n",
    "fig.suptitle('Impact of different parameters on ks_min')\n",
    "y = np.random.random(100) + 5\n",
    "x = np.linspace(0, 1, 100)\n",
    "\n",
    "rho1 = ks_min(y, rho=10.)\n",
    "rho10 = ks_min(y, rho=100.)\n",
    "rho100 = ks_min(y, rho=1000.)\n",
    "\n",
    "ax.plot(x, y, '.', label='y')\n",
    "ax.plot(x, rho1 * np.ones_like(x), label='ks_min(y, rho=10)')\n",
    "ax.plot(x, rho10 * np.ones_like(x), label='ks_min(y, rho=100)')\n",
    "ax.legend(ncol=1)\n",
    "ax.grid()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1ed36a21",
   "metadata": {},
   "source": [
    "## Getting derivatives from jax-composed functions\n",
    "\n",
    "If the user writes a function that is composed entirely using jax-based functions (from `jax.numpy`, etc.), then `jax` will in most cases be able to provide derivatives of those functions automatically.\n",
    "\n",
    "The library has several ways of doing this and the best approach will likely depend on the specific use-case at hand.\n",
    "Rather than provide a component to wrap a `jax` function and provide derivatives automatically, consider the following example as a template for how to utilize `jax` in combination with OpenMDAO components.\n",
    "\n",
    "The following component uses the `jax` library's numpy implementation to compute the root-mean-square (rms) of an array of data.  It then passes this data through the `openmdao.jax.act_tanh` activation function.\n",
    "\n",
    "The arguments to `act_tanh` are such that it will return a value of approximately 1.0 if the rms is greater than a threshold value of 0.5, or approximately 0.0 if the rms is less than this value.  This `act_tanh` function is an activation function that smoothly transitions from 0.0 to 1.0 such that it is differentiable. Near the threhold value it will return some value between 0.0 and 1.0.\n",
    "\n",
    "\\begin{align}\n",
    "\\mathrm{rms\\_switch} = \\mathrm{act\\_tanh}\\left(\\sqrt{\\frac{1}{n}\\sum{\\left(x^2\\right)}}\\right)\n",
    "\\end{align}\n",
    "\n",
    "\n",
    "### compute_primal\n",
    "\n",
    "If OpenMDAO sees a Component method with the name `compute_primal`, it assumes that the method takes \n",
    "the component's inputs as positional arguments and returns the component's outputs as a tuple. \n",
    "JaxExplicitComponent and JaxImplicitComponent both require a `compute_primal` method to be defined,\n",
    "but any OpenMDAO component may declare `compute_primal` instead of `compute`.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac28484d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import jax.numpy as jnp\n",
    "\n",
    "import openmdao.api as om\n",
    "from openmdao.jax_funcs import act_tanh\n",
    "\n",
    "\n",
    "class RootMeanSquareSwitchComp(om.JaxExplicitComponent):\n",
    "\n",
    "    def initialize(self):\n",
    "        self.options.declare('vec_size', types=(int,))\n",
    "        self.options.declare('mu', types=(float,), default=0.01)\n",
    "        self.options.declare('threshold', types=(float,), default=0.5)\n",
    "\n",
    "    def setup(self):\n",
    "        n = self.options['vec_size']\n",
    "        self.add_input('x', shape=(n,))\n",
    "        self.add_output('rms', shape=())\n",
    "        self.add_output('rms_switch', shape=())\n",
    "\n",
    "        # we can declare partials here if we know them, but in most cases it's best just to\n",
    "        # let the component determine them (and any sparsity) automatically.\n",
    "\n",
    "    # because our compute_primal references 'static' data, i.e. data that won't change during\n",
    "    # the execution of the component, we need to provide a way to let jax know about this data.\n",
    "    # This is important in order to cause jax to recompile the function if any of the static data\n",
    "    # changes, between runs for example.\n",
    "    def get_self_statics(self):\n",
    "        return (self.options['vec_size'], self.options['mu'], self.options['threshold'])\n",
    "\n",
    "    def compute_primal(self, x):\n",
    "        n = self.options['vec_size']\n",
    "        mu = self.options['mu']\n",
    "        z = self.options['threshold']\n",
    "        rms = jnp.sqrt(jnp.sum(x**2) / n)\n",
    "        return rms, act_tanh(rms, mu, z, 0.0, 1.0)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d19c963",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "N = 100\n",
    "np.random.seed(16)\n",
    "\n",
    "p = om.Problem()\n",
    "p.model.add_subsystem('counter', RootMeanSquareSwitchComp(vec_size=N),\n",
    "                      promotes_inputs=['x'], promotes_outputs=['rms', 'rms_switch'])\n",
    "\n",
    "\n",
    "p.setup(force_alloc_complex=True)\n",
    "p.set_val('x', np.random.random(N))\n",
    "p.run_model()\n",
    "\n",
    "print('Derivative method: {deriv_method}')\n",
    "print('rms = ', p.get_val('rms'))\n",
    "print('rms_switch = ', p.get_val('rms_switch'))\n",
    "\n",
    "print('\\nchecking partials')\n",
    "with np.printoptions(linewidth=1024):\n",
    "    cpd = p.check_partials(method='fd', compact_print=True)\n",
    "print()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aef56f31",
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "from openmdao.utils.assert_utils import assert_check_partials\n",
    "\n",
    "assert_check_partials(cpd)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dc8aae02",
   "metadata": {},
   "source": [
    "## Example 2: A component with vector inputs and outputs\n",
    "\n",
    "A common pattern is to have a vectorized input and a corresponding vectorized output.\n",
    "For a simple vectorized calculation this will typically result in a diagonal jacobian, where the\n",
    "n-th element of the input only impacts the n-th element of the output.  JaxExplicitComponent\n",
    "will automatically detect the sparsity of the partial jacobian if we don't declare any partials."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef4be7f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "class SinCosComp(om.JaxExplicitComponent):\n",
    "\n",
    "    def initialize(self):\n",
    "        self.options.declare('vec_size', types=(int,))\n",
    "\n",
    "    def setup(self):\n",
    "        n = self.options['vec_size']\n",
    "        self.add_input('x', shape=(n,))\n",
    "        self.add_output('sin_cos_x', shape=(n,))\n",
    "\n",
    "        # We'll let jax automatically detect our partials and their sparsity by not declaring\n",
    "        # any partials.\n",
    "\n",
    "    # because we don't reference any static data in our compute_primal, we don't need to provide\n",
    "    # a get_self_statics method.\n",
    "\n",
    "    def compute_primal(self, x):\n",
    "        return jnp.sin(jnp.cos(x))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55b8389e",
   "metadata": {},
   "outputs": [],
   "source": [
    "N = 8\n",
    "np.random.seed(16)\n",
    "\n",
    "p = om.Problem()\n",
    "scx = p.model.add_subsystem('scx', SinCosComp(vec_size=N),\n",
    "                             promotes_inputs=['x'], promotes_outputs=['sin_cos_x'])\n",
    "\n",
    "\n",
    "p.setup(force_alloc_complex=True)\n",
    "p.set_val('x', np.random.random(N))\n",
    "p.run_model()\n",
    "\n",
    "print('sin(cos(x)) = ', p.get_val('sin_cos_x'))\n",
    "\n",
    "print('\\nchecking partials')\n",
    "with np.printoptions(linewidth=1024):\n",
    "    cpd = p.check_partials(method='cs', compact_print=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c29a361b",
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "assert_check_partials(cpd, atol=1.0E-5, rtol=1.0E-5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dddcf7dc",
   "metadata": {},
   "source": [
    "## Example 3: A component with dynamically shaped inputs and outputs\n",
    "\n",
    "Jax can determine output shapes based on input shapes at runtime, so if no shape information is\n",
    "'hard wired' into your `compute_primal` method, you can use OpenMDAO's dynamic shaping capability\n",
    "to figure out the shapes automatically."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bdf9fcfc",
   "metadata": {},
   "outputs": [],
   "source": [
    "class SinCosDynamicComp(om.JaxExplicitComponent):\n",
    "\n",
    "    def setup(self):\n",
    "        self.add_input('x')\n",
    "        self.add_output('sin_cos_x')\n",
    "\n",
    "    def compute_primal(self, x):\n",
    "        return jnp.sin(jnp.cos(x))\n",
    "\n",
    "p = om.Problem()\n",
    "# by setting default_to_dyn_shapes to True, we tell OpenMDAO to use dynamic shapes by default\n",
    "# for any variables where we don't set a shape. Note that if you want to use this option and it is\n",
    "# not passed in when creating the component as it is in this case, you must set it BEFORE adding\n",
    "# any variables in your component's setup method. Otherwise it will be ignored.\n",
    "p.model.add_subsystem('scx', SinCosDynamicComp(default_to_dyn_shapes=True),\n",
    "                      promotes_inputs=['x'],\n",
    "                      promotes_outputs=['sin_cos_x'])\n",
    "\n",
    "p.setup(force_alloc_complex=True)\n",
    "p.set_val('x', np.random.random((3, 4)))\n",
    "p.run_model()\n",
    "\n",
    "print('sin(cos(x)) = ', p.get_val('sin_cos_x'))\n",
    "\n",
    "print('\\nchecking partials')\n",
    "with np.printoptions(linewidth=1024):\n",
    "    cpd = p.check_partials(method='cs', compact_print=True)\n",
    "print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15aafa71",
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "from openmdao.utils.assert_utils import assert_check_partials\n",
    "\n",
    "assert_check_partials(cpd)"
   ]
  }
 ],
 "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": 5
}
