{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# A function that is not jax\n",
    "\n",
    "cottax never looks inside a node's implementation: a run just calls it. So a function written\n",
    "in numpy, or a wrapper around a compiled code, is an ordinary node. What has to change is the\n",
    "function itself, because the solvers trace and differentiate whatever they are handed.\n",
    "\n",
    "This notebook goes through the steps:\n",
    "\n",
    "1. a plain numpy function, and where it breaks\n",
    "2. `jax.pure_callback`, so jax can call it\n",
    "3. finite differences, so jax can differentiate it\n",
    "4. the result inside the Sellar problem, solved by Newton and differentiated\n",
    "5. some notes on cost, and on codes that come with their own derivatives"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax\n",
    "jax.config.update('jax_enable_x64', True)   # needed for the finite differences, see section 3\n",
    "import numpy as np\n",
    "import jax.numpy as jnp\n",
    "from jax.flatten_util import ravel_pytree\n",
    "\n",
    "import cottax.assemblies.jax as cx\n",
    "from cottax.assemblies.jax import Inputs, Start\n",
    "flat = cx.flat"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2",
   "metadata": {},
   "source": [
    "## 1. A numpy function\n",
    "\n",
    "This stands in for something like a compiled solver. On concrete numbers it works fine, and for\n",
    "a graph without cycles that is run eagerly, that is all you need."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {},
   "outputs": [],
   "source": [
    "def disc2_host(y_1, z_1, z_2):\n",
    "    return np.sqrt(np.abs(y_1)) + z_1 + z_2\n",
    "\n",
    "disc2_host(4.0, 1.0, 2.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4",
   "metadata": {},
   "source": [
    "Hand it a tracer, though, and it fails:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def show(label, thunk):\n",
    "    try:\n",
    "        thunk(); print(f'{label:6s} ok')\n",
    "    except Exception as e:\n",
    "        print(f'{label:6s} {type(e).__name__}: {str(e).splitlines()[0][:80]}')\n",
    "\n",
    "a = (jnp.asarray(4.3), jnp.asarray(1.1), jnp.asarray(2.2))\n",
    "show('jit',  lambda: jax.jit(disc2_host)(*a))\n",
    "show('grad', lambda: jax.grad(disc2_host)(*a))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6",
   "metadata": {},
   "source": [
    "`np.sqrt` on a tracer calls `__array__`, which raises. Note that you meet this even if you never\n",
    "write `jit` yourself: a Newton solver differentiates the residual to take a step, so every\n",
    "function inside a solved cycle gets traced.\n",
    "\n",
    "## 2. `pure_callback`\n",
    "\n",
    "`jax.pure_callback` runs the function on the host with concrete numpy arrays, and only tells\n",
    "jax the shape of the result. That shape has to be given up front, since jax cannot run the\n",
    "function to find out."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "out_shape = jax.ShapeDtypeStruct((), jnp.float64)\n",
    "\n",
    "def disc2_cb(y_1, z_1, z_2):\n",
    "    return jax.pure_callback(disc2_host, out_shape, y_1, z_1, z_2, vmap_method='sequential')\n",
    "\n",
    "show('jit',  lambda: jax.jit(disc2_cb)(*a))\n",
    "show('grad', lambda: jax.grad(disc2_cb)(*a))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8",
   "metadata": {},
   "source": [
    "Tracing works now, differentiating does not: jax has no idea what the derivative of the callback\n",
    "is. That needs a `jax.custom_jvp`.\n",
    "\n",
    "## 3. Finite differences\n",
    "\n",
    "There are two ways to write a finite-difference JVP rule, and they behave differently:\n",
    "\n",
    "| mode | rule | calls per JVP | forward mode | reverse mode |\n",
    "|---|---|---|---|---|\n",
    "| `directional` | `(f(x + h v) - f(x)) / h` | 1 | yes | no |\n",
    "| `jacobian` | `J @ v`, `J` from one FD sweep | `n + 1` | yes | yes |\n",
    "\n",
    "Reverse mode needs jax to transpose the rule, which only works if the rule is linear in the\n",
    "tangent `v`. In `directional` mode `v` goes into the callback, so it is not; in `jacobian` mode\n",
    "it only appears in the matrix product, so it is."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Tally:\n",
    "    \"\"\"Counts how often the host function ran.\"\"\"\n",
    "    def __init__(self): self.calls = 0\n",
    "    def reset(self):    self.calls = 0; return self\n",
    "    def __repr__(self): return f'{self.calls} call(s)'\n",
    "\n",
    "\n",
    "def blackbox(fn, out_shape, *, mode='jacobian', tally=None):\n",
    "    \"\"\"`fn` made traceable (`pure_callback`) and differentiable (`custom_jvp`, finite differences).\"\"\"\n",
    "    step  = np.sqrt(np.finfo(np.dtype(out_shape.dtype)).eps)\n",
    "\n",
    "    def call(*args):\n",
    "        def host(*a):\n",
    "            if tally: tally.calls += 1\n",
    "            return fn(*a)\n",
    "        return jax.pure_callback(host, out_shape, *args, vmap_method='sequential')\n",
    "\n",
    "    @jax.custom_jvp\n",
    "    def f(*args):\n",
    "        return call(*args)\n",
    "\n",
    "    @f.defjvp\n",
    "    def _(primals, tangents):\n",
    "        x, unravel = ravel_pytree(primals)\n",
    "        v, _       = ravel_pytree(tangents)\n",
    "        y0         = call(*primals)\n",
    "\n",
    "        if mode == 'directional':\n",
    "            nv = jnp.linalg.norm(v)\n",
    "            h  = step * (1 + jnp.linalg.norm(x)) / jnp.where(nv == 0, 1, nv)\n",
    "            return y0, (call(*unravel(x + h * v)) - y0) / h\n",
    "\n",
    "        # one probe per input, so the rule is `J @ v`: linear in v, so jax can transpose it\n",
    "        h = step * (1 + jnp.abs(x))\n",
    "        Y = jax.vmap(lambda w: jnp.ravel(call(*unravel(w))))(x[None, :] + jnp.diag(h))\n",
    "        J = ((Y - jnp.ravel(y0)[None, :]) / h[:, None]).T\n",
    "        return y0, (J @ v).reshape(out_shape.shape)\n",
    "\n",
    "    return f"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10",
   "metadata": {},
   "source": [
    "Compared against the same function written in jax:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11",
   "metadata": {},
   "outputs": [],
   "source": [
    "def disc2_jax(y_1, z_1, z_2):\n",
    "    return jnp.sqrt(jnp.abs(y_1)) + z_1 + z_2\n",
    "\n",
    "exact = jax.jacfwd(disc2_jax, argnums=(0, 1, 2))(*a)\n",
    "\n",
    "for mode in ('directional', 'jacobian'):\n",
    "    tally = Tally()\n",
    "    f     = blackbox(disc2_host, out_shape, mode=mode, tally=tally)\n",
    "    print(mode)\n",
    "    for name, d in (('jacfwd', jax.jacfwd(f, argnums=(0, 1, 2))),\n",
    "                    ('jacrev', jax.jacrev(f, argnums=(0, 1, 2)))):\n",
    "        tally.reset()\n",
    "        try:\n",
    "            J   = d(*a)\n",
    "            err = max(float(abs(p - q)) for p, q in zip(J, exact))\n",
    "            print(f'   {name}   {tally}   max err {err:.1e}')\n",
    "        except Exception as e:\n",
    "            print(f'   {name}   {type(e).__name__}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12",
   "metadata": {},
   "source": [
    "Some things to note here:\n",
    "\n",
    "- `jacrev` with the directional rule fails, because the rule cannot be transposed. The error\n",
    "  is not very helpful, so it is good to recognise it.\n",
    "- A full Jacobian costs `n + 1` calls in both modes. The difference is only when they are paid:\n",
    "  once per Jacobian, or once per direction.\n",
    "- `n` is the number of inputs of this function, not of the whole model. Everything around it is\n",
    "  still differentiated exactly by jax, so a large model with one small black box only pays for\n",
    "  the box's own inputs. This is the main reason to wrap the node rather than finite-difference\n",
    "  the whole model.\n",
    "- Use double precision. Finite differences are accurate to about `sqrt(eps)`, which is `1e-8` in\n",
    "  double but only `5e-4` in single precision. The function itself has to compute in double as\n",
    "  well; casting its result afterwards does not help.\n",
    "\n",
    "## 4. In a cottax graph\n",
    "\n",
    "This is the Sellar problem from the [Sellar example](implicit_sellar.ipynb), without the\n",
    "optimisation: just the coupled analysis, with discipline 2 replaced by the black box. The same\n",
    "model is declared twice, once with each version of `D2`, so the two can be compared."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13",
   "metadata": {},
   "outputs": [],
   "source": [
    "TALLY = Tally()\n",
    "\n",
    "# built once: a node's function is part of the jit cache key\n",
    "disc2 = blackbox(disc2_host, out_shape, mode='jacobian', tally=TALLY)\n",
    "\n",
    "def sellar(d2):\n",
    "    @flat.function(\"y_1\")\n",
    "    def D1(x_1, z_1, z_2, y_2):\n",
    "        return z_1**2 + z_2 + x_1 - 0.2 * y_2\n",
    "\n",
    "    @flat.function(\"y_2\")\n",
    "    def D2(y_1, z_1, z_2):\n",
    "        return d2(y_1, z_1, z_2)\n",
    "\n",
    "    @flat.function(\"f\")\n",
    "    def obj(x_1, z_2, y_1, y_2):\n",
    "        return x_1**2 + z_2 + y_1 + jnp.exp(-y_2)\n",
    "\n",
    "    declared = flat.declare(D1, D2, obj)\n",
    "    # cut y_2 where D1 reads it, close the cycle with ^copy.y_2 = y_2 and solve that by Newton\n",
    "    return declared + flat.fixed_point_cut('y_2', 'D1') + flat.drive(cx.Newton(rtol=1e-10, atol=1e-10))\n",
    "\n",
    "boxed = sellar(disc2)\n",
    "plain = sellar(disc2_jax)\n",
    "cx.render_xdsm_svg(boxed, blocks=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14",
   "metadata": {},
   "source": [
    "The black box sits in the solved cycle like any other node. The inputs and the Newton's starting\n",
    "guess for `^copy.y_2` are given by name:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "metadata": {},
   "outputs": [],
   "source": [
    "def values(z_1=1.0, x_1=1.0):\n",
    "    return {Inputs: dict(z_1=z_1, z_2=2.0, x_1=x_1), Start: {cx.copy('y_2'): 1.0}}\n",
    "\n",
    "TALLY.reset()\n",
    "out = flat.run(boxed, values())\n",
    "ref = flat.run(plain, values())\n",
    "print('black box calls:', TALLY)\n",
    "for name in ('y_1', 'y_2', 'f'):\n",
    "    print(f'{name:4s} {float(out[name]):.10f}   jax: {float(ref[name]):.10f}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16",
   "metadata": {},
   "source": [
    "The Newton converged on a residual that runs through numpy, getting its Jacobian from the finite\n",
    "differences. With only `pure_callback`, the same solve fails at its first step:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17",
   "metadata": {},
   "outputs": [],
   "source": [
    "show('solve', lambda: flat.run(sellar(disc2_cb), values()))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18",
   "metadata": {},
   "source": [
    "The gradient through the solve also works. optimistix differentiates the converged solve\n",
    "implicitly, so the finite differences are used twice: in the Newton's Jacobian, and in the\n",
    "adjoint of the solved cycle."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19",
   "metadata": {},
   "outputs": [],
   "source": [
    "def objective(model):\n",
    "    return lambda z_1, x_1: flat.run(model, values(z_1, x_1))['f']\n",
    "\n",
    "TALLY.reset()\n",
    "g     = jax.grad(objective(boxed), argnums=(0, 1))(1.0, 1.0)\n",
    "calls = TALLY.calls\n",
    "g_ref = jax.grad(objective(plain), argnums=(0, 1))(1.0, 1.0)\n",
    "print('df/dz_1, df/dx_1 =', [float(v) for v in g], f'  ({calls} calls)')\n",
    "print('            jax  =', [float(v) for v in g_ref])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20",
   "metadata": {},
   "source": [
    "Nothing in cottax itself changed for any of this.\n",
    "\n",
    "## 5. Further notes\n",
    "\n",
    "### Reverse mode costs the same as forward mode\n",
    "\n",
    "`mode='jacobian'` already gives a VJP, and `n + 1` calls is the minimum. One evaluation of\n",
    "`f(x + h v)` gives a column of `J` (a JVP). A VJP needs a row, and no perturbation of the inputs\n",
    "gives you one, so you need all `n` columns. Automatic differentiation avoids this by reading the\n",
    "program; a black box cannot be read. So reverse mode loses its advantage through a black box.\n",
    "\n",
    "A hand-written `custom_vjp` costs the same, but then forward mode stops working. Writing the JVP\n",
    "and letting jax transpose it gives you both.\n",
    "\n",
    "### Not storing the Jacobian\n",
    "\n",
    "In reverse mode `J` (dense, `m x n`) is kept in memory until the backward pass. If that is too\n",
    "much, the sweep can be done in the backward pass instead, contracting each column with the\n",
    "cotangent as soon as it is computed:\n",
    "\n",
    "```python\n",
    "def bwd(res, ct):\n",
    "    x, y0 = res\n",
    "    def body(i, acc):\n",
    "        h   = step * (1 + jnp.abs(x[i]))\n",
    "        col = (f(x.at[i].add(h)) - y0) / h\n",
    "        return acc.at[i].set(jnp.vdot(ct, col))\n",
    "    return (jax.lax.fori_loop(0, x.size, body, jnp.zeros_like(x)),)\n",
    "```\n",
    "\n",
    "Same number of calls and the same result, with `O(n + m)` memory instead of `O(n m)`. This is a\n",
    "`custom_vjp`, so forward mode is lost.\n",
    "\n",
    "Which mode is best depends on the solver. A dense Newton builds the Jacobian anyway, so\n",
    "`jacobian` mode costs it nothing extra. A Krylov solver only asks for products `J @ v`, so there\n",
    "`directional` mode needs far fewer calls and less memory. `jacobian` mode inside a Krylov solver\n",
    "is the worst case: `n + 1` calls on every product.\n",
    "\n",
    "### `jax.ffi`: no Python in between\n",
    "\n",
    "All of the above can be jitted, but each callback still goes back to Python: the compiled program\n",
    "stops, hands buffers to CPython, takes the GIL and continues. That is on the order of a hundred\n",
    "microseconds per call, compared to a few for a direct call.\n",
    "\n",
    "If the function can be called through a C ABI, `jax.ffi` makes it a real XLA custom call, about as\n",
    "fast as a native jax operation, and a whole solved cycle then compiles into one XLA program:\n",
    "\n",
    "```python\n",
    "jax.ffi.register_ffi_target('disc2', jax.ffi.pycapsule(lib.Disc2), platform='cpu')\n",
    "jax.ffi.ffi_call('disc2', out_shape, vmap_method='sequential')(y_1, z_1, z_2)\n",
    "```\n",
    "\n",
    "This needs a compiled library (headers from `jax.ffi.include_dir()`) and a handler per platform.\n",
    "It does not make the function differentiable: `ffi_call` fails under `grad` just like\n",
    "`pure_callback`, and the finite-difference wrapper above goes on top unchanged.\n",
    "\n",
    "### Codes with their own derivatives\n",
    "\n",
    "Many older simulation codes come with a tangent-linear and/or adjoint routine. These cannot be\n",
    "attached with `custom_jvp` and `custom_vjp` together, since a function can only have one of the\n",
    "two. The way to do it is a custom primitive for the linear map `v -> J(x) v`, with\n",
    "`ad.primitive_transposes` giving reverse mode and `batching.primitive_batchers` giving the sweep\n",
    "`jacfwd` needs. `lineax` builds its `linear_solve_p` like this.\n",
    "\n",
    "| the code provides | one JVP | one gradient | forward | reverse |\n",
    "|---|---|---|---|---|\n",
    "| value only | 2 calls | `n + 1` calls | FD | FD |\n",
    "| value + tangent | 1 tangent | `n` tangents | exact | exact |\n",
    "| value + adjoint | 2 calls | 1 adjoint | FD | exact |\n",
    "| value + tangent + adjoint | 1 tangent | 1 adjoint | exact | exact |\n",
    "\n",
    "The same construction works for a fast implementation paired with a slower one written in jax:\n",
    "the fast one computes the values, and `jax.jvp` / `jax.vjp` of the jax one give the derivatives.\n",
    "The solve then runs at the speed of the fast implementation and the gradient is that of the jax\n",
    "one, as long as the two really compute the same function.\n",
    "\n",
    "### Summary\n",
    "\n",
    "- Give the output shape; jax cannot run the function to find it.\n",
    "- Build the wrapped function once, when declaring the model. It is part of the jit cache key, and\n",
    "  a new one on every call means retracing every time.\n",
    "- Every tangent a rule uses has to go through the linear part. Finite-differencing one argument\n",
    "  while passing another through the primitive makes the rule nonlinear again, and transposing it\n",
    "  fails.\n",
    "- Use double precision.\n",
    "- Jit the run: calling callbacks eagerly is much slower than the jitted call."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}