patrick-kidger / patrick-kidger/diffrax

Zero-order-hold inputs with fixed-step solvers in a single diffeqsolve

Open
#766 4 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Python
Stars
2.1k
Forks
189
Avg merge
3d 18h
Merged PRs (30d)
1

Description

Following the idea of forcing term in the docs, a natural extension is a zero-order-hold (piecewise-constant) input, which is very common for actuated systems:

import jax.numpy as jnp

ts = jnp.array([0.0, 1.0, 2.0, 3.0, 4.0])
us = jnp.array([1.0, 0.5, -0.5, 0.3, 0.8])
dt = 1.0
y0 = jnp.array(0.0)

def ode_zoh(t, y, args):
    ts_, us_ = args
    idx = jnp.clip(jnp.searchsorted(ts_, t, side="right") - 1, 0, us_.shape[0] - 2)
    return -y + us_[idx]

With a fixed-step higher-order solver stepping on ts, this is subtly wrong at the switching times:
the final stage of the step over [0, 1] is evaluated at exactly t=1.0 and sees u=0.5 instead of 1.0.
Using side="left" doesn't help, it just moves the error to the first stage of the following step.

Now this problem can be avoided completely if we write our integration as lax.scan of constant steps:

def ode(t, y, u):
    return -y + u

solver = diffrax.Tsit5()
n_intervals = ts.shape[0] - 1

def solve_scan(y0):
    def step(y, i):
        sol = diffrax.diffeqsolve(
            diffrax.ODETerm(ode),
            solver=solver,
            t0=ts[i],
            t1=ts[i + 1],
            dt0=dt,
            y0=y,
            args=us[i],
            saveat=diffrax.SaveAt(t1=True),
            stepsize_controller=diffrax.ConstantStepSize(),
        )
        y_next = sol.ys[0]
        return y_next, y_next

    _, ys_tail = jax.lax.scan(step, y0, jnp.arange(n_intervals))
    return jnp.concatenate([y0[None], ys_tail])

This is correct but gives up the single-call API: dense output, whole-trajectory SaveAt, events, and adjoint control over the full horizon (checkpointing is then dictated by the outer scan).
A single-call integration was possible in Diffrax 0.7.0 (and precisely before #608), because it was possible to have a step-size controller combining ClipStepSizeController(StepTo(step_ts=ts), step_ts=ts, jump_ts=ts).
My proposal would be to allow jump semantics with non-adaptive controllers, e.g. letting ClipStepSizeController inherit type from the inner controller. Happy to attempt a PR if this seems reasonable.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start by reading the forcing example in the documentation and the behavior of ClipStepSizeController, StepTo, and jump_ts, then compare with the pre-#608 behavior described in the issue. The change is complete when a single diffeqsolve call supports zero-order-hold switching with fixed-step controllers without losing dense output, full-horizon SaveAt, events, or adjoint control.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
backend-api-design
Issue type
Feature
Difficulty
4/5
Estimated time
3-5 days
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
52/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.