edmkit 0.0.14__tar.gz → 0.0.16__tar.gz

{edmkit-0.0.14 → edmkit-0.0.16}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: edmkit
-Version: 0.0.14
+Version: 0.0.16
 Summary: Simple EDM (Empirical Dynamic Modeling) library
 Author: FUJISHIGE TEMMA
 Author-email: FUJISHIGE TEMMA <tenma.x0@gmail.com>

{edmkit-0.0.14 → edmkit-0.0.16}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "edmkit"
-version = "0.0.14"
+version = "0.0.16"
 description = "Simple EDM (Empirical Dynamic Modeling) library"
 readme = "README.md"
 requires-python = ">= 3.13"
@@ -22,6 +22,7 @@ dev = [
     "hypothesis>=6.151.9",
     "matplotlib>=3.11.2",
     "pytest>=9.1.1",
+    "pywavelets>=1.10.0",
     "ruff>=0.15.20",
     "ty>=0.0.56",
 ]
@@ -42,5 +43,5 @@ markers = [
 ]
 
 [build-system]
-requires = ["uv_build>=0.10.9,<0.11.0"]
+requires = ["uv_build>=0.13.6,<0.14.0"]
 build-backend = "uv_build"

{edmkit-0.0.14 → edmkit-0.0.16}/pyproject.toml.orig

@@ -1,6 +1,6 @@
 [project]
 name = "edmkit"
-version = "0.0.14"
+version = "0.0.16"
 description = "Simple EDM (Empirical Dynamic Modeling) library"
 authors = [{ name = "FUJISHIGE TEMMA", email = "tenma.x0@gmail.com" }]
 readme = "README.md"
@@ -19,6 +19,7 @@ dev = [
     "hypothesis>=6.151.9",
     "matplotlib>=3.11.2",
     "pytest>=9.1.1",
+    "pywavelets>=1.10.0",
     "ruff>=0.15.20",
     "ty>=0.0.56",
 ]
