Repository: rtqichen/torchdiffeq

Issue:
Solvers for stiff problem
Hi,
First, thank you for your work to create a great package. I think it is your side project to contribute to a community but you made it well.
Although this package is good, I think it still lacks a solver for stiff problems, which is compatible to do backward when optimizing parameters in pytorch. So far until now, I see that Scipy solvers in your package only supports for forward process. If you can make a solver for stiff problems, or make Scipy solvers (BDF or LSODA) compatible with backward in pytorch, I believe that is a great improvement.
I see that your activity in this repo is less often, I hope that you have time to implement this new feature.
Thank you very much

Discussion:
@WUCHIL: I'm also having the same issue@timnotavailable: Me too~

Solution patch:
diff --git a/torchdiffeq/_impl/fixed_grid_implicit.py b/torchdiffeq/_impl/fixed_grid_implicit.py
new file mode 100644
index 0000000..7519efc
--- /dev/null
+++ b/torchdiffeq/_impl/fixed_grid_implicit.py
@@ -0,0 +1,140 @@
+import torch
+from .rk_common import FixedGridFIRKODESolver, FixedGridDIRKODESolver
+from .rk_common import _ButcherTableau
+
+_sqrt_2 = torch.sqrt(torch.tensor(2, dtype=torch.float64)).item()
+_sqrt_3 = torch.sqrt(torch.tensor(3, dtype=torch.float64)).item()
+_sqrt_6 = torch.sqrt(torch.tensor(6, dtype=torch.float64)).item()
+_sqrt_15 = torch.sqrt(torch.tensor(15, dtype=torch.float64)).item()
+
+_IMPLICIT_EULER_TABLEAU = _ButcherTableau(
+    alpha=torch.tensor([1], dtype=torch.float64),
+    beta=[
+        torch.tensor([1], dtype=torch.float64),
+    ],
+    c_sol=torch.tensor([1], dtype=torch.float64),
+    c_error=torch.tensor([], dtype=torch.float64),
+)
+
+class ImplicitEuler(FixedGridFIRKODESolver):
+    order = 1
+    tableau = _IMPLICIT_EULER_TABLEAU
+
+_IMPLICIT_MIDPOINT_TABLEAU = _ButcherTableau(
+    alpha=torch.tensor([1 / 2], dtype=torch.float64),
+    beta=[
+        torch.tensor([1 / 2], dtype=torch.float64),
+
+    ],
+    c_sol=torch.tensor([1], dtype=torch.float64),
+    c_error=torch.tensor([], dtype=torch.float64),
+)
+
+class ImplicitMidpoint(FixedGridFIRKODESolver):
+    order = 2
+    tableau = _IMPLICIT_MIDPOINT_TABLEAU
+
+_GAUSS_LEGENDRE_4_TABLEAU = _ButcherTableau(
+    alpha=torch.tensor([1 / 2 - _sqrt_3 / 6, 1 / 2 - _sqrt_3 / 6], dtype=torch.float64),
+    beta=[
+        torch.tensor([1 / 4, 1 / 4 - _sqrt_3 / 6], dtype=torch.float64),
+        torch.tensor([1 / 4 + _sqrt_3 / 6, 1 / 4], dtype=torch.float64),
+    ],
+    c_sol=torch.tensor([1 / 2, 1 / 2], dtype=torch.float64),
+    c_error=torch.tensor([], dtype=torch.float64),
+)
+
+_TRAPEZOID_TABLEAU = _ButcherTableau(
+    alpha=torch.tensor([0, 1], dtype=torch.float64),
+    beta=[
+        torch.tensor([0, 0], dtype=torch.float64),
+        torch.tensor([1 /2, 1 / 2], dtype=torch.float64),
+    ],
+    c_sol=torch.tensor([1 / 2, 1 / 2], dtype=torch.float64),
+    c_error=torch.tensor([], dtype=torch.float64),
+)
+
+class Trapezoid(FixedGridFIRKODESolver):
+    order = 2
+    tableau = _TRAPEZOID_TABLEAU
+
+
+class GaussLegendre4(FixedGridFIRKODESolver):
+    order = 4
+    tableau = _GAUSS_LEGENDRE_4_TABLEAU
+
+_GAUSS_LEGENDRE_6_TABLEAU = _ButcherTableau(
+    alpha=torch.tensor([1 / 2 - _sqrt_15 / 10, 1 / 2, 1 / 2 + _sqrt_15 / 10], dtype=torch.float64),
+    beta=[
+        torch.tensor([5 / 36                , 2 / 9 - _sqrt_15 / 15, 5 / 36 - _sqrt_15 / 30], dtype=torch.float64),
+        torch.tensor([5 / 36 + _sqrt_15 / 24, 2 / 9                , 5 / 36 - _sqrt_15 / 24], dtype=torch.float64),
+        torch.tensor([5 / 36 + _sqrt_15 / 30, 2 / 9 + _sqrt_15 / 15, 5 / 36                ], dtype=torch.float64),
+    ],
+    c_sol=torch.tensor([5 / 18, 4 / 9, 5 / 18], dtype=torch.float64),
+    c_error=torch.tensor([], dtype=torch.float64),
+)
+
+class GaussLegendre6(FixedGridFIRKODESolver):
+    order = 6
+    tableau = _GAUSS_LEGENDRE_6_TABLEAU
+
+_RADAU_IIA_3_TABLEAU = _ButcherTableau(
+    alpha=torch.tensor([1 / 3, 1], dtype=torch.float64),
+    beta=[
+        torch.tensor([5 / 12, -1 / 12], dtype=torch.float64),
+        torch.tensor([3 / 4, 1 / 4], dtype=torch.float64)
+    ],
+    c_sol=torch.tensor([3 / 4, 1 / 4], dtype=torch.float64),
+    c_error=torch.tensor([], dtype=torch.float64)
+)
+
+class RadauIIA3(FixedGridFIRKODESolver):
+    order = 3
+    tableau = _RADAU_IIA_3_TABLEAU
+
+_RADAU_IIA_5_TABLEAU = _ButcherTableau(
+    alpha=torch.tensor([2 / 5 - _sqrt_6 / 10, 2 / 5 + _sqrt_6 / 10, 1], dtype=torch.float64),
+    beta=[
+        torch.tensor([11 / 45 - 7 * _sqrt_6 / 360    , 37 / 225 - 169 * _sqrt_6 / 1800, -2 / 225 + _sqrt_6 / 75], dtype=torch.float64),
+        torch.tensor([37 / 225 + 169 * _sqrt_6 / 1800, 11 / 45 + 7 * _sqrt_6 / 360    , -2 / 225 - _sqrt_6 / 75], dtype=torch.float64),
+        torch.tensor([4 / 9 - _sqrt_6 / 36           , 4 / 9 + _sqrt_6 / 36           , 1 / 9], dtype=torch.float64)
+    ],
+    c_sol=torch.tensor([4 / 9 - _sqrt_6 / 36, 4 / 9 + _sqrt_6 / 36, 1 / 9], dtype=torch.float64),
+    c_error=torch.tensor([], dtype=torch.float64)
+)
+
+class RadauIIA5(FixedGridFIRKODESolver):
+    order = 5
+    tableau = _RADAU_IIA_5_TABLEAU
+
+gamma = (2. - _sqrt_2) / 2.
+_SDIRK_2_TABLEAU = _ButcherTableau(
+    alpha = torch.tensor([gamma, 1], dtype=torch.float64),
+    beta=[
+        torch.tensor([gamma], dtype=torch.float64),
+        torch.tensor([1 - gamma, gamma], dtype=torch.float64),
+    ],
+    c_sol=torch.tensor([1 - gamma, gamma], dtype=torch.float64),
+    c_error=torch.tensor([], dtype=torch.float64)
+)
+
+class SDIRK2(FixedGridDIRKODESolver):
+    order = 2
+    tableau = _SDIRK_2_TABLEAU
+
+gamma = 1. - _sqrt_2 / 2.
+beta = _sqrt_2 / 4.
+_TRBDF_2_TABLEAU = _ButcherTableau(
+    alpha = torch.tensor([0, 2 * gamma, 1], dtype=torch.float64),
+    beta=[
+        torch.tensor([0], dtype=torch.float64),
+        torch.tensor([gamma, gamma], dtype=torch.float64),
+        torch.tensor([beta, beta, gamma], dtype=torch.float64),
+    ],
+    c_sol=torch.tensor([beta, beta, gamma], dtype=torch.float64),
+    c_error=torch.tensor([], dtype=torch.float64)
+)
+
+class TRBDF2(FixedGridDIRKODESolver):
+    order = 2
+    tableau = _TRBDF_2_TABLEAU
diff --git a/torchdiffeq/_impl/odeint.py b/torchdiffeq/_impl/odeint.py
index 1514650..14a01ef 100644
--- a/torchdiffeq/_impl/odeint.py
+++ b/torchdiffeq/_impl/odeint.py
@@ -5,6 +5,10 @@ from .bosh3 import Bosh3Solver
 from .adaptive_heun import AdaptiveHeunSolver
 from .fehlberg2 import Fehlberg2
 from .fixed_grid import Euler, Midpoint, Heun2, Heun3, RK4
