patrick-kidger / patrick-kidger/diffrax
Zero-order-hold inputs with fixed-step solvers in a single diffeqsolve
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
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- 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