@@ -38,7 +39,7 @@ markers = [
 
 
 [build-system]
-requires = ["uv_build>=0.10.9,<0.11.0"]
+requires = ["uv_build>=0.13.6,<0.14.0"]
 build-backend = "uv_build"
 
 [tool.uv.build-backend]

edmkit-0.0.16/src/edmkit/generate/__init__.py

@@ -0,0 +1,40 @@
+from .chua import chua
+from .ctrnn import ctrnn
+from .double_pendulum import double_pendulum, to_xy
+from .duffing import duffing
+from .euler_spiral import euler_spiral
+from .fitzhugh_nagumo import fitzhugh_nagumo
+from .hindmarsh_rose import hindmarsh_rose
+from .kuramoto import kuramoto
+from .lorenz import lorenz
+from .lorenz96 import lorenz96
+from .lotka_volterra import lotka_volterra
+from .mackey_glass import mackey_glass
+from .morris_lecar import morris_lecar
+from .pyloric import pyloric
+from .rk45 import rk45
+from .rossler import rossler
+from .rulkov_map import rulkov_map
+from .three_scroll import three_scroll
+
+__all__ = [
+    "chua",
+    "ctrnn",
+    "double_pendulum",
+    "duffing",
+    "euler_spiral",
+    "fitzhugh_nagumo",
+    "hindmarsh_rose",
+    "kuramoto",
+    "lorenz",
+    "lorenz96",
+    "lotka_volterra",
+    "mackey_glass",
+    "morris_lecar",
+    "pyloric",
+    "rk45",
+    "rossler",
+    "rulkov_map",
+    "three_scroll",
+    "to_xy",
+]

edmkit-0.0.16/src/edmkit/generate/chua.py

@@ -0,0 +1,31 @@
+import numpy as np
+
+__all__ = ["chua"]
+
+
+def chua(alpha: float, beta: float, m0: float, m1: float):
+    """Vector field of Chua's circuit, with the piecewise-linear diode ``h(x) = m1 x + (m0 - m1) (|x + 1| - |x - 1|) / 2``.
+
+    Parameters
+    ----------
+    alpha : float
+        Ratio of capacitances (typical: 15.6).
+    beta : float
+        Ratio of inductance to resistance (typical: 28).
+    m0 : float
+        Diode slope for ``|x| < 1`` (typical: -8/7).
+    m1 : float
+        Diode slope for ``|x| > 1`` (typical: -5/7).
+
+    Returns
+    -------
+    f : Callable[[float, np.ndarray], np.ndarray]
+        Right-hand side ``f(t, x)`` for a state ``(x, y, z)`` of shape ``(3,)``, or ``(B, 3)`` for a batch.
+    """
+
+    def f(t: float, s: np.ndarray):
+        x, y, z = np.moveaxis(s, -1, 0)
+        h = m1 * x + (m0 - m1) * np.clip(x, -1, 1)  # (|x + 1| - |x - 1|) / 2 is clip(x, -1, 1)
+        return np.stack([alpha * (y - x - h), x - y + z, -beta * y], axis=-1)
+
+    return f

{edmkit-0.0.14 → edmkit-0.0.16}/src/edmkit/generate/ctrnn.py

@@ -16,10 +16,10 @@ def ctrnn(g: float, J: np.ndarray):
     Returns
     -------
     f : Callable[[float, np.ndarray], np.ndarray]
-        Right-hand side ``f(t, x)`` for a state of shape ``(N,)``.
+        Right-hand side ``f(t, x)`` for a state of shape ``(N,)``, or ``(B, N)`` for a batch.
     """
 
     def f(t: float, x: np.ndarray):
-        return -x + g * J @ np.tanh(x)
+        return -x + g * np.tanh(x) @ J.T
 
     return f

{edmkit-0.0.14 → edmkit-0.0.16}/src/edmkit/generate/double_pendulum.py

@@ -22,11 +22,11 @@ def double_pendulum(m1: float, m2: float, L1: float, L2: float, g: float):
     Returns
     -------
     f : Callable[[float, np.ndarray], np.ndarray]
-        Right-hand side ``f(t, x)`` for a state ``(theta1, theta2, omega1, omega2)`` of shape ``(4,)``.
+        Right-hand side ``f(t, x)`` for a state ``(theta1, theta2, omega1, omega2)`` of shape ``(4,)``, or ``(B, 4)`` for a batch.
     """
 
     def f(t: float, x: np.ndarray):
-        theta1, theta2, omega1, omega2 = x
+        theta1, theta2, omega1, omega2 = np.moveaxis(x, -1, 0)
         delta = theta1 - theta2
 
         denom = 2 * m1 + m2 - m2 * np.cos(2 * delta)
@@ -44,7 +44,7 @@ def double_pendulum(m1: float, m2: float, L1: float, L2: float, g: float):
             L2 * denom
         )
 
-        return np.array([dtheta1_dt, dtheta2_dt, domega1_dt, domega2_dt])
+        return np.stack([dtheta1_dt, dtheta2_dt, domega1_dt, domega2_dt], axis=-1)
 
     return f
 

{edmkit-0.0.14 → edmkit-0.0.16}/src/edmkit/generate/duffing.py

@@ -22,10 +22,11 @@ def duffing(alpha: float, beta: float, delta: float, gamma: float, omega: float)
     Returns
     -------
     f : Callable[[float, np.ndarray], np.ndarray]
