Moved some more functionality to the backends

Author: david-zwicker

pde/backends/numba/solvers.py

@@ -16,14 +16,60 @@
     AdaptiveStepperType,
     FixedStepperType,
     SolverBase,
+    _make_dt_adjuster,
 )
 from ...tools.math import OnlineStatistics
 from ...tools.numba import jit
-from ...tools.typing import NumericArray, TField
+from ...tools.typing import NumericArray, StepperHook, TField
 
 SingleStepType = Callable[[NumericArray, float], None]
 
 
+def _make_post_step_hook(solver: SolverBase, state: TField) -> StepperHook:
+    """Create a callable that executes the PDE's post-step hook.
+
+    The returned function calls the post-step hook provided by the PDE (if any)
+    after each completed time step. If the PDE implements make_post_step_hook,
+    this method attempts to obtain both the hook function and an initial value
+    for the hook's mutable data by calling
+    ``post_step_hook, post_step_data_init = self.pde.make_post_step_hook(state)``.
+    The initial data is stored on the solver instance as ``self._post_step_data_init``
+    (copied to ensure mutability) and will be passed to the hook when the stepper
+    is executed.
+
+    If no hook is provided by the PDE (i.e., ``make_post_step_hook`` raises
+    :class:`NotImplementedError`) or if the solver's ``_use_post_step_hook`` flag
+    is ``False``, a no-op hook is returned and ``self._post_step_data_init`` is set
+    to ``None``.
+
+    The hook returned by this method always conforms to the signature
+    ``(state_data: numpy.ndarray, t: float, post_step_data: numpy.ndarray) -> None``
+    and is suitable for JIT compilation where supported.
+
+    Args:
+        solver (:class:`~pde.solvers.base.SolverBase`):
+            The solver instance, which determines how the hook is constructed
+        state (:class:`~pde.fields.base.FieldBase`):
+            Example field providing the array shape and grid information required
+            by the PDE when constructing the post-step hook.
+
+    Returns:
+        callable:
+            A function that invokes the PDE's post-step hook (or a no-op) with the
+            signature described above.
+    """
+    # get uncompiled post_step_hook
+    post_step_hook = solver._make_post_step_hook(state)
+
+    # compile post_step_hook
+    post_step_data_type = nb.typeof(solver._post_step_data_init)
+    signature_hook = (nb.typeof(state.data), nb.float64, post_step_data_type)
+    post_step_hook = jit(signature_hook)(post_step_hook)
+
+    solver._logger.debug("Compiled post-step hook")
+    return post_step_hook  # type: ignore
+
+
 def _make_fixed_stepper(
     solver: SolverBase, state: TField, dt: float
 ) -> FixedStepperType:
@@ -42,7 +88,7 @@ def _make_fixed_stepper(
     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)
