Add Radau method for stiff problems.

torchdiffeq/_impl/odeint.py

@@ -2,6 +2,7 @@
 from torch.autograd.functional import vjp
 from .dopri5 import Dopri5Solver
 from .bosh3 import Bosh3Solver
+from .radau import RadauSolver
 from .adaptive_heun import AdaptiveHeunSolver
 from .fehlberg2 import Fehlberg2
 from .fixed_grid import Euler, Midpoint, Heun2, Heun3, RK4
@@ -14,6 +15,7 @@
 SOLVERS = {
     'dopri8': Dopri8Solver,
     'dopri5': Dopri5Solver,
+    'radau': RadauSolver,
     'bosh3': Bosh3Solver,
     'fehlberg2': Fehlberg2,
     'adaptive_heun': AdaptiveHeunSolver,

torchdiffeq/_impl/radau.py

@@ -0,0 +1,50 @@
+import torch
+from .rk_common import _ButcherTableau, RKAdaptiveStepsizeODESolver
+
+# Radau IIA coefficients (order 5, 3 stages)
+# Reference: E. Hairer, G. Wanner, "Solving Ordinary Differential Equations II: Stiff and Differential-Algebraic Problems"
+_RADAU_IIA_TABLEAU = _ButcherTableau(
+    alpha=torch.tensor([
+        (4 - torch.sqrt(torch.tensor(6.))) / 10,
+        (4 + torch.sqrt(torch.tensor(6.))) / 10,
+        1.
+    ], dtype=torch.float64),
+    beta=[
+        torch.tensor([(4 - torch.sqrt(torch.tensor(6.))) / 10], dtype=torch.float64),
+        torch.tensor([
+            (88 - 7 * torch.sqrt(torch.tensor(6.))) / 360,
+            (296 + 169 * torch.sqrt(torch.tensor(6.))) / 1800
+        ], dtype=torch.float64),
+        torch.tensor([
+            (296 - 169 * torch.sqrt(torch.tensor(6.))) / 1800,
+            (88 + 7 * torch.sqrt(torch.tensor(6.))) / 360,
+            (16 - torch.sqrt(torch.tensor(6.))) / 36
+        ], dtype=torch.float64),
+    ],
+    c_sol=torch.tensor([
+        (16 - torch.sqrt(torch.tensor(6.))) / 36,
+        (16 + torch.sqrt(torch.tensor(6.))) / 36,
+        1/9,
+        0.
+    ], dtype=torch.float64),
+    c_error=torch.tensor([
+        (16 - torch.sqrt(torch.tensor(6.))) / 36 - (1/9),
+        (16 + torch.sqrt(torch.tensor(6.))) / 36 - (1/9),
+        0.,
+        0.
+    ], dtype=torch.float64),
+)
+
+# Interpolation coefficients for dense output
+RADAU_C_MID = torch.tensor([
+    0.5 * ((16 - torch.sqrt(torch.tensor(6.))) / 36),
+    0.5 * ((16 + torch.sqrt(torch.tensor(6.))) / 36),
+    0.5 * (1/9),
+    0.
+], dtype=torch.float64)
+
+
+class RadauSolver(RKAdaptiveStepsizeODESolver):
+    order = 5
+    tableau = _RADAU_IIA_TABLEAU
+    mid = RADAU_C_MID

torchdiffeq/_impl/rk_common.py

@@ -281,6 +281,8 @@ def _adaptive_step(self, rk_state):
         ########################################################
         #                      Assertions                      #
         ########################################################
+        if (dt <= 0).any():
+            print('dt < 0')
         assert t0 + dt > t0, 'underflow in dt {}'.format(dt.item())
         assert torch.isfinite(y0).all(), 'non-finite values in state `y`: {}'.format(y0)