Comparison against the 'py-pde' package based on the example: Diffusion equation with spatial dependence

.gitignore

@@ -158,3 +158,6 @@ cython_debug/
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 .idea/
+
+# VS Code
+.vscode/

.zenodo.json

@@ -49,6 +49,22 @@
     {
       "affiliation": "Jagiellonian University, Kraków, Poland",
       "name": "Sadowski, Michał"
+    },
+    {
+      "affiliation": "AGH University of Krakow, Poland",
+      "name": "Jaśkowiec, Adrian"
+    },
+    {
+      "affiliation": "AGH University of Krakow, Poland",
+      "name": "Paterak, Arkadiusz"
+    },
+    {
+      "affiliation": "AGH University of Krakow, Poland",
+      "name": "Prosowicz, Wiktor"
+    },
+    {
+      "affiliation": "AGH University of Krakow, Poland",
+      "name": "Stryszewski, Jan"
     }
   ],
   "keywords": [

PyMPDATA/impl/formulae_laplacian.py

@@ -43,7 +43,56 @@ def apply(traversals_data, advector, advectee):
                 null_vecfield_bc,
                 *null_scalarfield,
                 null_scalarfield_bc,
-                traversals_data.buffer
+                traversals_data.buffer,
+            )
+
+    return apply
+
+
+def make_heterogeneous_laplacian(
+    non_unit_g_factor: bool, options: Options, traversals: Traversals
+):
+    """returns njit-ted function for heterogeneous diffusion with spatially varying diffusivity
+
+    Note: heterogeneous diffusion is only supported when options.non_zero_mu_coeff is enabled
+    """
+    if not options.non_zero_mu_coeff:
+        raise NotImplementedError(
+            "Heterogeneous diffusion requires options.non_zero_mu_coeff to be enabled"
+        )
+    elif not options.heterogeneous_diffusion:
+        raise NotImplementedError(
+            "Heterogeneous diffusion requires options.heterogeneous_diffusion to be enabled"
+        )
+
+    else:
+        idx = traversals.indexers[traversals.n_dims]
+        apply_vector = traversals.apply_vector()
+
+        formulae_heterogeneous = tuple(
+            (
+                __make_heterogeneous_laplacian(
+                    options.jit_flags, idx.ats[i], options.epsilon, non_unit_g_factor
+                )
+                if idx.ats[i] is not None
+                else None
+            )
+            for i in range(MAX_DIM_NUM)
+        )
+
+        @numba.njit(**options.jit_flags)
+        def apply(traversals_data, advector, advectee, diffusivity_field):
+            null_vecfield, null_vecfield_bc = traversals_data.null_vector_field
+            return apply_vector(
+                *formulae_heterogeneous,
+                *advector.field,
+                *advectee.field,
+                advectee.bc,
+                *null_vecfield,
+                null_vecfield_bc,
+                *diffusivity_field.field,
+                diffusivity_field.bc,
+                traversals_data.buffer,
             )
 
     return apply
@@ -62,3 +111,40 @@ def fun(advectee, _, __):
         )
 
     return fun
+
+
+def __make_heterogeneous_laplacian(jit_flags, ats, epsilon, non_unit_g_factor):
+    """Create heterogeneous Laplacian function that matches MPDATA's one-sided gradient pattern
+
+    Note: Diffusivity values are expected to be non-negative. Negative values will cause
+    an assertion error. Zero values are handled by setting a minimum threshold (epsilon).
+    """
+    if non_unit_g_factor:
+        raise NotImplementedError()
+
+    @numba.njit(**jit_flags)
+    def fun(advectee, _, diffusivity):
+        # Get concentration values (matching regular laplacian pattern)
+        c_curr = ats(*advectee, 0)
+        c_right = ats(*advectee, 1)
+
+        # Get diffusivity values
+        D_curr = ats(*diffusivity, 0)
+        D_right = ats(*diffusivity, 1)
+
+        # Input validation for diffusivity values
+        assert D_curr >= 0.0, "Diffusivity values must be non-negative"
+        assert D_right >= 0.0, "Diffusivity values must be non-negative"
+
+        # Handle near-zero diffusivity by setting minimum threshold
+        D_curr = max(D_curr, epsilon)
+        D_right = max(D_right, epsilon)
+
+        # Match the exact MPDATA pattern but with diffusivity weighting
+        # Regular: -2 * (c[i+1] - c[i]) / (c[i+1] + c[i] + epsilon)
+        # Heterogeneous: weight by diffusivity at face
+        D_face = 0.5 * (D_curr + D_right)
+
+        return -2 * D_face * (c_right - c_curr) / (c_right + c_curr + epsilon)
+
+    return fun

PyMPDATA/options.py

