Moved some code from solvers to backend

Split the ExplicitSolver into EulerSolver and RungeKuttaSolver

Author: david-zwicker

pde/backends/base.py

@@ -8,7 +8,7 @@
 import inspect
 import logging
 from collections import defaultdict
-from typing import Any, Callable, NamedTuple
+from typing import Any, Callable, Literal, NamedTuple
 
 from ..fields import DataFieldBase
 from ..grids import BoundariesBase, GridBase
@@ -342,9 +342,13 @@ def make_sde_rhs(
             f"SDE right hand side not defined for backend {self.name}"
         )
 
-    def make_stepper(
-        self, solver: SolverBase, state: TField, dt: float | None = None
-    ) -> Callable[[TField, float, float], float]:
+    def make_inner_stepper(
+        self,
+        solver: SolverBase,
+        stepper_style: Literal["fixed", "adaptive"],
+        state: TField,
+        dt: float,
+    ) -> Callable:
         """Return a stepper function using an explicit scheme.
 
         Args:
@@ -353,8 +357,8 @@ def make_stepper(
             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 used (Uses :attr:`SolverBase.dt_default` if `None`)
+            stepper_style (str):
+                Determines how the stepper is expected to work
 
         Returns:
             Function that can be called to advance the `state` from time `t_start` to

pde/backends/numba/backend.py

@@ -6,7 +6,7 @@
 from __future__ import annotations
 
 import functools
-from typing import Callable
+from typing import Callable, Literal
 
 import numba as nb
 import numpy as np
@@ -16,6 +16,7 @@
 from ...fields import DataFieldBase, VectorField
 from ...grids import BoundariesBase, DimensionError, GridBase
 from ...pdes import PDEBase
+from ...solvers import AdaptiveSolverBase, SolverBase
 from ...tools.numba import get_common_numba_dtype, jit, make_array_constructor
 from ...tools.typing import (
     DataSetter,
@@ -703,3 +704,38 @@ def make_sde_rhs(
             together with a noise realization.
         """
         return eq._make_sde_rhs_numba_cached(state, **kwargs)
+
+    def make_inner_stepper(
+        self,
+        solver: SolverBase,
+        stepper_style: Literal["fixed", "adaptive"],
+        state: TField,
+        dt: float,
+    ) -> Callable:
+        """Return a stepper function using an explicit scheme.
+
+        Args:
+            solver (:class:`~pde.solvers.base.SolverBase`):
+                The solver instance, which determines how the stepper is constructed
+            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 used (Uses :attr:`SolverBase.dt_default` if `None`)
+
+        Returns:
+            Function that can be called to advance the `state` from time `t_start` to
+            time `t_end`. The function call signature is `(state: numpy.ndarray,
+            t_start: float, t_end: float)`
+        """
+        assert solver.backend == "numba"
+
+        from .solvers import make_adaptive_stepper, make_fixed_stepper
+
+        if stepper_style == "fixed":
+            return make_fixed_stepper(solver, state, dt=dt)
+        elif stepper_style == "adaptive":
+            assert isinstance(solver, AdaptiveSolverBase)
+            return make_adaptive_stepper(solver, state)
+        else:
+            raise NotImplementedError

pde/backends/numba/solvers.py

