tinydiffeq 0.1.0__py3-none-any.whl

tinydiffeq/__init__.py

@@ -0,0 +1,37 @@
+"""Tiny differentiable ODE/SDE solvers for JAX.
+
+solve_ode integrates dx/dt = f(x, t, args, p) with fixed-step (Euler, RK4)
+or adaptive (Tsit5 + IController) explicit Runge-Kutta methods inside one
+bounded lax.scan of exactly max_steps iterations, so shapes are static and
+solves are differentiable in BOTH forward and reverse mode (including
+reverse-over-forward) with O(max_steps) memory. SaveAt picks the output:
+the endpoint, cubic-Hermite interpolation onto a fixed grid, or the raw
+padded step rows. solve_sde is fixed-step Euler-Maruyama with presampled
+diagonal noise. States are arrays (scalar or vector); pytree states,
+implicit/stiff solvers, PID control, events, dense output, and adjoint
+methods are non-goals — use diffrax for those.
+"""
+
+from tinydiffeq.controllers import ConstantStepSize, IController
+from tinydiffeq.interpolation import hermite_interpolate
+from tinydiffeq.ode import solve_ode
+from tinydiffeq.quadrature import cumulative_trapezoid
+from tinydiffeq.saveat import SaveAt
+from tinydiffeq.sde import solve_sde
+from tinydiffeq.solution import Solution
+from tinydiffeq.solvers import RK4, Euler, EulerMaruyama, Tsit5
+
+__all__ = [
+    "solve_ode",
+    "solve_sde",
+    "Euler",
+    "RK4",
+    "Tsit5",
+    "EulerMaruyama",
+    "ConstantStepSize",
+    "IController",
+    "SaveAt",
+    "Solution",
+    "hermite_interpolate",
+    "cumulative_trapezoid",
+]

tinydiffeq/controllers.py

@@ -0,0 +1,72 @@
+from dataclasses import dataclass
+
+import jax
+import jax.numpy as jnp
+
+# Controllers decide accept/reject and the next step size from the embedded
+# error estimate. Contract: adapt(x0, x1, err, dt_used, dt_prev, order) ->
+# (accept, dt_next), where dt_used is the horizon-clipped step the solver
+# actually took and dt_prev is the unclipped carried step size. All numeric
+# fields are pytree data leaves, so changing tolerances never recompiles.
+
+SAFETY, MIN_FACTOR, MAX_FACTOR = 0.9, 0.2, 5.0
+
+
+@jax.tree_util.register_dataclass
+@dataclass(frozen=True)
+class ConstantStepSize:
+    """Accept every step and keep the carried step size unchanged."""
+
+    uses_error_estimate = False
+
+    def adapt(self, x0, x1, err, dt_used, dt_prev, order):
+        return jnp.asarray(True), dt_prev
+
+
+@jax.tree_util.register_dataclass
+@dataclass(frozen=True)
+class IController:
+    """Integral step-size controller with max-norm error.
+
+    Accept iff ``E = max(|err| / (atol + rtol * max(|x0|, |x1|))) <= 1``
+    (forced accept once the step reaches ``dtmin``), and propose
+    ``dt_next = dt_used * clip(safety * E**(-1/order), factormin, factormax)``
+    clipped to ``[dtmin, dtmax]``. This is the classic integral controller —
+    equal to diffrax's ``PIDController`` at its default coefficients
+    (pcoeff=0, dcoeff=0); there is no proportional term, hence the name.
+    """
+
+    rtol: float
+    atol: float
+    dtmin: float = 1e-10
+    dtmax: float = float("inf")
+    safety: float = SAFETY
+    factormin: float = MIN_FACTOR
+    factormax: float = MAX_FACTOR
+
+    uses_error_estimate = True
+
+    def adapt(self, x0, x1, err, dt_used, dt_prev, order):
+        dtype = jnp.result_type(x0, float)
+        rtol = jnp.asarray(self.rtol, dtype)
+        atol = jnp.asarray(self.atol, dtype)
+        dtmin = jnp.asarray(self.dtmin, jnp.result_type(dt_used))
+        dtmax = jnp.asarray(self.dtmax, jnp.result_type(dt_used))
+        safety = jnp.asarray(self.safety, dtype)
+        factormin = jnp.asarray(self.factormin, dtype)
+        factormax = jnp.asarray(self.factormax, dtype)
+        scale = atol + rtol * jnp.maximum(jnp.abs(x0), jnp.abs(x1))
+        # The controller is wrapped in stop_gradient: accept/reject is a
+        # non-differentiable branch either way, the gradient of E**(-1/order)
+        # blows up at the exact-zero error of a flat-start policy, and the
+        # d(dt)/dtheta term only slides sample points along the visited
+        # trajectory -- irrelevant to a residual that must vanish at every
+        # state. The states themselves remain fully differentiable through
+        # the RK stages.
+        E = jax.lax.stop_gradient(jnp.max(jnp.abs(err) / scale))
+        accept = (E <= 1.0) | (dt_used <= dtmin)
+        factor = jnp.clip(
+            safety * jnp.maximum(E, 1e-12) ** (-1.0 / order), factormin, factormax
+        )
+        dt_next = jnp.clip(jax.lax.stop_gradient(dt_used) * factor, dtmin, dtmax)
+        return accept, dt_next