@@ -2,6 +2,8 @@
 MPDATA variants, iterations, data-type and jit-flags settings
 """
 
+from typing import Union
+
 import numpy as np
 from pystrict import strict
 
@@ -32,8 +34,9 @@ def __init__(
         DPDC: bool = False,  # pylint: disable=invalid-name
         epsilon: float = 1e-15,
         non_zero_mu_coeff: bool = False,
+        heterogeneous_diffusion: bool = False,
         dimensionally_split: bool = False,
-        dtype: [np.float32, np.float64] = np.float64
+        dtype: Union[np.float32, np.float64] = np.float64,
     ):
         self._values = HashableDict(
             {
@@ -47,6 +50,7 @@ def __init__(
                 "dimensionally_split": dimensionally_split,
                 "dtype": dtype,
                 "DPDC": DPDC,
+                "heterogeneous_diffusion": heterogeneous_diffusion,
             }
         )
 
@@ -136,6 +140,11 @@ def dimensionally_split(self) -> bool:
         """flag disabling cross-dimensional terms in antidiffusive velocities"""
         return self._values["dimensionally_split"]
 
+    @property
+    def heterogeneous_diffusion(self) -> bool:
+        """flag enabling spatially varying diffusivity support"""
+        return self._values["heterogeneous_diffusion"]
+
     def __str__(self):
         return str(self._values)
 

PyMPDATA/solver.py

@@ -50,7 +50,8 @@ def __init__(
         stepper: Stepper,
         advectee: ScalarField,
         advector: VectorField,
-        g_factor: [ScalarField, None] = None,
+        g_factor: Union[ScalarField, None] = None,
+        diffusivity_field: Union[ScalarField, None] = None,
     ):
         def scalar_field(dtype=None):
             return ScalarField.clone(advectee, dtype=dtype)
@@ -64,8 +65,10 @@ def vector_field():
         def null_vector_field():
             return VectorField.make_null(advector.n_dims, stepper.traversals)
 
-        for field in [advector, advectee] + (
-            [g_factor] if g_factor is not None else []
+        for field in (
+            [advector, advectee]
+            + ([g_factor] if g_factor is not None else [])
+            + ([diffusivity_field] if diffusivity_field is not None else [])
         ):
             assert field.halo == stepper.options.n_halo
             assert field.dtype == stepper.options.dtype
@@ -75,6 +78,7 @@ def null_vector_field():
             "advectee": advectee,
             "advector": advector,
             "g_factor": g_factor or null_scalar_field(),
+            "diffusivity_field": diffusivity_field or null_scalar_field(),
             "vectmp_a": vector_field(),
             "vectmp_b": vector_field(),
             "vectmp_c": (
@@ -122,6 +126,12 @@ def g_factor(self) -> ScalarField:
         e.g. the changing density of a fluid."""
         return self.__fields["g_factor"]
 
+    @property
+    def diffusivity_field(self) -> ScalarField:
+        """Diffusivity field (with halo), unmodified by advance(),
+        assumed to be constant-in-time. Used for heterogeneous diffusion."""
+        return self.__fields["diffusivity_field"]
+
     def advance(
         self,
         n_steps: int,
@@ -138,9 +148,12 @@ def advance(
             assert self.__stepper.options.non_zero_mu_coeff
         else:
             mu_coeff = (0.0, 0.0, 0.0)
+
+        # Check for heterogeneous diffusion
         if (
             self.__stepper.options.non_zero_mu_coeff
             and not self.__fields["g_factor"].meta[META_IS_NULL]
+            and not self.__stepper.options.heterogeneous_diffusion
         ):
             raise NotImplementedError()
 

PyMPDATA/stepper.py

@@ -13,7 +13,7 @@
 from .impl.formulae_antidiff import make_antidiff
 from .impl.formulae_axpy import make_axpy
 from .impl.formulae_flux import make_flux_first_pass, make_flux_subsequent
-from .impl.formulae_laplacian import make_laplacian
+from .impl.formulae_laplacian import make_heterogeneous_laplacian, make_laplacian
 from .impl.formulae_nonoscillatory import make_beta, make_correction, make_psi_extrema
 from .impl.formulae_upwind import make_upwind
 from .impl.meta import _Impl
@@ -34,7 +34,7 @@ def __init__(
         grid: (tuple, None) = None,
         n_threads: (int, None) = None,
         left_first: (tuple, None) = None,
-        buffer_size: int = 0
+        buffer_size: int = 0,
     ):
         if n_dims is not None and grid is not None:
             raise ValueError()
@@ -144,6 +144,7 @@ def make_step_impl(
     n_iters = options.n_iters
     non_zero_mu_coeff = options.non_zero_mu_coeff
     nonoscillatory = options.nonoscillatory
+    heterogeneous_diffusion = options.heterogeneous_diffusion
 
     upwind = make_upwind(options, non_unit_g_factor, traversals)
     flux_first_pass = make_flux_first_pass(options, traversals)
@@ -153,6 +154,9 @@ def make_step_impl(
         non_unit_g_factor, options, traversals, last_pass=True
     )
     laplacian = make_laplacian(non_unit_g_factor, options, traversals)
+    heterogeneous_laplacian = make_heterogeneous_laplacian(
+        non_unit_g_factor, options, traversals
+    )
     nonoscillatory_psi_extrema = make_psi_extrema(options, traversals)
     nonoscillatory_beta = make_beta(non_unit_g_factor, options, traversals)
     nonoscillatory_correction = make_correction(options, traversals)
@@ -168,6 +172,7 @@ def step(
         advectee,
         advector,
         g_factor,
+        diffusivity_field,
         vectmp_a,
         vectmp_b,
         vectmp_c,
@@ -185,7 +190,12 @@ def step(
                     if nonoscillatory:
                         nonoscillatory_psi_extrema(null_impl, psi_extrema, advectee)
                     if non_zero_mu_coeff:
-                        laplacian(null_impl, advector, advectee)
+                        if heterogeneous_diffusion:
+                            heterogeneous_laplacian(
+                                null_impl, advector, advectee, diffusivity_field
+                            )
+                        else:
+                            laplacian(null_impl, advector, advectee)
                         axpy(
                             *advector.field,
                             mu_coeff,

examples/PyMPDATA_examples/comparisons_against_pypde/diffusion_equation_with_spatial_dependence/__init__.py

@@ -0,0 +1,6 @@
+"""
+Comparison between [the *py-pde* example](https://py-pde.readthedocs.io/en/latest/examples_gallery/simple_pdes/pde_heterogeneous_diffusion.html) of d iffusion equation with spatial dependence and the reimplementation with PyMPDATA.
+
+diffusion_equation_with_spatial_dependence.ipynb:
+.. include:: ./diffusion_equation_with_spatial_dependence.ipynb
+"""