Events

[cholberg]

Updates to how events are handled in diffrax. The main changes are:

• There is now only one event class, Event.
• Multiple cond_fn are supported. An event is triggered whenever one of them changes sign.
• Supports differentiable event times if there are real-valued cond_fn and a root_finder is provided.
• New tests for the new event implementation.

Some things that still might require a little thinking about:

• Update documentation (and maybe examples, too).
• What to do with outdated tests.
• How to handle the case where the control is not smooth, e.g., Brownian motion.

[patrick-kidger]

This looks very clean! I've gone through and left some comments.

[patrick-kidger]

I don't think default arguments are needed here.

(This is also just a general good-practice-rule-of-thumb: try to avoid default arguments where possible, as they're a common source of unexpected behaviour. So in practice this usually means putting them on public interfaces.)

[patrick-kidger]

I think you can replace this whole function with just the built-in callable.

[patrick-kidger]

Don't use the Python builtin float (or int etc.) as a dtype. This will actually give different behaviour on different platforms (MacOS vs Windows etc.) for some strange reason.
I think in this case you're probably looking for dtype=jnp.result_type(old_event_result.dtype, jnp.float32)?

[patrick-kidger]

Nit: remove the not and flip teh branches.

[patrick-kidger]

Note that this will actually evaluate the function, I think essentially pointlessly (as it will be re-evaluated in the loop).
I think using jax.eval_shape or eqx.filter_eval_shape will allow you to get the structure/dtype/shape of dense_info_for_event without having to actually do any work at runtime. (And without having to compile solver.step an additional time during compilation time either, which is a nontrivial concern.)

[patrick-kidger]

(a) I think if you put this before init_state = ... then you can just pass in this information during initialisation, without tree_at.
(b) I think again you can use eval_shape to initialise arrays of the appropriate sort without having to actually evaluate each cond_fn.

[patrick-kidger]

Nit: prefer single-assignment form:

if event.root_finder is None:
    tevent = final_state.tprev
else:
    # all the rest of the branch below
    tevent = ...

one of these software best-practices for writing easy-to-read code.

[cholberg]

Thank you for the comments, Patrick. I have gone through and made changes accordingly. The only thing I would disagree with is your comment on not having to evaluate cond_fn. Correct me if I am wrong, but would we not need to know the sign of cond_fn evaluated at the intial condition. I.e., apriori you cannot know if an event happens when the condition function goes from negative to positive or vice versa. Furthermore, if we want cond_fn to also possibly take the state as an argument, event_result would need to be initialized after init_state is defined.

[patrick-kidger]

Just letting you know that I've not forgotten about this! I'm trying to focus on getting #344 in first, and then I'm hoping to return to this. They're both quite large changes so I don't want them to step on each other's toes.

[patrick-kidger]

Okay, #344 is in! I'd love to get this in next.

I appeciate that's rather a lot of merge conflicts. If you're able to rebase on to the latest dev branch then I'd be very happy to come back to this and start work on getting this in.

[cholberg]

Perfect, I rebased and squashed all the commits into a single big one. Quite a few tests are failing when I run it locally now, but I just wanted to update it so you could have a look.

Two thoughts since we last touched base:

1. We should probably narrow down exaclty what arguments should be passed to the functions in cond_fn.
2. We might want to make it so that the user has a way of passing arguments to the optimistix.root_find call that finds the exact event time.

[LuggiStruggi]

Hi! Maybe i am using it wrong but at the moment I can't get the root_finder to do smth. I get the same event times when i use the Newton Method as a root finder or just root_finder = None. Also the tests still pass when setting the root finder to None:

45 def test_continuous_event_time():
46     term = diffrax.ODETerm(lambda t, y, args: 1.0)
47     solver = diffrax.Tsit5()
48     t0 = 0
49     t1 = jnp.inf
50     dt0 = 1.0
51     y0 = -10.0
52 
53     def cond_fn(state, y, **kwargs):
54         assert isinstance(state.y, jax.Array)
55         return y
56 
57     #root_finder = optx.Newton(1e-5, 1e-5, optx.rms_norm)
58     root_finder = None
59     event = diffrax.Event(cond_fn, root_finder)
60     sol = diffrax.diffeqsolve(term, solver, t0, t1, dt0, y0, event=event)
61     assert jnp.all(jnp.isclose(cast(Array, sol.ts), 10.0, 1e-5))

Or in my code the resulting event times do not change if having a root finder or not:

  7 def bouncing_ball():
  8     g = 9.81
  9     damping = 0.8
 10     max_bounces = 10
 11     vx_0 = 5.0
 12 
 13     def dynamics(t, y, args):
 14         x, y, vx, vy = y
 15         dxdt = vx
 16         dydt = vy
 17         dvxdt = 0
 18         dvydt = -g
 19         return jnp.array([dxdt, dydt, dvxdt, dvydt])
 20 
 21     def cond_fn(state, y, **kwargs):
 22         return y[1] < 0
 23 
 24     y0 = jnp.array([0.0, 10.0, vx_0, 0.0])
 25     t0, t1 = 0, float('inf')
 26 
 27     times = []
 28     states = []
 29 
 30     for _ in range(max_bounces):
 31         root_finder = optx.Newton(1e-5, 1e-5)
 32         #root_finder = None
 33         event = Event(cond_fn, root_finder=root_finder)
 34         solver = Tsit5()
 35 
 36         sol = diffeqsolve(ODETerm(dynamics), solver, t0, t1, 0.01, y0, event=event)
 37 
 38         t0 = sol.ts[-1]
 39         last_y = sol.ys[-1]
 40         y0 = last_y * jnp.array([1, 0, 1, -damping])
 41         times.append(sol.ts)
 42         states.append(y0)
 43 
 44 
 45     return jnp.array(times), jnp.array(states)

Thanks for the help! (Also sorry If this is the wrong place to ask for this, just let me know where to write this) :)

