patrick-kidger / patrick-kidger/diffrax
sol.ts contains wrong values in some cases
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 2.1k
- Forks
- 189
- Avg merge
- 3d 18h
- Merged PRs (30d)
- 1
Description
When using saveat= and jump_ts=, the return values for sol.ts can be wrong/infinite. Here is a reproducible example:
import jax.numpy as jnp
import jax
import diffrax as dfx
Q = jnp.array([[-1., 1.], [0., 0.]])
y0 = jnp.array([0.5, 0.5])
def A(t, y, _):
return Q @ y
solver = dfx.Kvaerno3()
def f(t1):
saveat = dfx.SaveAt(t1=True, ts=[t1])
ssc = dfx.PIDController(atol=1e-6, rtol=1e-6, jump_ts=jnp.array([0., t1]))
res = dfx.diffeqsolve(
dfx.ODETerm(A),
solver=solver,
y0=y0,
t0=0.,
t1=t1,
dt0=0.01,
stepsize_controller=ssc,
saveat=saveat
)
return res.ts
for t1 in [1e0, 1e3]:
print(t1, f(t1))
# 1.0 [1. 1.]
# 1000.0 [1000. inf]
The expected output is [t1, t1] regardless of what t1 is. However, for certain values, the saved timepoints are erroneously [t1, inf]. The behavior seems to depend on the magnitude of t1, so I suspect maybe a jnp.nextafter-ish type of bug.
This occurs on HEAD as well as the latest release.
Might be related to #607.
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
Run the provided reproduction with diffeqsolve, SaveAt(t1=True, ts=[t1]), and PIDController(jump_ts=...) to confirm when sol.ts becomes inf. Trace the handling of saved timepoints and jump_ts, review the context of #607, and add a regression test showing that both returned values equal t1 for large and small endpoints.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- backend
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 42/100