Infinite final-step rejection loop due to conflicting endpoint clipping and adaptative RK step rejection

[Nicolas-Lepage]

Hi :)

I am using Diffrax through the Dynamiqs quantum simulation library, and I think I encountered an endpoint-handling issue in Diffrax.

This issue is similar but not the same mechanism as : Issue #703

Issue

I was investigating the stability of Tsit5() on a difficult quantum simulation where the PID controller converges to small time steps compared with the final integration time. The simulation runs fine for almost the whole interval, but at the very last step the solver gets stuck in what looks like an infinite rejection loop until max_steps is reached.

Diagnosis of the issue

Diffrax has a mechanism to avoid taking a very small last time step. If the proposed tnext is within the last 100 * ULP(t1) before the final time, then Diffrax treats this as close enough to the end and clips toward t1:

def _clip_to_end(tprev, tnext, t1, t1_clip_floor, keep_step):
    clip = tnext > t1_clip_floor
    tclip = jnp.where(keep_step, t1, tprev + 0.5 * (t1 - tprev))
    return jnp.where(clip, tclip, tnext)

I understand the motivation: if the solver is at

tprev = t1 - epsilon

with very small epsilon, then taking a tiny final step to exactly t1 may be numerically unstable, especially for dense output.

In my case, with float32 and t1 = 600.0:

ULP(t1) ≈ 6.1e-5
100 * ULP(t1) ≈ 6.1e-3
PID-selected dt ≈ 3e-3

So the PID controller wants steps smaller than the final clipping window.

For example, suppose we are at:

tprev ≈ 599.993
and the PID controller naturally proposes:

tnext ≈ 599.996
This is already inside the interval:

[t1 - 100 * ULP(t1), t1]
so Diffrax treats it as a near-final step. If it tries to go to t1, the effective final step becomes roughly:

599.993 -> 600.0
which is about 0.007, larger than what the PID controller wanted. The step is rejected.

Then Diffrax proposes the midpoint:

tnext = 0.5 * (tprev + t1)
which is around:

599.9965
but this is still inside the final 100-ULP window. So Diffrax again treats it as a near-final step, tries to finish, rejects, proposes another point still inside the window, and the loop repeats until max_steps is reached.

Suggested fix

A possible fix, if you agree with the diagnosis, would be to make endpoint handling an option. The current behavior could remain the default, but another mode could be “overshoot and interpolate”:

Instead of forcing the last numerical step to land exactly on t1, allow the solver to take one normal accepted step past t1, then evaluate the existing dense/local interpolation at t1.

This would treat the final time like any other saved intermediate time. For example:

9.9975 -> 9.9985 -> 9.9995 -> 10.005

and then return:

y(10.0)

by evaluating the local interpolation on the accepted step [9.9995, 10.005].

Would this endpoint mode make sense for Diffrax? If so, I would be happy to try opening a PR.

[jpbrodrick89]

Instead of overshooting I think we can drop the overly cautious keep_step=False now that we use ULP instead of hardcoded tolerances we can probably just get away with:

def clip_to_end(...):
    return jnp.where(keep_step & (tnext > t1_clip_floor), t1, tnext)

and trust the adaptive time stepper to handle.

[patrick-kidger]

Ah, this is an interesting problem.

I think I agree with @jpbrodrick89's solution. In the case that the stepsize controller's proposed tnext is close to the end and we have a step rejection, then we are necessarily in the case that the previous step was something large – and therefore also close to the end – but was rejected. In this case we are necessarily in the case the the ODE requires sub-100-ULP time steps, and hoping to get dense output in this case (the main reason to have the clip-to-end) might just be wishful thinking.

@Nicolas-Lepage great work figuring out what was going on, I can imagine that took some debugging! Do you have a simple reproduction of the issue? It would be great to be able to add this to our test suite, alongside @jpbrodrick89's fix.

[Nicolas-Lepage]

Thanks for both your responses, that makes sense.

Checking my understanding

