BUG: integrator reverses direction when diffrax.ClipStepSizeController and event detection collide

[FFroehlich]

When diffrax is set up to track roots with positive directions (this is one the dev branch) and starts in a root with negative direction with diffrax.ClipStepSizeController also tracking the location of the same root, it will integrate backwards in time and report that root as root with positive direction (notice 9.7999999999999765e-01 < 9.7999999999999987e-01). Code example could probably further reduced, but it was sufficient to identify diffrax.ClipStepSizeController as the culprit. Segment 1 runs simulation to the root, Segment 2 is where diffrax starts integrating backwards:

Segment 1: t_start = 0.0000000000000000e+00, t_end = 1.0000000000000000e+00
Segment 1: Event detected at t = 9.7999999999999987e-01
Segment 2: t_start = 9.7999999999999987e-01, t_end = 2.0000000000000000e+00
Segment 2: Event detected at t = 9.7999999999999765e-01

With

controller = (
        diffrax.ClipStepSizeController(
            controller,
            jump_ts=[0.98],
        )
    )

removed, all works fine:

Segment 1: t_start = 0.0000000000000000e+00, t_end = 1.0000000000000000e+00
Segment 1: Event detected at t = 9.7999999999999998e-01
Segment 2: t_start = 9.7999999999999998e-01, t_end = 2.0000000000000000e+00
Segment 2: No event detected

code:

import jax

from typing import Callable
from optimistix import AbstractRootFinder

# Enable 64-bit precision
jax.config.update("jax_enable_x64", True)

import jax.numpy as jnp
import diffrax
import equinox as eqx
import optimistix

def _run_segment(
    t_start: float,
    t_end: float,
    y0,
    solver: diffrax.AbstractSolver,
    controller: diffrax.AbstractStepSizeController,
    root_finder: AbstractRootFinder,
    max_steps: jnp.int_,
    adjoint: diffrax.AbstractAdjoint,
    cond_fns: list[Callable],
    cond_dirs: list[None | bool],
    saveat: diffrax.SaveAt,
    term: diffrax.ODETerm,
) -> diffrax.Solution:
    controller = (
        diffrax.ClipStepSizeController(
            controller,
            jump_ts=[0.98],
        )
    )

    # combine all discontinuity conditions into a single diffrax.Event
    event = diffrax.Event(
        cond_fn=cond_fns,
        root_finder=root_finder,
        direction=cond_dirs,
    )

    sol = diffrax.diffeqsolve(
        term,
        solver,
        t0=t_start,
        t1=t_end,
        dt0=None,
        y0=y0,
        stepsize_controller=controller,
        max_steps=max_steps,
        adjoint=adjoint,
        saveat=saveat,
        event=event,
        throw=False,
    )

    return sol


# ============================================================================
# Model 00979 functions (recreated from generated model code)
# ============================================================================

def _xdot(t, x, args):
    """ODE dynamics for model 00979 (all zero)."""

    damici_ydt = 0.0
    damici_xdt = 0.0
    damici_pdt = 0.0
    dqdt = 0.0

    return jnp.array([damici_ydt, damici_xdt, damici_pdt, dqdt])


def _root_cond_fn_combined(t, y, args, **_):
    """Combined root condition function that returns all roots as a vector."""

    eroot0 = t - 49/50      # t - 0.98
    eroot1 = 49/50 - t      # 0.98 - t

    return jnp.array([eroot0, eroot1])


def _x0(t, p):
    """Initial state."""
    x00 = 0.0
    x01 = 0.0
    x02 = 0.0
    x03 = 0.0
    return jnp.array([x00, x01, x02, x03])


# ============================================================================
# Main test
# ============================================================================

