{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Coupled ODE\n",
    "\n",
    "ODEs have no special status: they are just a conditional node as well. Let's show an example, the Van der Pol oscillator written as two first-order equations:\n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    "\\dot x &= y \\\\\n",
    "\\dot y &= \\mu (1 - x^2)\\, y - x\n",
    "\\end{aligned}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax\n",
    "jax.config.update('jax_enable_x64', True)\n",
    "import jax.numpy as jnp\n",
    "import cottax.assemblies.jax as cx\n",
    "from cottax.assemblies.jax import Start\n",
    "flat = cx.flat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {},
   "outputs": [],
   "source": [
    "@flat.rate(\"x\", initial=\"x(0)\", span=\"t_s\")\n",
    "def vel(y):\n",
    "    return y\n",
    "\n",
    "@flat.rate(\"y\", initial=\"y(0)\", span=\"t_s\")\n",
    "def accel(mu, x, y):\n",
    "    return mu * ( 1- x**2) * y - x"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3",
   "metadata": {},
   "source": [
    "Each rate declares both a functional node that computes the rate (which is trivial in the x case)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "declared = flat.declare(vel, accel)\n",
    "cx.render_html(declared, blocks=True, formatter = cx.LatexFormatter())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5",
   "metadata": {},
   "source": [
    "As is this graph is not executable:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [],
   "source": [
    "from cottax.core import ExecutableGraph\n",
    "from cottax.core.frozen import violations\n",
    "\n",
    "for v in violations(ExecutableGraph, declared.graph):\n",
    "    print(type(v).__name__, '--', v.message)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7",
   "metadata": {},
   "source": [
    "The easiest way to execute this is to define one integrator."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "combined = declared + flat.combine(\"vdp\", (vel.problem, accel.problem), cx.OdeCondition)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10",
   "metadata": {},
   "source": [
    "Now we drive it. A bare `drive` would integrate to the end of the span and report the final state only. To get the whole solution we assign the contract `Integration(trajectory=True)`: it additionally takes the `Times` to save at, and reports each state's `Trajectory`. `Diffrax` then implements it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = (combined\n",
    "         + flat.assign(\"vdp\", cx.Integration(trajectory=True))\n",
    "         + flat.drive(cx.Diffrax(rtol=1e-8, atol=1e-8)))\n",
    "cx.render_html(model, blocks=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12",
   "metadata": {},
   "source": [
    "The initial values and the span are part of the problem, read at the places the rates name (`x(0)`, `y(0)`, and one span `t_s` both share: two ODEs over different spans would not combine), so they are ordinary inputs. The save times are the driver's data, given by role at any one state of the problem."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13",
   "metadata": {},
   "outputs": [],
   "source": [
    "ts  = jnp.linspace(0.0, 20.0, 1001)\n",
    "out = flat.run(model, {cx.Inputs: {'mu': 1.0, 'x(0)': 2.0, 'y(0)': 0.0, 't_s': (0.0, 20.0)}, cx.Times: {'x': ts}})\n",
    "xs, ys = out[cx.Trajectory, 'x'], out[cx.Trajectory, 'y']\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "plt.plot(ts, xs, ts, ys)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14",
   "metadata": {},
   "source": [
    "A report is read by its role at a state, as the times were given: `out[cx.Trajectory, 'x']`, wherever the driver put it. `into` writes the run into a structure: the computed places it holds, or under a role, each state's value at the state's own place, so a history is keyed like the state:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "metadata": {},
   "outputs": [],
   "source": [
    "state   = out.into({'x': None, 'y': None})\n",
    "history = out.into({'x': None, 'y': None}, cx.Trajectory)\n",
    "state['x'], history['x'].shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16",
   "metadata": {},
   "source": [
    "The graph is structure only, so one model runs for every $\\mu$: only the inputs change."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17",
   "metadata": {},
   "source": [
    "And since the run is plain jax, it differentiates through the integrator (diffrax's adjoint) and jits once per graph:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18",
   "metadata": {},
   "outputs": [],
   "source": [
    "import equinox as eqx\n",
    "\n",
    "@eqx.filter_jit\n",
    "def x_end(mu):\n",
    "    return flat.run(model, {cx.Inputs: {'mu': mu, 'x(0)': 2.0, 'y(0)': 0.0, 't_s': (0.0, 5.0)},\n",
    "                            cx.Times: {'x': jnp.linspace(0.0, 5.0, 11)}})['x']\n",
    "\n",
    "print('x(5) =', x_end(1.0), ' dx(5)/dmu =', jax.grad(x_end)(1.0))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19",
   "metadata": {},
   "source": [
    "## An implicit model inside the ODE\n",
    "\n",
    "The conditional and relational conditional nodes can be combined in the same problem. For example, let's make the damping term computed by a residual (in this case trivial!)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20",
   "metadata": {},
   "outputs": [],
   "source": [
    "@flat.rate(\"y\", initial=\"y(0)\", span=\"t_s\")\n",
    "def accel_z(z, x):\n",
    "    return z - x\n",
    "\n",
    "@flat.residual(\"z\")\n",
    "def damping(z, mu, x, y):\n",
    "    return z - mu * (1 - x**2) * y\n",
    "\n",
    "implicit = (flat.declare(vel, accel_z, damping)\n",
    "            + flat.combine('vdp', (vel.problem, accel_z.problem), cx.OdeCondition)\n",
    "            + flat.mda())\n",
    "cx.render_html(implicit, blocks=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "21",
   "metadata": {},
   "source": [
    "Newton solves $z$ inside each rate evaluation, Diffrax integrates around it, and the history is the explicit one:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22",
   "metadata": {},
   "outputs": [],
   "source": [
    "implicit = (implicit + flat.assign('vdp', cx.Integration(trajectory=True))\n",
    "            + flat.drive(cx.Newton(), cx.Diffrax(rtol=1e-8, atol=1e-8)))\n",
    "out = flat.run(implicit, {cx.Inputs: {'mu': 1.0, 'x(0)': 2.0, 'y(0)': 0.0, 't_s': (0.0, 20.0)},\n",
    "                          cx.Times: {'x': jnp.linspace(0.0, 20.0, 1001)}, cx.Start: {'z': 0.0}})\n",
    "bool(jnp.allclose(out[cx.Trajectory, 'x'], history['x']))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "process_port",
   "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.12.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}