Comparison with PyPDE: Integrate Benchmarking and Compatibility Analysis

examples/PyMPDATA_examples/comparison_against_pypde_2025/diffusion_2d.py

@@ -0,0 +1,167 @@
+"""
+Test the similarity of solutions from MPDATA and PyPDE for 2D diffusion
+"""
+
+from dataclasses import dataclass
+
+import numpy as np
+import numpy.typing as npt
+from pde import CartesianGrid, DiffusionPDE
+from pde import ScalarField as PDEScalarField
+
+from PyMPDATA import Options, ScalarField, Solver, Stepper, VectorField
+from PyMPDATA.boundary_conditions import Periodic
+
+
+@dataclass
+class InitialConditions:
+    """
+    Initial conditions for the 2D diffusion problem.
+    """
+
+    def __init__(
+        self,
+        *,
+        diffusion_coefficient: float,
+        time_step: float,
+        time_end: float,
+        grid_shape: tuple[int, int],
+        grid_range_x: tuple[float, float],
+        grid_range_y: tuple[float, float],
+        pulse_position: tuple[float, float],
+    ) -> None:
+        self.diffusion_coefficient = diffusion_coefficient
+        self.time_step = time_step
+        self.time_end = time_end
+        self.grid_shape = grid_shape
+        self.grid_range_x = grid_range_x
+        self.grid_range_y = grid_range_y
+        self.pulse_position = pulse_position
+        self.nx, self.ny = grid_shape
+        self.min_x, self.max_x = grid_range_x
+        self.min_y, self.max_y = grid_range_y
+        self.pulse_x, self.pulse_y = pulse_position
+
+    def __repr__(self) -> str:
+        return (
+            f"InitialConditions(diffusion_coefficient={self.diffusion_coefficient}, "
+            f"time_step={self.time_step}, time_end={self.time_end}, "
+            f"grid_shape={self.grid_shape}, grid_range_x={self.grid_range_x}, "
+            f"grid_range_y={self.grid_range_y}, pulse_position={self.pulse_position})"
+        )
+
+    @property
+    def dx(self) -> float:
+        """Calculate the grid spacing in the x-direction."""
+        return (self.max_x - self.min_x) / self.nx
+
+    @property
+    def dy(self) -> float:
+        """Calculate the grid spacing in the y-direction."""
+        return (self.max_y - self.min_y) / self.ny
+
+    @property
+    def n_steps(self) -> int:
+        """Calculate the number of time steps based on the time range and time step."""
+        return int(self.time_end / self.time_step)
+
+
+type Two2DiffusionSolution = npt.NDArray[np.float64]
+
+
+def py_pde_solution(initial_conditions: InitialConditions) -> Two2DiffusionSolution:
+    """
+    Solve the 2D diffusion equation using PyPDE.
+    """
+    grid = CartesianGrid(
+        bounds=[
+            initial_conditions.grid_range_x,
+            initial_conditions.grid_range_y,
+        ],
+        shape=initial_conditions.grid_shape,
+    )
+    state = PDEScalarField(grid=grid)
+    state.insert(
+        point=np.array(initial_conditions.pulse_position),
+        amount=1,
+    )
+    eq = DiffusionPDE(diffusivity=initial_conditions.diffusion_coefficient)
+    result = eq.solve(
+        state=state,
+        t_range=1,
+        dt=initial_conditions.time_step,
+    )
+    return result.data
+
+
+def mpdata_solution(initial_conditions: InitialConditions) -> Two2DiffusionSolution:
+    """
+    Solve the 2D diffusion equation using PyMPDATA.
+    """
+    opt = Options(
+        n_iters=2,
+        non_zero_mu_coeff=True,
+    )
+    stepper = Stepper(
+        options=opt,
+        n_dims=2,
+    )
+
+    def create_pde_like_data(ic) -> npt.NDArray[np.float64]:
+        """
+        Create a 2D array with a pulse at the specified position.
+        """
+        x = np.linspace(ic.min_x + ic.dx / 2, ic.max_x - ic.dx / 2, ic.nx)
+        y = np.linspace(ic.min_y + ic.dy / 2, ic.max_y - ic.dy / 2, ic.ny)
+        result = np.zeros((ic.nx, ic.ny))
+        # Locate cell nearest (0, 1)
+        i = np.argmin(np.abs(x - ic.pulse_x))
+        j = np.argmin(np.abs(y - ic.pulse_y))
+        # Distribute mass over 2x2 cells (py-pde seems to do this internally)
+        mass_per_cell = 1.0 / (4 * ic.dx * ic.dy)
+        result[i, j] = mass_per_cell
+        result[i + 1, j] = mass_per_cell
+        result[i, j + 1] = mass_per_cell
+        result[i + 1, j + 1] = mass_per_cell
+        return result
+
+    data = create_pde_like_data(ic=initial_conditions)
+
+    advectee = ScalarField(
+        data=data, halo=opt.n_halo, boundary_conditions=(Periodic(), Periodic())
+    )
+
+    cx = np.zeros(
+        shape=(initial_conditions.nx + 1, initial_conditions.ny),
+        dtype=opt.dtype,
+    )
+    cy = np.zeros(
+        shape=(initial_conditions.nx, initial_conditions.ny + 1),
+        dtype=opt.dtype,
+    )
+    advector = VectorField(
+        data=(cx, cy),
+        halo=opt.n_halo,
+        boundary_conditions=(Periodic(), Periodic()),
+    )
+
+    solver = Solver(
+        stepper=stepper,
+        advector=advector,
+        advectee=advectee,
+    )
+    mu_x = (
+        initial_conditions.diffusion_coefficient
+        * initial_conditions.time_step
+        / initial_conditions.dx**2
+    )
+    mu_y = (
+        initial_conditions.diffusion_coefficient
+        * initial_conditions.time_step
+        / initial_conditions.dy**2
+    )
+    solver.advance(
+        n_steps=initial_conditions.n_steps,
+        mu_coeff=(mu_x, mu_y),
+    )
+    return solver.advectee.get()

