How to properly scale with multiple event conditions?

Hi diffrax team,

I am working on a problem, where the amount of cond_fn scales with my state dimension. Unfortunately, this currently scales pretty badly timewise even when jitting the code. Do you have any idea what could be done to improve here?

Here as an example:

\dot{y} + y = 0

y(0) \sim [0.9, 1.0]

event condition (any state):

y - 0.1 = 0

Solving with Euler and a fixed stepsize of 0.005ms between 0 and 30ms (or until first event), it scales really badly with the state dimension:

   Dim |   Time (s)
--------------------
     1 |   0.001022
     2 |   0.000625
     4 |   0.000657
     8 |   0.001162
    16 |   0.001870
    32 |   0.004110
    64 |   0.010566
   128 |   0.044464
   200 |   0.218422

see for comparison, the timing where i just used the cond_fn on the first element of the state y[0]:

   Dim |   Time (s)
--------------------
     1 |   0.000723
     2 |   0.000612
     4 |   0.000595
     8 |   0.000615
    16 |   0.000596
    32 |   0.000716
    64 |   0.000903
   128 |   0.000724
   200 |   0.001835

Do you see any way to speed this up a reasonable amount? Do you think it would be smh possible to additionally support vectorized condition functions for events? Like:

def cond_fn(t, y, args, **kwargs):
    return y - 0.1

Thanks for the help!

Here is the MWE i used (Single cond_fn commented out)

import jax
import jax.numpy as jnp
import diffrax as dfx
import optimistix as optx
from timeit import default_timer as timer

def solve_once(y0, t0, t1, dt0, solver, stepsize_controller, event):
    term = dfx.ODETerm(lambda t, y, args: -y)
    return dfx.diffeqsolve(
        terms=term,
        solver=solver,
        t0=t0,
        t1=t1,
        dt0=dt0,
        y0=y0,
        args=None,
        stepsize_controller=stepsize_controller,
        event=event,
        saveat=dfx.SaveAt(t0=False, t1=True, steps=False),
        max_steps=3000000,
        throw=True,
    )

solve_once = jax.jit(solve_once, static_argnames=["solver", "stepsize_controller", "event"])

def benchmark_diffeqsolve_with_event():
    key = jax.random.PRNGKey(0)
    t0, t1, dt0 = 0.0, 30.0, 0.005
    solver = dfx.Euler()
    controller = dfx.ConstantStepSize()
    dims = [1, 2, 4, 8, 16, 32, 64, 128, 200]

    print(f"{'Dim':>6} | {'Time (s)':>10}")
    print("-" * 20)

    for n in dims:
        key, subkey = jax.random.split(key)
        y0 = jax.random.uniform(subkey, shape=(n,), minval=0.9, maxval=1.1)

        cond_fns = [lambda t, y, *args, i=i, **kwargs: y[i] - 0.01 for i in range(n)]
        # cond_fns = lambda t, y, *args, **kwargs: y[0] - 0.01

        event = dfx.Event(cond_fns, root_finder=optx.Newton(rtol=1e-4, atol=1e-4))

        _ = solve_once(y0, t0, t1, dt0, solver, controller, event)  # warm-up

        start = timer()
        _ = solve_once(y0, t0, t1, dt0, solver, controller, event)
        end = timer()

        print(f"{n:6d} | {end - start:10.6f}")

benchmark_diffeqsolve_with_event()

Start with the supplied benchmark and the dfx.Event construction in solve_once, comparing the list of scalar cond_fns with the proposed vectorized function. Measure event-solving time across the listed state dimensions; done means the event API supports the intended condition shape or clearly documents the supported approach with improved scaling.