+    post_step_hook = _make_post_step_hook(solver, state)
 
     # provide compiled function doing all steps
     fixed_stepper_signature = (
@@ -86,7 +132,7 @@ def _make_adams_bashforth_stepper(
         raise NotImplementedError
 
     rhs_pde = solver._make_pde_rhs(state, backend=solver.backend)
-    post_step_hook = solver._make_post_step_hook(state)
+    post_step_hook = _make_post_step_hook(solver, state)
     sig_single_step = (nb.typeof(state.data), nb.double, nb.typeof(state.data))
 
     @jit(sig_single_step)
@@ -168,11 +214,12 @@ def _make_adaptive_stepper_general(
     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)
+    post_step_hook = _make_post_step_hook(solver, state)
     sync_errors = solver._make_error_synchronizer()
 
     # obtain auxiliary functions
-    adjust_dt = solver._make_dt_adjuster()
+    signature = (nb.double, nb.double)
+    adjust_dt = jit(signature)(_make_dt_adjuster(solver.dt_min, solver.dt_max))
     tolerance = solver.tolerance
     dt_min = solver.dt_min
 
@@ -231,7 +278,7 @@ def adaptive_stepper(
         adaptive_stepper = jit(signature_stepper)(adaptive_stepper)
 
     solver._logger.info("Initialized adaptive stepper")
-    return adaptive_stepper  # type: ignore
+    return adaptive_stepper
 
 
 def _make_adaptive_stepper_euler(
@@ -251,11 +298,21 @@ def _make_adaptive_stepper_euler(
         time `t_end`. The function call signature is `(state: numpy.ndarray,
         t_start: float, t_end: float)`
     """
-    stepper = solver._make_adaptive_stepper(state)
     if nb.config.DISABLE_JIT:
         # this can be useful to debug numba implementations and for test coverage checks
-        return stepper
+        return solver._make_adaptive_stepper(state)
     else:
+        # create compiled function for adjusting the time step
+        adjust_dt = _make_dt_adjuster(solver.dt_min, solver.dt_max)
+        adjust_signature = (nb.double, nb.double)
+        adjust_dt = jit(adjust_signature)(adjust_dt)
+
+        # create the adaptive stepper and compile it
+        stepper = solver._make_adaptive_stepper(
+            state,
+            post_step_hook=_make_post_step_hook(solver, state),
+            adjust_dt=adjust_dt,
+        )
         signature = (
             nb.typeof(state.data),
             nb.double,

pde/solvers/base.py

@@ -13,14 +13,13 @@
 from inspect import isabstract
 from typing import TYPE_CHECKING, Any, Callable, Union
 
-import numba as nb
 import numpy as np
 from numba.extending import register_jitable
 
 from ..pdes.base import PDEBase
 from ..tools.math import OnlineStatistics
 from ..tools.misc import classproperty
-from ..tools.numba import is_jitted, jit
+from ..tools.numba import is_jitted
 from ..tools.typing import BackendType, NumericArray, StepperHook, TField
 
 if TYPE_CHECKING:
@@ -205,12 +204,6 @@ def _backend_obj(self) -> BackendBase:
 
         return self.__backend_obj
 
-    @property
-    def _compiled(self) -> bool:
-        """bool: indicates whether functions need to be compiled"""
-        jit_enabled = not nb.config.DISABLE_JIT
-        return jit_enabled and self.backend != "numpy"
-
     def _make_error_synchronizer(
         self, operator: int | str = "MAX"
     ) -> Callable[[float], float]:
@@ -283,12 +276,6 @@ def post_step_hook(
             # ensure that the initial values is a mutable array
             self._post_step_data_init = np.array(self._post_step_data_init, copy=True)
 
-        self._post_step_data_type = nb.typeof(self._post_step_data_init)
-        if self._compiled:
-            sig_hook = (nb.typeof(state.data), nb.float64, self._post_step_data_type)
-            post_step_hook = jit(sig_hook)(post_step_hook)
-            self._logger.debug("Compiled post-step hook")
-
         return post_step_hook  # type: ignore
 
     def _check_backend(self, rhs: Callable) -> None:
@@ -498,54 +485,6 @@ def __init__(
         self.adaptive = adaptive
         self.tolerance = tolerance
 
-    def _make_dt_adjuster(self) -> Callable[[float, float], float]:
-        """Return a function that can be used to adjust time steps."""
-        dt_min = self.dt_min
-        dt_min_nan_err = f"Encountered NaN even though dt < {dt_min}"
-        dt_min_err = f"Time step below {dt_min}"
-        dt_max = self.dt_max
-
-        def adjust_dt(dt: float, error_rel: float) -> float:
-            """Helper function that adjust the time step.
-
-            The goal is to keep the relative error `error_rel` close to 1.
-
-            Args:
-                dt (float): Current time step
-                error_rel (float): Current (normalized) error estimate
-
-            Returns:
-                float: Time step of the next iteration
-            """
-            # adjust the time step
-            if error_rel < 0.00057665:
-                # error was very small => maximal increase in dt
-                # The constant on the right hand side of the comparison is chosen to
-                # agree with the equation for adjusting dt below
-                dt *= 4.0
-            elif np.isnan(error_rel):
-                # state contained NaN => decrease time step strongly
-                dt *= 0.25
-            else:
-                # otherwise, adjust time step according to error
-                dt *= max(0.9 * error_rel**-0.2, 0.1)
-
-            # limit time step to permissible bracket
-            if dt > dt_max:
-                dt = dt_max
-            elif dt < dt_min:
-                if np.isnan(error_rel):
-                    raise RuntimeError(dt_min_nan_err)
-                else:
-                    raise RuntimeError(dt_min_err)
-
-            return dt
-
-        if self._compiled:
-            adjust_dt = jit((nb.double, nb.double))(adjust_dt)
-
-        return adjust_dt
-
     def _make_single_step_variable_dt(
         self, state: TField
     ) -> Callable[[NumericArray, float, float], NumericArray]:
@@ -583,9 +522,6 @@ def _make_single_step_error_estimate(
             raise RuntimeError("Cannot use adaptive stepper with stochastic equation")
 
         single_step = self._make_single_step_variable_dt(state)
-        if self._compiled:
-            sig_single_step = (nb.typeof(state.data), nb.double, nb.double)
-            single_step = jit(sig_single_step)(single_step)
 
         def single_step_error_estimate(
             state_data: NumericArray, t: float, dt: float
@@ -605,13 +541,18 @@ def single_step_error_estimate(
 
         return single_step_error_estimate
 
-    def _make_adaptive_stepper(self, state: TField) -> AdaptiveStepperType:
+    def _make_adaptive_stepper(
+        self, state: TField, *, adjust_dt: Callable[[float, float], float] | None = None
+    ) -> AdaptiveStepperType:
         """Make an adaptive Euler stepper.
 
         Args:
             state (:class:`~pde.fields.base.FieldBase`):
                 An example for the state from which the grid and other information can
                 be extracted
+            adjust_dt (callable or None):
+                A function that is used to adjust the time step. The function takes the
+                current time step and a relative error and returns an adjusted time step.
 
         Returns:
             Function that can be called to advance the `state` from time `t_start` to
@@ -624,15 +565,11 @@ def _make_adaptive_stepper(self, state: TField) -> AdaptiveStepperType:
         sync_errors = self._make_error_synchronizer()
 
         # obtain auxiliary functions
-        adjust_dt = self._make_dt_adjuster()
+        if adjust_dt is None:
+            adjust_dt = _make_dt_adjuster(self.dt_min, self.dt_max)
         tolerance = self.tolerance
         dt_min = self.dt_min
 
-        if self._compiled:
-            # compile paired stepper
-            sig_stepper = (nb.typeof(state.data), nb.double, nb.double)
-            single_step_error = jit(sig_stepper)(single_step_error)
-
         def adaptive_stepper(
             state_data: NumericArray,
             t_start: float,
@@ -674,18 +611,6 @@ def adaptive_stepper(
 
             return t, dt_opt, steps
 
-        if self._compiled:
-            # compile inner stepper
-            sig_adaptive = (
-                nb.typeof(state.data),
-                nb.double,
-                nb.double,
-                nb.double,
-                nb.typeof(self.info["dt_statistics"]),
-                self._post_step_data_type,
-            )
-            adaptive_stepper = jit(sig_adaptive)(adaptive_stepper)
-
         self._logger.info("Initialized adaptive stepper")
         return adaptive_stepper
 
@@ -754,3 +679,68 @@ def wrapped_stepper(state: TField, t_start: float, t_end: float) -> float:
             return t_last
 
         return wrapped_stepper
+
+
+def _make_dt_adjuster(dt_min: float, dt_max: float) -> Callable[[float, float], float]:
+    """Return a function that can be used to adjust time steps.
+
+    The returned function adjust_dt(dt, error_rel) adjusts the current time step
+    `dt` based on the normalized error estimate `error_rel` with the goal of
+    keeping `error_rel` close to 1.
+
+    Behavior:
+    - If the error is very small the time step is increased (up to a factor 4).
+    - If the error is NaN the time step is reduced strongly.
+    - Otherwise the time step is scaled according to error_rel**-0.2 with a
+      conservative lower bound for the scaling factor.
+    - The adjusted time step is clamped to the interval [dt_min, dt_max].
+    - If the adjusted time step falls below dt_min a RuntimeError is raised.
+
+    Args:
+        dt_min (float): Minimal allowed time step.
+        dt_max (float): Maximal allowed time step.
+
+    Returns:
+        Callable[[float, float], float]:
+            Function that takes (dt, error_rel) and returns the adjusted dt.
+    """
+    dt_min_nan_err = f"Encountered NaN even though dt < {dt_min}"
+    dt_min_err = f"Time step below {dt_min}"
+
+    def adjust_dt(dt: float, error_rel: float) -> float:
+        """Helper function that adjust the time step.
+
+        The goal is to keep the relative error `error_rel` close to 1.
+
+        Args:
+            dt (float): Current time step
+            error_rel (float): Current (normalized) error estimate
+
+        Returns:
+            float: Time step of the next iteration
+        """
+        # adjust the time step
+        if error_rel < 0.00057665:
+            # error was very small => maximal increase in dt
+            # The constant on the right hand side of the comparison is chosen to
+            # agree with the equation for adjusting dt below
+            dt *= 4.0
+        elif np.isnan(error_rel):
+            # state contained NaN => decrease time step strongly
+            dt *= 0.25
+        else:
+            # otherwise, adjust time step according to error
+            dt *= max(0.9 * error_rel**-0.2, 0.1)
+
+        # limit time step to permissible bracket
+        if dt > dt_max:
+            dt = dt_max
+        elif dt < dt_min:
+            if np.isnan(error_rel):
+                raise RuntimeError(dt_min_nan_err)
+            else:
+                raise RuntimeError(dt_min_err)
+
+        return dt
+
+    return adjust_dt

pde/solvers/crank_nicolson.py

@@ -7,7 +7,6 @@
 
 from typing import Callable
 
-import numba as nb
 import numpy as np
 
 from ..pdes.base import PDEBase