Add dedicated interpolate function

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 = [
@@ -43,4 +43,5 @@ tests = [
     "pytest-mpi",
     "pytest-timeout",
     "matplotlib",
+    "pympdata-examples"
 ]

PyMPDATA/impl/formulae_divide.py

@@ -0,0 +1,61 @@
+"""basic a*x+y operation logic for use in Fickian term handling"""
+
+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

@@ -31,6 +31,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 +78,30 @@ 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(i):
+                dim, _ii, _kk = INNER, int(i), int(k - 0.5)
+            else:
+                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
+
+        @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)
+            else:
+                dim, _ii, _kk = OUTER, int(i - 0.5), int(k)
+            return (
+                arrs[dim][focus[OUTER] + _ii, focus[INNER] + _kk]
+                + arrs[dim][focus[OUTER] + _ii, focus[INNER] + _kk + 1]
+            ) / 2
+
         @staticmethod
         @numba.njit(**jit_flags)
         def set(arr, i, _, k, value):
@@ -140,24 +172,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"""

PyMPDATA/solver.py

@@ -3,7 +3,7 @@ class grouping user-supplied stepper, fields and post-step/post-iter hooks,
 as well as self-initialised temporary storage
 """
 
-from typing import Union
+from typing import Hashable, Iterable, Mapping, Union
 
 import numba
 
@@ -13,14 +13,37 @@ class grouping user-supplied stepper, fields and post-step/post-iter hooks,
 from .vector_field import VectorField
 
 
+@numba.experimental.jitclass([])
+class AnteStepNull:  # pylint: disable=too-few-public-methods
+    """do-nothing version of the ante-step hook"""
+
+    def __init__(self):
+        pass
+
+    def call(
+        self,
+        traversals_data,
+        advectee,
+        advector,
+        step,
+        index,
+        todo_outer,
+        todo_mid3d,
+        todo_inner,
+    ):  # pylint: disable-next=unused-argument,disable=too-many-arguments
+        """think of it as a `__call__` method (which Numba does not allow)"""
+
+
 @numba.experimental.jitclass([])
 class PostStepNull:  # pylint: disable=too-few-public-methods
     """do-nothing version of the post-step hook"""
 
     def __init__(self):
         pass
 
-    def call(self, psi, step):  # pylint: disable-next=unused-argument
+    def call(
+        self, traversals_data, psi, step, index
+    ):  # pylint: disable-next=unused-argument
         """think of it as a `__call__` method (which Numba does not allow)"""
 
 
@@ -31,7 +54,9 @@ class PostIterNull:  # pylint: disable=too-few-public-methods
     def __init__(self):
         pass
 
-    def call(self, flux, g_factor, step, iteration):  # pylint: disable=unused-argument
+    def call(
+        self, traversals_data, flux, g_factor, step, iteration
+    ):  # pylint: disable=unused-argument,disable=too-many-arguments
         """think of it as a `__call__` method (which Numba does not allow)"""
 
 
@@ -48,66 +73,88 @@ class Solver:
     def __init__(
         self,
         stepper: Stepper,
-        advectee: ScalarField,
+        advectee: [ScalarField, Iterable[ScalarField], Mapping[Hashable, ScalarField]],
         advector: VectorField,
         g_factor: [ScalarField, None] = None,
     ):
+        self.advectee = advectee
+        n_dims = advector.n_dims
+        if isinstance(advectee, ScalarField):
+            self.__fields = {"advectee": (advectee,)}
+
+        elif isinstance(advectee, Mapping):
+            self.__fields = {"advectee": tuple(advectee.values())}
+
+        elif isinstance(advectee, Iterable):
+            self.__fields = {"advectee": tuple(advectee)}
+
         def scalar_field(dtype=None):
-            return ScalarField.clone(advectee, dtype=dtype)
+            return ScalarField.clone(self.__fields["advectee"][0], dtype=dtype)
 
         def null_scalar_field():
-            return ScalarField.make_null(advectee.n_dims, stepper.traversals)
+            return ScalarField.make_null(n_dims, stepper.traversals)
 
         def vector_field():
             return VectorField.clone(advector)
 
         def null_vector_field():
