Refactor backend classes to use generic type parameters and improve type hints across various backends

docs/source/_static/requirements_main.csv

@@ -1,7 +1,7 @@
 Package,Minimal version,Usage
 matplotlib,3.1,Visualizing results
 numba,0.59,Just-in-time compilation to accelerate numerics
-numpy,1.22,Handling numerical data
+numpy,2,Handling numerical data
 scipy,1.10,Miscellaneous scientific functions
 sympy,1.9,Dealing with user-defined mathematical expressions
 tqdm,4.66,Display progress bars during calculations

docs/source/create_performance_plots.py → docs/source/create_performance_plots_ops.py

@@ -1,5 +1,5 @@
 #!/usr/bin/env python3
-"""Code for creating performance plots.
+"""Code for creating plots showing performance of operators.
 
 .. codeauthor:: David Zwicker <david.zwicker@ds.mpg.de>
 """
@@ -44,7 +44,7 @@
         cv2.Laplacian, ddepth=cv2.CV_64F, borderType=cv2.BORDER_REFLECT
     )
 # determine path of the cache for the ground truth of simulations
-GROUND_TRUTH_CACHE = Path(__file__).resolve().parent / "_cache" / "performance_data"
+RESULT_CACHE = Path(__file__).resolve().parent / "_cache" / "performance_ops"
 
 
 def time_function(func, arg, repeat: int = 3, use_out: bool = False) -> float:
@@ -111,7 +111,7 @@ def get_single_run(backend: str, size: int, periodic=False) -> float:
         periodic (bool):
             Whether the grid is periodic or not
     """
-    cache_file = GROUND_TRUTH_CACHE / f"{backend}_{size}_{periodic}.json"
+    cache_file = RESULT_CACHE / f"{backend}_{size}_{periodic}.json"
 
     if cache_file.exists():
         # read performance data
@@ -130,7 +130,7 @@ def get_single_run(backend: str, size: int, periodic=False) -> float:
         # write  performance data
         cache_file.parent.mkdir(parents=True, exist_ok=True)
         with cache_file.open("w") as f:
-            return json.dump(data, f)
+            json.dump(data, f)
 
     return runtime
 

docs/source/create_performance_plots_pde.py

