patrick-kidger / patrick-kidger/diffrax
Documentation and implementation of ConstantStepSize do not quite align
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 2.1k
- Forks
- 189
- Avg merge
- 3d 18h
- Merged PRs (30d)
- 1
Description
The step size in ConstantStepSize is currently computed as: (t1-t0)/(step/n_steps (see https://github.com/patrick-kidger/diffrax/pull/666). I understand the rationale and advantages of this, but IMO, the documentation is unclear on this, and the results can be somewhat unexpected for. Specifically:
- If
t1-t0is not divisible bydt0, the step size will not be equal todt0, unlike what is stated in the docs: use a constant step size, equal to the dt0 argument of diffrax.diffeqsolve - When run using
dfx.SaveAt(t0=True, steps=True), the step fromt0to the first solver step is not necessarily the same as the rest of the steps. This might be unexpected for a class calledConstantStepSize. See example below. - The documentation for the
dt0argument ofdiffeqsolvecurrently states: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. This seems outdated, as with the current implementation ofConstantStepSize, step size is computed such that the last step occurs exactly att1, without clipping.
Consider this simple example:
import diffrax as dfx
import jax.numpy as jnp
import jax.random as jr
t0, t1, dt0 = 0.0, 1.05, 0.1
key = jr.PRNGKey(0)
class TestModel:
@property
def initial(self):
return jnp.array(0.0)
def drift(self, t, x, args):
return jnp.array(0.0)
def diffusion(self, t, x, args):
return jnp.array(0.0)
def terms(self, key):
process_noise = dfx.UnsafeBrownianPath(
shape=self.initial.shape, key=key, levy_area=dfx.SpaceTimeLevyArea
)
return dfx.MultiTerm(
dfx.ODETerm(self.drift), dfx.ControlTerm(self.diffusion, process_noise)
)
model = TestModel()
terms = model.terms(jr.PRNGKey(0))
sol = dfx.diffeqsolve(
terms,
dfx.Euler(),
t0=t0,
t1=t1,
dt0=dt0,
y0=model.initial,
args={},
saveat= dfx.SaveAt(t0=True, steps=True),
adjoint=dfx.ForwardMode(),
stepsize_controller=dfx.ConstantStepSize(),
)
print("Timesteps from diffrax solution:", sol.ts[jnp.isfinite(sol.ys)])
# prints: Timesteps from diffrax solution: [0. 0.1 0.19090909 0.28636363 0.38181818 0.47727272 0.57272726 0.6681818 0.76363635 0.8590909 0.95454544 1.05 ]
As you can see, the steps are neither constant (if you include the t0) nor are they equal to dt0. Of course, for simulations with many steps, this differences will be minimal, so it is likely not a significant problem in practice, but it might be confusing to some Diffrax beginners such as myself. Curious to hear your thoughts on this!
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with the ConstantStepSize implementation and the stepsize-controller documentation, then compare them with the dt0 description in diffeqsolve. Reproduce the provided t0=0.0, t1=1.05, dt0=0.1 example and decide whether the behavior or wording should change. Done means the documented constant-step semantics and implementation agree, including the first and final steps.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- backend, documentation
- Issue type
- Documentation
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 52/100