+from .fixed_grid_implicit import ImplicitEuler, ImplicitMidpoint, Trapezoid
+from .fixed_grid_implicit import GaussLegendre4, GaussLegendre6
+from .fixed_grid_implicit import RadauIIA3, RadauIIA5
+from .fixed_grid_implicit import SDIRK2, TRBDF2
 from .fixed_adams import AdamsBashforth, AdamsBashforthMoulton
 from .dopri8 import Dopri8Solver
 from .tsit5 import Tsit5Solver
@@ -26,6 +30,15 @@ SOLVERS = {
     'rk4': RK4,
     'explicit_adams': AdamsBashforth,
     'implicit_adams': AdamsBashforthMoulton,
+    'implicit_euler': ImplicitEuler,
+    'implicit_midpoint': ImplicitMidpoint,
+    'trapezoid': Trapezoid,
+    'radauIIA3': RadauIIA3,
+    'gl4': GaussLegendre4,
+    'radauIIA5': RadauIIA5,
+    'gl6': GaussLegendre6,
+    'sdirk2': SDIRK2,
+    'trbdf2': TRBDF2,
     # Backward compatibility: use the same name as before
     'fixed_adams': AdamsBashforthMoulton,
     # ~Backwards compatibility
diff --git a/torchdiffeq/_impl/rk_common.py b/torchdiffeq/_impl/rk_common.py
index 3b07877..f0050db 100644
--- a/torchdiffeq/_impl/rk_common.py
+++ b/torchdiffeq/_impl/rk_common.py
@@ -5,9 +5,11 @@ from .event_handling import find_event
 from .interp import _interp_evaluate, _interp_fit
 from .misc import (_compute_error_ratio,
                    _select_initial_step,
-                   _optimal_step_size)
+                   _optimal_step_size,
+                   _handle_unused_kwargs)
 from .misc import Perturb