@@ -0,0 +1,164 @@
+#!/usr/bin/env python3
+"""Code for creating plots showing performance of PDEs.
+
+.. codeauthor:: David Zwicker <david.zwicker@ds.mpg.de>
+"""
+
+import sys
+from pathlib import Path
+
+# determine path of the `py-pde` package
+PACKAGE_PATH = Path(__file__).resolve().parents[2]
+sys.path.insert(0, str(PACKAGE_PATH))
+
+# disable multithreading
+import os
+
+from pde import config
+
+os.environ["NUMBA_NUM_THREADS"] = "1"
+config["backend.numba.multithreading"] = "never"
+
+# import remaining packages
+import json
+import time
+from datetime import datetime, timezone
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+import pde
+from pde import CahnHilliardPDE, ScalarField, TrackerBase, UnitGrid
+
+# determine path of the cache for the ground truth of simulations
+RESULT_CACHE = Path(__file__).resolve().parent / "_cache" / "performance_pde"
+
+
+class RuntimeTacker(TrackerBase):
+    def __init__(self, interrupts):
+        super().__init__(interrupts=interrupts)
+        self.start_time = time.monotonic()
+        self.data: list[tuple[float, float]] = []
+
+    def initialize(self, field, info):
+        return super().initialize(field, info)
+
+    def handle(self, field, t) -> None:
+        """Handle data supplied to this tracker.
+
+        Args:
+            field (:class:`~pde.fields.FieldBase`):
+                The current state of the simulation
+            t (float):
+                The associated time
+        """
+        self.data.append((t, time.monotonic() - self.start_time))
+
+
+def calculate_single_run(backend: str, size: int) -> list[tuple[float, float]]:
+    """Calculate performance data for a particular configuration.
+
+    Args:
+        backend (str):
+            Backend to test
+        size (int):
+            Dimension of the grid
+    """
+    print(f"Run simulation for {backend=} and {size=}")
+    rng = np.random.default_rng(0)
+    grid = UnitGrid([size] * 2, periodic=True)
+    field = ScalarField.random_uniform(grid, rng=rng)
+    eq = CahnHilliardPDE()
+    runtime = RuntimeTacker(pde.ConstantInterrupts(1e2, 1))
+    eq.solve(field, t_range=1e3 + 1, dt=1e-2, backend=backend, tracker=runtime)
+    return runtime.data
+
+
+def get_single_run(backend: str, size: int) -> list[tuple[float, float]]:
+    """Get performance data for a particular configuration.
+
+    Args:
+        backend (str):
+            Backend to test
+        size (int):
+            Dimension of the grid
+    """
+    cache_file = RESULT_CACHE / f"{backend}_{size}.json"
+
+    if cache_file.exists():
+        # read performance data
+        with cache_file.open() as f:
+            runtime = json.load(f)["runtime"]
+
+    else:
+        # calculate  performance data
+        runtime = calculate_single_run(backend, size)
+        data = {
+            "runtime": runtime,
+            "version": pde.__version__,
+            "date": datetime.now(timezone.utc).isoformat(),
+        }
+
+        # write  performance data
+        cache_file.parent.mkdir(parents=True, exist_ok=True)
+        with cache_file.open("w") as f:
+            json.dump(data, f)
+
+    return runtime
+
+
+def collect_performance_data(size: int) -> dict[list[tuple[float, float]]]:
+    """Obtain the data used in the performance plot.
+
+    Args:
+        size (int):
+            Dimension of the grid
+
+    Returns:
+        dict: The durations of calculating the Laplacian on different grids
+        using different methods
+    """
+    return {
+        backend: get_single_run(backend, size)
+        for backend in ["numpy", "numba", "torch", "jax"]
+    }
+
+
+def plot_performance(performance_data):
+    """Plot the performance data.
+
+    Args:
+        performance_data:
+            The data obtained from calling :func:`get_performance_data`.
+    """
+    plt.figure(figsize=[4, 3])
+
+    PLOT_DATA = [
+        {"key": "numpy", "label": "numpy", "fmt": "C0-"},
+        {"key": "numba", "label": "numba", "fmt": "C1-"},
+        {"key": "torch", "label": "torch:cpu", "fmt": "C2-"},
+        {"key": "jax", "label": "jax", "fmt": "C3-"},
+    ]
+
+    for plot in PLOT_DATA:
+        data = np.asarray(performance_data[plot["key"]])
+        plt.plot(data[:, 0], data[:, 1], plot["fmt"], label=plot["label"])
+
+    plt.xlim(0, data[-1, 0])
+    plt.xlabel("Simulation time")
+    plt.ylabel("Runtime [s]")
+    plt.legend(loc="best", fontsize=8)
+    plt.tight_layout()
+
+
+def main():
+    """Run main scripts."""
+    data = collect_performance_data(size=256)
+    plot_performance(data)
+    plt.savefig("performance_cahn_hilliard.pdf", transparent=True)
+    plt.savefig("performance_cahn_hilliard.png", transparent=True, dpi=200)
+    plt.close()
+
+
+if __name__ == "__main__":
+    main()

docs/source/manual/advanced_usage.rst

@@ -444,7 +444,7 @@ introduced above:
             diffusivity = self.diffusivity
 
             # create operators with correct attributes
-            args = {"backend": backend, "native": True, "dtype": state.dtype}
+            args = {"backend": backend, "dtype": state.dtype}
             laplace_u = state.grid.make_operator("laplace", bc=self.bc, **args)
             gradient_u = state.grid.make_operator("gradient", bc=self.bc, **args)
             laplace2_u = state.grid.make_operator("laplace", bc=self.bc_laplace, **args)

docs/source/manual/contributing.rst

@@ -48,15 +48,49 @@ where the arguments denote the grid class, the name of the operator, and the fac
 function, respectively.
 
 
