diff --git a/.github/workflows/pdoc.yml b/.github/workflows/pdoc.yml index f39d3ca8..a441140f 100644 --- a/.github/workflows/pdoc.yml +++ b/.github/workflows/pdoc.yml @@ -25,7 +25,7 @@ jobs: persist-credentials: false - uses: actions/setup-python@v5.2.0 with: - python-version: 3.9 + python-version: "3.13" - env: JUPYTER_PLATFORM_DIRS: 1 run: | diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 0b2fffe9..c078eb2e 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -38,7 +38,7 @@ jobs: - name: Set up Python uses: actions/setup-python@v5.2.0 with: - python-version: 3.11 + python-version: "3.13" - name: Install dependencies run: | python -m pip install --upgrade pip diff --git a/examples/PyMPDATA_examples/comparison_against_pypde_2025/__init__.py b/examples/PyMPDATA_examples/comparison_against_pypde_2025/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/examples/PyMPDATA_examples/comparison_against_pypde_2025/cartesian_grid-1.ipynb b/examples/PyMPDATA_examples/comparison_against_pypde_2025/cartesian_grid-1.ipynb new file mode 100644 index 00000000..81cad4f8 --- /dev/null +++ b/examples/PyMPDATA_examples/comparison_against_pypde_2025/cartesian_grid-1.ipynb @@ -0,0 +1,256 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "[![preview notebook](https://img.shields.io/static/v1?label=render%20on&logo=github&color=87ce3e&message=GitHub)](https://github.com/open-atmos/PyMPDATA/blob/main/examples/PyMPDATA_examples/comparison_against_pypde_2025/cartesian_grid-1.ipynb)\n", + "[![launch on Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1an38vReLVS7q1Wd-qE_CiDtjmPJK-lU1?usp=sharing)\n", + " " + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "license: GPL v3 \n", + "authors: Piotr Karaś, Norbert Klockiewicz, Jakub Kot, Kacper Majchrzak \n", + "Diffusion on a Cartesian grid \n", + "This example shows how to solve the diffusion equation on a Cartesian grid using both py-pde and PyMPDATA." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import sys\n", + "if 'google.colab' in sys.modules:\n", + " !pip --quiet install open-atmos-jupyter-utils\n", + " from open_atmos_jupyter_utils import pip_install_on_colab\n", + " pip_install_on_colab('PyMPDATA-examples')" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from examples.PyMPDATA_examples.comparison_against_pypde_2025.diffusion_2d import InitialConditions\n", + "\n", + "initial_conditions = InitialConditions(\n", + " diffusion_coefficient=0.1,\n", + " time_step=0.001,\n", + " time_end=1,\n", + " grid_shape=(30, 18),\n", + " grid_range_x=(-1.0, 1.0),\n", + " grid_range_y=(0.0, 2.0),\n", + " pulse_position=(0.0, 1.0),\n", + " pulse_shape=(2, 2),\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "ExecuteTime": { + "end_time": "2025-06-20T18:48:41.195420Z", + "start_time": "2025-06-20T18:48:41.055293Z" + } + }, + "outputs": [], + "source": [ + "from examples.PyMPDATA_examples.comparison_against_pypde_2025.diffusion_2d import create_pde_state\n", + "\n", + "from open_atmos_jupyter_utils import show_plot\n", + "from pde import DiffusionPDE\n", + "\n", + "state = create_pde_state(initial_conditions=initial_conditions)\n", + "state.plot(title=\"Initial condition\", cmap=\"magma\")\n", + "show_plot(filename=\"Initial_condition\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "eq = DiffusionPDE(0.1) # define the pde\n", + "result_pdepde = eq.solve(state, t_range=1, dt=0.001)\n", + "result_pdepde.plot(cmap=\"magma\")\n", + "result_pdepde = result_pdepde.data" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "ExecuteTime": { + "end_time": "2025-06-20T18:48:45.524859Z", + "start_time": "2025-06-20T18:48:44.260427Z" + } + }, + "outputs": [], + "source": [ + "from examples.PyMPDATA_examples.comparison_against_pypde_2025.diffusion_2d import py_pde_solution\n", + "\n", + "\n", + "result_pdepde2 = py_pde_solution(initial_conditions=initial_conditions)\n", + "assert (result_pdepde == result_pdepde2).all()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "ExecuteTime": { + "end_time": "2025-06-20T18:50:59.449433Z", + "start_time": "2025-06-20T18:50:58.089985Z" + } + }, + "outputs": [], + "source": [ + "import numpy as np\n", + "from matplotlib import pyplot as plt\n", + "from examples.PyMPDATA_examples.comparison_against_pypde_2025.diffusion_2d import (\n", + " mpdata_solution,\n", + ")\n", + "\n", + "result_mpdata = mpdata_solution(initial_conditions=initial_conditions)\n", + "\n", + "# Plot final result\n", + "plt.figure(figsize=(6, 4))\n", + "plt.imshow(\n", + " result_mpdata,\n", + " cmap=\"magma\",\n", + " origin=\"lower\",\n", + " extent=[\n", + " initial_conditions.min_x,\n", + " initial_conditions.max_x,\n", + " initial_conditions.min_y,\n", + " initial_conditions.max_y,\n", + " ],\n", + ")\n", + "plt.title(\"PyMPDATA Diffusion Result\")\n", + "plt.xlabel(\"x\")\n", + "plt.ylabel(\"y\")\n", + "plt.colorbar()\n", + "plt.tight_layout()\n", + "show_plot(filename=\"PyMPDATA_diffusion_result\")\n", + "\n", + "# Total mass check\n", + "mass = np.sum(result_mpdata) * initial_conditions.dx * initial_conditions.dy\n", + "print(f\"Total mass after diffusion: {mass:.6f} (should be close to 1)\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "ExecuteTime": { + "end_time": "2025-06-20T18:51:09.076355Z", + "start_time": "2025-06-20T18:51:09.072679Z" + } + }, + "outputs": [], + "source": [ + "diff = result_mpdata - result_pdepde\n", + "rmse = np.sqrt(np.mean(diff**2))\n", + "l1 = np.sum(np.abs(diff)) * initial_conditions.dx * initial_conditions.dy\n", + "mass_py_pde = result_pdepde.sum() * initial_conditions.dx * initial_conditions.dy\n", + "mass_mpdata = result_mpdata.sum() * initial_conditions.dx * initial_conditions.dy\n", + "print(\"Total mass:\")\n", + "print(f\"py-pde:{mass_py_pde:>12.6f}\")\n", + "print(f\"PyMPDATA:{mass_mpdata:>10.6f}\")\n", + "\n", + "print(f\"RMSE: {rmse:.6f}, L1 Error: {l1:.6f}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "ExecuteTime": { + "end_time": "2025-06-20T18:53:49.719836Z", + "start_time": "2025-06-20T18:53:49.672914Z" + } + }, + "outputs": [], + "source": [ + "def analytical_solution(x, y, t, D=0.1, center=(0.0, 1.0)):\n", + " x0, y0 = center\n", + " r2 = (x - x0) ** 2 + (y - y0) ** 2\n", + " return (1 / (4 * np.pi * D * t)) * np.exp(-r2 / (4 * D * t))\n", + "\n", + "\n", + "x = np.linspace(initial_conditions.min_x + initial_conditions.dx / 2, initial_conditions.max_x - initial_conditions.dx / 2, initial_conditions.nx)\n", + "y = np.linspace(initial_conditions.min_y + initial_conditions.dy / 2, initial_conditions.max_y - initial_conditions.dy / 2, initial_conditions.ny)\n", + "X, Y = np.meshgrid(x, y, indexing=\"ij\")\n", + "u_exact = analytical_solution(X, Y, initial_conditions.time_end, D=initial_conditions.diffusion_coefficient)\n", + "\n", + "plt.figure(figsize=(6, 4))\n", + "plt.imshow(u_exact.T, origin=\"lower\", extent=[\n", + " initial_conditions.min_x,\n", + " initial_conditions.max_x,\n", + " initial_conditions.min_y,\n", + " initial_conditions.max_y,\n", + "], cmap=\"magma\")\n", + "plt.title(\"Analytical Diffusion Solution\")\n", + "plt.colorbar()\n", + "show_plot(filename=\"Analytical_diffusion_solution\")\n", + "\n", + "diff_pde = result_pdepde - u_exact\n", + "rmse_pde = np.sqrt(np.mean((diff_pde) ** 2))\n", + "l1_error_pde = np.sum(np.abs(diff_pde)) * initial_conditions.dx * initial_conditions.dy\n", + "\n", + "COL_WIDTH_METHOD = 10\n", + "COL_WIDTH_METRIC = 18\n", + "\n", + "print(\"=\" * 50)\n", + "print(\"Comparison with Analytical Solution\")\n", + "print(\"=\" * 50)\n", + "print(\n", + " f\"{'Method':<{COL_WIDTH_METHOD}} {'RMSE':>{COL_WIDTH_METRIC}} {'L1 Error':>{COL_WIDTH_METRIC}}\"\n", + ")\n", + "print(\"-\" * 50)\n", + "print(\n", + " f\"{'PDE':<{COL_WIDTH_METHOD}} {rmse_pde:>{COL_WIDTH_METRIC}.6f} {l1_error_pde:>{COL_WIDTH_METRIC}.6f}\"\n", + ")\n", + "\n", + "diff_mpdata = result_mpdata - u_exact\n", + "rmse_mpdata = np.sqrt(np.mean((diff_mpdata) ** 2))\n", + "l1_error_mpdata = np.sum(np.abs(diff_mpdata)) * initial_conditions.dx * initial_conditions.dy\n", + "\n", + "print(\n", + " f\"{'MPDATA':<{COL_WIDTH_METHOD}} {rmse_mpdata:>{COL_WIDTH_METRIC}.6f} {l1_error_mpdata:>{COL_WIDTH_METRIC}.6f}\"\n", + ")\n", + "print(\"=\" * 50)" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": ".venv", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.11" + } + }, + "nbformat": 4, + "nbformat_minor": 0 +} diff --git a/examples/PyMPDATA_examples/comparison_against_pypde_2025/diffusion_2d.py b/examples/PyMPDATA_examples/comparison_against_pypde_2025/diffusion_2d.py new file mode 100644 index 00000000..afcf92c4 --- /dev/null +++ b/examples/PyMPDATA_examples/comparison_against_pypde_2025/diffusion_2d.py @@ -0,0 +1,174 @@ +""" +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, + DataFieldBase, + 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. + """ + + N_DIMS = 2 + + 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] + pulse_shape: tuple[int, int] + + @property + def nx(self) -> int: + return self.grid_shape[0] + + @property + def ny(self) -> int: + return self.grid_shape[1] + + @property + def min_x(self) -> float: + return self.grid_range_x[0] + + @property + def max_x(self) -> float: + return self.grid_range_x[1] + + @property + def min_y(self) -> float: + return self.grid_range_y[0] + + @property + def max_y(self) -> float: + return self.grid_range_y[1] + + @property + def pulse_x(self) -> float: + return self.pulse_position[0] + + @property + def pulse_y(self) -> float: + return self.pulse_position[1] + + @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) + + +def create_pde_state(initial_conditions: InitialConditions) -> DataFieldBase: + 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, + ) + return state + + +Grid = npt.NDArray[np.float64] + + +def py_pde_solution(initial_conditions: InitialConditions) -> Grid: + """ + Solve the 2D diffusion equation using PyPDE. + """ + state = create_pde_state(initial_conditions=initial_conditions) + eq = DiffusionPDE(diffusivity=initial_conditions.diffusion_coefficient) + result = eq.solve( + state=state, + t_range=1, + dt=initial_conditions.time_step, + ) + if result is not None and hasattr(result, "data"): + return result.data + raise RuntimeError( + "PyPDE solve did not return a valid result with a 'data' attribute." + ) + + +def mpdata_solution(initial_conditions: InitialConditions) -> Grid: + """ + Solve the 2D diffusion equation using PyMPDATA. + """ + opt = Options( + n_iters=2, + non_zero_mu_coeff=True, + ) + stepper = Stepper( + options=opt, + n_dims=initial_conditions.N_DIMS, + ) + + data = create_pde_state(initial_conditions=initial_conditions).data + + 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() diff --git a/examples/setup.py b/examples/setup.py index e1ed3658..63be2284 100644 --- a/examples/setup.py +++ b/examples/setup.py @@ -2,6 +2,7 @@ import os import re +import sys from setuptools import find_packages, setup @@ -37,6 +38,7 @@ def get_long_description(): "imageio", "meshio", "numdifftools", + *(["py-pde==0.45.0"] if CI and sys.version_info.minor >= 12 else []), "pandas", ], author="https://github.com/open-atmos/PyMPDATA/graphs/contributors", diff --git a/setup.py b/setup.py index 3d8633d4..d48ba4b9 100644 --- a/setup.py +++ b/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 and sys.version_info.minor >= 12 else []), ] }, author="https://github.com/open-atmos/PyMPDATA/graphs/contributors", diff --git a/tests/smoke_tests/comparison_against_pypde_2025/__init__.py b/tests/smoke_tests/comparison_against_pypde_2025/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/smoke_tests/comparison_against_pypde_2025/test_comparison.py b/tests/smoke_tests/comparison_against_pypde_2025/test_comparison.py new file mode 100644 index 00000000..f623769e --- /dev/null +++ b/tests/smoke_tests/comparison_against_pypde_2025/test_comparison.py @@ -0,0 +1,122 @@ +""" +Test the similarity of solutions from MPDATA and PyPDE for 2D diffusion +""" + +import time + +import numba +import numpy as np +import pytest + +from examples.PyMPDATA_examples.comparison_against_pypde_2025.diffusion_2d import ( + Grid, + InitialConditions, + mpdata_solution, + py_pde_solution, +) + +if numba.__version__ < "0.59.1": + pytest.skip("Skipping whole file: needs numba >= 0.59.1", allow_module_level=True) + + +@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), + pulse_shape=(2, 2), + ) + + +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 + assert isinstance(initial_conditions.pulse_shape, tuple) + assert len(initial_conditions.pulse_shape) == 2 + + +def test_similarity_of_solutions(initial_conditions: InitialConditions) -> None: + """Test that the solutions from PyPDE and MPDATA for 2D diffusion are similar.""" + + max_elapsed_time = 10 # seconds + + # initial solutions and timing for PyPDE + py_pde_start = time.perf_counter() + py_pde_result: Grid = py_pde_solution( + initial_conditions=initial_conditions, + ) + py_pde_elapsed = time.perf_counter() - py_pde_start + assert py_pde_elapsed < max_elapsed_time, "PyPDE solution took too long to compute" + # ensure consistency across runs + assert np.all( + py_pde_result + == py_pde_solution( + initial_conditions=initial_conditions, + ) + ), "PyPDE results are not consistent across runs" + + # initial solutions and timing for MPDATA + mpdata_start = time.perf_counter() + mpdata_result: Grid = mpdata_solution( + initial_conditions=initial_conditions, + ) + mpdata_elapsed = time.perf_counter() - mpdata_start + assert mpdata_elapsed < max_elapsed_time, "MPDATA solution took too long to compute" + # ensure consistency across runs + assert np.all( + mpdata_result + == mpdata_solution( + initial_conditions=initial_conditions, + ) + ), "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 np.isclose(py_pde_result.sum(), mpdata_result.sum(), rtol=1e-5, atol=1e-10) + + 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 + l1_error = np.sum(np.abs(diff)) * initial_conditions.dx * initial_conditions.dy + l2_norm = np.linalg.norm(diff) + rmse = l2_norm / np.sqrt(py_pde_result.size) + 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