def main():
    # Model parameters
    x0 = _x0(0.0, None)

    # Create solver and controller
    atol = 1e-4
    rtol = 1e-4
    tol_factor = 1e2
    solver = diffrax.Kvaerno5()
    controller = diffrax.PIDController(
        rtol=rtol / tol_factor,
        atol=atol / tol_factor,
        pcoeff=0.0,
        icoeff=0.3,
        dcoeff=0.4,
    )
    root_finder = optimistix.Newton(atol=atol, rtol=rtol)

    # Create ODE term
    term = diffrax.ODETerm(_xdot)

    # Create individual root condition functions from the combined one
    def _root_cond_fn_event(
        ie: int,
        t: float,
        y,
        args: tuple,
        **_
    ):
        rval = _root_cond_fn_combined(t, y, args, **_)
        return rval.at[ie].get()

    # ========================================================================
    # Segment 1: t in [0.0, 1.0]
    # ========================================================================

    print(f"Segment 1: t_start = {0.0:.16e}, t_end = {1.0:.16e}")

    saveat_1 = diffrax.SaveAt(t1=True)
    sol1 = _run_segment(
        t_start=0.0,
        t_end=1.0,
        y0=x0,
        solver=solver,
        controller=controller,
        root_finder=root_finder,
        max_steps=10,
        adjoint=diffrax.DirectAdjoint(),
        cond_fns=eqx.Partial(_root_cond_fn_event, 0),
        cond_dirs=True,
        saveat=saveat_1,
        term=term,
    )

    if diffrax.is_event(sol1.result):
        t1_event = float(sol1.ts[-1]) if len(sol1.ts) > 0 else sol1.t1
        print(f"Segment 1: Event detected at t = {t1_event:.16e}")
    else:
        print(f"Segment 1: No event detected")

    # ========================================================================
    # Segment 2: t in [t0_next, 2.0]
    # ========================================================================

    t0_next = float(sol1.ts[-1]) if len(sol1.ts) > 0 else sol1.t1
    print(f"Segment 2: t_start = {t0_next:.16e}, t_end = {2.0:.16e}")

    saveat_2 = diffrax.SaveAt(t1=True)
    sol2 = _run_segment(
        t_start=sol1.ts[-1],
        t_end=2.0,
        y0=sol1.ys[-1],
        solver=solver,
        controller=controller,
        root_finder=root_finder,
        max_steps=10,
        adjoint=diffrax.DirectAdjoint(),
        cond_fns=eqx.Partial(_root_cond_fn_event, 1),
        cond_dirs=True,
        saveat=saveat_2,
        term=term,
    )

    if diffrax.is_event(sol2.result):
        t2_event = float(sol2.ts[-1]) if len(sol2.ts) > 0 else sol2.t1
        print(f"Segment 2: Event detected at t = {t2_event:.16e}")
    else:
        print(f"Segment 2: No event detected")

    return sol1, sol2


if __name__ == '__main__':
    main()

[FFroehlich]

as a follow up: why would you do such a crazy thing as doing both event tracking and step size clipping in the controller? Turns out gradient calculations with respect to state variables are off if you don't:

output with controller (correct):

Gradient final state: 1.0000152392863118e+00
Gradient event time: 1.0000000000000000e+00
Segment 1: Event detected at t = 9.7999999999999987e-01
Segment 2: Event detected at t = 9.7999999999999765e-01
sol2.ys[-1, 0] = 9.8000748277170480e-01

output without controller (incorrect for state)

Gradient final state: 3.0751597615596221e-01
Gradient event time: 1.0000000000000000e+00
Segment 1: Event detected at t = 9.7999999999999998e-01
Segment 2: No event detected
sol2.ys[-1, 0] = 9.7995867587552554e-01

Note that the gradient for the event time is correct in both cases, so I can't quite make sense of as to whats going on.

Code:

import jax

from typing import Callable
from optimistix import AbstractRootFinder

# Enable 64-bit precision
jax.config.update("jax_enable_x64", True)

import jax.numpy as jnp
import diffrax
import equinox as eqx
import optimistix

