PID controller tolerances sharp bits

[ParticularlyPythonicBS]

Hi,
I recently noticed that the default tolerances suggested for the PID controller can give erroneous results even for very simple systems.
Here is a minimum working example:
Lets start with a simple pendulum with a small angle kick:

class SimplePendulum(eqx.Module):
    length: float
    def __call__(self, t, state, args=None):
        theta, omega = state
        dtheta = omega
        domega = - self.length * jnp.sin(theta)
        return jnp.array([dtheta, domega])
dynamics = SimplePendulum(length=1)
init_angle = 0.1
init_omega = 0
state0 = jnp.array([init_angle, init_omega])

Then we can specify an integration function:

def integrator(dynamics, state0, stepsize_controller):
    term = diffrax.ODETerm(dynamics)
    solver = diffrax.Dopri8()
    t0 = 0.0
    dt0 = 0.1
    t1 = 1e4 * dt0
    ts = jnp.arange(t0, t1, dt0)
    stepsize_controller = stepsize_controller
    solution = diffrax.diffeqsolve(
        term,
        solver,
        t0,
        t1,
        dt0,
        state0,
        saveat=diffrax.SaveAt(ts=ts),
        stepsize_controller=stepsize_controller,
        max_steps=2**20,
    )
    return solution

Using this function we can compare a constant step size controller to PID with different tolerances:

step_size_controller = diffrax.PIDController(rtol=1e-3, atol=1e-6)
sol_incorrect_pid = integrator(dynamics, state0, step_size_controller)
step_size_controller = diffrax.PIDController(rtol=1e-7, atol=1e-9)
sol_correct_pid = integrator(dynamics, state0, step_size_controller)
step_size_controller = diffrax.ConstantStepSize()
sol_const = integrator(dynamics, state0, step_size_controller)

And then we can plot the phase portraits by:

fig, axs = plt.subplots(1, 3, figsize=(15, 5), sharey=True)
axs[0].plot(np.sin(sol_const.ys[:,0]), sol_const.ys[:,1])
axs[0].set_title('Constant step size')
axs[1].plot(np.sin(sol_incorrect_pid.ys[:,0]), sol_incorrect_pid.ys[:,1])
axs[1].set_title('PID (incorrect) rtol=1e-3, atol=1e-6')
axs[2].plot(np.sin(sol_correct_pid.ys[:,0]), sol_correct_pid.ys[:,1])
axs[2].set_title('PID (correct) rtol=1e-7, atol=1e-9')

And the results make the variation very clear:

This may be obvious to someone with more experience with adaptive stepsize controllers but I failed to notice it when I was looking at less trivial time series and it may be helpful for someone else.

If there is something incorrect in my implementation please let me know.
If not I think having an example like this in the documentation could be useful.

Thanks for the great package!

[patrick-kidger]

Right! This is actually the reason that rtol and atol have no default value, to force an end user to make a particular choice about how accurate the solution should be.

I'd be happy to take a PR adding to the documentation here, to help make clear the role of rtol and atol / what happens if they're too loose. :)

[ParticularlyPythonicBS]

Sounds good!
Should I add it as a separate example or extend the PID controller choosing tolerances section?

[patrick-kidger]

Either I think! (Assuming that "separate example" means an !!! Example admonition.) It's already quite a long docstring so I think the only important bit is that it be concise. :)

[ParticularlyPythonicBS]

I am not sure what !!! Example means, I meant adding a notebook to the https://docs.kidger.site/diffrax/examples/ section.

[patrick-kidger]

Ah, right! Yeah I definitely wouldn't add an example there -- we have a lot of those already so easy to get lost amongst them I think.

As for the !!!: our docstrings use an extended version of markdown. In particular it is possible to write

!!! Example

    Lorem ipsum...

for am admonition, or for

??? Example

    Lorem ipsum...

for an expandable admonition.

Some examples of the latter are in the existing docstring for PIDController.

(Here is the documentation on admonitions.)

[ParticularlyPythonicBS]

Thanks! I have submitted a PR. I added it under the choosing tolerances tip. Hopefully it's not too obtrusive

[ParticularlyPythonicBS]

Solution PR merged