Repository: google/jax

Issue:
Add odeint max step size optional argument
It'd be nice if `odeint` had an optional argument analogous to scipy `solve_ivp`'s `max_step` argument, which controls the maximum step size. I'll copy a code example below, but I'm running into an issue with `odeint` in which it steps over a feature in the ODE `rhs` after a long-ish period for which `rhs=0.`

One minor issue is that `max_step` unfortunately has a naming collision with `odeint`'s optional arg `mxstep`, which is an integer controlling the total number of iterations `odeint` can take before exiting, so this name may not be usable without changing `mxstep` which may not be desirable.

Aside from choosing the argument name, I think implementing this is nearly trivial: just need to add the argument, and I believe change [line 199](https://github.com/google/jax/blob/bab8520d0c6c40aa3799b897a7f2629ba746c2ad/jax/experimental/ode.py#L199) to take the minimum of the max step size and the computed optimal step. The example I post below could be rolled into a unittest validating this behaviour (perhaps with comparison directly with scipy).

Thoughts? Given how simple it is I'm happy to implement this myself, but would definitely defer to the core JAX developers for what argument name to use (or for any insight into why this might be an issue with other parts of `odeint`/JAX control flow).

## Motivating example

```
import jax
jax.config.update("jax_enable_x64", True)
jax.config.update('jax_platform_name', 'cpu')

import jax.numpy as jnp
from jax.experimental.ode import odeint

dt = jnp.array(1.)
k = 5
def deriv(t):
    t = jnp.array(t)
    return jnp.piecewise(
        t,
        [t <= 2 * dt, (t >= k * dt) & (t <= (k + 2) * dt)],
        [lambda s: jnp.array(1.), lambda s: jnp.array(-1.), lambda s: jnp.array(0.)]
    )

def rhs(y, t):
    return deriv(t)

odeint(
    func=rhs, 
    y0=jnp.array(0.), 
    t=jnp.array([0., k*dt, 2*k*dt]),
    atol=1e-8, rtol=1e-8
)
```
Using `odeint`, the above code is essentially integrating a function which is `1.` over the interval `[0., 2*dt]`, `-1.` over the interval `[k*dt, (k+1)*dt]`, and zero elsewhere. If integrated over an interval containing both non-zero intervals, the correct result should be `0.`. However , the above code outputs:
```
DeviceArray([0.        , 1.99999565, 1.99999565], dtype=float64)
```
and changing the tolerance doesn't change this (it just yields better or worse approximations of `[0., 2., 2.]`)

The analogous code in scipy using `solve_ivp` with `method='RK45'` yields the same results, but with `solve_ivp` this can be corrected by setting `max_step=dt`. Shortening the intermediate `0.` interval by either reducing `dt` or `k` both solve the problem as well, so I'm pretty confident it's an issue of step size overshoot.

### Slightly more context

The above is a minimal example, but the similarities to the situation in which I'm using `odeint` and the above are: 
- I don't know ahead of time if a period of `0.` will be present or for how long. 
- I do however have something analogous to `dt` - basically a natural time scale over which to restrict the step so that features aren't missed, but that also isn't so short that it will totally mess with the efficiency of the solver.


Discussion:
Have a look at [Diffrax](https://github.com/patrick-kidger/diffrax), which supports this as:
```python
diffeqsolve(..., stepsize_controller=PIDController(..., dtmax=...))
```
@patrick-kidger we are actually in the process of adding `diffrax` as a dependency to the project in which this came up :) . 

In any case it would still be a nice feature to have in `odeint`

Resolving patch:
diff --git a/jax/experimental/ode.py b/jax/experimental/ode.py
--- a/jax/experimental/ode.py
+++ b/jax/experimental/ode.py
@@ -142,7 +142,7 @@ def optimal_step_size(last_step, mean_error_ratio, safety=0.9, ifactor=10.0,
                       jnp.maximum(mean_error_ratio**(-1.0 / order) * safety, dfactor))
   return jnp.where(mean_error_ratio == 0, last_step * ifactor, last_step * factor)
 
-def odeint(func, y0, t, *args, rtol=1.4e-8, atol=1.4e-8, mxstep=jnp.inf):
+def odeint(func, y0, t, *args, rtol=1.4e-8, atol=1.4e-8, mxstep=jnp.inf, hmax=jnp.inf):
   """Adaptive stepsize (Dormand-Prince) Runge-Kutta odeint implementation.
 
   Args:
@@ -157,6 +157,7 @@ def odeint(func, y0, t, *args, rtol=1.4e-8, atol=1.4e-8, mxstep=jnp.inf):
     rtol: float, relative local error tolerance for solver (optional).
     atol: float, absolute local error tolerance for solver (optional).
     mxstep: int, maximum number of steps to take for each timepoint (optional).
+    hmax: float, maximum step size allowed (optional).
 
   Returns:
     Values of the solution `y` (i.e. integrated system values) at each time
@@ -171,17 +172,17 @@ def odeint(func, y0, t, *args, rtol=1.4e-8, atol=1.4e-8, mxstep=jnp.inf):
     raise TypeError(f"t must be an array of floats, but got {t}.")
 
   converted, consts = custom_derivatives.closure_convert(func, y0, t[0], *args)
-  return _odeint_wrapper(converted, rtol, atol, mxstep, y0, t, *args, *consts)
+  return _odeint_wrapper(converted, rtol, atol, mxstep, hmax, y0, t, *args, *consts)
 
-@partial(jax.jit, static_argnums=(0, 1, 2, 3))
-def _odeint_wrapper(func, rtol, atol, mxstep, y0, ts, *args):
+@partial(jax.jit, static_argnums=(0, 1, 2, 3, 4))
+def _odeint_wrapper(func, rtol, atol, mxstep, hmax, y0, ts, *args):
   y0, unravel = ravel_pytree(y0)
   func = ravel_first_arg(func, unravel)
-  out = _odeint(func, rtol, atol, mxstep, y0, ts, *args)
+  out = _odeint(func, rtol, atol, mxstep, hmax, y0, ts, *args)
   return jax.vmap(unravel)(out)
 
-@partial(jax.custom_vjp, nondiff_argnums=(0, 1, 2, 3))
-def _odeint(func, rtol, atol, mxstep, y0, ts, *args):
+@partial(jax.custom_vjp, nondiff_argnums=(0, 1, 2, 3, 4))
+def _odeint(func, rtol, atol, mxstep, hmax, y0, ts, *args):
   func_ = lambda y, t: func(y, t, *args)
 
   def scan_fun(carry, target_t):
@@ -196,7 +197,7 @@ def body_fun(state):
       next_t = t + dt
       error_ratio = mean_error_ratio(next_y_error, rtol, atol, y, next_y)
       new_interp_coeff = interp_fit_dopri(y, next_y, k, dt)
-      dt = optimal_step_size(dt, error_ratio)
+      dt = jnp.clip(optimal_step_size(dt, error_ratio), a_min=0., a_max=hmax)
 
       new = [i + 1, next_y, next_f, next_t, dt,      t, new_interp_coeff]
       old = [i + 1,      y,      f,      t, dt, last_t,     interp_coeff]
@@ -209,17 +210,17 @@ def body_fun(state):
     return carry, y_target
 
   f0 = func_(y0, ts[0])
-  dt = initial_step_size(func_, ts[0], y0, 4, rtol, atol, f0)
+  dt = jnp.clip(initial_step_size(func_, ts[0], y0, 4, rtol, atol, f0), a_min=0., a_max=hmax)
   interp_coeff = jnp.array([y0] * 5)
   init_carry = [y0, f0, ts[0], dt, ts[0], interp_coeff]
   _, ys = lax.scan(scan_fun, init_carry, ts[1:])
   return jnp.concatenate((y0[None], ys))
 
-def _odeint_fwd(func, rtol, atol, mxstep, y0, ts, *args):
-  ys = _odeint(func, rtol, atol, mxstep, y0, ts, *args)
+def _odeint_fwd(func, rtol, atol, mxstep, hmax, y0, ts, *args):
+  ys = _odeint(func, rtol, atol, mxstep, hmax, y0, ts, *args)
   return ys, (ys, ts, args)
 
-def _odeint_rev(func, rtol, atol, mxstep, res, g):
+def _odeint_rev(func, rtol, atol, mxstep, hmax, res, g):
   ys, ts, args = res
 
   def aug_dynamics(augmented_state, t, *args):
@@ -244,7 +245,7 @@ def scan_fun(carry, i):
     _, y_bar, t0_bar, args_bar = odeint(
         aug_dynamics, (ys[i], y_bar, t0_bar, args_bar),
         jnp.array([-ts[i], -ts[i - 1]]),
-        *args, rtol=rtol, atol=atol, mxstep=mxstep)
+        *args, rtol=rtol, atol=atol, mxstep=mxstep, hmax=hmax)
     y_bar, t0_bar, args_bar = tree_map(op.itemgetter(1), (y_bar, t0_bar, args_bar))
     # Add gradient from current output
     y_bar = y_bar + g[i - 1]


Test patch:
diff --git a/tests/ode_test.py b/tests/ode_test.py
--- a/tests/ode_test.py
+++ b/tests/ode_test.py
@@ -250,6 +250,21 @@ def f(y0, ts, alpha):
 
     jtu.check_grads(f, (y0, ts, alpha), modes=["rev"], order=2, atol=tol, rtol=tol)
 
+  @jtu.skip_on_devices("tpu", "gpu")
+  def test_hmax(self):
+    """Test max step size control."""
+
+    def rhs(y, t):
+      return jnp.piecewise(
+        t,
+        [t <= 2., (t >= 5.) & (t <= 7.)],
+        [lambda s: jnp.array(1.), lambda s: jnp.array(-1.), lambda s: jnp.array(0.)]
+      )
+    ys = odeint(func=rhs, y0=jnp.array(0.), t=jnp.array([0., 5., 10.]), hmax=1.)
+
+    self.assertTrue(jnp.abs(ys[1] - 2.) < 1e-4)
+    self.assertTrue(jnp.abs(ys[2]) < 1e-4)
+
 
 if __name__ == '__main__':
   absltest.main(testLoader=jtu.JaxTestLoader())