tinydiffeq/interpolation.py

@@ -0,0 +1,47 @@
+import jax.numpy as jnp
+
+# Cubic Hermite interpolation over the raw attempt rows of the bounded scan.
+# The knot times are nondecreasing with exact duplicates (rejected attempts
+# and the post-horizon frozen tail repeat the previous row), so no compaction
+# is needed: searchsorted(side="right") - 1 lands on the LAST duplicate at or
+# before each query, giving a positive-width bracket everywhere except the
+# frozen tail. Degenerate (zero-width) brackets return the left knot -- flat
+# extrapolation beyond the reached time when the budget was exhausted. The
+# bracketing indices are integer and non-differentiable, consistent with the
+# stop-gradiented step controller (the sliding-knot d(dt)/dtheta term is
+# deliberately excluded); values differentiate fully through xL, xR, fL, fR.
+
+
+def hermite_interpolate(ts_query, knot_ts, knot_xs, knot_fs):
+    """Cubic Hermite interpolation of ``(knot_ts, knot_xs, knot_fs)`` at
+    ``ts_query``, where ``knot_fs`` holds the time derivatives at the knots.
+
+    ``knot_ts`` must be nondecreasing; duplicate knots are allowed and
+    queries falling on a zero-width bracket return the left knot value.
+    Queries outside the knot span clamp to the boundary knot values (flat
+    extrapolation) rather than evaluating the cubic outside its bracket.
+    """
+    n = knot_ts.shape[0]
+    idx = jnp.clip(jnp.searchsorted(knot_ts, ts_query, side="right") - 1, 0, n - 2)
+    t_left, t_right = knot_ts[idx], knot_ts[idx + 1]
+    x_left, x_right = knot_xs[idx], knot_xs[idx + 1]
+    f_left, f_right = knot_fs[idx], knot_fs[idx + 1]
+    width = t_right - t_left
+    degenerate = width <= 0.0
+    # double-where: divide by the safe width BEFORE branching so neither the
+    # primal nor its jvp/vjp ever sees a 0/0.
+    width_safe = jnp.where(degenerate, 1.0, width)
+    s = jnp.clip((ts_query - t_left) / width_safe, 0.0, 1.0)
+
+    extra = knot_xs.ndim - 1
+
+    def bc(a):
+        return a.reshape(a.shape + (1,) * extra)
+
+    s_, w_, deg_ = bc(s), bc(width_safe), bc(degenerate)
+    h00 = (1.0 + 2.0 * s_) * (1.0 - s_) ** 2
+    h10 = s_ * (1.0 - s_) ** 2
+    h01 = s_**2 * (3.0 - 2.0 * s_)
+    h11 = s_**2 * (s_ - 1.0)
+    value = h00 * x_left + h10 * w_ * f_left + h01 * x_right + h11 * w_ * f_right
+    return jnp.where(deg_, x_left, value)

tinydiffeq/ode.py