-Design choices
-""""""""""""""
-The data layout of field classes (subclasses of
-:class:`~pde.fields.base.FieldBase`) was chosen to allow for a simple
-decomposition of different fields and tensor components. Consequently, the data
-is laid out in memory such that spatial indices are last. For instance, the data
-of a vector field ``field`` defined on a 2d Cartesian grid will have three
-dimensions and can be accessed as ``field.data[vector_component, x, y]``,
+Data layout
+"""""""""""
+
+..  figure:: /_images/data_layout.*
+    :figwidth: 50%
+    :align: right
+    :alt: Schematic of the data flow during a simulation.
+
+    Schematic of the data flow during a simulation.
+
+Since the package supports several different backends, data can reside in various
+places (e.g., CPU or GPU).
+To avoid confusion, we adhere to the following principles: Data that the user
+manipulates directly should always be stored in numpy arrays on the CPU.
+In contrast, inner loops for solving PDEs can move memory to GPU or any other device.
+Consequently, operators typically assume that data is stored in the native version of
+the respective backend.
+
+The data layout of field classes (subclasses of :class:`~pde.fields.base.FieldBase`) was
+chosen to allow for a simple decomposition of different fields and tensor components.
+Consequently, data is laid out in memory such that spatial indices are last.
+For instance, the data of a vector field ``field`` defined on a 2d Cartesian grid will
+have three dimensions and can be accessed as ``field.data[vector_component, x, y]``,
 where ``vector_component`` is either 0 or 1.
+Note that :class:`~pde.fields.collection.FieldCollection` linearizes all vector- and
+tensor-components, to force all data in a unified array format.
+How this works is shown in the following example:
+
+.. code-block:: python
+
+    grid = pde.UnitGrid([7, 9])
+    scalar = pde.ScalarField.random_uniform(grid)
+    assert scalar.data.shape == (7, 9)
+    tensor = pde.Tensor2Field.random_uniform(grid)
+    assert tensor.data.shape == (2, 2, 7, 9)
+    collection = pde.FieldCollection([scalar, tensor])
+    assert collection.data.shape == (5, 7, 9)
+
+    assert np.all(collection.data[0] == scalar.data)
+    assert np.all(collection.data[1] == tensor.data[0, 0])
+    assert np.all(collection.data[2] == tensor.data[0, 1])
+    assert np.all(collection.data[3] == tensor.data[1, 0])
+    assert np.all(collection.data[4] == tensor.data[1, 1])
 
 
 Coding style

docs/source/manual/performance.rst

@@ -3,13 +3,29 @@ Performance
 
 Measuring performance
 """""""""""""""""""""
+..  figure:: /_images/performance_cahn_hilliard.*
+    :figwidth: 50%
+    :align: right
+    :alt: Runtime as a function of simulation time for various backends.
+
+    Performance of solving the Cahn-Hilliard equation with various backends.
 
 The performance of the `py-pde` package depends on many details and general 
 statements are thus difficult to make.
 However, since the core operators can be compiled using :mod:`numba` or :mod:`torch`,
 many operations of the package proceed at performances close to most compiled
 languages.
-For instance, a simple Laplace operator applied to fields defined on a Cartesian
+For a simple test case, consider the plot on the right, which shows how much real time
+passed until a certain simulation time was reached.
+Here, the slope of the curve quantifies the raw speed of the solver, whereas the
+intercept with the y-axis indicates the start-up time, which is mainly dominated by the
+compilation of the code.
+The trade-off between compilation time and run time varies between backends and will
+also depend on details of the simulation.
+
+A more microscopic performance measurement concerns the actual operators, which can also
+be compiled using various backends.
+Generally, a simple Laplace operator applied to fields defined on a Cartesian
 grid has performance that is similar to the operators supplied by the
 `OpenCV <https://opencv.org>`_ package.
 The following figures illustrate this by showing the duration of evaluating the

examples/advanced_pdes/pde_brusselator_class.py

@@ -53,7 +53,7 @@ def make_evolution_rate(self, state, backend):
         d0, d1 = self.diffusivity
         a, b = self.a, self.b
         laplace = state.grid.make_operator(
-            "laplace", bc=self.bc, backend=backend, native=True, dtype=state.dtype
+            "laplace", bc=self.bc, backend=backend, dtype=state.dtype
         )
 
         def pde_rhs(state_data, t):

examples/advanced_pdes/pde_custom_numba.py

@@ -31,14 +31,10 @@ def evolution_rate(self, state, t=0):
     def make_evolution_rate(self, state, backend):
         """Compilable implementation of the PDE."""
         gradient_squared = state.grid.make_operator(
-            "gradient_squared",
-            bc=self.bc,
-            backend=backend,
-            native=True,
-            dtype=state.dtype,
+            "gradient_squared", bc=self.bc, backend=backend, dtype=state.dtype
         )
         laplace = state.grid.make_operator(
-            "laplace", bc=self.bc, backend=backend, native=True, dtype=state.dtype
+            "laplace", bc=self.bc, backend=backend, dtype=state.dtype
         )
 
         def pde_rhs(data, t):