@@ -0,0 +1,283 @@
+"""Implements numba-accelerated solvers.
+
+.. codeauthor:: David Zwicker <[email protected]>
+"""
+
+from __future__ import annotations
+
+from typing import Callable
+
+import numba as nb
+import numpy as np
+
+from ...solvers import *
+from ...solvers.base import (
+    AdaptiveSolverBase,
+    AdaptiveStepperType,
+    FixedStepperType,
+    SolverBase,
+)
+from ...tools.math import OnlineStatistics
+from ...tools.numba import jit
+from ...tools.typing import NumericArray, TField
+
+SingleStepType = Callable[[NumericArray, float], None]
+
+
+def _make_fixed_stepper(
+    solver: SolverBase, state: TField, dt: float
+) -> FixedStepperType:
+    """Return a stepper function using an explicit scheme with fixed time steps.
+
+    Args:
+        solver (:class:`~pde.solvers.base.SolverBase`):
+            The solver instance, which determines how the stepper is constructed
+        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 explicit stepping.
+    """
+    # get compiled version of a single step
+    single_step = solver._make_single_step_fixed_dt(state, dt)
+    single_step_signature = (nb.typeof(state.data), nb.double)
+    single_step = jit(single_step_signature)(single_step)
+    post_step_hook = solver._make_post_step_hook(state)
+
+    # provide compiled function doing all steps
+    fixed_stepper_signature = (
+        nb.typeof(state.data),
+        nb.double,
+        nb.int_,
+        nb.typeof(solver._post_step_data_init),
+    )
+
+    @jit(fixed_stepper_signature)
+    def fixed_stepper(
+        state_data: NumericArray, t_start: float, steps: int, post_step_data
+    ) -> float:
+        """Perform `steps` steps with fixed time steps."""
+        for i in range(steps):
+            # calculate the right hand side
+            t = t_start + i * dt
+            single_step(state_data, t)
+            post_step_hook(state_data, t, post_step_data)
+
+        return t + dt
+
+    return fixed_stepper  # type: ignore
+
+
+def _make_adams_bashforth_stepper(
+    solver: AdamsBashforthSolver, state: TField, dt: float
+) -> FixedStepperType:
+    """Return a stepper function using an explicit scheme with fixed time steps.
+
+    Args:
+        solver (:class:`~pde.solvers.adams_bashforth.AdamsBashforthSolver`):
+            The solver instance, which determines how the stepper is constructed
+        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 explicit stepping.
+    """
+    if solver.pde.is_sde:
+        raise NotImplementedError
+
+    rhs_pde = solver._make_pde_rhs(state, backend=solver.backend)
+    post_step_hook = solver._make_post_step_hook(state)
+    sig_single_step = (nb.typeof(state.data), nb.double, nb.typeof(state.data))
+
+    @jit(sig_single_step)
+    def single_step(
+        state_data: NumericArray, t: float, state_prev: NumericArray
+    ) -> None:
+        """Perform a single Adams-Bashforth step."""
+        rhs_prev = rhs_pde(state_prev, t - dt).copy()
+        rhs_cur = rhs_pde(state_data, t)
+        state_prev[:] = state_data  # save the previous state
+        state_data += dt * (1.5 * rhs_cur - 0.5 * rhs_prev)
+
+    # allocate memory to store the state of the previous time step
+    state_prev = np.empty_like(state.data)
+    init_state_prev = True
+
+    def fixed_stepper(
+        state_data: NumericArray, t_start: float, steps: int, post_step_data
+    ) -> float:
+        """Perform `steps` steps with fixed time steps."""
+        nonlocal state_prev, init_state_prev
+
+        if init_state_prev:
+            # initialize the state_prev with an estimate of the previous step
+            state_prev[:] = state_data - dt * rhs_pde(state_data, t_start)
+            init_state_prev = False
+
+        for i in range(steps):
+            # calculate the right hand side
+            t = t_start + i * dt
+            single_step(state_data, t, state_prev)
+            post_step_hook(state_data, t, post_step_data=post_step_data)
+
+        return t + dt
+
+    solver._logger.info("Init explicit Adams-Bashforth stepper with dt=%g", dt)
+
+    return fixed_stepper
+
+
+def make_fixed_stepper(
+    solver: SolverBase, state: TField, dt: float
+) -> FixedStepperType:
+    """Return a stepper function using an explicit scheme with fixed time steps.
+
+    Args:
+        solver (:class:`~pde.solvers.base.SolverBase`):
+            The solver instance, which determines how the stepper is constructed
+        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 explicit stepping.
+    """
+    if isinstance(solver, AdamsBashforthSolver):
+        return _make_adams_bashforth_stepper(solver, state, dt)
+    else:
+        return _make_fixed_stepper(solver, state, dt)
+
+
+def _make_adaptive_stepper_general(
+    solver: AdaptiveSolverBase, state: TField
+) -> AdaptiveStepperType:
+    """Return a stepper function using an explicit scheme.
+
+    Args:
+        solver (:class:`~pde.solvers.base.AdaptiveSolverBase`):
+            The solver instance, which determines how the stepper is constructed
+        state (:class:`~pde.fields.base.FieldBase`):
+            An example for the state from which the grid and other information can
+            be extracted
+
+    Returns:
+        Function that can be called to advance the `state` from time `t_start` to
+        time `t_end`. The function call signature is `(state: numpy.ndarray,
+        t_start: float, t_end: float)`
+    """
+    # obtain functions determining how the PDE is evolved
+    single_step_error = solver._make_single_step_error_estimate(state)
+    signature_single_step = (nb.typeof(state.data), nb.double, nb.double)
+    single_step_error = jit(signature_single_step)(single_step_error)
+    post_step_hook = solver._make_post_step_hook(state)
+    sync_errors = solver._make_error_synchronizer()
+
+    # obtain auxiliary functions
+    adjust_dt = solver._make_dt_adjuster()
+    tolerance = solver.tolerance
+    dt_min = solver.dt_min
+
+    signature_stepper = (
+        nb.typeof(state.data),
+        nb.double,
+        nb.double,
+        nb.double,
+        nb.typeof(solver.info["dt_statistics"]),
+        nb.typeof(solver._post_step_data_init),
+    )
+
+    @jit(signature_stepper)
+    def adaptive_stepper(
+        state_data: NumericArray,
+        t_start: float,
+        t_end: float,
+        dt_init: float,
+        dt_stats: OnlineStatistics | None = None,
+        post_step_data=None,
+    ) -> tuple[float, float, int]:
+        """Adaptive stepper that advances the state in time."""
+        dt_opt = dt_init
+        t = t_start
+        steps = 0
+        while True:
+            # use a smaller (but not too small) time step if close to t_end
+            dt_step = max(min(dt_opt, t_end - t), dt_min)
+
+            # try two different step sizes to estimate errors
+            new_state, error = single_step_error(state_data, t, dt_step)
+
+            error_rel = error / tolerance  # normalize error to given tolerance
+            # synchronize the error between all processes (necessary for MPI)
+            error_rel = sync_errors(error_rel)
+
+            # do the step if the error is sufficiently small
+            if error_rel <= 1:
+                steps += 1
+                t += dt_step
+                state_data[...] = new_state
+                post_step_hook(state_data, t, post_step_data)
+
+                if dt_stats is not None:
+                    dt_stats.add(dt_step)
+
+            if t < t_end:
+                # adjust the time step and continue (happens in every MPI process)
+                dt_opt = adjust_dt(dt_step, error_rel)
+            else:
+                break  # return to the controller
+
+        return t, dt_opt, steps
+
+    solver._logger.info("Initialized adaptive stepper")
+    return adaptive_stepper  # type: ignore
+
+
+def _make_adaptive_stepper_euler(
+    solver: EulerSolver, state: TField
+) -> AdaptiveStepperType:
+    """Return a stepper function using an explicit scheme.
+
+    Args:
+        solver (:class:`~pde.solvers.explicit.EulerSolver`):
+            The solver instance, which determines how the stepper is constructed
+        state (:class:`~pde.fields.base.FieldBase`):
+            An example for the state from which the grid and other information can
+            be extracted
+
+    Returns:
+        Function that can be called to advance the `state` from time `t_start` to
+        time `t_end`. The function call signature is `(state: numpy.ndarray,
+        t_start: float, t_end: float)`
+    """
+    stepper = solver._make_adaptive_stepper(state)
+    signature = (
+        nb.typeof(state.data),
+        nb.double,
+        nb.double,
+        nb.double,
+        nb.typeof(solver.info["dt_statistics"]),
+        nb.typeof(solver._post_step_data_init),
+    )
+    return jit(signature)(stepper)  # type: ignore
+
+
+def make_adaptive_stepper(
+    solver: AdaptiveSolverBase, state: TField
+) -> AdaptiveStepperType:
+    """Return a stepper function using an explicit scheme.
+
+    Args:
+        solver (:class:`~pde.solvers.base.AdaptiveSolverBase`):
+            The solver instance, which determines how the stepper is constructed
+        state (:class:`~pde.fields.base.FieldBase`):
+            An example for the state from which the grid and other information can
+            be extracted
+
+    Returns:
+        Function that can be called to advance the `state` from time `t_start` to
+        time `t_end`. The function call signature is `(state: numpy.ndarray,
+        t_start: float, t_end: float)`
+    """
+    if isinstance(solver, EulerSolver):
+        return _make_adaptive_stepper_euler(solver, state)
+    else:
+        return _make_adaptive_stepper_general(solver, state)