@@ -0,0 +1,176 @@
+import inspect
+
+import jax
+import jax.numpy as jnp
+
+from tinydiffeq.controllers import ConstantStepSize
+from tinydiffeq.interpolation import hermite_interpolate
+from tinydiffeq.saveat import SaveAt
+from tinydiffeq.solution import Solution
+
+
+def identity_project(x):
+    return x
+
+
+def canonicalize_field(f, name="f"):
+    # The vector field may take (x), (x, t), (x, t, args), or (x, t, args, p),
+    # always in that order; it is wrapped into the canonical 4-arg form here
+    # so the compiled code is identical for all four. Uninspectable
+    # signatures (or *args) are assumed 4-arg.
+    try:
+        signature = inspect.signature(f)
+    except (TypeError, ValueError):
+        arity = 4
+    else:
+        arity = 0
+        for parameter in signature.parameters.values():
+            if parameter.kind in (
+                inspect.Parameter.POSITIONAL_ONLY,
+                inspect.Parameter.POSITIONAL_OR_KEYWORD,
+            ):
+                arity += 1
+            elif parameter.kind == inspect.Parameter.VAR_POSITIONAL:
+                arity = 4
+                break
+        if arity < 1 or arity > 4:
+            raise ValueError(
+                f"{name} must take 1 to 4 positional arguments: "
+                "(x), (x, t), (x, t, args), or (x, t, args, p)"
+            )
+    if arity == 1:
+        return lambda x, t, args, p: f(x)
+    if arity == 2:
+        return lambda x, t, args, p: f(x, t)
+    if arity == 3:
+        return lambda x, t, args, p: f(x, t, args)
+    return f
+
+
+def solve_ode(
+    f,
+    solver,
+    t0,
+    t1,
+    x0,
+    *,
+    p=None,
+    args=None,
+    dt0=None,
+    saveat=None,
+    controller=None,
+    max_steps=4096,
+    project=None,
+):
+    """Integrate ``dx/dt = f(x, t, args, p)`` from ``t0`` to ``t1 > t0``.
+
+    The vector field may be declared ``f(x)``, ``f(x, t)``, ``f(x, t, args)``,
+    or ``f(x, t, args, p)`` — always in that order. ``x`` is an array state
+    (scalar or vector), ``args`` is pass-through data that is by convention
+    not an AD target, and ``p`` holds differentiable parameters (any pytree);
+    jvp/vjp with respect to ``p`` and ``x0`` are first-class.
+
+    Fixed and adaptive stepping share one bounded ``lax.scan`` of exactly
+    ``max_steps`` iterations, so shapes are static and curvature-dependent
+    step counts never retrace. ``dt0`` is required (no auto-initial-step
+    heuristic). Each attempt is clipped to the remaining horizon; the clipped
+    step also feeds the controller's next-step proposal, which doubles as the
+    growth guard — near-flat fields otherwise grow steps into quarter-horizon
+    leaps. ``project`` (e.g. a positivity clamp, assumed idempotent) is
+    applied at every point where ``f`` is evaluated and to every accepted
+    state. Returns a :class:`Solution`; ``sol.ok`` is False if the budget ran
+    out before ``t1`` (outputs are then the reached prefix, never poisoned).
+
+    The time dtype follows ``jnp.result_type(x0, float)``; the library never
+    sets ``jax_enable_x64`` — do that in your application.
+    """
+    if dt0 is None:
+        raise ValueError("dt0 is required (tinydiffeq has no initial-step heuristic)")
+    if saveat is None:
+        saveat = SaveAt(t1=True)
+    if controller is None:
+        controller = ConstantStepSize()
+    if project is None:
+        project = identity_project
+    if controller.uses_error_estimate and not solver.has_error_estimate:
+        raise ValueError(
+            f"{type(controller).__name__} needs an embedded error estimate, "
+            f"which {type(solver).__name__} does not provide"
+        )
+    f = canonicalize_field(f)
+
+    x0 = jnp.asarray(x0)
+    tdt = jnp.result_type(x0, float)
+    t0 = jnp.asarray(t0, tdt)
+    t1 = jnp.asarray(t1, tdt)
+    dt0 = jnp.asarray(dt0, tdt)
+    t_eps = 4.0 * jnp.finfo(tdt).eps * jnp.maximum(1.0, jnp.abs(t1))
+    # Summing ~max_steps rounded steps can leave t short of t1 by up to
+    # ~max_steps * eps, so a step whose remaining horizon is within that
+    # slack of the desired dt is stretched to land on t1 exactly; otherwise
+    # dt0 = (t1 - t0)/n with max_steps = n would strand a one-ulp sliver.
+    t_slack = max_steps * t_eps
+
+    def g(x, t):
+        return f(project(x), t, args, p)
+
+    need_f = solver.fsal or (saveat.ts is not None)
+    f_init = g(x0, t0) if need_f else jnp.zeros_like(x0)
+
+    def body(carry, _):
+        t, x, dt, f_cur, done, num_accepted = carry
+        remaining = t1 - t
+        h = jnp.where(remaining <= dt + t_slack, jnp.maximum(remaining, 1e-12), dt)
+        x1, f1, err = solver.step(g, t, x, h, f_cur if need_f else None, project)
+        if need_f and f1 is None:
+            f1 = g(x1, t + h)
+        accept, dt_next = controller.adapt(x, x1, err, h, dt, solver.order)
+        advance = accept & ~done
+        x_new = jnp.where(advance, x1, x)
+        t_new = jnp.where(advance, t + h, t)
+        f_new = jnp.where(advance, f1, f_cur) if need_f else f_cur
+        dt_new = jnp.where(done, dt, dt_next)
+        done_new = done | (t_new >= t1 - t_eps)
+        num_new = num_accepted + advance.astype(jnp.int32)
+        carry_new = (t_new, x_new, dt_new, f_new, done_new, num_new)
+        if saveat.t1:
+            out = None
+        elif saveat.steps:
+            out = (t_new, x_new, advance)
+        else:
+            out = (t_new, x_new, f_new, advance)
+        return carry_new, out
+
+    carry0 = (t0, x0, dt0, f_init, jnp.asarray(False), jnp.asarray(0, jnp.int32))
+    (t_final, x_final, _, _, done, num_accepted), rows = jax.lax.scan(
+        body, carry0, None, length=max_steps
+    )
+
+    if saveat.t1:
+        return Solution(ts=t_final, xs=x_final, ok=done, num_accepted=num_accepted)
+
+    if saveat.steps:
+        ts_s, xs_s, adv_s = rows
+    else:
+        ts_s, xs_s, fs_s, adv_s = rows
+    ts_all = jnp.concatenate([t0[None], ts_s])
+    xs_all = jnp.concatenate([x0[None], xs_s])
+    accepted = jnp.concatenate([jnp.ones((1,), bool), adv_s])
+
+    if saveat.steps:
+        if saveat.fill == "inf":
+            ts_all = jnp.where(accepted, ts_all, jnp.inf)
+            keep = accepted.reshape(accepted.shape + (1,) * (xs_all.ndim - 1))
+            xs_all = jnp.where(keep, xs_all, jnp.inf)
+        return Solution(
+            ts=ts_all,
+            xs=xs_all,
+            ok=done,
+            num_accepted=num_accepted,
+            accepted=accepted,
+        )
+
+    fs_all = jnp.concatenate([f_init[None], fs_s])
+    ts_query = jnp.asarray(saveat.ts, tdt)
+    xs_query = hermite_interpolate(ts_query, ts_all, xs_all, fs_all)
+    return Solution(ts=ts_query, xs=xs_query, ok=done, num_accepted=num_accepted)