[cholberg]

Hi! Maybe i am using it wrong but at the moment I can't get the root_finder to do smth. I get the same event times when i use the Newton Method as a root finder or just root_finder = None. Also the tests still pass when setting the root finder to None:

Ah, thanks for mentioning this. This is essentially due to the fact that the solution to the ODE is linear and the fact that dt0 divides 10.0 which is exactly the time at which the solution crosses 0. In other words, it just so happens that the root is exactly at the end point of the last step of the solver which is also the event time that is returned when no root finder is provided. I have tweaked the test so this is no longer the case.

Or in my code the resulting event times do not change if having a root finder or not:

In your example the cond_fn returns a boolean. In this case the returned event time is exactly the first step of the solver for which cond_fn switches sign, i.e., your event time will be n * dt0 where n is the number of steps taken until and including the point at which y[1] >= 0. Note that this is the same behaviour as when root_finder=None.

If you want continuous event times, you should specify a real-valued condition function. In your bouncing ball example, this would simply correspond to setting:

def cond_fn(state, y, **kwargs):
	return y[1]

[LuggiStruggi]

Ah i see! Was just a bit confusing that both boolean and comparison with 0 is possible. I tried a real valued cond_fn initially but then it was every time directly triggered at t=0 bc i reset the state to 0 at every jump. I guess setting it to a small value should work. Thank you! :)

[LuggiStruggi]

Regarding my confusion: Maybe this is a bad idea but maybe it would make sense to have a warning if both a boolean cond_fn and a root finder are passed to the event since this should never make sense then?

[LuggiStruggi]

If I pass multiple cond_fn, what would you suggest to determine which of them caused the Event? :)
Thanks for the help!

[cholberg]

Regarding my confusion: Maybe this is a bad idea but maybe it would make sense to have a warning if both a boolean cond_fn and a root finder are passed to the event since this should never make sense then?

Hmm, I'm not too sure about this. There might be cases where you have one real-valued event function and one boolean. E.g., in the bouncing ball example we might want to add an extra function to cond_fn stopping the solve when the velocity is low enough (i.e., the ball has reached a steady state and stopped bouncing). But happy to hear your thoughts as well @patrick-kidger.

If I pass multiple cond_fn, what would you suggest to determine which of them caused the Event? :)
Thanks for the help!

