From b5015a3b5bef3efccd08bb2061434f86e04989cd Mon Sep 17 00:00:00 2001 From: David Zwicker Date: Mon, 29 Jan 2024 10:58:41 +0100 Subject: [PATCH] Added a semi-implicit Euler-Maruyama solver (#530) This solver could improve stiff stochastic partial differential equations. --- pde/grids/coordinates/__init__.py | 12 ++-- pde/solvers/explicit.py | 2 + pde/solvers/implicit.py | 96 +++++++++++++++++++++++++- tests/solvers/test_implicit_solvers.py | 52 ++++++++++++++ 4 files changed, 154 insertions(+), 8 deletions(-) create mode 100644 tests/solvers/test_implicit_solvers.py diff --git a/pde/grids/coordinates/__init__.py b/pde/grids/coordinates/__init__.py index 5cb40cb2..16bf7add 100644 --- a/pde/grids/coordinates/__init__.py +++ b/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 diff --git a/pde/solvers/explicit.py b/pde/solvers/explicit.py index 78fb3a1d..ac529978 100644 --- a/pde/solvers/explicit.py +++ b/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: diff --git a/pde/solvers/implicit.py b/pde/solvers/implicit.py index 60a7740c..adc127cf 100644 --- a/pde/solvers/implicit.py +++ b/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) diff --git a/tests/solvers/test_implicit_solvers.py b/tests/solvers/test_implicit_solvers.py new file mode 100644 index 00000000..c6ee4dbb --- /dev/null +++ b/tests/solvers/test_implicit_solvers.py @@ -0,0 +1,52 @@ +""" +.. codeauthor:: David Zwicker +""" + +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)