tinydiffeq/quadrature.py

@@ -0,0 +1,31 @@
+import jax
+import jax.numpy as jnp
+
+
+def cumulative_trapezoid(g, ts, *, substeps=1):
+    """Cumulative composite-trapezoid integral of a time-only ``g(t)`` on the
+    (possibly nonuniform) sorted grid ``ts``.
+
+    Each grid interval is subdivided into ``substeps`` uniform panels, so the
+    quadrature error shrinks with ``substeps`` without changing the output
+    grid. Returns ``(integral, values)`` where ``integral[k]`` approximates
+    the integral of ``g`` from ``ts[0]`` to ``ts[k]`` (``integral[0] = 0``)
+    and ``values = g(ts)``; ``g`` may return any array shape, which is
+    appended to the leading grid axis.
+    """
+    if substeps < 1:
+        raise ValueError("substeps must be at least 1")
+    values = jax.vmap(g)(ts)
+
+    def interval_increment(left, right):
+        nodes = jnp.linspace(left, right, substeps + 1)
+        node_values = jax.vmap(g)(nodes)
+        widths = nodes[1:] - nodes[:-1]
+        widths = widths.reshape(widths.shape + (1,) * (node_values.ndim - 1))
+        return jnp.sum(0.5 * widths * (node_values[:-1] + node_values[1:]), axis=0)
+
+    increments = jax.vmap(interval_increment)(ts[:-1], ts[1:])
+    integral = jnp.concatenate(
+        [jnp.zeros_like(increments[:1]), jnp.cumsum(increments, axis=0)]
+    )
+    return integral, values

tinydiffeq/saveat.py

@@ -0,0 +1,39 @@
+from dataclasses import dataclass, field
+
+import jax
+
+
+@jax.tree_util.register_dataclass
+@dataclass(frozen=True)
+class SaveAt:
+    """What ``solve_ode``/``solve_sde`` return. Exactly one mode must be set.
+
+    - ``t1=True``: the endpoint only (the default in the solve functions).
+    - ``ts=grid``: cubic-Hermite interpolation of the internal steps onto a
+      fixed, sorted query grid in ``[t0, t1]``. Output shape is
+      ``(len(ts), ...)`` regardless of how many internal steps the controller
+      takes, so changing curvature never changes shapes or recompiles.
+      ``ts`` is a data leaf; a different grid of the same length retraces
+      nothing. ODE only.
+    - ``steps=True``: the raw padded attempt rows, ``max_steps + 1`` of them
+      including the initial state. Rejected attempts and post-horizon frozen
+      iterations duplicate the previous row. ``fill="last"`` (default) keeps
+      those duplicates; ``fill="inf"`` overwrites every non-accepted row
+      (including mid-trajectory rejection rows) with ``inf``, diffrax-style.
+
+    ``fill`` only applies to ``steps=True``.
+    """
+
+    t1: bool = field(default=False, metadata=dict(static=True))
+    ts: jax.Array | None = None
+    steps: bool = field(default=False, metadata=dict(static=True))
+    fill: str = field(default="last", metadata=dict(static=True))
+
+    def __post_init__(self):
+        modes = int(bool(self.t1)) + int(self.ts is not None) + int(bool(self.steps))
+        if modes != 1:
+            raise ValueError(
+                "SaveAt requires exactly one of t1=True, ts=..., steps=True"
+            )
+        if self.fill not in ("last", "inf"):
+            raise ValueError('SaveAt fill must be "last" or "inf"')