setup.py

@@ -82,6 +82,7 @@ def get_long_description():
             "joblib" + ("==1.4.0" if CI else ""),
             "imageio",
             "nbformat",
+            "py-pde" + ("==0.45.0" if CI else ""),
         ]
     },
     author="https://github.com/open-atmos/PyMPDATA/graphs/contributors",

tests/smoke_tests/comparison_against_pypde_2025/test_comparison.py

@@ -0,0 +1,115 @@
+"""
+Test the similarity of solutions from MPDATA and PyPDE for 2D diffusion
+"""
+
+import numpy as np
+import pytest
+import time
+from numpy.ma.testutils import assert_almost_equal
+
+from examples.PyMPDATA_examples.comparison_against_pypde_2025.diffusion_2d import (
+    InitialConditions,
+    Two2DiffusionSolution,
+    mpdata_solution,
+    py_pde_solution,
+)
+
+
+@pytest.fixture(name="initial_conditions")
+def _initial_conditions() -> InitialConditions:
+    """Fixture providing initial conditions for the diffusion problem."""
+
+    return InitialConditions(
+        diffusion_coefficient=0.1,
+        time_step=0.001,
+        time_end=1,
+        grid_shape=(30, 18),
+        grid_range_x=(-1.0, 1.0),
+        grid_range_y=(0.0, 2.0),
+        pulse_position=(0.0, 1.0),
+    )
+
+
+def test_initial_conditions(initial_conditions: InitialConditions) -> None:
+    """Test that the initial conditions are set up correctly."""
+
+    assert initial_conditions.diffusion_coefficient > 0
+    assert initial_conditions.time_step > 0
+    assert initial_conditions.time_end > 0
+    assert isinstance(initial_conditions.grid_shape, tuple)
+    assert len(initial_conditions.grid_shape) == 2
+    assert all(isinstance(dim, int) for dim in initial_conditions.grid_shape)
+    assert all(dim > 0 for dim in initial_conditions.grid_shape)
+    assert isinstance(initial_conditions.grid_range_x, tuple)
+    assert len(initial_conditions.grid_range_x) == 2
+    assert isinstance(initial_conditions.grid_range_y, tuple)
+    assert len(initial_conditions.grid_range_y) == 2
+    assert isinstance(initial_conditions.pulse_position, tuple)
+    assert len(initial_conditions.pulse_position) == 2
+
+
+def test_similarity_of_solutions(initial_conditions: InitialConditions) -> None:
+    """Test that the solutions from PyPDE and MPDATA for 2D diffusion are similar."""
+
+    # initial solutions
+    py_pde_result: Two2DiffusionSolution = py_pde_solution(
+        initial_conditions=initial_conditions,
+    )
+
+    mpdata_result: Two2DiffusionSolution = mpdata_solution(
+        initial_conditions=initial_conditions,
+    )
+
+    # calculate solutions again to time them and ensure they are consistent across runs
+    py_pde_start = time.perf_counter()
+    py_pde_result2: Two2DiffusionSolution = py_pde_solution(
+        initial_conditions=initial_conditions,
+    )
+    assert np.all(py_pde_result == py_pde_result2), (
+        "PyPDE results are not consistent across runs"
+    )
+
+    py_pde_elapsed = time.perf_counter() - py_pde_start
+    assert py_pde_elapsed < 10, "PyPDE solution took too long to compute"
+
+    mpdata_start = time.perf_counter()
+    mpdata_result2: Two2DiffusionSolution = mpdata_solution(
+        initial_conditions=initial_conditions,
+    )
+    mpdata_elapsed = time.perf_counter() - mpdata_start
+    assert mpdata_elapsed < 10, "MPDATA solution took too long to compute"
+
+    assert np.all(mpdata_result == mpdata_result2), (
+        "MPDATA results are not consistent across runs"
+    )
+
+    # sanity checks
+    assert py_pde_result.shape == mpdata_result.shape
+    assert np.all(np.isfinite(py_pde_result))
+    assert np.all(np.isfinite(mpdata_result))
+    assert np.all(py_pde_result >= 0)
+    assert np.all(mpdata_result >= 0)
+
+    # total mass check
+    assert_almost_equal(py_pde_result.sum(), mpdata_result.sum(), decimal=5)
+
+    py_pde_total_mass = (
+        py_pde_result.sum() * initial_conditions.dx * initial_conditions.dy
+    )
+    mpdata_total_mass = (
+        mpdata_result.sum() * initial_conditions.dx * initial_conditions.dy
+    )
+    assert np.isclose(py_pde_total_mass, 1.0, rtol=1e-3)
+    assert np.isclose(mpdata_total_mass, 1.0, rtol=1e-3)
+
+    # compare results
+    corr = np.corrcoef(py_pde_result.ravel(), mpdata_result.ravel())[0, 1]
+    diff = py_pde_result - mpdata_result
+    rmse = np.sqrt(np.mean(diff**2))
+    l1_error = np.sum(np.abs(diff)) * initial_conditions.dx * initial_conditions.dy
+    l2_norm = np.linalg.norm(py_pde_result - mpdata_result)
+    assert np.allclose(py_pde_result, mpdata_result, rtol=1e-2, atol=1e-1)
+    assert corr > 0.99
+    assert rmse < 0.02
+    assert l1_error < 0.05
+    assert l2_norm < 0.5