The solution returned by diffeqsolve has an attribute called event_mask which is a PyTree of the same structure as your cond_fn where each leaf is False if the corresponding condition function did not trigger an event and True otherwise. (Note: for now, only one leaf can be True). If you're interested in a more involved example of how this all works with multiple event handling you might want to check out this repo (specifically snn.py). This is a very much a work in progress, but it should serve as an example of how it all works. Hope that helps!

[patrick-kidger]

Agreed on the first point. Happy to add a warning if all events are Boolean, though -- no strong feelings.

[dkweiss31]

Hi folks! Very excited about this PR, as I'm thinking about quantum jump applications for dynamiqs. I'm unfortunately running into an error If I try to pass an option different from saveat = SaveAt(t1=True) :

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

term = dx.ODETerm(lambda t, y, args: y)
solver = dx.Tsit5()
t0 = 0
t1 = 100.0
dt0 = 1.0
y0 = 1.0
ts = jnp.arange(t0, t1, dt0)

def cond_fn(state, y, **kwargs):
    assert isinstance(state.y, jax.Array)
    return y - jnp.exp(1.0)

fn = lambda t, y, args: y

subsaveat_a = dx.SubSaveAt(ts=ts, fn=fn)  # save solution regularly
subsaveat_b = dx.SubSaveAt(t1=True)  # save last state
saveat = dx.SaveAt(subs=[subsaveat_a, subsaveat_b])

root_finder = optx.Newton(1e-5, 1e-5, optx.rms_norm)
event = dx.Event(cond_fn, root_finder)
sol = dx.diffeqsolve(
    term, solver, t0, t1, dt0, y0, saveat=saveat, event=event
)

This runs into an error on line 1204 of _integrate.py

    ys = jtu.tree_map(lambda _y, _yevent: _y.at[-1].set(_yevent), ys, yevent)
ValueError: Expected list, got Traced<ShapedArray(float32[1])>with<DynamicJaxprTrace(level=1/0)>.

Thanks!

[cholberg]

Ah, yes I forgot to handle the case where multiple SubSaveAt are passed. Should be fixed now with the latest commit. Let me know if that works for you.

[dkweiss31]

Thanks for the quick response! So that fixed the example I posted, however I am still running into issues on slightly more complicated examples more in line with how dynamiqs actually calls diffeqsolve and more specifically how it saves data as the simulation progresses. Here is a MWE where y saves the state and y2 saves "expectation values".

import diffrax as dx
import optimistix as optx
import jax
import jax.numpy as jnp
import equinox as eqx
from jax import Array

term = dx.ODETerm(lambda t, y, args: y + t)
solver = dx.Tsit5()
t0 = 0
t1 = 100.0
dt0 = 1.0
y0 = jnp.array([[1.0], [0.0]])
ts = jnp.arange(t0, t1, dt0)


def cond_fn(state, y, **kwargs):
    assert isinstance(state.y, jax.Array)
    norm = jnp.einsum("ij,ij->", y, y)
    return norm - jnp.exp(1.0)


class Saved(eqx.Module):
    y: Array
    y2: Array

def save_fn(t, y, args):
    ynorm = jnp.einsum("ij,ij->", y, y)
    return Saved(y, jnp.array([ynorm, 3 * ynorm]))

subsaveat_a = dx.SubSaveAt(ts=ts, fn=save_fn)  # save solution regularly
subsaveat_b = dx.SubSaveAt(t1=True)  # save last state
saveat = dx.SaveAt(subs=[subsaveat_a, subsaveat_b])

root_finder = optx.Newton(1e-5, 1e-5, optx.rms_norm)
event = dx.Event(cond_fn, root_finder)
sol = dx.diffeqsolve(
    term, solver, t0, t1, dt0, y0, saveat=saveat, event=event
)

This runs into

ValueError: Cannot broadcast to shape with fewer dimensions: arr_shape=(2, 2) shape=(2,)

Interestingly, the code runs without errors if I replace return Saved(y, jnp.array([ynorm, 3 * ynorm])) with return Saved(y, jnp.array([y, 3 * y]))

[cholberg]

You're right I did not account for the fact that SubSaveAt.fn could return a PyTree. Should be fixed now. At least your MWE works with the latest commit.

[dkweiss31]