tinydiffeq/sde.py

@@ -0,0 +1,86 @@
+import jax
+import jax.numpy as jnp
+
+from tinydiffeq.ode import canonicalize_field, identity_project
+from tinydiffeq.saveat import SaveAt
+from tinydiffeq.solution import Solution
+
+
+def solve_sde(
+    drift,
+    diffusion,
+    solver,
+    t0,
+    t1,
+    x0,
+    *,
+    key,
+    n_steps,
+    p=None,
+    args=None,
+    saveat=None,
+    project=None,
+):
+    """Integrate the Ito SDE ``dx = drift dt + diffusion dW`` (diagonal noise)
+    on the fixed grid of ``n_steps`` uniform steps from ``t0`` to ``t1 > t0``.
+
+    ``drift`` and ``diffusion`` follow the same signature convention as
+    ``solve_ode`` — ``(x)``, ``(x, t)``, ``(x, t, args)``, or
+    ``(x, t, args, p)``. ``n_steps`` must be a static Python int (the honest
+    static-shape contract; there is no adaptive SDE stepping in v1). The
+    Brownian increments are presampled from ``key`` as
+    ``sqrt(dt) * normal(key, (n_steps,) + x0.shape)``, so a fixed key gives a
+    fixed noise process: reproducible across calls and differentiable with
+    respect to ``x0`` and ``p`` (not ``key``).
+
+    ``SaveAt(ts=...)`` raises — cubic Hermite interpolation is wrong for
+    rough paths; use ``t1`` (default) or ``steps`` (here ``n_steps + 1``
+    rows, all accepted).
+    """
+    if not isinstance(n_steps, int):
+        raise TypeError("n_steps must be a static Python int")
+    if n_steps < 1:
+        raise ValueError("n_steps must be at least 1")
+    if saveat is None:
+        saveat = SaveAt(t1=True)
+    if saveat.ts is not None:
+        raise ValueError(
+            "SaveAt(ts=...) is not supported for SDEs: Hermite interpolation "
+            "is wrong for rough paths; use SaveAt(t1=True) or SaveAt(steps=True)"
+        )
+    if project is None:
+        project = identity_project
+    drift = canonicalize_field(drift, name="drift")
+    diffusion = canonicalize_field(diffusion, name="diffusion")
+
+    x0 = jnp.asarray(x0)
+    tdt = jnp.result_type(x0, float)
+    t0 = jnp.asarray(t0, tdt)
+    t1 = jnp.asarray(t1, tdt)
+    dt = (t1 - t0) / n_steps
+    dW = jnp.sqrt(dt) * jax.random.normal(key, (n_steps,) + x0.shape, dtype=tdt)
+    ts_grid = jnp.linspace(t0, t1, n_steps + 1)
+
+    def g_drift(x, t):
+        return drift(project(x), t, args, p)
+
+    def g_diffusion(x, t):
+        return diffusion(project(x), t, args, p)
+
+    def body(x, inputs):
+        t, dW_step = inputs
+        x1 = solver.step(g_drift, g_diffusion, t, x, dt, dW_step, project)
+        return x1, x1 if saveat.steps else None
+
+    x_final, xs_s = jax.lax.scan(body, x0, (ts_grid[:-1], dW))
+    num_accepted = jnp.asarray(n_steps, jnp.int32)
+    ok = jnp.asarray(True)
+
+    if saveat.t1:
+        return Solution(ts=t1, xs=x_final, ok=ok, num_accepted=num_accepted)
+
+    xs_all = jnp.concatenate([x0[None], xs_s])
+    accepted = jnp.ones((n_steps + 1,), bool)
+    return Solution(
+        ts=ts_grid, xs=xs_all, ok=ok, num_accepted=num_accepted, accepted=accepted
+    )

tinydiffeq/solution.py

