Skip to content

API Reference

Solve functions

tinydiffeq.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.

Source code in src/tinydiffeq/ode.py 
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
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.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).

Source code in src/tinydiffeq/sde.py 
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
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
    )

Solvers

tinydiffeq.Euler dataclass

Explicit Euler. Fixed-step only: no embedded error estimate.

Source code in src/tinydiffeq/solvers.py 
55
56
57
58
59
60
61
62
63
64
65
66
67
@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

tinydiffeq.RK4 dataclass

Classic fourth-order Runge-Kutta. Fixed-step only: no error estimate.

Source code in src/tinydiffeq/solvers.py 
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
@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

tinydiffeq.Tsit5 dataclass

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.

Source code in src/tinydiffeq/solvers.py 
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
@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

tinydiffeq.EulerMaruyama dataclass

Euler-Maruyama for Ito SDEs with diagonal noise. Fixed-step only.

Source code in src/tinydiffeq/solvers.py 
121
122
123
124
125
126
127
128
129
@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)

Step-size controllers

tinydiffeq.ConstantStepSize dataclass

Accept every step and keep the carried step size unchanged.

Source code in src/tinydiffeq/controllers.py 
15
16
17
18
19
20
21
22
23
@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

tinydiffeq.IController dataclass

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.

Source code in src/tinydiffeq/controllers.py 
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
@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

Output selection and results

tinydiffeq.SaveAt dataclass

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.

Source code in src/tinydiffeq/saveat.py 
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
@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.Solution dataclass

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.
Source code in src/tinydiffeq/solution.py 
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
@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

Utilities

tinydiffeq.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.

Source code in src/tinydiffeq/interpolation.py 
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
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.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.

Source code in src/tinydiffeq/quadrature.py 
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
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