"""dfun codegen + the ring-buffer step(carry, _) factory.
Adapted from ``vbi/simulator/backend/jax_/codegen.py`` and
``vbi/simulator/backend/jax_/simulator.py`` - see NOTICE.md in this directory
for provenance and what was dropped (stochastic noise, state clipping,
sigmoidal/kuramoto coupling, monitors/stimuli).
This is the piece that unifies ODE and DDE: ``has_delays=False`` reads
coupling from the instantaneous state; ``has_delays=True`` reads it from a
ring buffer of past coupling-variable states via a flat-index gather. Both
paths produce the same ``step(carry, _) -> (new_carry, new_state)`` shape
consumed by ``jax.lax.scan``, and - critically for the Lyapunov engines - both
are plain, differentiable JAX functions of ``carry``.
"""
from __future__ import annotations
from typing import Callable
import jax.numpy as jnp
from ..integrators import rk6_combine
from .model_spec import ModelSpec
[docs]
def build_jax_dfun(spec: ModelSpec) -> Callable:
"""
Compile dfun_str expressions into a JAX-traceable function.
Generated signature:
fn(state, coupling, params) -> jnp.ndarray shape (n_sv, n_nodes)
where:
state : (n_sv, n_nodes)
coupling : (n_cvar, n_nodes) - one row per coupling variable
params : dict[str, scalar | jnp.ndarray]
Compiled once via exec() on the spec's own strings, same as the vbi
original - nothing about the generated function blocks jax.jacfwd /
jax.jvp.
.. warning::
The ``dfun_str`` expressions are code-generated and executed via
``exec()`` with **no sanitization**. This is fine for specs a
caller writes themselves, but never build a ``ModelSpec`` from
untrusted input (user uploads, network data, config files from
unknown sources) - the strings can execute arbitrary code.
"""
sv = spec.sv_names
param_names = spec.param_names
cvar_names = spec.cvar
lines = [
"import jax.numpy as _jnp",
"from jax.numpy import pi, exp, log, sin, cos, tanh, sqrt",
"def _dfun_jax(state, coupling, params):",
]
for i, name in enumerate(sv):
lines.append(f" {name} = state[{i}]")
for name in param_names:
lines.append(f" {name} = params['{name}']")
for i, cname in enumerate(cvar_names):
lines.append(f" c_{cname} = coupling[{i}]")
lines.append(" c = coupling[0] if coupling.shape[0] > 0 else 0.0")
sv_exprs = ", ".join(f"({spec.dfun_str[name]})" for name in sv)
lines.append(f" return _jnp.stack([{sv_exprs}])")
src = "\n".join(lines)
globs: dict = {}
exec(compile(src, "<dfun_jax>", "exec"), globs)
return globs["_dfun_jax"]
# ---------------------------------------------------------------------------
# Coupling primitives
# ---------------------------------------------------------------------------
def _instant_coupling(cvar_state: jnp.ndarray, weights: jnp.ndarray,
G: float, a: float, b: float) -> jnp.ndarray:
"""No-delay path. cvar_state: (n_cvar, N) -> coupling: (n_cvar, N)."""
return G * a * jnp.einsum("ts,cs->ct", weights, cvar_state) + b
def _write_ring(buf: jnp.ndarray, step: int, cvar_state: jnp.ndarray,
horizon: int) -> jnp.ndarray:
"""buf: (horizon, n_cvar, n_nodes) -> updated buf.
Time convention: slot ``k % horizon`` holds the
coupling-variable value at physical time ``k * dt``. ``step()``'s
caller must pass the physical time index of ``cvar_state`` itself as
``step`` -- e.g. the newly integrated state advances the carry's own
step counter ``t`` to ``t + 1``, so it is written at ``step=t + 1``,
not ``t``. Writing under the wrong index by one silently shifts every
delayed read by one step (previously the case here -- fixed)."""
return buf.at[step % horizon].set(cvar_state)
def _read_delayed_coupling(
buf: jnp.ndarray, step: int,
delay_steps: jnp.ndarray, weights: jnp.ndarray,
G: float, a: float, b: float,
horizon: int, n_nodes: int) -> jnp.ndarray:
"""
Returns coupling (n_cvar, n_nodes), reading each source node's delayed
coupling-variable state out of the ring buffer via a flat-index gather
(avoids a Python loop / dynamic 2-D gather inside traced code).
Per-edge delay (delay_steps is a full (N, N) matrix): the coupling
formula (G*a*sum+b) is baked in here rather than exposed as a separate
"gather" step, because different edges read different time offsets --
there's no single (n_cvar, n_nodes) "the delayed cvar state" independent
of which edge is asking, so a plain coupling_fn(cvar_state, weights,
params) signature (one reading per node) can't express this in general.
Left untouched for M5 (per-edge delay + coupling_fn needs an edge-aware
signature, a separate design fork) -- see _read_uniform_delayed_cvar
below for the M4 case where this collapses back to one reading per node.
"""
idx_time = (step - delay_steps) % horizon # (N, N)
src_idx = jnp.arange(n_nodes, dtype=jnp.int32)
flat_idx = idx_time * n_nodes + src_idx[None, :] # (N, N)
n_cvar = buf.shape[1]
cvars = []
for cv in range(n_cvar):
buf_cv = buf[:, cv, :] # (horizon, N)
delayed = buf_cv.reshape(-1)[flat_idx] # (N, N)
cvars.append(G * a * jnp.sum(weights * delayed, axis=1) + b)
return jnp.stack(cvars) # (n_cvar, N)
def _read_uniform_delayed_cvar(
buf: jnp.ndarray, step: int, tau_steps: int, horizon: int) -> jnp.ndarray:
"""M4: a single global delay (not per-edge), so the delayed reading is
the same for every node -- one O(1) modular ring-buffer read, no
flat-index gather needed. Returns cvar_state (n_cvar, n_nodes), the same
shape lyapax.coupling's plain-callable coupling_fn expects for the
zero-delay case, so any existing coupling_fn works unmodified against
either instantaneous or uniformly-delayed cvar_state. ``step`` is the
physical time index of *this* read (i.e. the caller's current time in
units of dt) -- see ``_write_ring``'s docstring for the matching write
convention this read relies on."""
return buf[(step - tau_steps) % horizon]
def _write_ring_interp(
buf: jnp.ndarray, step: int, cvar_value: jnp.ndarray,
cvar_deriv: jnp.ndarray, horizon: int) -> jnp.ndarray:
"""Interpolated-history counterpart of ``_write_ring``: ``buf`` gains a
length-2 axis (value, derivative) so later reads can reconstruct a
delayed value *between* grid points via Hermite interpolation instead
of only ever reading an exact stored sample -- see
``_read_uniform_delayed_cvar_interp``. ``buf``: ``(horizon, 2, n_cvar,
n_nodes) -> updated buf``. Same ``step`` time convention as
``_write_ring``."""
idx = step % horizon
buf = buf.at[idx, 0].set(cvar_value)
return buf.at[idx, 1].set(cvar_deriv)
def _read_uniform_delayed_cvar_interp(
buf: jnp.ndarray, step, tau_steps: float, horizon: int, dt: float) -> jnp.ndarray:
"""Cubic Hermite-interpolated counterpart of
``_read_uniform_delayed_cvar``: reconstructs the delayed
coupling-variable value at the exact (possibly non-integer) time
``step*dt - tau_steps*dt``, from the stored value *and derivative* at
the two grid points bracketing it, rather than snapping ``tau_steps``
to the nearest integer. ``tau_steps`` is a plain (generally
non-integer) float here.
``buf``: ``(horizon, 2, n_cvar, n_nodes)``, axis 1 = (value,
derivative), as written by ``_write_ring_interp``.
"""
t_delayed = step - tau_steps
i0 = jnp.floor(t_delayed).astype(jnp.int32)
theta = t_delayed - i0
idx0 = i0 % horizon
idx1 = (i0 + 1) % horizon
y0, dy0 = buf[idx0, 0], buf[idx0, 1]
y1, dy1 = buf[idx1, 0], buf[idx1, 1]
# Standard cubic Hermite basis on theta in [0, 1); dy0/dy1 are d(y)/dt,
# so they're scaled by dt (the spacing between grid points i0 and i1).
h00 = 2 * theta ** 3 - 3 * theta ** 2 + 1
h10 = theta ** 3 - 2 * theta ** 2 + theta
h01 = -2 * theta ** 3 + 3 * theta ** 2
h11 = theta ** 3 - theta ** 2
return h00 * y0 + h10 * dt * dy0 + h01 * y1 + h11 * dt * dy1
# ---------------------------------------------------------------------------
# Integrators (pure JAX, traceable)
# ---------------------------------------------------------------------------
#
# Each takes ``coupling_at(y_stage, c) -> coupling`` -- a callable, not a
# precomputed array -- and calls it fresh at every stage, passing that
# stage's own intra-step state estimate ``y_stage`` and Butcher node ``c``
# (0 at the first stage, 1 at the last), instead of freezing coupling at
# the value read once at the step's start. Freezing it, the earlier
# design, silently capped every one of these methods at O(dt) whenever
# coupling was actually nonzero, regardless of the base method's nominal
# order (rk4 and rk6 gave *identical* errors on a coupled network).
# ``make_step_fn``'s ``coupling_at``
# exploits this in two places: for the zero-delay case, coupling depends
# on ``y_stage`` (the evolving state), evaluated fresh at each stage; for
# ``has_delays=True`` with ``interpolate=True``, the delayed lookup
# depends on *time*, so it is re-read via the Hermite interpolant at each
# stage's own intra-step time ``step + c``. (An earlier attempt at that
# per-stage delayed read appeared not to help -- but it was made while
# the ring-buffer write index was off by one full step, itself an O(dt)
# bias that masked the improvement; with the write convention fixed, the
# per-stage read restores the integrator's real order, capped at ~4 by
# the cubic Hermite history.)
# The grid-snapped delayed path (``interpolate=False``) still ignores
# ``c``: without an interpolant there is no sub-step history to read, so
# it stays O(dt) by construction. A caller that has no coupling at all
# (``coupling_at`` returns a constant, or is never dependent on ``y``/``c``)
# pays a few extra, cheap function calls for no behavior change.
_EULER_STAGE_C = (0.0,)
_HEUN_STAGE_C = (0.0, 1.0)
_RK4_STAGE_C = (0.0, 0.5, 0.5, 1.0)
def _euler(state, dfun, coupling_at, dt, params):
c = _EULER_STAGE_C
k1 = dfun(state, coupling_at(state, c[0]), params)
return state + dt * k1
def _heun(state, dfun, coupling_at, dt, params):
c = _HEUN_STAGE_C
k1 = dfun(state, coupling_at(state, c[0]), params)
y2 = state + dt * k1
k2 = dfun(y2, coupling_at(y2, c[1]), params)
return state + 0.5 * dt * (k1 + k2)
def _rk4(state, dfun, coupling_at, dt, params):
c = _RK4_STAGE_C
k1 = dfun(state, coupling_at(state, c[0]), params)
y2 = state + 0.5 * dt * k1
k2 = dfun(y2, coupling_at(y2, c[1]), params)
y3 = state + 0.5 * dt * k2
k3 = dfun(y3, coupling_at(y3, c[2]), params)
y4 = state + dt * k3
k4 = dfun(y4, coupling_at(y4, c[3]), params)
return state + (dt / 6.0) * (k1 + 2.0 * k2 + 2.0 * k3 + k4)
def _rk6(state, dfun, coupling_at, dt, params):
"""See ``lyapax.integrators.rk6_combine`` for the tableau, provenance,
and correctness checks -- this just plugs the coupled-network
derivative ``dfun(y, coupling_at(y, c), params)``, recomputed fresh at
each of the 8 stages, into that shared combination."""
return rk6_combine(lambda y, c: dfun(y, coupling_at(y, c), params), state, dt)
_STEP_INTEGRATORS: dict[str, Callable] = {
"euler": _euler,
"heun": _heun,
"rk4": _rk4,
"rk6": _rk6,
}
# ---------------------------------------------------------------------------
# step(carry, _) factory
# ---------------------------------------------------------------------------
[docs]
def make_step_fn(
dfun: Callable,
weights: jnp.ndarray,
has_delays: bool,
horizon: int,
n_nodes: int,
cvar_indices: tuple[int, ...],
dt: float,
delay_steps: jnp.ndarray | None = None,
G_default: float = 1.0,
coup_a: float = 1.0,
coup_b: float = 0.0,
integrator: str | Callable = "heun",
coupling_fn: Callable | None = None,
tau_steps: int | float | None = None,
interpolate: bool = False,
) -> Callable:
"""
Returns step(carry, _) -> (new_carry, new_state).
carry = (state, buf, step_int32, params)
``params`` travels inside the carry (not closed over) so that, later,
jax.vmap over swept parameters and jax.jacfwd/jvp w.r.t. state both see
it as data rather than a baked-in constant.
``buf`` is a no-op array when ``has_delays=False`` - the ODE and DDE
cases share this exact function; only the coupling branch differs.
integrator : ``"euler"``, ``"heun"``, ``"rk4"``, ``"rk6"``, or a callable
``(state, dfun, coupling_at, dt, params) -> new_state``, where
``coupling_at(y_stage, c) -> coupling`` is called fresh at each
internal stage (``y_stage``: that stage's intra-step state
estimate; ``c``: its Butcher node, 0 at the first stage, 1 at the
last) rather than being frozen once per step -- see the
"Integrators" section's module comment, and
:ref:`dde-rk-stage-order` for why freezing it instead
silently caps the whole scheme at O(dt) whenever coupling is
nonzero. The zero-delay case reads coupling from ``y_stage`` (the
evolving state) fresh at each stage; the interpolated-delay case
(``interpolate=True``) reads the delayed history at each stage's
own intra-step time ``step + c`` via the Hermite interpolant, so
both realize the integrator's nominal order (the delayed case up
to the cubic Hermite history's own ~4th-order ceiling). Only the
grid-snapped delayed default (``interpolate=False``) still ignores
``c`` and reads once per step -- without an interpolant there is
no sub-step history to read -- and remains O(dt); pass
``interpolate=True`` for higher-order DDE stepping.
coupling_fn : optional ``lyapax.coupling``-style plain callable,
``(cvar_state, weights, params) -> coupling`` (see
``lyapax.coupling.CouplingFn``). When given, replaces the
hardcoded ``G_default``/``coup_a``/``coup_b`` linear formula for
*both* the zero-delay and (uniform-delay-only, see ``tau_steps``)
delayed branches, unifying this step with ``lyapax.coupling``'s
plain-callable design. ``G_default``/``coup_a``/``coup_b`` are
ignored when ``coupling_fn`` is given. Default ``None`` preserves
the exact original vendored behavior.
tau_steps : required alongside ``coupling_fn`` when ``has_delays=True``
-- a single global delay (in steps), read via
``_read_uniform_delayed_cvar`` (O(1) ring-buffer read) or, when
``interpolate=True``, ``_read_uniform_delayed_cvar_interp`` (in
which case ``tau_steps`` need not be an integer -- see
``interpolate``). Per-edge heterogeneous delays (the general
``delay_steps`` matrix) are not supported together with a custom
``coupling_fn`` yet -- that combination needs an edge-aware
coupling_fn signature, which does not exist yet; use the
legacy ``coupling_fn=None`` path (which does support per-edge
``delay_steps``, just with the hardcoded linear formula) until
then.
interpolate : reconstruct the delayed coupling-variable value via cubic
Hermite interpolation (using the stored value *and* derivative at
the two grid points bracketing the delayed time) instead of
snapping to the nearest stored grid sample -- see
:ref:`dde-history-interpolation`. Requires ``has_delays=True``,
``coupling_fn`` given, and ``tau_steps`` given (the uniform-delay
path only; not supported with the legacy per-edge ``delay_steps``
path). ``tau_steps`` may then be a non-integer float: the physical
delay no longer needs rounding to a whole number of ``dt`` steps.
Requires ``tau_steps >= 1`` (``tau >= dt``): the per-stage delayed
reads at intra-step times up to ``step + 1`` (and the derivative
stored at write time) must only ever land on ring-buffer slots
already finalized by earlier steps, which fails once the delay is
shorter than one step. Costs one extra ``dfun`` evaluation per
step (to compute the derivative to store) and doubles ``buf``'s
per-slot size (value and derivative, instead of just value).
"""
if interpolate and not (has_delays and coupling_fn is not None and tau_steps is not None):
raise ValueError(
"interpolate=True requires has_delays=True, a coupling_fn, and "
"tau_steps (the uniform-delay path) -- interpolated reads are "
"not supported with has_delays=False or the legacy per-edge "
"delay_steps path."
)
if interpolate and tau_steps < 1:
raise ValueError(
f"interpolate=True requires tau_steps >= 1 (tau >= dt); got "
f"tau_steps={tau_steps}. A sub-step delay would make the "
"per-stage Hermite history reads (and the stored-derivative "
"computation) depend on ring-buffer slots not yet written at "
"read time -- decrease dt below tau instead."
)
if getattr(integrator, "_lyapax_adaptive_ode_only", False):
raise ValueError(
"adaptive ODE integrators (lyapax.adaptive.diffrax_adaptive_step) "
"are only supported for a single, uncoupled ode_problem -- not "
"for network_problem (even without delays) or DDEs. Networks "
"and DDEs go through this shared step-function machinery, whose "
"per-stage coupling callback (coupling_at) has no adaptive-"
"integrator equivalent; DDEs additionally can't use diffrax at "
"all (diffrax has no DDE support, patrick-kidger/diffrax#406). "
"Both stay fixed-step only."
)
cvar_idx = jnp.array(list(cvar_indices), dtype=jnp.int32)
integrate = _STEP_INTEGRATORS[integrator] if isinstance(integrator, str) else integrator
# coupling_at(y_stage, c, buf, step, params) -> coupling, called fresh
# at each integrator stage (see the "Integrators" section above): for
# has_delays=False, at that stage's own intra-step state y_stage (c is
# unused -- coupling only depends on the current state, not on time).
# This measurably fixes the zero-delay case: a coupled linear network's error
# dropped roughly five orders of magnitude at the same dt, and rk4/rk6
# stopped giving identical (order-1) errors.
#
# For has_delays=True with interpolate=True, the delayed lookup is
# likewise re-read at each stage's own intra-step time (step + c) via
# the Hermite interpolant -- a delayed lookup depends on *time*, not
# y_stage. Reading it once per step and freezing it across stages is
# an O(dt) error in the delayed argument, which caps the whole scheme
# at first order exactly like frozen zero-delay coupling did. (This
# same change was once tried and looked ineffective -- but only
# because the ring-buffer write was then off by one slot, a second,
# independent O(dt) bias hiding the gain. With both fixed, the
# scalar linear DDE's measured order is ~2.0 under heun and ~4 under
# rk4/rk6, the cubic Hermite history's ceiling.) The stage times
# step + c never exceed step + 1, so with tau_steps >= 1 (validated
# above) every read lands on slots finalized by earlier steps.
#
# The grid-snapped default (interpolate=False) still ignores c: there
# is no sub-step history to read without an interpolant, so it stays
# O(dt) by construction.
if coupling_fn is not None:
def _coupling_at(y_stage, c, buf, step, params):
if has_delays and interpolate:
cvar_state = _read_uniform_delayed_cvar_interp(
buf, step + c, tau_steps, horizon, dt)
elif has_delays:
cvar_state = _read_uniform_delayed_cvar(buf, step, tau_steps, horizon)
else:
cvar_state = y_stage[cvar_idx]
return coupling_fn(cvar_state, weights, params)
else:
def _coupling_at(y_stage, c, buf, step, params):
G = params.get("G", G_default)
if has_delays:
return _read_delayed_coupling(
buf, step, delay_steps, weights, G, coup_a, coup_b,
horizon, n_nodes)
cvar_state = y_stage[cvar_idx]
return _instant_coupling(cvar_state, weights, G, coup_a, coup_b)
def step(carry, _):
state, buf, t, params = carry
coupling_at = lambda y_stage, c: _coupling_at(y_stage, c, buf, t, params) # noqa: E731
new_state = integrate(state, dfun, coupling_at, dt, params)
if has_delays and interpolate:
# The derivative to store is the true d(cvar)/dt at new_state,
# which needs the coupling *at the new time* -- re-read it
# (interpolated) at t+1 rather than reusing this step's coupling:
# since tau_steps >= 1, this only touches ring-buffer slots
# already finalized by earlier steps, not the slot this step
# is about to write.
coup_new = _coupling_at(new_state, 0.0, buf, t + 1, params)
new_deriv = dfun(new_state, coup_new, params)[cvar_idx]
new_buf = _write_ring_interp(buf, t + 1, new_state[cvar_idx], new_deriv, horizon)
elif has_delays:
new_buf = _write_ring(buf, t + 1, new_state[cvar_idx], horizon)
else:
new_buf = buf
return (new_state, new_buf, t + 1, params), new_state
return step