-from .solvers import AdaptiveStepsizeEventODESolver
+from .solvers import AdaptiveStepsizeEventODESolver, FixedGridODESolver
+import warnings
 
 
 _ButcherTableau = collections.namedtuple('_ButcherTableau', 'alpha, beta, c_sol, c_error')
@@ -371,3 +373,186 @@ def _sort_tvals(tvals, t0):
     # TODO: add warning if tvals come before t0?
     tvals = tvals[tvals >= t0]
     return torch.sort(tvals).values
+
+
+class FixedGridFIRKODESolver(FixedGridODESolver):
+    order: int
+    tableau: _ButcherTableau
+
+    def __init__(self, func, y0, step_size=None, grid_constructor=None, interp='linear', perturb=False, max_iters=100, **unused_kwargs):
+
+        self.max_iters = max_iters
+        self.atol = unused_kwargs.pop('atol')
+        unused_kwargs.pop('rtol', None)
+        unused_kwargs.pop('norm', None)
+        _handle_unused_kwargs(self, unused_kwargs)
+        del unused_kwargs
+
+        self.func = func
+        self.y0 = y0
+        self.dtype = y0.dtype
+        self.device = y0.device
+        self.step_size = step_size
+        self.interp = interp
+        self.perturb = perturb
+
+        if step_size is None:
+            if grid_constructor is None:
+                self.grid_constructor = lambda f, y0, t: t
+            else:
+                self.grid_constructor = grid_constructor
+        else:
+            if grid_constructor is None:
+                self.grid_constructor = self._grid_constructor_from_step_size(step_size)
+            else:
+                raise ValueError("step_size and grid_constructor are mutually exclusive arguments.")
+            
+        self.tableau = _ButcherTableau(alpha=self.tableau.alpha.to(device=self.device, dtype=y0.dtype),
+                                       beta=[b.to(device=self.device, dtype=y0.dtype) for b in self.tableau.beta],
+                                       c_sol=self.tableau.c_sol.to(device=self.device, dtype=y0.dtype),
+                                       c_error=self.tableau.c_error.to(device=self.device, dtype=y0.dtype))
+        
+    def _step_func(self, func, t0, dt, t1, y0):
+        if not isinstance(t0, torch.Tensor):
+            t0 = torch.tensor(t0)
+        if not isinstance(dt, torch.Tensor):
+            dt = torch.tensor(dt)
+        if not isinstance(t1, torch.Tensor):
+            t1 = torch.tensor(t1)
+        f0 = func(t0, y0, perturb=Perturb.NEXT if self.perturb else Perturb.NONE)
+        
+        t_dtype = y0.abs().dtype
+        tol = 1e-8
+        if t_dtype == torch.float64:
+            tol = 1e-8
+        if t_dtype == torch.float32:
+            tol = 1e-6
+
+        t0 = t0.to(t_dtype)
+        dt = dt.to(t_dtype)
+        t1 = t1.to(t_dtype)
+
+        k = f0.clone().unsqueeze(-1).tile(len(self.tableau.alpha))
+        beta = torch.stack(self.tableau.beta, -1)
+
+        # Broyden's Method to solve the system of nonlinear equations
+        y = torch.matmul(k, beta * dt).add(y0.unsqueeze(-1)).movedim(-1, 0)
+        f = self._residual(func, k, y, t0, dt, t1)
+        J = torch.ones_like(f).diag()
+        converged = False
+        for _ in range(self.max_iters):
+            if torch.linalg.norm(f, 2) < tol:
+                converged = True
+                break
+
+            # If the matrix becomes singular, just stop and return the last value
+            try:
+                s = -torch.linalg.solve(J, f)
+            except torch._C._LinAlgError:
+                break
+
+            k = k + s.reshape_as(k)
+            y = torch.matmul(k, beta * dt).add(y0.unsqueeze(-1)).movedim(-1, 0)
+            newf = self._residual(func, k, y, t0, dt, t1)
+            z = newf - f
+            f = newf
+            J = J + (torch.outer ((z - torch.linalg.vecdot(J,s)),s)) / (torch.dot(s,s))
+
+        if not converged:
+            warnings.warn('Functional iteration did not converge. Solution may be incorrect.')
+
+        dy = torch.matmul(k, dt * self.tableau.c_sol)
+
+        return dy, f0
+    
+    def _residual(self, func, K, y, t0, dt, t1):
+        res = torch.zeros_like(K)
+        for i, (y_i, alpha_i) in enumerate(zip(y, self.tableau.alpha)):
+            perturb = Perturb.NONE
+            if alpha_i == 1.:
+                ti = t1
+                perturb = Perturb.PREV
+            elif alpha_i == 0.:
+                if not torch.all(self.tableau.beta[i]):
+                    # Same slope as stored so skip
+                    continue
+                ti = t0
+            else:
+                ti = t0 + alpha_i * dt
+            res[...,i] = K[...,i] - func(ti, y_i, perturb=perturb)
+        return res.flatten()
+    
+
+class FixedGridDIRKODESolver(FixedGridFIRKODESolver):
+
+    def _step_func(self, func, t0, dt, t1, y0):
+        if not isinstance(t0, torch.Tensor):
+            t0 = torch.tensor(t0)
+        if not isinstance(dt, torch.Tensor):
+            dt = torch.tensor(dt)
+        if not isinstance(t1, torch.Tensor):
+            t1 = torch.tensor(t1)
+        f0 = func(t0, y0, perturb=Perturb.NEXT if self.perturb else Perturb.NONE)
+        
+        t_dtype = y0.abs().dtype
+        tol = 1e-8
+        if t_dtype == torch.float64:
+            tol = 1e-8
+        if t_dtype == torch.float32:
+            tol = 1e-6
+
+        t0 = t0.to(t_dtype)
+        dt = dt.to(t_dtype)
+        t1 = t1.to(t_dtype)
+
+        k = [f0.clone()] * len(self.tableau.alpha)
+
+        for i, (alpha_i, beta_i) in enumerate(zip(self.tableau.alpha, self.tableau.beta)):
+            perturb = Perturb.NONE
+            if alpha_i == 1.:
+                ti = t1
+                perturb = Perturb.PREV
+            elif alpha_i == 0.:
+                if not torch.all(self.tableau.beta[i]):
+                    # Same slope as stored so skip
+                    continue
+                ti = t0
+            else:
+                ti = t0 + alpha_i * dt
+
+            k_i = torch.stack(k[:i+1], -1)
+
+            # Broyden's Method to solve the system of nonlinear equations
+            y_i = torch.matmul(k_i, beta_i * dt).add(y0)
+            f = self._residual(func, k_i, y_i, ti, perturb)
+            J = torch.ones_like(f).diag()
+            converged = False
+            for _ in range(self.max_iters):
+                if torch.linalg.norm(f, 2) < tol:
+                    converged = True
+                    break
+
+                # If the matrix becomes singular, just stop and return the last value
+                try:
+                    s = -torch.linalg.solve(J, f)
+                except torch._C._LinAlgError:
+                    break
+
+                k[i] = k[i] + s.reshape_as(k[i])
+                k_i = torch.stack(k[:i+1], -1)
+                y_i = torch.matmul(k_i, beta_i * dt).add(y0)
+                newf = self._residual(func, k_i, y_i, ti, perturb)
+                z = newf - f
+                f = newf
+                J = J + (torch.outer ((z - torch.linalg.vecdot(J,s)),s)) / (torch.dot(s,s))
+
+            if not converged:
+                warnings.warn('Functional iteration did not converge. Solution may be incorrect.')
+
+        dy = torch.matmul(torch.stack(k, -1), dt * self.tableau.c_sol)
+
+        return dy, f0
+    
+    def _residual(self, func, K, y, t, perturb):
+        res = K[...,-1] - func(t, y, perturb=perturb)
+        return res.flatten()