@@ -0,0 +1,28 @@
+from dataclasses import dataclass
+
+import jax
+
+
+@jax.tree_util.register_dataclass
+@dataclass(frozen=True)
+class Solution:
+    """Result of ``solve_ode``/``solve_sde``.
+
+    - ``ts``/``xs``: times and states in the shape dictated by ``SaveAt`` —
+      scalar/endpoint for ``t1``, ``(len(ts), ...)`` for ``ts``,
+      ``(max_steps + 1, ...)`` for ``steps``.
+    - ``ok``: scalar bool, True iff the integration reached ``t1`` within the
+      attempt budget. The package never poisons outputs; callers that want
+      diverging residuals do ``jnp.where(sol.ok, sol.xs, jnp.inf)``.
+    - ``num_accepted``: number of accepted steps (excluding the initial
+      state).
+    - ``accepted``: ``steps`` mode only (otherwise None): per-row bool mask,
+      True for rows holding a newly computed state. Row 0 (the initial state)
+      is always True, so ``accepted.sum() == num_accepted + 1``.
+    """
+
+    ts: jax.Array
+    xs: jax.Array
+    ok: jax.Array
+    num_accepted: jax.Array
+    accepted: jax.Array | None = None

tinydiffeq/solvers.py

@@ -0,0 +1,129 @@
+from dataclasses import dataclass
+
+import jax
+
+# Solvers are stateless frozen dataclasses registered as pytrees so they pass
+# through jit/vmap as ordinary arguments. `step` receives `g(x, t)` -- the
+# user vector field already wrapped so every evaluation goes through
+# `project` -- plus `f0`, the loop-carried value of `g` at (x, t) when the
+# loop guarantees it is current (FSAL, or interpolation output requested);
+# otherwise `f0` is None and the solver evaluates its own first stage.
+# The step contract is `step(g, t, x, dt, f0, project) -> (x1, f1, err)`:
+# `x1` is the projected accepted candidate, `f1` is `g(x1, t + dt)` when the
+# solver produces it for free (FSAL) and None otherwise, and `err` is the
+# embedded error estimate or None. `project` is assumed idempotent (a clamp).
+
+# Tsitouras (2011) 5(4) coefficients (FSAL: k7 = f(x1) is the next step's k1).
+A21 = 0.161
+A31, A32 = -0.008480655492356989, 0.335480655492357
+A41, A42, A43 = 2.8971530571054935, -6.359448489975075, 4.3622954328695815
+A51, A52, A53, A54 = (
+    5.325864828439257,
+    -11.748883564062828,
+    7.4955393428898365,
+    -0.09249506636175525,
+)
+A61, A62, A63, A64, A65 = (
+    5.86145544294642,
+    -12.92096931784711,
+    8.159367898576159,
+    -0.071584973281401,
+    -0.028269050394068383,
+)
+B1, B2, B3, B4, B5, B6 = (
+    0.09646076681806523,
+    0.01,
+    0.4798896504144996,
+    1.379008574103742,
+    -3.290069515436081,
+    2.324710524099774,
+)
+# embedded 4th-order error coefficients (b - bhat)
+E1, E2, E3, E4, E5, E6, E7 = (
+    -0.00178001105222577714,
+    -0.0008164344596567469,
+    0.007880878010261995,
+    -0.1447110071732629,
+    0.5823571654525552,
+    -0.45808210592918697,
+    1.0 / 66.0,
+)
+# stage times c_i (needed for non-autonomous fields)
+C2, C3, C4, C5, C6, C7 = 0.161, 0.327, 0.9, 0.9800255409045097, 1.0, 1.0
+
+
+@jax.tree_util.register_dataclass
+@dataclass(frozen=True)
+class Euler:
+    """Explicit Euler. Fixed-step only: no embedded error estimate."""
+
+    order = 1
+    fsal = False
+    has_error_estimate = False
+
+    def step(self, g, t, x, dt, f0, project):
+        k1 = g(x, t) if f0 is None else f0
+        x1 = project(x + dt * k1)
+        return x1, None, None
+
+
+@jax.tree_util.register_dataclass
+@dataclass(frozen=True)
+class RK4:
+    """Classic fourth-order Runge-Kutta. Fixed-step only: no error estimate."""
+
+    order = 4
+    fsal = False
+    has_error_estimate = False
+
+    def step(self, g, t, x, dt, f0, project):
+        k1 = g(x, t) if f0 is None else f0
+        k2 = g(x + 0.5 * dt * k1, t + 0.5 * dt)
+        k3 = g(x + 0.5 * dt * k2, t + 0.5 * dt)
+        k4 = g(x + dt * k3, t + dt)
+        x1 = project(x + (dt / 6.0) * (k1 + 2.0 * k2 + 2.0 * k3 + k4))
+        return x1, None, None
+
+
+@jax.tree_util.register_dataclass
+@dataclass(frozen=True)
+class Tsit5:
+    """Tsitouras 5(4) explicit Runge-Kutta with embedded error estimate.
+
+    FSAL: the last stage k7 = g(x1, t + dt) is the next step's first stage,
+    so an accepted adaptive step costs six fresh evaluations. Note k7 is
+    evaluated at the *projected* accepted state, so the FSAL cache stays
+    consistent with the state actually carried forward when `project` binds.
+    """
+
+    order = 5
+    fsal = True
+    has_error_estimate = True
+
+    def step(self, g, t, x, dt, f0, project):
+        k1 = g(x, t) if f0 is None else f0
+        k2 = g(x + dt * (A21 * k1), t + C2 * dt)
+        k3 = g(x + dt * (A31 * k1 + A32 * k2), t + C3 * dt)
+        k4 = g(x + dt * (A41 * k1 + A42 * k2 + A43 * k3), t + C4 * dt)
+        k5 = g(x + dt * (A51 * k1 + A52 * k2 + A53 * k3 + A54 * k4), t + C5 * dt)
+        k6 = g(
+            x + dt * (A61 * k1 + A62 * k2 + A63 * k3 + A64 * k4 + A65 * k5),
+            t + C6 * dt,
+        )
+        x1 = project(
+            x + dt * (B1 * k1 + B2 * k2 + B3 * k3 + B4 * k4 + B5 * k5 + B6 * k6)
+        )
+        k7 = g(x1, t + C7 * dt)
+        err = dt * (E1 * k1 + E2 * k2 + E3 * k3 + E4 * k4 + E5 * k5 + E6 * k6 + E7 * k7)
+        return x1, k7, err
+
+
+@jax.tree_util.register_dataclass
+@dataclass(frozen=True)
+class EulerMaruyama:
+    """Euler-Maruyama for Ito SDEs with diagonal noise. Fixed-step only."""
+
+    order = 1
+
+    def step(self, g_drift, g_diffusion, t, x, dt, dW, project):
+        return project(x + dt * g_drift(x, t) + g_diffusion(x, t) * dW)