Indeed that fixed my MWE! I hate to be such a pain but I am now running into another issue, here is an example that is now much closer to the actual code I am interested in running.

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

L_op = 0.1 * jnp.array([[0.0, 1.0],
                        [0.0, 0.0]], dtype=complex)
H = 0.0 * L_op
t0 = 0
t1 = 100.0
dt0 = 1.0
y0 = jnp.array([[0.0], [1.0]], dtype=complex)


def vector_field(t, state, _args):
    L_d_L = jnp.transpose(L_op) @ L_op
    new_state = -1j * (H - 1j * 0.5 * L_d_L) @ state
    return new_state


def cond_fn(state, **kwargs):
    psi = state.y
    prob = jnp.abs(jnp.einsum("id,id->", jnp.conj(psi), psi))
    return prob - 0.95


term = dx.ODETerm(vector_field)
root_finder = optx.Newton(1e-5, 1e-5, optx.rms_norm)
event = dx.Event(cond_fn, root_finder)

sol = dx.diffeqsolve(
    term, dx.Tsit5(), t0, t1, dt0, y0, event=event
)

This runs into

equinox._errors.EqxRuntimeError: The linear solver returned non-finite (NaN or inf) output. This usually means that the
operator was not well-posed, and that the solver does not support this.

It's possible this could be due to my use of complex numbers, which as I understand are only partly supported in diffrax? However with the previous DiscreteTerminatingEvent I did not run into such errors. Note that if I change 0.95 to 0.0 I don't see this error (likely because the event is not triggered)

[patrick-kidger]

Complex number support is definitely still iffy. Can you try reproducing this without using them? You can still solve the same equation, mathematically speaking, just by splitting things into separate real and imaginary components.

[dkweiss31]

Right, here is the same example using the complex->real isomorphism described e.g. here (see Eq. (9)). I am getting the same error as before, so it seems then this is not a complex number issue

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


def mat_cmp_to_real(matrix):
    re_matrix = jnp.real(matrix)
    im_matrix = jnp.imag(matrix)
    top_row = jnp.hstack((re_matrix, -im_matrix))
    bottom_row = jnp.hstack((im_matrix, re_matrix))
    return jnp.vstack((top_row, bottom_row))


def vec_cmp_to_real(vector):
    re_vec = jnp.real(vector)
    im_vec = jnp.imag(vector)
    return jnp.vstack((re_vec, im_vec))


L_op = 0.1 * jnp.array([[0.0, 1.0],
                        [0.0, 0.0]], dtype=complex)
L_d_L = jnp.transpose(L_op) @ L_op
H = 0.0 * L_op
_prop = -1j * (H - 1j * 0.5 * L_d_L)
_y0 = jnp.array([[0.0], [1.0]], dtype=complex)
prop = mat_cmp_to_real(_prop)
y0 = vec_cmp_to_real(_y0)
t0 = 0
t1 = 100.0
dt0 = 1.0


def vector_field(t, state, _args):
    new_state = prop @ state
    return new_state


def cond_fn(state, **kwargs):
    psi = state.y
    prob = jnp.abs(jnp.einsum("id,id->", jnp.conj(psi), psi))
    return prob - 0.95


term = dx.ODETerm(vector_field)
root_finder = optx.Newton(1e-5, 1e-5, optx.rms_norm)
event = dx.Event(cond_fn, root_finder)

sol = dx.diffeqsolve(
    term, dx.Tsit5(), t0, t1, dt0, y0, event=event
)

[patrick-kidger]

Thankyou @dkweiss31!
@cholberg, are you able to have a look at this example?

Anyway, as promised! Getting this in is my next priority. As such I've gone through and submitted a PR against this branch here. I don't claim that everything I've done is necessarily correct, so @cholberg I'd appreciate a review! :D

[cholberg]

Perfect! Rebased on the latest dev now :)

[johannahaffner]

Hi All!

I just opened #448, I get the same behavior on diffrax:main and cholberg/diffrax:dev.
When solver tolerances are atol=1e-06, rtol=1e-03 on cholberg/diffrax:dev, the solver terminates without a single step (?), reported t1 is 0.0.

There also appears to be a difference with respect to how the results are reported, as I get