-        Right-hand side ``f(t, x)`` for a state ``(x, v)`` of shape ``(2,)``.
+        Right-hand side ``f(t, x)`` for a state ``(x, v)`` of shape ``(2,)``, or ``(B, 2)`` for a batch.
     """
 
-    def f(t: float, x: np.ndarray):
-        return np.array([[0, 1], [-alpha - beta * x[0] ** 2, -delta]]) @ x + np.array([0, gamma * np.cos(omega * t)])
+    def f(t: float, s: np.ndarray):
+        x, v = np.moveaxis(s, -1, 0)
+        return np.stack([v, (-alpha - beta * x**2) * x - delta * v + gamma * np.cos(omega * t)], axis=-1)
 
     return f

edmkit-0.0.16/src/edmkit/generate/euler_spiral.py

@@ -0,0 +1,24 @@
+import numpy as np
+
+__all__ = ["euler_spiral"]
+
+
+def euler_spiral(a: float):
+    """Vector field of the Euler spiral (clothoid).
+
+    Parameters
+    ----------
+    a : float
+        Rate at which the curvature grows (typical: 1).
+
+    Returns
+    -------
+    f : Callable[[float, np.ndarray], np.ndarray]
+        Right-hand side ``f(t, x)`` for a state ``(x, y, theta)`` of shape ``(3,)``, or ``(B, 3)`` for a batch.
+    """
+
+    def f(t: float, x: np.ndarray):
+        theta = x[..., 2]
+        return np.stack([np.cos(theta), np.sin(theta), np.full_like(theta, a * t)], axis=-1)
+
+    return f

edmkit-0.0.16/src/edmkit/generate/fitzhugh_nagumo.py

@@ -0,0 +1,30 @@
+import numpy as np
+
+__all__ = ["fitzhugh_nagumo"]
+
+
+def fitzhugh_nagumo(a: float, b: float, epsilon: float, I: float):
+    """Vector field of the FitzHugh-Nagumo neuron model.
+
+    Parameters
+    ----------
+    a : float
+        Offset of the recovery variable (typical: 0.7).
+    b : float
+        Damping of the recovery variable (typical: 0.8).
+    epsilon : float
+        Time scale of the recovery variable (typical: 0.08).
+    I : float
+        Applied current (typical: 0.5 for spiking).
+
+    Returns
+    -------
+    f : Callable[[float, np.ndarray], np.ndarray]
+        Right-hand side ``f(t, x)`` for a state ``(v, w)`` of shape ``(2,)``, or ``(B, 2)`` for a batch.
+    """
+
+    def f(t: float, x: np.ndarray):
+        v, w = np.moveaxis(x, -1, 0)
+        return np.stack([v - v**3 / 3 - w + I, epsilon * (v + a - b * w)], axis=-1)
+
+    return f

edmkit-0.0.16/src/edmkit/generate/hindmarsh_rose.py

@@ -0,0 +1,38 @@
+import numpy as np
+
+__all__ = ["hindmarsh_rose"]
+
+
+def hindmarsh_rose(a: float, b: float, c: float, d: float, r: float, s: float, x_R: float, I: float):
+    """Vector field of the Hindmarsh-Rose neuron model.
+
+    Parameters
+    ----------
+    a : float
+        Cubic coefficient of the membrane potential (typical: 1).
+    b : float
+        Quadratic coefficient of the membrane potential (typical: 3).
+    c : float
+        Constant drive of the recovery variable (typical: 1).
+    d : float
+        Quadratic coupling of the potential into the recovery variable (typical: 5).
+    r : float
+        Time scale of the slow adaptation current (typical: 0.006).
+    s : float
+        Adaptation strength (typical: 4).
+    x_R : float
+        Resting potential (typical: -1.6).
+    I : float
+        Applied current (typical: 3.25 for chaotic bursting).
+
+    Returns
+    -------
+    f : Callable[[float, np.ndarray], np.ndarray]
+        Right-hand side ``f(t, x)`` for a state ``(x, y, z)`` of shape ``(3,)``, or ``(B, 3)`` for a batch: membrane potential, recovery variable and adaptation current.
+    """
+
+    def f(t: float, x: np.ndarray):
+        v, w, u = np.moveaxis(x, -1, 0)
+        return np.stack([w - a * v**3 + b * v**2 - u + I, c - d * v**2 - w, r * (s * (v - x_R) - u)], axis=-1)
+
+    return f

edmkit-0.0.16/src/edmkit/generate/kuramoto.py

@@ -0,0 +1,26 @@
+import numpy as np
+
+__all__ = ["kuramoto"]
+
+
+def kuramoto(omega: np.ndarray, K: np.ndarray):
+    """Vector field of the Kuramoto model ``dθ_i/dt = ω_i + Σ_j K_ij sin(θ_j - θ_i)``.
+
+    Parameters
+    ----------
+    omega : np.ndarray
+        Natural frequencies of shape ``(N,)``.
+    K : np.ndarray
+        Coupling of shape ``(N, N)``; ``K[i, j]`` is the strength with which oscillator ``j`` drives oscillator ``i`` (typical: ``k / N`` everywhere).
+
+    Returns
+    -------
+    f : Callable[[float, np.ndarray], np.ndarray]
+        Right-hand side ``f(t, x)`` for unwrapped phases of shape ``(N,)``, or ``(B, N)`` for a batch.
+    """
+
+    def f(t: float, x: np.ndarray):
+        s, c = np.sin(x), np.cos(x)
+        return omega + c * (s @ K.T) - s * (c @ K.T)
+
+    return f

{edmkit-0.0.14 → edmkit-0.0.16}/src/edmkit/generate/lorenz.py

@@ -18,10 +18,11 @@ def lorenz(sigma: float, rho: float, beta: float):
     Returns
     -------
     f : Callable[[float, np.ndarray], np.ndarray]
-        Right-hand side ``f(t, x)`` for a state ``(x, y, z)`` of shape ``(3,)``.
+        Right-hand side ``f(t, x)`` for a state ``(x, y, z)`` of shape ``(3,)``, or ``(B, 3)`` for a batch.
     """
 
