patrick-kidger / patrick-kidger/diffrax
Inconsistency between constant and adaptive step size solvers with Discrete terminating events
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 2.1k
- Forks
- 189
- Avg merge
- 3d 18h
- Merged PRs (30d)
- 1
Description
When using the adaptive step size solver to integrate an ODE with a discrete terminating event, the detected event time is significantly after the detected event time computed with a constant step size solver. This error is significant for my application and above all tolerances I set. Is this expected behavior?
I am attaching a plot to illustrate the problem. You can see that the adaptive step solver continues integrating long after the event should be triggered (computed analytically).
And a minimum working example:
import diffrax as de
import jax.numpy as jnp
n = 2
m = 1
mass = 1.0
l = 1.0
gravity = 9.81
damping = 0.0
cost_threshold = 1e0
def f(x):
theta, theta_dot = x[..., 0], x[..., 1]
damp = (-theta_dot * damping) / (mass * l ** 2)
f_1 = theta_dot
f_2 = damp + (gravity/l)*jnp.sin(theta)
return jnp.stack([f_1, f_2], axis=-1)
def g(x):
return jnp.array([[0.0], [1.0/(mass * l**2)]])
def cost(x):
return jnp.sum(x**2)
def vf(t, y_rc, args):
y, rc = y_rc
y_dot = f(y) + g(y) @ jnp.array([10.0])
rc_dot = cost(y)
return y_dot, rc_dot
def terminating_event(state, **kwargs):
y_rc = state.y
y, rc = y_rc
return rc > cost_threshold
ts = jnp.linspace(0, 10, 2**13)
adaptive_solution = de.diffeqsolve(
de.ODETerm(vf),
solver=de.Tsit5(),
t0=ts[0],
t1=ts[-1],
dt0=ts[1]-ts[0],
saveat=de.SaveAt(ts=ts),
y0=(jnp.array([-jnp.pi, 0.0]), jnp.array(0.0)),
stepsize_controller=de.PIDController(rtol=1e-6, atol=1e-12, dtmin=1e-5),
discrete_terminating_event=de.DiscreteTerminatingEvent(terminating_event),
max_steps=8192
)
constant_solution = de.diffeqsolve(
de.ODETerm(vf),
solver=de.Tsit5(),
t0=ts[0],
t1=ts[-1],
dt0=ts[1]-ts[0],
saveat=de.SaveAt(ts=ts),
y0=(jnp.array([-jnp.pi, 0.0]), jnp.array(0.0)),
stepsize_controller=de.ConstantStepSize(),
discrete_terminating_event=de.DiscreteTerminatingEvent(terminating_event),
max_steps=8192
)
finite_xor = jnp.isfinite(constant_solution.ys[1]) ^ jnp.isfinite(adaptive_solution.ys[1])
adaptive_fail_ts = finite_xor * ts
non_zeros = jnp.nonzero(adaptive_fail_ts)[0]
fail_ts = adaptive_fail_ts[non_zeros]
first_event = ts[non_zeros[0]-1]
second_event = ts[non_zeros[-1]]
print(f'First event: {first_event}')
print(f'Second event: {second_event}')
print(f"Event time difference: {second_event - first_event}")
P.S.: Thanks so much for your work on diffrax. It has been invaluable for my research.
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
Run the provided minimum working example and compare Tsit5 with PIDController against ConstantStepSize, both using DiscreteTerminatingEvent. Inspect the event-time difference and determine whether adaptive stepping detects the terminating event incorrectly or whether the behavior is expected; done means the behavior is corrected or clearly documented.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- backend
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100