-            return VectorField.make_null(advector.n_dims, stepper.traversals)
+            return VectorField.make_null(n_dims, stepper.traversals)
 
-        for field in [advector, advectee] + (
+        for field in [advector, *self.__fields["advectee"]] + (
             [g_factor] if g_factor is not None else []
         ):
             assert field.halo == stepper.options.n_halo
             assert field.dtype == stepper.options.dtype
             assert field.grid == advector.grid
 
-        self.__fields = {
-            "advectee": advectee,
-            "advector": advector,
-            "g_factor": g_factor or null_scalar_field(),
-            "vectmp_a": vector_field(),
-            "vectmp_b": vector_field(),
-            "vectmp_c": (
-                vector_field()
-                if stepper.options.non_zero_mu_coeff
-                else null_vector_field()
-            ),
-            "nonosc_xtrm": (
-                scalar_field(dtype=complex)
-                if stepper.options.nonoscillatory
-                else null_scalar_field()
-            ),
-            "nonosc_beta": (
-                scalar_field(dtype=complex)
-                if stepper.options.nonoscillatory
-                else null_scalar_field()
-            ),
-        }
-        for field in self.__fields.values():
-            field.assemble(stepper.traversals)
+        self.__fields.update(
+            {
+                "advector": advector,
+                "g_factor": g_factor or null_scalar_field(),
+                "vectmp_a": vector_field(),
+                "vectmp_b": vector_field(),
+                "vectmp_c": (
+                    vector_field()
+                    if stepper.options.non_zero_mu_coeff
+                    else null_vector_field()
+                ),
+                "todo_outer": (
+                    scalar_field()
+                    if stepper.options.dynamic_advector and n_dims > 1
+                    else null_scalar_field()
+                ),
+                "todo_mid3d": (
+                    scalar_field()
+                    if stepper.options.dynamic_advector and n_dims > 2
+                    else null_scalar_field()
+                ),
+                "todo_inner": (
+                    scalar_field()
+                    if stepper.options.dynamic_advector
+                    else null_scalar_field()
+                ),
+                "nonosc_xtrm": (
+                    scalar_field(dtype=complex)
+                    if stepper.options.nonoscillatory
+                    else null_scalar_field()
+                ),
+                "nonosc_beta": (
+                    scalar_field(dtype=complex)
+                    if stepper.options.nonoscillatory
+                    else null_scalar_field()
+                ),
+            }
+        )
+        for key, value in self.__fields.items():
+            for field in (value,) if key != "advectee" else value:
+                field.assemble(stepper.traversals)
 
         self.__stepper = stepper
 
-    @property
-    def advectee(self) -> ScalarField:
-        """advectee scalar field (with halo), modified by advance(),
-        may be modified from user code (e.g., source-term handling)"""
-        return self.__fields["advectee"]
-
     @property
     def advector(self) -> VectorField:
-        """advector vector field (with halo), unmodified by advance(),
-        may be modified from user code"""
+        """advector vector field , todo_outer, todo_mid3d, todo_inner(with halo),
+        unmodified by advance(), may be modified from user code"""
         return self.__fields["advector"]
 
     @property
@@ -126,9 +173,10 @@ def advance(
         self,
         n_steps: int,
         mu_coeff: Union[tuple, None] = None,
+        ante_step=None,
         post_step=None,
         post_iter=None,
-    ):
+    ):  # pylint: disable=too-many-arguments
         """advances solution by `n_steps` steps, optionally accepts: a tuple of diffusion
         coefficients (one value per dimension) as well as `post_iter` and `post_step`
         callbacks expected to be `numba.jitclass`es with a `call` method, for
@@ -144,12 +192,14 @@ def advance(
         ):
             raise NotImplementedError()
 
+        ante_step = ante_step or AnteStepNull()
         post_step = post_step or PostStepNull()
         post_iter = post_iter or PostIterNull()
 
         return self.__stepper(
             n_steps=n_steps,
             mu_coeff=mu_coeff,
+            ante_step=ante_step,
             post_step=post_step,
             post_iter=post_iter,
             fields=self.__fields,