---> 28 assert solution_a.result == dfx.RESULTS.discrete_terminating_event_occurred
(...)
--> 158 raise ValueError(
    159     "Can only compare equality between enumerations of the same type."
    160 )

ValueError: Can only compare equality between enumerations of the same type.

which is not raised on diffrax:main.

[johannahaffner]

Alright, I was merely being obtuse and didn't think about the most obvious thing: how the absolute and relative tolerances actually affect the terminating condition, and what that is based on. I assumed that it would compare rates of change $dy$ for consecutive steps. If you think that others might make the same mistake, I suggest adding a half-sentence to the documentation, something like

class SteadyStateEvent(AbstractDiscreteTerminatingEvent):
    """Terminates the solve once it reaches a steady state, defined as 
    norm(dy) < atol + rtol * norm(y).
    """

[patrick-kidger]

Aaaaand... merged! 🎉

Thank you for all your hard work @cholberg! This was a massive PR, so getting this in is no small feat. Especially in an large existing codebase like Diffrax.

This is now on the dev branch, which I'll be merging into main and doing a new release for shortly. :)

[patrick-kidger]

By the way, @dkweiss31 -- since I've not seen any movement on jax-ml/jax#22011, my guess is that we should probably work around this by registering a transpose rule for lax.stop_gradient ourselves. This transpose rule should just be the identity function.

I'd be happy to take a PR on this!

[dkweiss31]

Nice I'd love to take a swing at that! This would probably want to be in equinox rather than diffrax?

[patrick-kidger]

Hmm, I think probably Optimistix actually, now you mention it.
This is what's doing the manual linearisation and transposition, so it's there that we fall afoul of this bug.

[dkweiss31]

Hey @patrick-kidger, sorry for the delay I was on vacation last week :).

Please forgive and apologies in advance for my likely elementary jax mistakes: as a test I've tried doing the simplest possible thing, and added

from jax.interpreters import ad
from jax._src.ad_util import stop_gradient_p


def stop_gradient_transpose(ct, x):
    return ct

ad.primitive_transposes[stop_gradient_p] = stop_gradient_transpose

to the top of your example in jax-ml/jax#22011 as well as in my above example. Interestingly, your example now runs, while my example fails on line 720 of _integrate.py

line 720, in loop
    tfinal, yfinal, result = lax.cond(
jax._src.source_info_util.JaxStackTraceBeforeTransformation: AssertionError

Any idea what might be happening? (Also happy to move this conversation elsewhere, but just thought I'd keep it on this thread for continuities sake :))

[patrick-kidger]

Happy to keep the discussion here!
The fix is just three characters: you want return (ct,). :)

With this, your example works for me!

[dkweiss31]

Ha! I knew it was something silly. Thank you! Just opened a PR in optimistix :)

[dkweiss31]

I am, unfortunately, still encountering errors after the transpose fix. See my latest MWE, which fails with

EqxRuntimeError: Terminating differential equation solve because an event occurred.

after a few times through the for loop of the optimizer. Such a bizarre error, since I expect the differential equation to terminate!

from functools import partial
import optax
import jax
import jax.numpy as jnp
import diffrax as dx
import optimistix as optx


def run_traj(H, jump_ops, psi0):
    def vector_field(t, state, _args):
        Ls = jnp.stack([L for L in jump_ops])
        LdL = (Ls.mT.conj() @ Ls).sum(axis=0)
        new_state = -1j * (H - 1j * 0.5 * LdL) @ state
        return new_state

    def cond_fn(t, y, *args, **kwargs):
        return jnp.abs(jnp.conj(y) @ y) - no_jump_prob

    event = dx.Event(cond_fn, optx.Newton(1e-5, 1e-5, optx.rms_norm))
    solution = dx.diffeqsolve(
        dx.ODETerm(vector_field), dx.Tsit5(), t0=ts[0], t1=ts[-1], dt0=dt0, y0=psi0, event=event,
    )
    return solution.ys


def run(H, jump_ops, psi0):
    f = jax.vmap(run_traj, in_axes=(None, None, 0))
    return f(H, jump_ops, psi0)


