Skip to content

Commit

Permalink
Example for adaptive step size choice (#372)
Browse files Browse the repository at this point in the history
* Example for adaptive step size choice

* Implemented PR feedback
  • Loading branch information
ParticularlyPythonicBS committed Feb 20, 2024
1 parent 553ecea commit d97ba20
Show file tree
Hide file tree
Showing 2 changed files with 34 additions and 0 deletions.
34 changes: 34 additions & 0 deletions diffrax/_step_size_controller/adaptive.py
Original file line number Diff line number Diff line change
Expand Up @@ -179,6 +179,39 @@ class PIDController(
common to refer to solving an equation to specific tolerances, without
necessarily stating which solver was used.)
??? Example
The choice of `rtol` and `atol` can have a significant impact on the
accuracy of even simple systems.
Consider a simple pendulum with a small angle kick:
```python
import diffrax as dfx
def dynamics(t, y, args):
dtheta = y["omega"]
domega = - jnp.sin(y["theta"])
return dict(theta=dtheta, omega=domega)
y0 = dict(theta=0.1, omega=0)
term = dfx.ODETerm(dynamics)
sol = dfx.diffeqsolve(
term, solver, t0=0, t1=1000, dt0=0.1, y0,
saveat=dfx.SaveAts(ts=jnp.linspace(0, 1000, 10000),
max_steps=2**20,
stepsize_controller=...
)
```
to compare the effect of different tolerances:
```python
PID_controller_incorrect = diffrax.PIDController(rtol=1e-3, atol=1e-6)
PID_controller_correct = diffrax.PIDController(rtol=1e-7, atol=1e-9)
Constant_controller = diffrax.ConstantStepSize()
```
The phase portraits of the pendulum from the different tolerances clearly
illustrate the impact of the choice of `rtol` and `atol` on the accuracy of
the solution.
![Phase portrait of pendulum](../imgs/pendulum_adaptive_steps.png)
??? tip "Choosing PID coefficients"
This controller can be reduced to any special case (e.g. just a PI controller,
Expand Down Expand Up @@ -217,6 +250,7 @@ class PIDController(
print(sol.stats["num_steps"])
```
??? cite "References"
Both the initial step size selection algorithm for ODEs, and the use of
Expand Down
Binary file added imgs/pendulum_adaptive_steps.png
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.

0 comments on commit d97ba20

Please sign in to comment.