Add documentation for second-order ODEs with JointEq

Co-authored-by: Routhleck <[email protected]>

Author: Copilot

docs_classic/tutorial_toolbox/joint_equations.ipynb

@@ -24,7 +24,7 @@
    "id": "c9df7780",
    "metadata": {},
    "source": [
-    "In a [dynamical system](../tutorial_building/dynamical_systems.ipynb), there may be multiple variables that change dynamically over time. Sometimes these variables are interconnected, and updating one variable requires others as the input. For example, in the widely known Hodgkin–Huxley model, the variables $V$, $m$, $h$, and $n$ are updated synchronously and interdependently (please refer to [Building Neuron Models](../tutorial_building/neuron_models.ipynb)for details). To achieve higher integral accuracy, it is recommended to use ``brainpy.JointEq`` to jointly solving interconnected differential equations."
+    "In a [dynamical system](../tutorial_building/dynamical_systems.ipynb), there may be multiple variables that change dynamically over time. Sometimes these variables are interconnected, and updating one variable requires others as the input. For example, in the widely known Hodgkin\u2013Huxley model, the variables $V$, $m$, $h$, and $n$ are updated synchronously and interdependently (please refer to [Building Neuron Models](../tutorial_building/neuron_models.ipynb)for details). To achieve higher integral accuracy, it is recommended to use ``brainpy.JointEq`` to jointly solving interconnected differential equations."
    ]
   },
   {
@@ -304,6 +304,181 @@
     "It is shown in this output code that second differential values of $v$ and $u$ are calculated by using the updated values (`k2_V_arg` and `k2_u_arg`) at the same time. This will result in a more accurate integral."
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "second_order_ode_title",
+   "metadata": {},
+   "source": [
+    "## Second-Order ODEs with `brainpy.JointEq`\n",
+    "\n",
+    "A common use case for `JointEq` is solving second-order ordinary differential equations (ODEs). ",
+    "Second-order ODEs appear in many physical systems, such as the harmonic oscillator, pendulum, ",
+    "or neural mass models like the Jansen-Rit model.\n",
+    "\n",
+    "When using `JointEq` for second-order ODEs, it's important to follow the correct function signature pattern."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "second_order_example",
+   "metadata": {},
+   "source": [
+    "### Example: Harmonic Oscillator\n",
+    "\n",
+    "Consider a damped harmonic oscillator described by:\n",
+    "\n",
+    "$$\\frac{d^2x}{dt^2} = -kx - c\\frac{dx}{dt}$$\n",
+    "\n",
+    "To solve this with `JointEq`, we split it into two first-order ODEs:\n",
+    "\n",
+    "$$\\frac{dx}{dt} = v$$\n",
+    "$$\\frac{dv}{dt} = -kx - cv$$\n",
+    "\n",
+    "Where $x$ is position and $v$ is velocity."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "id": "second_order_code",
+   "metadata": {},
+   "source": [
+    "import brainpy as bp\n",
+    "import brainpy.math as bm\n",
+    "\n",
+    "# Parameters\n",
+    "k = 1.0  # spring constant\n",
+    "c = 0.1  # damping coefficient\n",
+    "\n",
+    "# Define derivative functions\n",
+    "# IMPORTANT: Each state variable appears as the FIRST parameter before 't'\n",
+    "# Other state variables appear AFTER 't' as dependencies\n",
+    "def dx(x, t, v):\n",
+    "    \"\"\"dx/dt = v\"\"\"\n",
+    "    return v\n",
+    "\n",
+    "def dv(v, t, x):\n",
+    "    \"\"\"dv/dt = -k*x - c*v\"\"\"\n",
+    "    return -k * x - c * v\n",
+    "\n",
+    "# Create joint equation\n",
+    "joint_eq = bp.JointEq(dx, dv)\n",
+    "print(f\"Joint equation signature: {joint_eq.__signature__}\")"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "cell_type": "markdown",
+   "id": "second_order_signature",
+   "metadata": {},
+   "source": [
+    "### Important: Function Signature Pattern\n",
+    "\n",
+    "When defining derivative functions for `JointEq`, follow this pattern:\n",
+    "\n",
+    "**Correct:**\n",
+    "```python\n",
+    "def dx(x, t, v):  # x is the state variable, v is a dependency\n",
+    "    return v\n",
+    "\n",
+    "def dv(v, t, x):  # v is the state variable, x is a dependency\n",
+    "    return -k * x - c * v\n",
+    "```\n",
+    "\n",
+    "**Incorrect:**\n",
+    "```python\n",
+    "def dx(x, v, t):  # WRONG: Both x and v before t\n",
+    "    return v\n",
+    "```\n",
+    "\n",
+    "**Rule:** Each state variable should appear as the **first parameter before** `t` in exactly one derivative function. ",
+    "If a variable is needed as a dependency in another function, it should be placed **after** `t`.\n",
+    "\n",
+    "This ensures that `JointEq` knows which variable each function is differentiating and which variables are dependencies."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "second_order_jansen_rit",
+   "metadata": {},
+   "source": [
+    "### Example: Jansen-Rit Model\n",
+    "\n",
+    "The Jansen-Rit model is a neural mass model with three coupled second-order ODEs. ",
+    "Here's how to implement it correctly with `JointEq`:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "id": "jansen_rit_code",
+   "metadata": {},
+   "source": [
+    "class JansenRitModel(bp.dyn.NeuDyn):\n",
+    "    def __init__(self, size=1, A=3.25, te=10, B=22, ti=20, C=135, \n",
+    "                 e0=2.5, r=0.56, v0=6, method='rk4', **kwargs):\n",
+    "        super().__init__(size=size, **kwargs)\n",
+    "        self.A, self.te = A, te\n",
+    "        self.B, self.ti = B, ti\n",
+    "        self.C = C\n",
+    "        self.e0, self.r, self.v0 = e0, r, v0\n",
+    "        \n",
+    "        # State variables: positions (y0, y1, y2) and velocities (y3, y4, y5)\n",
+    "        self.y0 = bm.Variable(bm.zeros(self.num))\n",
+    "        self.y1 = bm.Variable(bm.zeros(self.num))\n",
+    "        self.y2 = bm.Variable(bm.zeros(self.num))\n",
+    "        self.y3 = bm.Variable(bm.zeros(self.num))  # velocity for y0\n",
+    "        self.y4 = bm.Variable(bm.zeros(self.num))  # velocity for y1\n",
+    "        self.y5 = bm.Variable(bm.zeros(self.num))  # velocity for y2\n",
+    "        \n",
+    "        self.integral = bp.odeint(f=self.derivative, method=method)\n",
+    "    \n",
+    "    # Position derivatives: dx/dt = v\n",
+    "    def dy0(self, y0, t, y3):  # y0 is state, y3 is dependency\n",
+    "        return y3 / 1000\n",
+    "    \n",
+    "    def dy1(self, y1, t, y4):  # y1 is state, y4 is dependency\n",
+    "        return y4 / 1000\n",
+    "    \n",
+    "    def dy2(self, y2, t, y5):  # y2 is state, y5 is dependency\n",
+    "        return y5 / 1000\n",
+    "    \n",
+    "    # Velocity derivatives: dv/dt = ...\n",
+    "    def dy3(self, y3, t, y0, y1, y2):  # y3 is state, others are dependencies\n",
+    "        Sp = 2 * self.e0 / (1 + bm.exp(self.r * (self.v0 - y1 + y2)))\n",
+    "        return (self.A * Sp - 2 * y3 - y0 / self.te * 1000) / self.te\n",
+    "    \n",
+    "    def dy4(self, y4, t, y0, y1, inp=0.):  # y4 is state, others are dependencies\n",
+    "        Se = 2 * self.e0 / (1 + bm.exp(self.r * (self.v0 - self.C * y0)))\n",
+    "        return (self.A * (inp + 0.8 * self.C * Se) - 2 * y4 - y1 / self.te * 1000) / self.te\n",
+    "    \n",
+    "    def dy5(self, y5, t, y0, y2):  # y5 is state, others are dependencies\n",
+    "        Si = 2 * self.e0 / (1 + bm.exp(self.r * (self.v0 - 0.25 * self.C * y0)))\n",
+    "        return (self.B * 0.25 * self.C * Si - 2 * y5 - y2 / self.ti * 1000) / self.ti\n",
+    "    \n",
+    "    @property\n",
+    "    def derivative(self):\n",
+    "        # Join all derivatives - order matches the state variables\n",
+    "        return bp.JointEq([self.dy0, self.dy1, self.dy2, self.dy3, self.dy4, self.dy5])\n",
+    "    \n",
+    "    def update(self, inp=0.):\n",
+    "        y0, y1, y2, y3, y4, y5 = self.integral(\n",
+    "            self.y0, self.y1, self.y2, self.y3, self.y4, self.y5,\n",
+    "            bp.share['t'], inp, bp.share['dt']\n",
+    "        )\n",
+    "        self.y0.value = y0\n",
+    "        self.y1.value = y1\n",
+    "        self.y2.value = y2\n",
+    "        self.y3.value = y3\n",
+    "        self.y4.value = y4\n",
+    "        self.y5.value = y5\n",
+    "\n",
+    "# Create and test the model\n",
+    "model = JansenRitModel(size=1)\n",
+    "print(\"Jansen-Rit model created successfully!\")"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
   {
    "cell_type": "markdown",
    "id": "73051bec",
@@ -336,4 +511,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 5
-}
+}