diff --git a/diffrax/_integrate.py b/diffrax/_integrate.py index 5441e0a9..15d57366 100644 --- a/diffrax/_integrate.py +++ b/diffrax/_integrate.py @@ -929,10 +929,11 @@ def diffeqsolve( choose a solver](../usage/how-to-choose-a-solver.md). - `t0`: The start of the region of integration. - `t1`: The end of the region of integration. - - `dt0`: The step size to use for the first step. If using fixed step sizes then - this will also be the step size for all other steps. (Except the last one, - which may be slightly smaller and clipped to `t1`.) If set as `None` then the - initial step size will be determined automatically. + - `dt0`: The step size to use for the first step. If using + [`diffrax.ConstantStepSize`][] on a finite interval, let + `n = ceil((t1 - t0) / dt0)`. The first step targets `t0 + dt0`; subsequent + steps target `t0 + k * (t1 - t0) / n` for `k = 2, ..., n`. If set as `None` + then the initial step size will be determined automatically. - `y0`: The initial value. This can be any PyTree of JAX arrays. (Or types that can be coerced to JAX arrays, like Python floats.) - `args`: Any additional arguments to pass to the vector field. diff --git a/diffrax/_step_size_controller/constant.py b/diffrax/_step_size_controller/constant.py index 02252869..75bc44bd 100644 --- a/diffrax/_step_size_controller/constant.py +++ b/diffrax/_step_size_controller/constant.py @@ -20,8 +20,12 @@ class ConstantStepSize( AbstractStepSizeController[_ConstantStepSizeState, RealScalarLike] ): - """Use a constant step size, equal to the `dt0` argument of + """Use a fixed number of constant steps determined by the `dt0` argument of [`diffrax.diffeqsolve`][]. + + On a finite interval, let `n = ceil((t1 - t0) / dt0)`. The first step targets + `t0 + dt0`; subsequent steps target `t0 + k * (t1 - t0) / n` for + `k = 2, ..., n`. """ def wrap(self, direction: IntScalarLike): diff --git a/docs/usage/getting-started.md b/docs/usage/getting-started.md index 92a4d7eb..567ddad8 100644 --- a/docs/usage/getting-started.md +++ b/docs/usage/getting-started.md @@ -53,7 +53,7 @@ print(sol.ys) # DeviceArray([1. , 0.368, 0.135, 0.0498]) - Where to save the result (e.g. to obtain dense output) can be adjusted by changing [`diffrax.SaveAt`][]. - Step sizes and locations can be changed. - The initial step size can be selected adaptively by setting `dt0=None`. - - A constant step size can be used by setting `stepsize_controller = ConstantStepSize()`. (This is also the default choice for `stepsize_controller` if you do not pass one at all.) + - Fixed steps can be used by setting `stepsize_controller = ConstantStepSize()`. On finite intervals the first step ends at `t0 + dt0`; later step targets lie on a uniform grid over `[t0, t1]`, ending exactly at `t1`. (This is also the default choice for `stepsize_controller` if you do not pass one at all.) - Things like solver tolerances, jumps in the vector field, etc. can be passed as arguments to the step size controller. - See the page on [Step size controllers](../api/stepsize_controller.md). - Any static arguments (that do not change during the integration) for the `vector_field` can be passed as `diffeqsolve(..., args=...)`. diff --git a/test/test_saveat_solution.py b/test/test_saveat_solution.py index d1aa6037..1818842c 100644 --- a/test/test_saveat_solution.py +++ b/test/test_saveat_solution.py @@ -4,6 +4,7 @@ import diffrax import equinox as eqx +import equinox.internal as eqxi import jax import jax.numpy as jnp import optimistix as optx @@ -247,6 +248,46 @@ def _step_integrate(saveat: diffrax.SaveAt, with_7: bool): assert jnp.allclose(ts, jnp.array([0.0, 3.0, 6.0])) +def test_constant_stepsize_rescales_finite_interval_steps(): + t0 = 0.0 + t1 = 1.05 + dt0 = 0.1 + term = diffrax.ODETerm(lambda t, y, args: 0.0) + sol = diffrax.diffeqsolve( + term, + solver=diffrax.Euler(), + t0=t0, + t1=t1, + dt0=dt0, + y0=0.0, + saveat=diffrax.SaveAt(t0=True, steps=True), + stepsize_controller=diffrax.ConstantStepSize(), + ) + + assert sol.ts is not None + ts = sol.ts[jnp.isfinite(sol.ts)] + num_steps = int(jnp.ceil((t1 - t0) / eqxi.nextafter(dt0))) + expected_ts = jnp.concatenate( + [ + jnp.array([t0, t0 + dt0]), + t0 + jnp.arange(2, num_steps + 1) * (t1 - t0) / num_steps, + ] + ) + assert jnp.allclose(ts, expected_ts) + + +def test_constant_stepsize_docs_describe_rescaled_finite_steps(): + constant_doc = diffrax.ConstantStepSize.__doc__ + solve_doc = diffrax.diffeqsolve.__doc__ + + assert constant_doc is not None + assert solve_doc is not None + constant_doc = " ".join(constant_doc.split()) + solve_doc = " ".join(solve_doc.split()) + assert "subsequent steps target `t0 + k * (t1 - t0) / n`" in constant_doc + assert "steps target `t0 + k * (t1 - t0) / n`" in solve_doc + + def test_saveat_solution_skip_vs_saveat(): ts = jnp.array([0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0]) n = 2