def optimize(params_to_optimize):
    opt_state = optimizer.init(params_to_optimize)
    for epoch in range(10):
        params_to_optimize, opt_state = step(params_to_optimize, opt_state)
    return params_to_optimize


@partial(jax.jit)
def step(params_to_optimize, opt_state,):
    grads = jax.grad(loss)(params_to_optimize,)
    updates, opt_state = optimizer.update(grads, opt_state)
    params_to_optimize = optax.apply_updates(params_to_optimize, updates)
    return params_to_optimize, opt_state


def loss(params_to_optimize):
    res = run(params_to_optimize * H1, [a, ], initial_states)
    return jnp.log(jnp.sum(jnp.abs(res)))


dim = 4
a = jnp.diag(jnp.sqrt(jnp.arange(1, stop=dim, dtype=complex)), k=1)
H1 = a + a.mT.conj()
ket0 = jnp.array([1.0, 0.0, 0.0, 0.0], dtype=complex)
ket1 = jnp.array([0.0, 1.0, 0.0, 0.0], dtype=complex)
initial_states = jnp.asarray([ket0, ket1])
dt0 = 1.0
ts = jnp.arange(0.0, 100.0, dt0)
no_jump_prob = 0.9
optimizer = optax.adam(learning_rate=0.9)
init_drive_params = 0.1
opt_params = optimize(
    params_to_optimize=init_drive_params,
)

[dkweiss31]

Sorry, the previous example was not exactly minimal. Here is a more minimal example that doesn't need the optimization layer. The error is now

equinox._errors.EqxRuntimeError: The maximum number of steps was reached in the nonlinear solver. The problem may not be solveable (e.g., a root-find on a function that has no roots), or you may need to increase `max_steps`.

I've tried increasing the number of max steps to no avail. Interestingly, choosing init_drive_params = -1.1 or 1.2 runs without an error

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


def run_traj(H, jump_op, psi0):
    def vector_field(t, state, _args):
        LdL = jump_op.mT.conj() @ jump_op
        new_state = -1j * (H - 1j * 0.5 * LdL) @ state
        return new_state

    def cond_fn(t, y, *args, **kwargs):
        return jnp.abs(jnp.conj(y) @ y) - no_jump_prob

    event = dx.Event(cond_fn, optx.Newton(1e-5, 1e-5, optx.rms_norm))
    solution = dx.diffeqsolve(
        dx.ODETerm(vector_field), dx.Tsit5(), t0=ts[0], t1=ts[-1], dt0=dt0, y0=psi0, event=event
    )
    return solution.ys


a = jnp.diag(jnp.sqrt(jnp.arange(1, stop=4, dtype=complex)), k=1)
ket1 = jnp.array([0.0, 1.0, 0.0, 0.0], dtype=complex)
dt0 = 1.0
ts = jnp.arange(0.0, 100.0, dt0)
no_jump_prob = 0.9
init_drive_params = -1.19
res = run_traj(init_drive_params * (a + a.mT.conj()), a, ket1)

[patrick-kidger]

Hmm! I think a big part of what this is speaking to is the need for better debug tooling.

There is one big thing I've been wanting to do here, which is to stop using throw=False on all of our calls to Optimistix. Instead, we should respect the value of diffeqsolve(..., throw=...). Then if this is True, the error will be localized to exactly where it occured. If this is False, we still get the desired behaviour of not throwing. I'd be happy to take a PR on this. (I think it might need a bit of threading through to all the places that we use it.)

Tweaking this locally and setting EQX_ON_ERROR=breakpoint, JAX_DISABLE_JIT=1, then I can see that this error is thrown from a root find to locate the event, seemingly on the very first step. Which seems accurate:

> event.cond_fn(0, _interpolator.evaluate(0))
Array(0.10000002, dtype=float32)
> event.cond_fn(1, _interpolator.evaluate(1))
Array(-0.47552538, dtype=float32)

Can you add some jax.debug.breakpoint (or some of the tools from eqx.debug.* inside the iteration to try debugging this?

[patrick-kidger]

FWIW you've spurred me to try and improve things from the Equinox side, at least.

patrick-kidger/equinox#785

This should now give much more clearly informative error messages, at least once the error_if is actually triggered.