def _run_segment(
    t_start: float,
    t_end: float,
    y0,
    args,
    solver: diffrax.AbstractSolver,
    controller: diffrax.AbstractStepSizeController,
    root_finder: AbstractRootFinder,
    max_steps: jnp.int_,
    adjoint: diffrax.AbstractAdjoint,
    cond_fns: list[Callable],
    cond_dirs: list[None | bool],
    saveat: diffrax.SaveAt,
    term: diffrax.ODETerm,
) -> diffrax.Solution:
    # controller = (
    #     diffrax.ClipStepSizeController(
    #         controller,
    #         jump_ts=[args],
    #     )
    # )

    event = diffrax.Event(
        cond_fn=cond_fns,
        root_finder=root_finder,
        direction=cond_dirs,
    )

    sol = diffrax.diffeqsolve(
        term,
        solver,
        args=args,
        t0=t_start,
        t1=t_end,
        dt0=None,
        y0=y0,
        stepsize_controller=controller,
        max_steps=max_steps,
        adjoint=adjoint,
        saveat=saveat,
        event=event,
        throw=False,
    )

    return sol


# ============================================================================
# Model
# ============================================================================

def _xdot(t, x, args):

    dqdt = jnp.select([jnp.less(t,args),True], [1.0, 0.0])

    return jnp.array([dqdt])


def _root_cond_fn_combined(t, y, args, **_):
    event_time = args
    eroot0 = t - event_time
    eroot1 = event_time - t

    return jnp.array([eroot0, eroot1])


def _x0(t, p):
    """Initial state."""
    x03 = 0.0
    return jnp.array([x03])


# ============================================================================
# Main test
# ============================================================================

def run_both_segments(event_time, solver, controller, root_finder, term, x0):
    """
    Run both segments (0.0 -> 1.0 -> 2.0) with event detection.

    This function supports gradient computation through JAX autodiff.

    Args:
        event_time: The time at which to detect events
        solver: ODE solver instance
        controller: Step size controller instance
        root_finder: Root finder instance
        term: ODE term
        x0: Initial state

    Returns:
        Tuple of (sol1, sol2, t1_event, t2_event)
    """
    # Create root condition function wrapper
    def _root_cond_fn_event(ie: int, t: float, y, args: tuple, **_):
        rval = _root_cond_fn_combined(t, y, args, **_)
        return rval.at[ie].get()

    # ========================================================================
    # Segment 1: t in [0.0, 1.0]
    # ========================================================================
    saveat_1 = diffrax.SaveAt(t1=True)
    sol1 = _run_segment(
        t_start=0.0,
        t_end=1.0,
        y0=x0,
        args=event_time,
        solver=solver,
        controller=controller,
        root_finder=root_finder,
        max_steps=100,
        adjoint=diffrax.DirectAdjoint(),
        cond_fns=eqx.Partial(_root_cond_fn_event, 0),
        cond_dirs=True,
        saveat=saveat_1,
        term=term,
    )

    # Extract event time from segment 1
    t1_event = sol1.ts[-1] if len(sol1.ts) > 0 else sol1.t1

    # ========================================================================
    # Segment 2: t in [t1_event, 2.0]
    # ========================================================================
    saveat_2 = diffrax.SaveAt(t1=True)
    sol2 = _run_segment(
        t_start=t1_event,
        t_end=2.0,
        y0=sol1.ys[-1],
        args=event_time,
        solver=solver,
        controller=controller,
        root_finder=root_finder,
        max_steps=100,
        adjoint=diffrax.DirectAdjoint(),
        cond_fns=eqx.Partial(_root_cond_fn_event, 1),
        cond_dirs=True,
        saveat=saveat_2,
        term=term,
    )

    # Extract event time from segment 2
    t2_event = sol2.ts[-1] if len(sol2.ts) > 0 else sol2.t1

    return sol1, sol2, t1_event, t2_event