Test patch:
diff --git a/tests/event_tests.py b/tests/event_tests.py
index 9f0c88c..176e3fc 100644
--- a/tests/event_tests.py
+++ b/tests/event_tests.py
@@ -25,8 +25,10 @@ class TestEventHandling(unittest.TestCase):
                             with self.subTest(reverse=reverse, dtype=dtype, device=device, ode=ode, method=method):
                                 if method == "explicit_adams":
                                     tol = 7e-2
-                                elif method == "euler":
+                                elif method == "euler" or method == "implicit_euler":
                                     tol = 5e-3
+                                elif method == "gl6":
+                                    tol = 2e-3
                                 else:
                                     tol = 1e-4
 
diff --git a/tests/gradient_tests.py b/tests/gradient_tests.py
index bf419a2..bbb7464 100644
--- a/tests/gradient_tests.py
+++ b/tests/gradient_tests.py
@@ -44,6 +44,8 @@ class TestGradient(unittest.TestCase):
                         eps = 1e-5
                     elif ode == 'sine':
                         eps = 5e-3
+                    elif ode == 'exp':
+                        eps = 1e-2
                     else:
                         raise RuntimeError
 
diff --git a/tests/norm_tests.py b/tests/norm_tests.py
index b053a6d..30abbda 100644
--- a/tests/norm_tests.py
+++ b/tests/norm_tests.py
@@ -273,6 +273,12 @@ class TestNorms(unittest.TestCase):
         for dtype in DTYPES:
             for device in DEVICES:
                 for method in ADAPTIVE_METHODS:
