Added a semi-implicit Euler-Maruyama solver (#530)

This solver could improve stiff stochastic partial differential
equations.

pde/grids/coordinates/__init__.py

@@ -4,12 +4,12 @@
 .. autosummary::
    :nosignatures:
 
-    bipolar.BipolarCoordinates
-    bispherical.BisphericalCoordinates
-    cartesian.CartesianCoordinates
-    cylindrical.CylindricalCoordinates
-    polar.PolarCoordinates
-    spherical.SphericalCoordinates
+    ~bipolar.BipolarCoordinates
+    ~bispherical.BisphericalCoordinates
+    ~cartesian.CartesianCoordinates
+    ~cylindrical.CylindricalCoordinates
+    ~polar.PolarCoordinates
+    ~spherical.SphericalCoordinates
 """
 
 from .base import CoordinatesBase, DimensionError

pde/solvers/explicit.py

@@ -74,6 +74,7 @@ def _make_single_step_fixed_euler(
         """
         if self.pde.is_sde:
             # handle stochastic version of the pde
+            self.info["scheme"] = "euler-maruyama"
             rhs_sde = self._make_sde_rhs(state, backend=self.backend)
 
             def stepper(state_data: np.ndarray, t: float) -> None:
@@ -87,6 +88,7 @@ def stepper(state_data: np.ndarray, t: float) -> None:
 
         else:
             # handle deterministic version of the pde
+            self.info["scheme"] = "euler"
             rhs_pde = self._make_pde_rhs(state, backend=self.backend)
 
             def stepper(state_data: np.ndarray, t: float) -> None:

pde/solvers/implicit.py

@@ -46,10 +46,10 @@ def __init__(
         self.maxiter = maxiter
         self.maxerror = maxerror
 
-    def _make_single_step_fixed_dt(
+    def _make_single_step_fixed_dt_deterministic(
         self, state: FieldBase, dt: float
     ) -> Callable[[np.ndarray, float], None]:
-        """return a function doing a single step with an implicit Euler scheme
+        """return a function doing a deterministic step with an implicit Euler scheme
 
         Args:
             state (:class:`~pde.fields.base.FieldBase`):
@@ -110,4 +110,96 @@ def implicit_step(state_data: np.ndarray, t: float) -> None:
                 raise ConvergenceError("Implicit Euler step did not converge.")
             nfev += n + 1
 
+        self._logger.info("Init implicit Euler stepper with dt=%g", dt)
+        return implicit_step
+
+    def _make_single_step_fixed_dt_stochastic(
+        self, state: FieldBase, dt: float
+    ) -> Callable[[np.ndarray, float], None]:
+        """return a function doing a step for a SDE with an implicit Euler scheme
+
+        Args:
+            state (:class:`~pde.fields.base.FieldBase`):
+                An example for the state from which the grid and other information can
+                be extracted
+            dt (float):
+                Time step of the implicit step
+        """
+        self.info["function_evaluations"] = 0
+        self.info["scheme"] = "implicit-euler-maruyama"
+        self.info["stochastic"] = True
+        self.info["dt_adaptive"] = False
+
+        rhs = self.pde.make_pde_rhs(state, backend=self.backend)  # type: ignore
+        rhs_sde = self._make_sde_rhs(state, backend=self.backend)
+        maxiter = int(self.maxiter)
+        maxerror2 = self.maxerror**2
+
+        # handle deterministic version of the pde
+        def implicit_step(state_data: np.ndarray, t: float) -> None:
+            """compiled inner loop for speed"""
+            nfev = 0  # count function evaluations
+
+            # save state at current time point t for stepping
+            state_t = state_data.copy()
+
+            # estimate state at next time point
+            evolution_rate, noise_realization = rhs_sde(state_data, t)
+            state_data[:] = state_t + dt * evolution_rate  # estimated state
+
+            if noise_realization is not None:
+                # add the noise to the reference state at the current time point and
+                # adept the state at the next time point iteratively below
+                state_t += np.sqrt(dt) * noise_realization
+
+            state_prev = np.empty_like(state_data)
+
+            # fixed point iteration for improving state after dt
+            for n in range(maxiter):
+                state_prev[:] = state_data  # keep previous state to judge convergence
+                # another interation to improve estimate
+                state_data[:] = state_t + dt * rhs(state_data, t + dt)
+
+                # calculate mean squared error to judge convergence
+                err = 0.0
+                for j in range(state_data.size):
+                    diff = state_data.flat[j] - state_prev.flat[j]
+                    err += (diff.conjugate() * diff).real
+                err /= state_data.size
+
+                if err < maxerror2:
+                    # fix point iteration converged
+                    break
+            else:
+                with nb.objmode:
+                    self._logger.warning(
+                        "Semi-implicit Euler-Maruyama step did not converge after %d "
+                        "iterations at t=%g (error=%g)",
+                        maxiter,
+                        t,
+                        err,
+                    )
+                raise ConvergenceError(
+                    "Semi-implicit Euler-Maruyama step did not converge."
+                )
+            nfev += n + 1
+
+        self._logger.info("Init semi-implicit Euler-Maruyama stepper with dt=%g", dt)
         return implicit_step
+
+    def _make_single_step_fixed_dt(
+        self, state: FieldBase, dt: float
+    ) -> Callable[[np.ndarray, float], None]:
+        """return a function doing a single step with an implicit Euler scheme
+
+        Args:
+            state (:class:`~pde.fields.base.FieldBase`):
+                An example for the state from which the grid and other information can
+                be extracted
+            dt (float):
+                Time step of the implicit step
+        """
+        if self.pde.is_sde:
+            return self._make_single_step_fixed_dt_stochastic(state, dt)
+        else:
+            return self._make_single_step_fixed_dt_deterministic(state, dt)

tests/solvers/test_implicit_solvers.py

@@ -0,0 +1,52 @@
+"""
+.. codeauthor:: David Zwicker <[email protected]>
+"""
+
+import numpy as np
+import pytest
+
+from pde import PDE, DiffusionPDE, ScalarField, UnitGrid
+from pde.solvers import Controller, ImplicitSolver
+from pde.tools import mpi
+
+
+@pytest.mark.parametrize("backend", ["numpy", "numba"])
+def test_implicit_solvers_simple_fixed(backend):
+    """test implicit solvers"""
+    grid = UnitGrid([4])
+    xs = grid.axes_coords[0]
+    field = ScalarField.from_expression(grid, "x")
+    eq = PDE({"c": "c"})
+
+    dt = 0.01
+    solver = ImplicitSolver(eq, backend=backend)
+    controller = Controller(solver, t_range=10.0, tracker=None)
+    res = controller.run(field, dt=dt)
+
+    if mpi.is_main:
+        np.testing.assert_allclose(res.data, xs * np.exp(10), rtol=0.1)
+        assert solver.info["steps"] == pytest.approx(10 / dt, abs=1)
+        assert not solver.info.get("dt_adaptive", False)
+
+
+@pytest.mark.parametrize("backend", ["numpy", "numba"])
+def test_implicit_stochastic_solvers(backend, rng):
+    """test simple version of the stochastic implicit solver"""
+    field = ScalarField.random_uniform(UnitGrid([16]), -1, 1, rng=rng)
+    eq = DiffusionPDE()
+    seq = DiffusionPDE(noise=1e-10)
+
+    solver1 = ImplicitSolver(eq, backend=backend)
+    c1 = Controller(solver1, t_range=1, tracker=None)
+    s1 = c1.run(field, dt=1e-3)
+
+    solver2 = ImplicitSolver(seq, backend=backend)
+    c2 = Controller(solver2, t_range=1, tracker=None)
+    s2 = c2.run(field, dt=1e-3)
+
+    np.testing.assert_allclose(s1.data, s2.data, rtol=1e-4, atol=1e-4)
+    assert not solver1.info["stochastic"]
+    assert solver2.info["stochastic"]
+
+    assert not solver1.info.get("dt_adaptive", False)
+    assert not solver2.info.get("dt_adaptive", False)