If I understand correctly, the proposed fix is to keep the accepted-step clipping behavior, since this is still a useful safeguard against tiny final intervals, but to remove the special rejected-step midpoint rule.

So if a near-final step is accepted, it still gets clipped to t1. But if a near-final step is rejected, Diffrax no longer overrides the step-size controller with the midpoint rule, and instead trusts the controller’s proposed next tnext.

That does seem like it directly fixes the loop I was seeing.

One point of caution

One small caveat I wanted to check: this does not completely guarantee the original intent of _clip_to_end: preventing a very small final interval, right?

For example, suppose t1 = 600.0, the 100-ULP window is roughly 6e-3, and the suggested step size around the end is around 3e-3( and max acceptable step size is very close, something like 3.39e-3).  I could imagine a sequence where the controller accepts a few small steps near the boundary of the final window, e.g. something schematically like:

599.9933 -> 599.9966 -> 599.9999

where the attempted step

599.9966 -> 600.0

was just too large and got rejected.

Then the remaining final interval is very small:

600.0 - 599.9999

which is the kind of final interval that the original clipping logic was trying to avoid for dense output.

So my understanding is that the proposed fix is not a complete endpoint-policy solution in the same sense as “overshoot and interpolate” would be. An overshoot/interpolate mode would avoid both forcing a larger final step and taking a tiny final interval, by treating t1 like any other intermediate saved time inside an accepted step.

That said, this remaining case seems much less likely/pathological than the rejection loop, and the proposed fix is much smaller and less invasive.

I’ll leave it to you both to decide whether this tradeoff is acceptable.

[Nicolas-Lepage]

Here is a stripped-down version of the problem I was simulating. It is still complex-valued and has a matrix-valued unknown, so I am not sure whether this is ideal for your test suite, but it reproduces the issue on my side.

It is a transient simulation of the von Neumann equation for a closed Kerr oscillator with two-photon drive.

import diffrax as dfx
import jax
import jax.numpy as jnp
import numpy as np


jax.config.update("jax_enable_x64", False)


def destroy(n):
    data = np.zeros((n, n), dtype=np.complex64)
    for k in range(1, n):
        data[k - 1, k] = np.sqrt(k)
    return jnp.asarray(data)


def coherent(n, alpha):
    coeffs = np.zeros(n, dtype=np.complex64)
    coeffs[0] = np.exp(-0.5 * abs(alpha) ** 2)
    for k in range(1, n):
        coeffs[k] = coeffs[k - 1] * alpha / np.sqrt(k)
    psi = jnp.asarray(coeffs)
    return psi / jnp.linalg.norm(psi)


n = 64
K = np.float32(0.1)
P = np.float32(0.4)
T_END = np.float32(600.0)

a = destroy(n)
adag = a.conj().T

# Kerr oscillator with two-photon drive:
# H = -0.5 K a†² a² + 0.5 P (a² + a†²)
H = -0.5 * K * (adag @ adag @ a @ a) + 0.5 * P * (a @ a + adag @ adag)
H = H.astype(jnp.complex64)

# Initial cat-state density matrix.
psi0 = coherent(n, -2.0) + coherent(n, 2.0)
psi0 = psi0 / jnp.linalg.norm(psi0)
rho0 = jnp.outer(psi0, psi0.conj()).astype(jnp.complex64)


def vf(t, rho, H):
    # von Neumann equation: dρ/dt = -i[H, ρ]
    return -1j * (H @ rho - rho @ H)


sol = dfx.diffeqsolve(
    dfx.ODETerm(vf),
    dfx.Tsit5(),
    t0=jnp.float32(0.0),
    t1=jnp.float32(T_END),
    dt0=None,
    y0=rho0,
    args=H,
    saveat=dfx.SaveAt(t1=True),
    stepsize_controller=dfx.PIDController(rtol=1e-6, atol=1e-6, safety=0.9),
    max_steps=200_000,
    throw=False,
    progress_meter=dfx.TextProgressMeter(),
)

print("result:", sol.result)
print("stats:", sol.stats)