patrick-kidger / patrick-kidger/diffrax

Inconsistency between constant and adaptive step size solvers with Discrete terminating events

Open
#416 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

question
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).

image

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

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

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.