patrick-kidger / patrick-kidger/diffrax
How to use reversible integrators?
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 2.1k
- Forks
- 189
- Avg merge
- 3d 18h
- Merged PRs (30d)
- 1
Description
I wanted a time-reversible integrator, and according to the documentation SemiImplicitEuler() should work.
However, I ran the following:
def f(t, p, args):
return -p
def g(t, q, args):
return -q
term = (diffrax.ODETerm(f), diffrax.ODETerm(g))
solver = diffrax.SemiImplicitEuler()
q0 = jnp.array([12., 3.])
p0 = jnp.array([1., 2.])
# Forward pass
y0 = (q0, p0)
solution = diffrax.diffeqsolve(term, solver, t0=0, t1=1, dt0=0.01, y0=y0,
saveat=diffrax.SaveAt(t1=True, solver_state=True))
y1 = jax.tree_map(lambda arr: arr[0], solution.ys)
# Backward pass
solution = diffrax.diffeqsolve(term, solver, t0=1, t1=0, dt0=-0.01, y0=y1,
solver_state=solution.solver_state)
y0_computed = jax.tree_map(lambda arr: arr[0], solution.ys)
I see that y0_computed is:
(Array([11.797033 , 2.9734273]),
Array([0.85248 , 1.9949753])
which is close to, but different from y0. Am I using the wrong arguments?
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 by running the Python reproduction in the issue with SemiImplicitEuler(), SaveAt(t1=True, solver_state=True), and the forward and backward diffeqsolve calls. Compare the returned y0_computed with y0 and determine whether the arguments are incorrect or the reversible-integrator behavior is inconsistent; the issue is done when the expected behavior is established and documented or corrected.
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
- 30/100