Add examples as PyMPDATA-MPI dependency, dynamic_advector option w/ stashes, fully JIT-able SWE code, interpolated indexers logic (for 1D & 2D), multiple advectees in solver, AnteStep class (#612)

Co-authored-by: Sylwester Arabas <[email protected]>

Author: Sfonxu

.github/workflows/mpi.yml

@@ -88,7 +88,7 @@ jobs:
         - if: steps.cache.outputs.cache-hit != 'true'
           run: |
             HDF5_MPI="ON" CC=mpicc pip install --no-binary=h5py h5py==3.13.0
-            pip install -e MPI[tests]
+            pip install -e .[tests]  -e ./examples[tests] -e ./MPI[tests]
         - run: pip show numpy
         - id: cache-save
           if: steps.cache.outputs.cache-hit != 'true'
@@ -101,15 +101,18 @@ jobs:
     needs: [tests_setup]
     strategy:
       matrix:
-        platform: [macos-15-intel, macos-14, ubuntu-latest]
+        # Due to reported performance issues, runs on macos-15-intel are skipped
+        # https://github.com/actions/runner-images/issues/12545
+        # platform: [macos-15-intel, macos-14, ubuntu-latest]
+        platform: [macos-14, ubuntu-latest]
         mpi: [ 'mpich', 'openmpi', 'intelmpi']
         python-version: ["3.10"]
         disable-jit: [1, 0]
         mpi-np: [1, 2, 3]
         exclude:
           # as of time of writing, mpi4py/setup-mpi does not support it
-          - platform: macos-15-intel
-            mpi: intelmpi
+          # - platform: macos-15-intel
+          #  mpi: intelmpi
           - platform: macos-14
             mpi: intelmpi
 
@@ -118,14 +121,14 @@ jobs:
             mpi: intelmpi
 
           # https://github.com/Homebrew/homebrew-core/issues/26974
-          - platform: macos-15-intel
-            mpi: mpich
+          # - platform: macos-15-intel
+          #  mpi: mpich
           - platform: macos-14
             mpi: mpich
 
           # HDF5_MPI expects arm architecture with python 3.12 on macos
-          - platform: macos-15-intel
-            python-version: '3.12'
+          # - platform: macos-15-intel
+          #  python-version: '3.12'
 
       fail-fast: false
 

MPI/pyproject.toml

@@ -14,7 +14,7 @@ build-backend = "setuptools.build_meta"
 name = "pympdata-mpi"
 description = "PyMPDATA + numba-mpi coupler sandbox"
 readme = "README.md"
-requires-python = ">=3.10"
+requires-python = ">=3.9"
 keywords = ["MPI", "MPDATA", "Numba", "PyMPDATA"]
 license = {text = "GPL-3.0"}
 classifiers = [
@@ -32,6 +32,7 @@ dependencies = [
     "numpy<1.25.0",
     "numba_mpi>=0.30",
     "PyMPDATA",
+    "PyMPDATA-examples",
     "mpi4py==4.0.3",
     "h5py",
 ]

PyMPDATA/impl/formulae_divide.py

@@ -0,0 +1,66 @@
+"""Operation logic for dividing the field by
+a set divisor table and saving the result
+to a seperate field. Requires 'dynmaic_advector'
+option to be enabled.
+Scalar field inputs are named after dimensional
+components of the VectorField, not to be confused
+with internal enumerations on axis indexing"""
+
+import numba
+import numpy as np
+
+from .enumerations import INNER, MID3D, OUTER
+from .meta import META_HALO_VALID
+
+
+def make_divide_or_zero(options, traversals):
+    """returns njit-ted function for use with given traversals"""
+
+    n_dims = traversals.n_dims
+
+    @numba.njit(**options.jit_flags)
+    # pylint: disable=too-many-arguments
+    def divide_or_zero(
+        out_outer_meta,
+        out_outer_data,
+        out_mid3d_meta,
+        out_mid3d_data,
+        out_inner_meta,
+        out_inner_data,
+        _,
+        dividend_outer,
+        __,
+        dividend_mid3d,
+        ___,
+        dividend_inner,
+        ____,
+        divisor,
+        time_step,
+        grid_step,
+    ):
+        eps = 1e-7
+        for i in np.ndindex(out_inner_data.shape):
+            if n_dims > 1:
+                out_outer_data[i] = (
+                    dividend_outer[i] / divisor[i] * time_step / grid_step[OUTER]
+                    if divisor[i] > eps
+                    else 0
+                )
+                if n_dims > 2:
+                    out_mid3d_data[i] = (
+                        dividend_mid3d[i] / divisor[i] * time_step / grid_step[MID3D]
+                        if divisor[i] > eps
+                        else 0
+                    )
+            out_inner_data[i] = (
+                dividend_inner[i] / divisor[i] * time_step / grid_step[INNER]
+                if divisor[i] > eps
+                else 0
+            )
+        if n_dims > 1:
+            out_outer_meta[META_HALO_VALID] = False
+            if n_dims > 2:
+                out_mid3d_meta[META_HALO_VALID] = False
+        out_inner_meta[META_HALO_VALID] = False
+
+    return divide_or_zero

PyMPDATA/impl/indexers.py

@@ -7,6 +7,7 @@
 from pathlib import Path
 
 import numba
+import numpy as np
 
 from .enumerations import INNER, INVALID_INDEX, MID3D, OUTER
 
@@ -31,6 +32,14 @@ def ats_1d(focus, arr, k, _=INVALID_INDEX, __=INVALID_INDEX):
         def atv_1d(focus, arrs, k, _=INVALID_INDEX, __=INVALID_INDEX):
             return arrs[INNER][focus[INNER] + int(k - 0.5)]
 
+        @staticmethod
+        @numba.njit(**jit_flags)
+        def ati_1d(focus, arrs, k, _=INVALID_INDEX, __=INVALID_INDEX):
+            return (
+                arrs[INNER][focus[INNER] + int(k - 0.5)]
+                + arrs[INNER][focus[INNER] + int(k + 0.5)]
+            ) / 2
+
         @staticmethod
         @numba.njit(**jit_flags)
         def set(arr, _, __, k, value):
@@ -70,6 +79,28 @@ def atv_axis1(focus, arrs, k, i=0, _=INVALID_INDEX):
                 dim, _ii, _kk = OUTER, int(i - 0.5), int(k)
             return arrs[dim][focus[OUTER] + _ii, focus[INNER] + _kk]
 
+        @staticmethod
+        @numba.njit(**jit_flags)
+        def ati_axis0(focus, arrs, i, k=0, _=INVALID_INDEX):
+            if _is_integral(k):
+                dim, _ii, _kk = OUTER, int(i - 0.5), int(k)
+                return (
+                    arrs[dim][focus[OUTER] + _ii, focus[INNER] + _kk]
+                    + arrs[dim][focus[OUTER] + _ii + 1, focus[INNER] + _kk]
+                ) / 2
+            return np.nan
+
+        @staticmethod
+        @numba.njit(**jit_flags)
+        def ati_axis1(focus, arrs, k, i=0, _=INVALID_INDEX):
+            if _is_integral(i):
+                dim, _ii, _kk = INNER, int(i), int(k - 0.5)
+                return (
+                    arrs[dim][focus[OUTER] + _ii, focus[INNER] + _kk]
+                    + arrs[dim][focus[OUTER] + _ii, focus[INNER] + _kk + 1]
+                ) / 2
+            return np.nan
+
         @staticmethod
         @numba.njit(**jit_flags)
         def set(arr, i, _, k, value):
@@ -140,24 +171,31 @@ def get(arr, i, j, k):
             return arr[i, j, k]
 
     Indexers = namedtuple(  # pylint: disable=invalid-name
-        Path(__file__).stem + "_Indexers", ("ats", "atv", "set", "get", "n_dims")
+        Path(__file__).stem + "_Indexers", ("ats", "atv", "ati", "set", "get", "n_dims")
     )
 
     indexers = (
         None,
         Indexers(
-            (None, None, _1D.ats_1d), (None, None, _1D.atv_1d), _1D.set, _1D.get, 1
+            (None, None, _1D.ats_1d),
+            (None, None, _1D.atv_1d),
+            (None, None, _1D.ati_1d),
+            _1D.set,
+            _1D.get,
+            1,
         ),
         Indexers(
             (_2D.ats_axis0, None, _2D.ats_axis1),
             (_2D.atv_axis0, None, _2D.atv_axis1),
+            (_2D.ati_axis0, None, _2D.ati_axis1),
             _2D.set,
             _2D.get,
             2,
         ),
         Indexers(
             (_3D.ats_axis0, _3D.ats_axis1, _3D.ats_axis2),
             (_3D.atv_axis0, _3D.atv_axis1, _3D.atv_axis2),
+            (None, None, None),
             _3D.set,
             _3D.get,
             3,

PyMPDATA/options.py

@@ -32,6 +32,7 @@ def __init__(
         DPDC: bool = False,  # pylint: disable=invalid-name
         epsilon: float = 1e-15,
         non_zero_mu_coeff: bool = False,
+        dynamic_advector: bool = False,
         dimensionally_split: bool = False,
         dtype: [np.float32, np.float64] = np.float64
     ):
@@ -44,6 +45,7 @@ def __init__(
                 "nonoscillatory": nonoscillatory,
                 "third_order_terms": third_order_terms,
                 "non_zero_mu_coeff": non_zero_mu_coeff,
+                "dynamic_advector": dynamic_advector,
                 "dimensionally_split": dimensionally_split,
                 "dtype": dtype,
                 "DPDC": DPDC,
@@ -131,6 +133,11 @@ def non_zero_mu_coeff(self) -> bool:
         """flag enabling handling of Fickian diffusion term"""
         return self._values["non_zero_mu_coeff"]
 
+    @property
+    def dynamic_advector(self) -> bool:
+        """flag enabling (todo desc)"""
+        return self._values["dynamic_advector"]
+
     @property
     def dimensionally_split(self) -> bool:
         """flag disabling cross-dimensional terms in antidiffusive velocities"""