Add traversals_data as argument to stepper loop, disable mpi_dim=INNER for ShallowWaterScenario

Author: Sfonxu

PyMPDATA/solver.py

@@ -15,13 +15,21 @@ class grouping user-supplied stepper, fields and post-step/post-iter hooks,
 
 @numba.experimental.jitclass([])
 class AnteStepNull:  # pylint: disable=too-few-public-methods
-    """do-nothing version of the post-step hook"""
+    """do-nothing version of the ante-step hook"""
 
     def __init__(self):
         pass
 
     def call(
-        self, advectee, advector, step, index, todo_outer, todo_mid3d, todo_inner
+        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)"""
 
@@ -33,7 +41,9 @@ class PostStepNull:  # pylint: disable=too-few-public-methods
     def __init__(self):
         pass
 
-    def call(self, psi, step, index):  # 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)"""
 
 
@@ -44,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
         """think of it as a `__call__` method (which Numba does not allow)"""
 
 

PyMPDATA/stepper.py

@@ -193,7 +193,14 @@ def step(
         for step in range(n_steps):
             for index, advectee in enumerate(advectees):
                 ante_step.call(
-                    advectees, advector, step, index, todo_outer, todo_mid3d, todo_inner
+                    traversals_data,
+                    advectees,
+                    advector,
+                    step,
+                    index,
+                    todo_outer,
+                    todo_mid3d,
+                    todo_inner,
                 )
                 if non_zero_mu_coeff:
                     advector_orig = advector
@@ -261,11 +268,12 @@ def step(
                                 traversals_data, flux, advectee, advector_nonos
                             )
                     upwind(traversals_data, advectee, flux, g_factor)
-                    post_iter.call(flux.field, g_factor.field, step, iteration)
+                    post_iter.call(
+                        traversals_data, flux.field, g_factor.field, step, iteration
+                    )
                 if non_zero_mu_coeff:
                     advector = advector_orig
-
-                post_step.call(advectees, step, index)
+                post_step.call(traversals_data, advectees, step, index)
         return (clock() - time) / n_steps if n_steps > 0 else np.nan
 
     return step, traversals

examples/PyMPDATA_examples/Jarecka_et_al_2015/simulation.py

@@ -100,8 +100,8 @@ def apply(traversals_data, momentum_x, momentum_y, advector):
 
     return apply
 
+
 def make_hooks(*, traversals, options, grid_step, time_step):
-    traversals_data = traversals.data
 
     divide_or_zero = make_divide_or_zero(options, traversals)
     interpolate = make_interpolate(options, traversals)
@@ -115,6 +115,7 @@ def __init__(self):
 
         def call(
             self,
+            traversals_data,
             advectees,
             advector,
             step,
@@ -145,7 +146,8 @@ def call(
     class PostStep:
         def __init__(self):
             pass
-        def call(self, advectees, step, index):
+
+        def call(self, traversals_data, advectees, step, index):
             if index == 0:
                 pass
             if index == 1:
@@ -189,7 +191,7 @@ def __init__(self, settings):
             traversals=stepper.traversals,
             options=settings.options,
             grid_step=(s.dx, None, s.dy),
-            time_step=s.dt
+            time_step=s.dt,
         )
 
         self.solver = Solver(stepper, self.advectees, self.advector)

scenarios_mpi/shallow_water.py

@@ -37,7 +37,7 @@ def initial_condition(x, y, lx, ly):
             """returns advectee array for a given grid indices"""
             # pylint: disable=invalid-name
             A = 1 / lx / ly
-            h = A * (1 - (x / lx) ** 2 - (y / ly) ** 2) * 6.25
+            h = A * (1 - (x / lx) ** 2 - (y / ly) ** 2)
             return np.where(h > 0, h, 0)
 
         # pylint: disable=invalid-name
@@ -124,7 +124,7 @@ def data(self, key):
 
     def _solver_advance(self, n_steps):
         for _ in range(n_steps):
-            self.solvers.advance(1, ante_step=self.ante_step)
+            self.solvers.advance(1, ante_step=self.ante_step, post_step=self.post_step)
         return -1
 
     @staticmethod

tests_mpi/contract_tests/test_single_vs_multi_node.py

@@ -22,9 +22,9 @@
 
 OPTIONS_KWARGS = (
     {"n_iters": 1},
-    # {"n_iters": 2, "third_order_terms": True},
-    # {"n_iters": 2, "nonoscillatory": True, "infinite_gauge": True},
-    # {"n_iters": 3},
+    {"n_iters": 2, "third_order_terms": True},
+    {"n_iters": 2, "nonoscillatory": True, "infinite_gauge": True},
+    {"n_iters": 3},
 )
 
 COURANT_FIELD_MULTIPLIER = (
@@ -37,7 +37,7 @@
 
 CARTESIAN_OUTPUT_STEPS = range(0, 24, 2)
 
-SHALLOW_WATER_OUTPUT_STEPS = range(0, 2, 1)
+SHALLOW_WATER_OUTPUT_STEPS = range(0, 48, 4)
 
 SPHERICAL_OUTPUT_STEPS = range(0, 2000, 100)
 
@@ -103,7 +103,9 @@ def test_single_vs_multi_node(  # pylint: disable=too-many-arguments,too-many-br
         pytest.skip("threading requires Numba JIT to be enabled")
 
     if scenario_class is ShallowWaterScenario and (
-        options_kwargs["n_iters"] == 3 or options_kwargs.get("third_order_terms", False)
+        options_kwargs["n_iters"] == 3
+        or options_kwargs.get("third_order_terms", False)
+        or mpi_dim == INNER
     ):
         pytest.skip("Unsupported method for simulation")
     # pylint: disable=too-many-boolean-expressions