tinydiffeq-0.1.0.dist-info/METADATA

@@ -0,0 +1,116 @@
+Metadata-Version: 2.4
+Name: tinydiffeq
+Version: 0.1.0
+Summary: Tiny differentiable ODE/SDE solvers for JAX — bounded-scan adaptive stepping, static shapes, jvp/vjp-safe
+Project-URL: Homepage, https://github.com/HighDimensionalEconLab/tinydiffeq
+Project-URL: Repository, https://github.com/HighDimensionalEconLab/tinydiffeq
+Project-URL: Documentation, https://highdimensionaleconlab.github.io/tinydiffeq/
+Project-URL: Issues, https://github.com/HighDimensionalEconLab/tinydiffeq/issues
+Author-email: Jesse Perla <jesseperla@gmail.com>
+License-Expression: MIT
+License-File: LICENSE
+Keywords: autodiff,differential-equations,jax,ode,runge-kutta,sde
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Science/Research
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Topic :: Scientific/Engineering :: Mathematics
+Requires-Python: >=3.11
+Requires-Dist: jax>=0.7.0
+Description-Content-Type: text/markdown
+
+# tinydiffeq
+
+[![CI](https://github.com/HighDimensionalEconLab/tinydiffeq/actions/workflows/ci.yml/badge.svg)](https://github.com/HighDimensionalEconLab/tinydiffeq/actions/workflows/ci.yml)
+[![Docs](https://github.com/HighDimensionalEconLab/tinydiffeq/actions/workflows/docs.yml/badge.svg)](https://highdimensionaleconlab.github.io/tinydiffeq/)
+[![PyPI](https://img.shields.io/pypi/v/tinydiffeq.svg)](https://pypi.org/project/tinydiffeq/)
+[![Python versions](https://img.shields.io/pypi/pyversions/tinydiffeq.svg)](https://pypi.org/project/tinydiffeq/)
+[![License: MIT](https://img.shields.io/github/license/HighDimensionalEconLab/tinydiffeq)](https://github.com/HighDimensionalEconLab/tinydiffeq/blob/main/LICENSE)
+[![Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff)
+
+Tiny differentiable ODE/SDE solvers for JAX: fixed-step Euler/RK4, adaptive
+Tsit5 with an integral step-size controller, and Euler–Maruyama for Itô SDEs.
+One bounded `lax.scan` of exactly `max_steps` iterations serves fixed and
+adaptive stepping, so shapes are static, nothing recompiles as tolerances or
+curvature change, and every solve is differentiable in **both** forward and
+reverse mode — including reverse-over-forward, the pattern a
+Levenberg–Marquardt optimizer with geodesic acceleration needs when it
+differentiates through a rollout.
+
+This is a deliberately small, jvp/vjp-friendly subset of
+[diffrax](https://docs.kidger.site/diffrax/). Use diffrax if you need pytree
+states, stiff/implicit solvers, PID control, events, dense output, or
+checkpointed/backsolve adjoints. tinydiffeq's single runtime dependency is
+`jax`.
+
+## Install
+
+```bash
+uv add tinydiffeq
+```
+
+For GPU use, install the JAX accelerator build that matches your hardware,
+for example:
+
+```bash
+uv add tinydiffeq "jax[cuda13]"
+```
+
+## Minimal example
+
+The vector field may take `(x)`, `(x, t)`, `(x, t, args)`, or
+`(x, t, args, p)` — always in that order. `args` is pass-through data (not an
+AD target by convention); `p` holds differentiable parameters (any pytree).
+
+```python
+import jax
+import jax.numpy as jnp
+from tinydiffeq import solve_ode, Tsit5, IController, SaveAt
+
+jax.config.update("jax_enable_x64", True)  # your call — the library never sets it
+
+
+def f(x, t, args, p):
+    return -p * x
+
+
+sol = solve_ode(
+    f, Tsit5(), 0.0, 2.0, jnp.asarray(1.0),
+    p=jnp.asarray(1.3),
+    dt0=0.1,
+    controller=IController(rtol=1e-8, atol=1e-10),
+    max_steps=512,
+    saveat=SaveAt(ts=jnp.linspace(0.0, 2.0, 21)),  # fixed output shape,
+)                                                  # however many steps adapt
+print(sol.xs)   # states on the grid
+print(sol.ok)   # reached t1 within the max_steps budget?
+```
+
+## Gradients through the solve
+
+```python
+def endpoint(p):
+    return solve_ode(
+        f, Tsit5(), 0.0, 2.0, jnp.asarray(1.0), p=p,
+        dt0=0.1, controller=IController(rtol=1e-10, atol=1e-12),
+        max_steps=512,
+    ).xs
+
+jax.grad(endpoint)(jnp.asarray(1.3))                      # reverse mode
+jax.jvp(endpoint, (jnp.asarray(1.3),), (jnp.asarray(1.0),))  # forward mode
+jax.grad(lambda p: jax.jvp(endpoint, (p,), (jnp.asarray(1.0),))[1])(
+    jnp.asarray(1.3)
+)                                                          # reverse-over-forward
+```
+
+The step-size controller is wrapped in `stop_gradient` (accept/reject is
+non-differentiable either way, and the error-ratio power blows up at exactly
+zero error); the states differentiate fully through the RK stages. See the
+[docs](https://highdimensionaleconlab.github.io/tinydiffeq/) for the design
+contracts: static shapes and `SaveAt`, AD through adaptive stepping, SDE key
+semantics, and migration recipes from hand-rolled RK4/Tsit5 loops.
+
+## License
+
+MIT

tinydiffeq-0.1.0.dist-info/RECORD

@@ -0,0 +1,13 @@
+tinydiffeq/__init__.py,sha256=OhMHS5DJX26rPrnTBIM4p6vCPYg8rrlB9aGXyw-e_7w,1383
+tinydiffeq/controllers.py,sha256=Umq0d3hku0TeKhqsOgxMgdAa_bDstx-o5E-fb8jx10k,2956
+tinydiffeq/interpolation.py,sha256=uYHO58eAIfw32ueG6SYGtPqv3p9LvWibzpTtvOhBpgk,2280
+tinydiffeq/ode.py,sha256=ruG2EGs-8HjPZQn5ogGAvvTWr-Kp_2l0BDPNuJT68X0,6584
+tinydiffeq/quadrature.py,sha256=0YglgE1bjhMfeXsIAGUzQAQKsOsHbHoBxtxu_bm42Vw,1283
+tinydiffeq/saveat.py,sha256=iFh04a8J9qJy8FY8Dvu4iJfcO8Q-e52l9Uekxo2siUs,1707
+tinydiffeq/sde.py,sha256=-pkSvcte1VxagUhoyfh1sJb_Gx1Lf2z3U3Qs_D7Z6ck,2912
+tinydiffeq/solution.py,sha256=9i2hV_s4Mhc4hk7K2pEWACMvAXxsIoQlFjcScY4wPFw,1015
+tinydiffeq/solvers.py,sha256=a8BJymEoHAcihl6J9KtB1Du2PjzpGq75GCVwRwU_zZQ,4291
+tinydiffeq-0.1.0.dist-info/METADATA,sha256=DoxN0K77HZyTDMVWEffde585ev9cc12IHaQkCgGGz_0,4851
+tinydiffeq-0.1.0.dist-info/WHEEL,sha256=lCkmxWfQsSc9CfIClYeavTdQeEX2toPqufh9gI35EQA,87
+tinydiffeq-0.1.0.dist-info/licenses/LICENSE,sha256=6ecMv6jswIubE8SkuH00m6qcbn45ZDuvG-fzsAYNFek,1079
+tinydiffeq-0.1.0.dist-info/RECORD,,

tinydiffeq-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.31.0
+Root-Is-Purelib: true
+Tag: py3-none-any

tinydiffeq-0.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 HighDimensionalEconLab
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.