diff --git a/diffrax/__init__.py b/diffrax/__init__.py index b207173a..6ba2c41e 100644 --- a/diffrax/__init__.py +++ b/diffrax/__init__.py @@ -81,4 +81,4 @@ ) -__version__ = "0.0.5" +__version__ = "0.0.6" diff --git a/diffrax/integrate.py b/diffrax/integrate.py index 5f9337b8..05d42568 100644 --- a/diffrax/integrate.py +++ b/diffrax/integrate.py @@ -88,8 +88,14 @@ def _save(state: _State, t: Scalar) -> _State: ) -def _clip_to_end(tnext, t1): - return jnp.where(tnext > t1 - 1e-6, t1, tnext) +def _clip_to_end(tprev, tnext, t1, keep_step): + if tnext.dtype is jnp.dtype("float64"): + tol = 1e-10 + else: + tol = 1e-6 + clip = tnext > t1 - tol + tclip = jnp.where(keep_step, t1, tprev + 0.5 * (t1 - tprev)) + return jnp.where(clip, tclip, tnext) def loop( @@ -161,8 +167,8 @@ def body_fun(state, inplace): # The 1e-6 tolerance means that we don't end up with too-small intervals for # dense output, which then gives numerically unstable answers due to floating # point errors. - tnext = _clip_to_end(tnext, t1) tprev = jnp.minimum(tprev, t1) + tnext = _clip_to_end(tprev, tnext, t1, keep_step) # The other parts of the mutable state are kept/not-kept (based on whether the # step was accepted) by the stepsize controller. But it doesn't get access to @@ -403,7 +409,7 @@ def _cond_fun(_state): def _body_fun(_state): _step, _t = _state - return _step + 1, _clip_to_end(_t + dt0, t1) + return _step + 1, _clip_to_end(_t, _t + dt0, t1, True) compiled_num_steps, _ = lax.while_loop( _cond_fun, _body_fun, (0, t0) diff --git a/test/test_detest.py b/test/test_detest.py index 78e40bef..9ca511f1 100644 --- a/test/test_detest.py +++ b/test/test_detest.py @@ -404,7 +404,7 @@ def _test(solver_ctr, problems, higher): # build up by t=20. # Teeny-tiny steps fix this. dt0 = 0.000001 - max_steps = 20_000_000 + max_steps = 20_000_001 stepsize_controller = diffrax.ConstantStepSize() elif solver_ctr is diffrax.ReversibleHeun and problem is _a1: # ReversibleHeun is a bit like LeapfrogMidpoint, and therefore bad over