From 99885422785590c16c057934cc20f9ccd3a599b6 Mon Sep 17 00:00:00 2001 From: whning0513 <3267069960@qq.com> Date: Wed, 1 Jul 2026 17:27:23 +0800 Subject: [PATCH 1/2] Clarify ConstantStepSize finite-interval docs --- diffrax/_integrate.py | 9 ++++--- diffrax/_step_size_controller/constant.py | 5 +++- docs/usage/getting-started.md | 2 +- test/test_saveat_solution.py | 33 +++++++++++++++++++++++ 4 files changed, 43 insertions(+), 6 deletions(-) diff --git a/diffrax/_integrate.py b/diffrax/_integrate.py index 5441e0a9..4d7b456f 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 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. diff --git a/diffrax/_step_size_controller/constant.py b/diffrax/_step_size_controller/constant.py index 02252869..c2b0484b 100644 --- a/diffrax/_step_size_controller/constant.py +++ b/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): diff --git a/docs/usage/getting-started.md b/docs/usage/getting-started.md index 92a4d7eb..f82ad44d 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 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=...)`. diff --git a/test/test_saveat_solution.py b/test/test_saveat_solution.py index d1aa6037..a0519ffe 100644 --- a/test/test_saveat_solution.py +++ b/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 From aad4faac3639d369a0c8da0d34373ede070452d5 Mon Sep 17 00:00:00 2001 From: whning Date: Fri, 4 Sep 2026 19:58:58 +0800 Subject: [PATCH 2/2] Clarify finite constant-step targets --- diffrax/_integrate.py | 8 ++++---- diffrax/_step_size_controller/constant.py | 5 +++-- docs/usage/getting-started.md | 2 +- test/test_saveat_solution.py | 20 ++++++++++++++------ 4 files changed, 22 insertions(+), 13 deletions(-) diff --git a/diffrax/_integrate.py b/diffrax/_integrate.py index 4d7b456f..15d57366 100644 --- a/diffrax/_integrate.py +++ b/diffrax/_integrate.py @@ -930,10 +930,10 @@ def diffeqsolve( - `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 - [`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. + [`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 c2b0484b..75bc44bd 100644 --- a/diffrax/_step_size_controller/constant.py +++ b/diffrax/_step_size_controller/constant.py @@ -23,8 +23,9 @@ class ConstantStepSize( """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. + 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 f82ad44d..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`. - - 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.) + - 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 a0519ffe..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 @@ -248,13 +249,16 @@ def _step_integrate(saveat: diffrax.SaveAt, with_7: bool): 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=0.0, - t1=1.05, - dt0=0.1, + t0=t0, + t1=t1, + dt0=dt0, y0=0.0, saveat=diffrax.SaveAt(t0=True, steps=True), stepsize_controller=diffrax.ConstantStepSize(), @@ -262,8 +266,12 @@ def test_constant_stepsize_rescales_finite_interval_steps(): 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([0.0, 0.1]), jnp.linspace(2 * 1.05 / 11, 1.05, 10)] + [ + jnp.array([t0, t0 + dt0]), + t0 + jnp.arange(2, num_steps + 1) * (t1 - t0) / num_steps, + ] ) assert jnp.allclose(ts, expected_ts) @@ -276,8 +284,8 @@ def test_constant_stepsize_docs_describe_rescaled_finite_steps(): 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 + 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():