-    def f(t: float, x: np.ndarray):
-        return np.array([[-sigma, sigma, 0], [rho, -1, -x[0]], [0, x[0], -beta]]) @ x
+    def f(t: float, s: np.ndarray):
+        x, y, z = np.moveaxis(s, -1, 0)
+        return np.stack([sigma * (y - x), rho * x - y - x * z, x * y - beta * z], axis=-1)
 
     return f

{edmkit-0.0.14 → edmkit-0.0.16}/src/edmkit/generate/lorenz96.py

@@ -14,12 +14,12 @@ def lorenz96(F: float):
     Returns
     -------
     f : Callable[[float, np.ndarray], np.ndarray]
-        Right-hand side ``f(t, x)`` for a state of shape ``(D,)`` with ``D >= 4``.
+        Right-hand side ``f(t, x)`` for a state of shape ``(D,)``, or ``(B, D)`` for a batch, with ``D >= 4``.
     """
 
     def f(t: float, x: np.ndarray):
         # p[i + 2] is x[i] on the ring, so this is (x[i + 1] - x[i - 2]) * x[i - 1] - x[i] + F
-        p = np.concatenate([x[-2:], x, x[:1]])
-        return (p[3:] - p[:-3]) * p[1:-2] - x + F
+        p = np.concatenate([x[..., -2:], x, x[..., :1]], axis=-1)
+        return (p[..., 3:] - p[..., :-3]) * p[..., 1:-2] - x + F
 
     return f

edmkit-0.0.16/src/edmkit/generate/lotka_volterra.py

@@ -0,0 +1,27 @@
+import numpy as np
+
+__all__ = ["lotka_volterra"]
+
+
+def lotka_volterra(rho: np.ndarray, sigma: float | np.ndarray, S: float | np.ndarray):
+    """Vector field of the generalized Lotka-Volterra system ``dx/dt = x (sigma - rho x) + S``.
+
+    Parameters
+    ----------
+    rho : np.ndarray
+        Interaction matrix of shape ``(N, N)``; ``rho[i, j]`` is the effect of species ``j`` on species ``i``.
+    sigma : float | np.ndarray
+        Growth rate, scalar or of shape ``(N,)`` (typical: 1).
+    S : float | np.ndarray
+        Constant input, scalar or of shape ``(N,)`` (typical: 0; small positive values keep winnerless competition switching).
+
+    Returns
+    -------
+    f : Callable[[float, np.ndarray], np.ndarray]
+        Right-hand side ``f(t, x)`` for a state of shape ``(N,)``, or ``(B, N)`` for a batch.
+    """
+
+    def f(t: float, x: np.ndarray):
+        return x * (sigma - x @ rho.T) + S
+
+    return f

edmkit-0.0.16/src/edmkit/generate/morris_lecar.py

@@ -0,0 +1,70 @@
+import numpy as np
+
+__all__ = ["morris_lecar"]
+
+
+def morris_lecar(
+    C: float,
+    g_L: float,
+    g_Ca: float,
+    g_K: float,
+    V_L: float,
+    V_Ca: float,
+    V_K: float,
+    V1: float,
+    V2: float,
+    V3: float,
+    V4: float,
+    phi: float,
+    I: float,
+):
+    """Vector field of the Morris-Lecar neuron model.
+
+    Parameters
+    ----------
+    C : float
+        Membrane capacitance (typical: 20).
+    g_L : float
+        Leak conductance (typical: 2).
+    g_Ca : float
+        Calcium conductance (typical: 4.4).
+    g_K : float
+        Potassium conductance (typical: 8).
+    V_L : float
+        Leak reversal potential (typical: -60).
+    V_Ca : float
+        Calcium reversal potential (typical: 120).
+    V_K : float
+        Potassium reversal potential (typical: -84).
+    V1 : float
+        Half-activation voltage of the calcium channel (typical: -1.2).
+    V2 : float
+        Slope of the calcium activation (typical: 18).
+    V3 : float
+        Half-activation voltage of the potassium channel (typical: 2).
+    V4 : float
+        Slope of the potassium activation (typical: 30).
+    phi : float
+        Rate of the potassium activation (typical: 0.04).
+    I : float
+        Applied current (typical: 100 for spiking).
+
+    Returns
+    -------
+    f : Callable[[float, np.ndarray], np.ndarray]
+        Right-hand side ``f(t, x)`` for a state ``(V, n)`` of shape ``(2,)``, or ``(B, 2)`` for a batch: membrane potential and potassium activation.
+    """
+
+    def f(t: float, x: np.ndarray):
+        V, n = np.moveaxis(x, -1, 0)
+        m_inf = 0.5 * (1 + np.tanh((V - V1) / V2))
+        n_inf = 0.5 * (1 + np.tanh((V - V3) / V4))
+        return np.stack(
+            [
+                (I - g_L * (V - V_L) - g_Ca * m_inf * (V - V_Ca) - g_K * n * (V - V_K)) / C,
+                phi * (n_inf - n) * np.cosh((V - V3) / (2 * V4)),
+            ],
+            axis=-1,
+        )
+
+    return f

edmkit-0.0.16/src/edmkit/generate/pyloric.py

@@ -0,0 +1,105 @@
+import numpy as np
+
+__all__ = ["pyloric"]
+
+AREA = 0.628e-3  # membrane area in cm^2 (specific capacitance 1 uF/cm^2, i.e. C = 0.628 nF)
+RT_2F = 12.2  # RT / 2F in mV at 10 C, for the Nernst potential of Ca
+E_REV = np.array([50.0, 0.0, 0.0, -80.0, -80.0, -80.0, -20.0, -50.0])  # Na, CaT, CaS, A, KCa, Kd, H, leak; Ca is replaced by Nernst
+POWER = np.array([3, 1, 3, 1, 3, 1, 3, 1, 4, 4, 1])[:, None]  # exponents of the 11 gates, see the state layout
+E_SYN = np.array([-70.0, -80.0])[:, None]  # reversal potentials of the glutamatergic and cholinergic synapses
+K_SYN = np.array([1 / 40, 1 / 100])[:, None]  # decay rates of the glutamatergic and cholinergic synapses
+
+
+def pyloric(g_mem: np.ndarray, g_glu: np.ndarray, g_ach: np.ndarray):
+    """Vector field of the three-neuron pyloric circuit of the crustacean stomatogastric ganglion (Prinz, Bucher & Marder 2004), in ms and mV.
+
+    Parameters
+    ----------
+    g_mem : np.ndarray
+        Maximal conductances in mS/cm^2 of shape ``(3, 8)``: rows are AB/PD, LP, PY and columns are Na, CaT, CaS, A, KCa, Kd, H, leak
+        (typical: ``[[400, 2.5, 6, 50, 10, 100, 0.01, 0], [100, 0, 8, 40, 5, 75, 0.05, 0.02], [100, 2.5, 2, 50, 0, 125, 0.05, 0.01]]``).
+    g_glu : np.ndarray
+        Fast glutamatergic synapse strengths in nS of shape ``(3, 3)``; ``g_glu[i, j]`` is the synapse from neuron ``j`` onto neuron ``i``
+        (typical: ``g_glu[1, 0] = 30`` (AB to LP), ``g_glu[2, 0] = 3`` (AB to PY), ``g_glu[0, 1] = 10`` (LP to AB/PD), ``g_glu[2, 1] = 3`` (LP to PY),
+        ``g_glu[1, 2] = 3`` (PY to LP), the rest 0, for a PD, LP, PY rhythm of period about 1.7 s).
+    g_ach : np.ndarray
+        Slow cholinergic synapse strengths in nS of shape ``(3, 3)``, indexed like ``g_glu`` (typical: ``g_ach[1, 0] = 3`` (PD to LP), ``g_ach[2, 0] = 30`` (PD to PY), the rest 0).
+
+    Returns
+    -------
+    f : Callable[[float, np.ndarray], np.ndarray]
+        Right-hand side ``f(t, x)`` for a state of shape ``(45,)``, or ``(B, 45)`` for a batch; ``x.reshape(15, 3)`` has one column per neuron and rows: membrane potential,
+        intracellular [Ca] in uM, the gates m and h of Na, m and h of CaT, m and h of CaS, m and h of A, m of KCa, n of Kd, m of H, and the glutamatergic
+        and cholinergic synaptic activations of the neuron as presynaptic cell
+        (typical initial state: V = -50, [Ca] = 0.05, all gates and synaptic activations 0, discard the first 3 s as transient).
+    """
+    g_mem = np.asarray(g_mem, dtype=float)
+    g_syn = np.stack([g_glu, g_ach]) * 1e-6 / AREA  # nS to uA/cm^2 per mV
+
+    def f(t: float, x: np.ndarray):
+        batch = x.shape[:-1]
+        V, Ca = x[..., :3], x[..., 3:6]
+        gate = x[..., 6:39].reshape(*batch, 11, 3)
+        s = x[..., 39:].reshape(*batch, 2, 3)
+
+        # steady states and time constants of the 11 gates
+        def sig(a, b):
+            return 1 / (1 + np.exp((V + a) / b))
+
+        def tau(a, b, c, d):
+            return a - b / (1 + np.exp((V + c) / -d))
+
+        inf = np.stack(
+            [
+                sig(25.5, -5.29),
+                sig(48.9, 5.18),
+                sig(27.1, -7.2),
+                sig(32.1, 5.5),
+                sig(33.0, -8.1),
+                sig(60.0, 6.2),
+                sig(27.2, -8.7),
+                sig(56.9, 4.9),
+                Ca / (Ca + 3) * sig(28.3, -12.6),
+                sig(12.3, -11.8),
+                sig(75.0, 5.5),
+            ],
+            axis=-2,
+        )
+        tau_g = np.stack(
+            [
+                tau(2.64, 2.52, 120, 25),
+                1.34 / (1 + np.exp((V + 62.9) / -10)) * (1.5 + 1 / (1 + np.exp((V + 34.9) / 3.6))),
+                tau(43.4, 42.6, 68.1, 20.5),
+                tau(210, 179.6, 55, 16.9),
+                2.8 + 14 / (np.exp((V + 27) / 10) + np.exp((V + 70) / -13)),
+                120 + 300 / (np.exp((V + 55) / 9) + np.exp((V + 65) / -16)),
+                tau(23.2, 20.8, 32.9, 15.2),
+                tau(77.2, 58.4, 38.9, 26.5),
+                tau(180.6, 150.2, 46, 22.7),
+                tau(14.4, 12.8, 28.3, 19.2),
+                2 / (np.exp((V + 169.7) / -11.6) + np.exp((V - 26.7) / 14.3)),
+            ],
+            axis=-2,
+        )
+
+        # membrane currents in uA/cm^2
+        m_Na, h_Na, m_CaT, h_CaT, m_CaS, h_CaS, m_A, h_A, m_KCa, n_Kd, m_H = np.moveaxis(gate**POWER, -2, 0)
+        activation = np.stack([m_Na * h_Na, m_CaT * h_CaT, m_CaS * h_CaS, m_A * h_A, m_KCa, n_Kd, m_H, np.ones_like(V)], axis=-2)
+        drive = V[..., None, :] - E_REV[:, None]
+        drive[..., 1:3, :] = (V - RT_2F * np.log(3000 / Ca))[..., None, :]
+        I_ion = g_mem.T * activation * drive
+        # sum over presynaptic cells first: g_s[p, i] = sum_j g_syn[p, i, j] s[p, j]
+        g_s = (g_syn @ s[..., None])[..., 0]
+        I_syn = (g_s * (V[..., None, :] - E_SYN)).sum(axis=-2)
+
+        # I_ion[1:3] in uA/cm^2 times area is uA, times 1e3 is nA, times f = 14.96 uM/nA gives the Ca influx
+        dCa = (0.05 - Ca - 14.96 * 1e3 * AREA * I_ion[..., 1:3, :].sum(axis=-2)) / 200
+        s_inf = sig(35.0, -5.0)[..., None, :]
+        # tau_s = (1 - s_inf) / k, floored at 1e-4 / k to keep the spikes from stiffening the system
+        ds = (s_inf - s) * K_SYN / np.maximum(1 - s_inf, 1e-4)
+
+        dV = -(I_ion.sum(axis=-2) + I_syn)
+        dgate = (inf - gate) / tau_g
+        return np.concatenate([dV, dCa, dgate.reshape(*batch, -1), ds.reshape(*batch, -1)], axis=-1)
+
+    return f

edmkit-0.0.16/src/edmkit/generate/rk45.py

@@ -0,0 +1,106 @@
+from collections.abc import Callable
+
+import numpy as np
+
+__all__ = ["rk45"]
+
+# Dormand-Prince 5(4)
+A = np.array(
+    [
+        [0, 0, 0, 0, 0, 0],
+        [1 / 5, 0, 0, 0, 0, 0],
+        [3 / 40, 9 / 40, 0, 0, 0, 0],
+        [44 / 45, -56 / 15, 32 / 9, 0, 0, 0],
+        [19372 / 6561, -25360 / 2187, 64448 / 6561, -212 / 729, 0, 0],
+        [9017 / 3168, -355 / 33, 46732 / 5247, 49 / 176, -5103 / 18656, 0],
+        [35 / 384, 0, 500 / 1113, 125 / 192, -2187 / 6784, 11 / 84],
+    ]
+)
+C = A.sum(axis=1)
+E = np.array([71 / 57600, 0, -71 / 16695, 71 / 1920, -17253 / 339200, 22 / 525, -1 / 40])
+# quartic dense output (Shampine 1986): x(s + theta h) = x + h sum_j K_j sum_k P[j, k] theta^(k + 1)
+P = np.array(
+    [
+        [1, -8048581381 / 2820520608, 8663915743 / 2820520608, -12715105075 / 11282082432],
+        [0, 0, 0, 0],
+        [0, 131558114200 / 32700410799, -68118460800 / 10900136933, 87487479700 / 32700410799],
+        [0, -1754552775 / 470086768, 14199869525 / 1410260304, -10690763975 / 1880347072],
+        [0, 127303824393 / 49829197408, -318862633887 / 49829197408, 701980252875 / 199316789632],
+        [0, -282668133 / 205662961, 2019193451 / 616988883, -1453857185 / 822651844],
+        [0, 40617522 / 29380423, -110615467 / 29380423, 69997945 / 29380423],
+    ]
+)
+THETA_POWERS = np.arange(1, 5)
+
+
+def rk45(
+    f: Callable[[float, np.ndarray], np.ndarray],
+    X0: np.ndarray,
+    dt: float,
+    t_max: float,
+    rtol: float = 1e-6,
+    atol: float = 1e-9,
+):
+    """Integrate ``dx/dt = f(t, x)`` via the adaptive Dormand-Prince RK45 method.
+
+    Steps are chosen by the error alone, and the uniform grid of spacing ``dt`` is filled in by the method's dense output.
+    A batch of initial conditions is integrated together, with steps that keep every member within the tolerance.
+
+    Parameters
+    ----------
+    f : Callable[[float, np.ndarray], np.ndarray]
+        Right-hand side of the system, called with states of the shape of ``X0``.
+    X0 : np.ndarray
+        Initial condition of shape ``(D,)``, or ``(B, D)`` for a batch.
+    dt : float
+        Sampling time step.
+    t_max : float
+        End time, excluded from the grid.
+    rtol : float
+        Relative tolerance of the local error.
+    atol : float
+        Absolute tolerance of the local error.
+
+    Returns
+    -------
+    t : np.ndarray
+        Time array of shape ``(N,)``.
+    X : np.ndarray
+        Trajectory of shape ``(N, D)``, or ``(B, N, D)`` for a batch.
+
+    Raises
+    ------
+    ValueError
+        If ``X0`` is not of shape ``(D,)`` or ``(B, D)``.
+    """
+    X0 = np.asarray(X0, dtype=float)
+    if X0.ndim not in (1, 2):
+        raise ValueError(f"X0 must be of shape (D,) or (B, D), got X0.shape={X0.shape}")
+
+    t = np.arange(0, t_max, dt)
+    X = np.empty((len(t), *X0.shape))
+    X[0] = X0
+
+    K = np.empty((7, *X0.shape))
+    Kf = K.reshape(7, -1)  # flat view of K, so the stage sums are plain matrix products
+    K[0] = f(0.0, X0)
+    s, x, h, i = 0.0, X0, dt, 1
+
+    while i < len(t):
+        step = min(h, t[-1] - s)
+        for j in range(1, 7):
+            # the last stage is evaluated at x_new, so an accepted step hands K[6] on as the next K[0]
+            x_new = x + step * (A[j, :j] @ Kf[:j]).reshape(x.shape)
+            K[j] = f(s + C[j] * step, x_new)
+        error = np.max(np.abs(step * (E @ Kf).reshape(x.shape)) / (atol + rtol * np.maximum(np.abs(x), np.abs(x_new))))
+
+        if error <= 1:
+            n = np.searchsorted(t, s + step, side="right")
+            if n > i:
+                theta = (t[i:n] - s) / step
+                X[i:n] = x + step * (theta[:, None] ** THETA_POWERS @ (P.T @ Kf)).reshape(n - i, *x.shape)
+            s, x, i = s + step, x_new, n
+            K[0] = K[6]
+        h = step * (10 if error == 0 else min(10, max(0.2, 0.9 * error**-0.2)))
+
+    return t, np.moveaxis(X, 0, -2)

{edmkit-0.0.14 → edmkit-0.0.16}/src/edmkit/generate/rossler.py

@@ -18,10 +18,11 @@ def rossler(a: float, b: float, c: float):
     Returns
     -------
     f : Callable[[float, np.ndarray], np.ndarray]
-        Right-hand side ``f(t, x)`` for a state ``(x, y, z)`` of shape ``(3,)``.
+        Right-hand side ``f(t, x)`` for a state ``(x, y, z)`` of shape ``(3,)``, or ``(B, 3)`` for a batch.
     """
 
-    def f(t: float, x: np.ndarray):
-        return np.array([[0, -1, -1], [1, a, 0], [0, 0, x[0] - c]]) @ x + np.array([0, 0, b])
+    def f(t: float, s: np.ndarray):
+        x, y, z = np.moveaxis(s, -1, 0)
+        return np.stack([-y - z, x + a * y, b + (x - c) * z], axis=-1)
 
     return f