+                    # Tests with known failures
+                    if (
+                        dtype in [torch.float32] and
+                        method in ['tsit5']
+                    ):
+                        continue
 
                     with self.subTest(dtype=dtype, device=device, method=method):
 
diff --git a/tests/odeint_tests.py b/tests/odeint_tests.py
index 7d775a3..620839e 100644
--- a/tests/odeint_tests.py
+++ b/tests/odeint_tests.py
@@ -5,7 +5,7 @@ import warnings
 import torch
 import torchdiffeq
 
-from problems import (construct_problem, PROBLEMS, DTYPES, DEVICES, METHODS, ADAPTIVE_METHODS, FIXED_METHODS, SCIPY_METHODS)
+from problems import (construct_problem, PROBLEMS, DTYPES, DEVICES, METHODS, ADAPTIVE_METHODS, FIXED_METHODS, SCIPY_METHODS, IMPLICIT_METHODS)
 
 
 def rel_error(true, estimate):
@@ -31,12 +31,23 @@ class TestSolverError(unittest.TestCase):
                         if method == 'dopri8' and dtype == torch.float32:
                             kwargs = dict(rtol=1e-7, atol=1e-7)
 
-                        problems = PROBLEMS if method in ADAPTIVE_METHODS else ('constant',)
+                        if method in ADAPTIVE_METHODS:
+                            if method in IMPLICIT_METHODS:
+                                problems = PROBLEMS
+                            else:
+                                problems = tuple(problem for problem in PROBLEMS)
+                        elif method in IMPLICIT_METHODS:
+                            problems = ('constant', 'exp')
+                        else:
+                            problems = ('constant',)
+
                         for ode in problems:
                             if method in ['adaptive_heun', 'bosh3']:
                                 eps = 4e-3
                             elif ode == 'linear':
                                 eps = 2e-3
