Clarify ConstantStepSize finite-interval step semantics

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.

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):

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=...)`.

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