def main():
    # Model parameters
    x0 = _x0(0.0, None)
    event_time = jnp.asarray(0.98)  # Parameter for event detection

    # Create solver and controller
    atol = 1e-4
    rtol = 1e-4
    tol_factor = 1e2
    solver = diffrax.Kvaerno5()
    controller = diffrax.PIDController(
        rtol=rtol / tol_factor,
        atol=atol / tol_factor,
        pcoeff=0.0,
        icoeff=0.3,
        dcoeff=0.4,
    )
    root_finder = optimistix.Newton(atol=atol, rtol=rtol)

    # Create ODE term
    term = diffrax.ODETerm(_xdot)

    # Define a wrapper function for gradient computation
    def compute_both_segments_state(event_time):
        sol1, sol2, t1_event, t2_event = run_both_segments(
            event_time, solver, controller, root_finder, term, x0
        )
        return sol2.ys[-1, 0]

    # Compute gradients with respect to event_time
    grad_state = jax.grad(compute_both_segments_state)(event_time)
    print(f"Gradient final state: {float(grad_state):.16e}")

    # Define a wrapper function for gradient computation
    def compute_both_segments_time(event_time):
        sol1, sol2, t1_event, t2_event = run_both_segments(
            event_time, solver, controller, root_finder, term, x0
        )
        return sol2.ys[-1, 0]

    # Define a wrapper function for gradient computation
    def compute_both_segments_time(event_time):
        sol1, sol2, t1_event, t2_event = run_both_segments(
            event_time, solver, controller, root_finder, term, x0
        )
        return sol1.ts[-1]

    # Compute gradients with respect to event_time
    grad_event_time = jax.grad(compute_both_segments_time)(event_time)
    print(f"Gradient event time: {float(grad_event_time):.16e}")
    print()



    # Run the actual computation
    sol1, sol2, t1_event, t2_event = run_both_segments(
        event_time, solver, controller, root_finder, term, x0
    )

    # Print results
    if diffrax.is_event(sol1.result):
        print(f"Segment 1: Event detected at t = {float(t1_event):.16e}")
    else:
        print(f"Segment 1: No event detected")

    if diffrax.is_event(sol2.result):
        print(f"Segment 2: Event detected at t = {float(t2_event):.16e}")
    else:
        print(f"Segment 2: No event detected")

    # Print final state
    print(f"sol2.ys[-1, 0] = {float(sol2.ys[-1, 3]):.16e}")

    return sol1, sol2


if __name__ == '__main__':
    main()

[patrick-kidger]

Thanks for the report! Sorry that this took a while to get around to. I think I should have fixed this in #714. Can you give that branch a try and see if it fixes your overall problem?

Meta-comment, the MWE here was quite large, so I had to spent some time trimming out the unused components, down to just this:

import jax.numpy as jnp
import optimistix as optx
import diffrax

jump_time = 0.98
controller = diffrax.PIDController(rtol=1e-6, atol=1e-6)
controller = diffrax.ClipStepSizeController(controller, jump_ts=[jump_time])
sol = diffrax.diffeqsolve(
    diffrax.ODETerm(lambda t, y, args: jnp.zeros_like(y)),
    diffrax.Heun(),
    t0=jnp.nextafter(jnp.asarray(jump_time), -1),
    t1=1.2,
    dt0=None,
    y0=jnp.array([0, 0, 0, 0.0]),
    stepsize_controller=controller,
    event=diffrax.Event(
        cond_fn=lambda t, y, args, **kw: jump_time - t,
        root_finder=optx.Newton(atol=1e-4, rtol=1e-4),
        direction=True,
    ),
    max_steps=100,
)
assert sol.result == diffrax.RESULTS.successful

which is much easier to understand :)

I don't completely understand your follow-up comment I'm afraid, do you mean that there is an issue with gradient calculations? If so I'd be very happy to take a look at a MWE (of similar size to the one I've added above!)