+                            elif ode == 'exp':
+                                eps = 5e-2
                             else:
                                 eps = 3e-4
 
@@ -155,6 +166,11 @@ class TestDiscontinuities(unittest.TestCase):
             for dtype in DTYPES:
                 for device in DEVICES:
                     for method in FIXED_METHODS:
+
+                        # Singluar matrix error with float32 and implicit_euler
+                        if dtype == torch.float32 and method == 'implicit_euler':
+                            continue
+
                         for perturb in (True, False):
                             with self.subTest(adjoint=adjoint, dtype=dtype, device=device, method=method,
                                               perturb=perturb):
diff --git a/tests/problems.py b/tests/problems.py
index 9825203..da945d4 100644
--- a/tests/problems.py
+++ b/tests/problems.py
@@ -53,15 +53,26 @@ class LinearODE(torch.nn.Module):
         return torch.stack([torch.tensor(ans_) for ans_ in ans]).reshape(len(t_numpy), self.dim).to(t)
 
 
-PROBLEMS = {'constant': ConstantODE, 'linear': LinearODE, 'sine': SineODE}
+class ExpODE(torch.nn.Module):
+    def forward(self, t, y):
+        return -0.1 * self.y_exact(t)
+    
+    def y_exact(self, t):
+        return torch.exp(-0.1 * t)
+
+
+PROBLEMS = {'constant': ConstantODE, 'linear': LinearODE, 'sine': SineODE, 'exp': ExpODE}
 DTYPES = (torch.float32, torch.float64)
 DEVICES = ['cpu']
 if torch.cuda.is_available():
     DEVICES.append('cuda')
-FIXED_METHODS = ('euler', 'midpoint', 'heun2', 'heun3', 'rk4', 'explicit_adams', 'implicit_adams')
+FIXED_EXPLICIT_METHODS = ('euler', 'midpoint', 'heun2', 'heun3', 'rk4', 'explicit_adams', 'implicit_adams')
+FIXED_IMPLICIT_METHODS = ('implicit_euler', 'implicit_midpoint', 'trapezoid', 'radauIIA3', 'gl4', 'radauIIA5', 'gl6', 'sdirk2', 'trbdf2')
+FIXED_METHODS = FIXED_EXPLICIT_METHODS + FIXED_IMPLICIT_METHODS
 ADAMS_METHODS = ('explicit_adams', 'implicit_adams')
 ADAPTIVE_METHODS = ('adaptive_heun', 'fehlberg2', 'bosh3', 'tsit5', 'dopri5', 'dopri8')
 SCIPY_METHODS = ('scipy_solver',)
+IMPLICIT_METHODS = FIXED_IMPLICIT_METHODS
 METHODS = FIXED_METHODS + ADAPTIVE_METHODS + SCIPY_METHODS