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 then this determines the
+        number of fixed steps via `ceil((t1 - t0) / dt0)`, and rescales the remaining
+        fixed steps to land exactly on `t1`. 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,11 @@
 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 finite intervals this chooses `ceil((t1 - t0) / dt0)` steps and rescales the
+    later fixed steps to hit `t1` exactly.
     """
 
     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 this uses `dt0` to determine how many steps to take, then spaces those steps so the solve lands exactly on `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

@@ -247,6 +247,39 @@ 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():
+    term = diffrax.ODETerm(lambda t, y, args: 0.0)
+    sol = diffrax.diffeqsolve(
+        term,
+        solver=diffrax.Euler(),
+        t0=0.0,
+        t1=1.05,
+        dt0=0.1,
+        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)]
+    expected_ts = jnp.concatenate(
+        [jnp.array([0.0, 0.1]), jnp.linspace(2 * 1.05 / 11, 1.05, 10)]
+    )
+    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 "hit `t1` exactly" in constant_doc
+    assert "rescales the remaining fixed steps to land exactly on `t1`" 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