add visualisation of threading domain decomposition to cartesian scen…

…ario quick_look (incl. README animations) (#110)
open-atmos · Mar 29, 2024 · 38e0706 · 38e0706
1 parent f0c3589
commit 38e0706
Show file tree

Hide file tree

Showing 4 changed files with 59 additions and 30 deletions.
diff --git a/README.md b/README.md
@@ -43,15 +43,15 @@ The inner dimension uses the [`MPIPolar`](https://open-atmos.github.io/PyMPDATA-
 Note that the spherical animations below depict simulations without MPDATA corrective iterations,
   i.e. only plain first-order upwind scheme is used (FIX ME).
 
-### 1 worker
+### 1 worker (n_threads = 1)
 <p align="middle">
-  <img src="https://github.com/open-atmos/PyMPDATA-MPI/releases/download/latest-generated-plots/n_iters.1_rank_0_size_1_c_field_.0.5.0.25._mpi_dim_0-SphericalScenario-anim.gif" width="49%" /> 
+  <img src="https://github.com/open-atmos/PyMPDATA-MPI/releases/download/latest-generated-plots/n_iters.1_rank_0_size_1_c_field_.0.5.0.25._mpi_dim_0_n_threads_1-SphericalScenario-anim.gif" width="49%" /> 
 </p>
 
-### 2 workers
+### 2 workers (MPI_DIM = 0, n_threads = 1)
 <p align="middle">
-  <img src="https://github.com/open-atmos/PyMPDATA-MPI/releases/download/latest-generated-plots/n_iters.1_rank_1_size_2_c_field_.0.5.0.25._mpi_dim_0-SphericalScenario-anim.gif" width="49%" /> 
-  <img src="https://github.com/open-atmos/PyMPDATA-MPI/releases/download/latest-generated-plots/n_iters.1_rank_0_size_2_c_field_.0.5.0.25._mpi_dim_0-SphericalScenario-anim.gif" width="49%" />
+  <img src="https://github.com/open-atmos/PyMPDATA-MPI/releases/download/latest-generated-plots/n_iters.1_rank_1_size_2_c_field_.0.5.0.25._mpi_dim_0_n_threads_1-SphericalScenario-anim.gif" width="49%" /> 
+  <img src="https://github.com/open-atmos/PyMPDATA-MPI/releases/download/latest-generated-plots/n_iters.1_rank_0_size_2_c_field_.0.5.0.25._mpi_dim_0_n_threads_1-SphericalScenario-anim.gif" width="49%" />
 </p>
 
 ### Cartesian scenario (2D)
@@ -62,38 +62,39 @@ MPI (Message Passing Interface) is used
   for handling data transfers and synchronisation with the domain decomposition
   across MPI workers done in either inner or in the outer dimension (user setting).
 Multi-threading (using, e.g., OpenMP via Numba) is used for shared-memory parallelisation 
-  within subdomains with further subdomain split across the inner dimension (PyMPDATA logic).
+  within subdomains (indicated by dotted lines in the animations below) with threading subdomain
+  split done across the inner dimension (internal PyMPDATA logic).
 In this example, two corrective MPDATA iterations are employed.
 
-### 1 worker
+### 1 worker (n_threads=3)
 <p align="middle">
-  <img src="https://github.com/open-atmos/PyMPDATA-MPI/releases/download/latest-generated-plots/n_iters.3_rank_0_size_1_c_field_.0.5.0.25._mpi_dim_0-CartesianScenario-anim.gif" width="49%" /> 
+  <img src="https://github.com/open-atmos/PyMPDATA-MPI/releases/download/latest-generated-plots/n_iters.3_rank_0_size_1_c_field_.0.5.0.25._mpi_dim_0_n_threads_3-CartesianScenario-anim.gif" width="49%" /> 
 </p>
 
-### 2 workers (MPI_DIM = 0)
+### 2 workers (MPI_DIM = 0, n_threads = 3)
 <p align="middle">
-  <img src="https://github.com/open-atmos/PyMPDATA-MPI/releases/download/latest-generated-plots/n_iters.3_rank_0_size_2_c_field_.0.5.0.25._mpi_dim_0-CartesianScenario-anim.gif" width="49%" />
-  <img src="https://github.com/open-atmos/PyMPDATA-MPI/releases/download/latest-generated-plots/n_iters.3_rank_1_size_2_c_field_.0.5.0.25._mpi_dim_0-CartesianScenario-anim.gif" width="49%" /> 
+  <img src="https://github.com/open-atmos/PyMPDATA-MPI/releases/download/latest-generated-plots/n_iters.3_rank_0_size_2_c_field_.0.5.0.25._mpi_dim_0_n_threads_3-CartesianScenario-anim.gif" width="49%" />
+  <img src="https://github.com/open-atmos/PyMPDATA-MPI/releases/download/latest-generated-plots/n_iters.3_rank_1_size_2_c_field_.0.5.0.25._mpi_dim_0_n_threads_3-CartesianScenario-anim.gif" width="49%" /> 
 </p>
 
-### 2 workers (MPI_DIM = -1)
+### 2 workers (MPI_DIM = -1, n_threads = 3)
 <p align="middle">
-  <img src="https://github.com/open-atmos/PyMPDATA-MPI/releases/download/latest-generated-plots/n_iters.3_rank_0_size_2_c_field_.0.5.0.25._mpi_dim_-1-CartesianScenario-anim.gif" width="49%" />
-  <img src="https://github.com/open-atmos/PyMPDATA-MPI/releases/download/latest-generated-plots/n_iters.3_rank_1_size_2_c_field_.0.5.0.25._mpi_dim_-1-CartesianScenario-anim.gif" width="49%" /> 
+  <img src="https://github.com/open-atmos/PyMPDATA-MPI/releases/download/latest-generated-plots/n_iters.3_rank_0_size_2_c_field_.0.5.0.25._mpi_dim_-1_n_threads_3-CartesianScenario-anim.gif" width="49%" />
+  <img src="https://github.com/open-atmos/PyMPDATA-MPI/releases/download/latest-generated-plots/n_iters.3_rank_1_size_2_c_field_.0.5.0.25._mpi_dim_-1_n_threads_3-CartesianScenario-anim.gif" width="49%" /> 
 </p>
 
-### 3 workers (MPI_DIM = 0)
+### 3 workers (MPI_DIM = 0, n_threads = 3)
 <p align="middle">
-  <img src="https://github.com/open-atmos/PyMPDATA-MPI/releases/download/latest-generated-plots/n_iters.3_rank_0_size_3_c_field_.0.5.0.25._mpi_dim_0-CartesianScenario-anim.gif" width="32%" />
-  <img src="https://github.com/open-atmos/PyMPDATA-MPI/releases/download/latest-generated-plots/n_iters.3_rank_1_size_3_c_field_.0.5.0.25._mpi_dim_0-CartesianScenario-anim.gif" width="32%" />
-  <img src="https://github.com/open-atmos/PyMPDATA-MPI/releases/download/latest-generated-plots/n_iters.3_rank_2_size_3_c_field_.0.5.0.25._mpi_dim_0-CartesianScenario-anim.gif" width="32%" />
+  <img src="https://github.com/open-atmos/PyMPDATA-MPI/releases/download/latest-generated-plots/n_iters.3_rank_0_size_3_c_field_.0.5.0.25._mpi_dim_0_n_threads_3-CartesianScenario-anim.gif" width="32%" />
+  <img src="https://github.com/open-atmos/PyMPDATA-MPI/releases/download/latest-generated-plots/n_iters.3_rank_1_size_3_c_field_.0.5.0.25._mpi_dim_0_n_threads_3-CartesianScenario-anim.gif" width="32%" />
+  <img src="https://github.com/open-atmos/PyMPDATA-MPI/releases/download/latest-generated-plots/n_iters.3_rank_2_size_3_c_field_.0.5.0.25._mpi_dim_0_n_threads_3-CartesianScenario-anim.gif" width="32%" />
 </p>
 
-### 3 workers (MPI_DIM = -1)
+### 3 workers (MPI_DIM = -1, n_threads = 3)
 <p align="middle">
-  <img src="https://github.com/open-atmos/PyMPDATA-MPI/releases/download/latest-generated-plots/n_iters.3_rank_0_size_3_c_field_.0.5.0.25._mpi_dim_-1-CartesianScenario-anim.gif" width="32%" />
-  <img src="https://github.com/open-atmos/PyMPDATA-MPI/releases/download/latest-generated-plots/n_iters.3_rank_1_size_3_c_field_.0.5.0.25._mpi_dim_-1-CartesianScenario-anim.gif" width="32%" />
-  <img src="https://github.com/open-atmos/PyMPDATA-MPI/releases/download/latest-generated-plots/n_iters.3_rank_2_size_3_c_field_.0.5.0.25._mpi_dim_-1-CartesianScenario-anim.gif" width="32%" />
+  <img src="https://github.com/open-atmos/PyMPDATA-MPI/releases/download/latest-generated-plots/n_iters.3_rank_0_size_3_c_field_.0.5.0.25._mpi_dim_-1_n_threads_3-CartesianScenario-anim.gif" width="32%" />
+  <img src="https://github.com/open-atmos/PyMPDATA-MPI/releases/download/latest-generated-plots/n_iters.3_rank_1_size_3_c_field_.0.5.0.25._mpi_dim_-1_n_threads_3-CartesianScenario-anim.gif" width="32%" />
+  <img src="https://github.com/open-atmos/PyMPDATA-MPI/releases/download/latest-generated-plots/n_iters.3_rank_2_size_3_c_field_.0.5.0.25._mpi_dim_-1_n_threads_3-CartesianScenario-anim.gif" width="32%" />
 </p>
 
 ## Package architecture

diff --git a/scenarios/cartesian.py b/scenarios/cartesian.py
@@ -1,15 +1,19 @@
 """ 2D constant-advector carthesian example """
 
+import numba
 import numpy as np
 from matplotlib import pyplot
 from PyMPDATA import ScalarField, Stepper, VectorField
 from PyMPDATA.boundary_conditions import Periodic
+from PyMPDATA.impl.domain_decomposition import make_subdomain
 from PyMPDATA.impl.enumerations import INNER, OUTER
 
 from PyMPDATA_MPI.domain_decomposition import mpi_indices
 from PyMPDATA_MPI.mpi_periodic import MPIPeriodic
 from scenarios._scenario import _Scenario
 
+subdomain = make_subdomain(jit_flags={})
+
 
 class CartesianScenario(_Scenario):
     """class representation of a test case from
@@ -85,9 +89,9 @@ def initial_condition(xi, yi, grid):
         return psi
 
     @staticmethod
-    def quick_look(psi, zlim=(-1, 1), norm=None):
+    def quick_look(psi, n_threads, zlim=(-1, 1), norm=None):
         """plots the passed advectee field"""
-        # pylint: disable=invalid-name
+        # pylint: disable=invalid-name,too-many-locals
         xi, yi = np.indices(psi.shape)
         _, ax = pyplot.subplots(subplot_kw={"projection": "3d"})
         pyplot.gca().plot_wireframe(xi + 0.5, yi + 0.5, psi, color="red", linewidth=0.5)
@@ -101,6 +105,24 @@ def quick_look(psi, zlim=(-1, 1), norm=None):
         ax.set_xlabel("x/dx")
         ax.set_ylabel("y/dy")
         ax.set_proj_type("ortho")
+
+        if n_threads > 1 and not numba.config.DISABLE_JIT:  # pylint: disable=no-member
+            first_i_with_finite_values = -1
+            for i in range(psi.shape[0]):
+                if sum(np.isfinite(psi[i, :])) > 0:
+                    first_i_with_finite_values = i
+            finite_slice = np.isfinite(psi[first_i_with_finite_values, :])
+            span = sum(finite_slice)
+            assert span != 0
+            zero = np.argmax(finite_slice > 0)
+            for i in range(n_threads):
+                start, stop = subdomain(span, i, n_threads)
+                kwargs = {"zs": -1, "zdir": "z", "color": "black", "linestyle": ":"}
+                x = [0, psi.shape[0] - 1]
+                ax.plot(x, [zero + start] * 2, **kwargs)
+                if i == n_threads - 1:
+                    ax.plot(x, [zero + stop] * 2, **kwargs)
+
         cnt = ax.contourf(
             xi + 0.5,
             yi + 0.5,
@@ -109,6 +131,8 @@ def quick_look(psi, zlim=(-1, 1), norm=None):
             offset=-1,
             norm=norm,
             levels=np.linspace(*zlim, 11),
+            alpha=0.75,
         )
         cbar = pyplot.colorbar(cnt, pad=0.1, aspect=10, fraction=0.04)
+
         return cbar.norm
diff --git a/scenarios/spherical.py b/scenarios/spherical.py
@@ -194,7 +194,7 @@ def __init__(  # pylint: disable=too-many-arguments
             g_factor=g_factor,
         )
 
-    def quick_look(self, state):
+    def quick_look(self, state, _):
         """plots the passed advectee field in spherical geometry"""
         # pylint: disable=invalid-name
         theta = np.linspace(0, 1, self.settings.nlat + 1, endpoint=True) * np.pi

diff --git a/tests/local/contract_tests/test_single_vs_multi_node.py b/tests/local/contract_tests/test_single_vs_multi_node.py
@@ -82,7 +82,7 @@ def test_single_vs_multi_node(  # pylint: disable=too-many-arguments,too-many-br
     if n_threads > 1 and numba.config.DISABLE_JIT:  # pylint: disable=no-member
         pytest.skip("threading requires Numba JIT to be enabled")
 
-    plot = True and (
+    plot = (
         "CI_PLOTS_PATH" in os.environ
         and courant_field_multiplier == COURANT_FIELD_MULTIPLIER[0]
         and (
@@ -127,8 +127,12 @@ def test_single_vs_multi_node(  # pylint: disable=too-many-arguments,too-many-br
                 Path(os.environ["CI_PLOTS_PATH"])
                 / Path(scenario_class.__name__)
                 / Path(
-                    f"{options_str}_rank_{mpi.rank()}_size_{truncated_size}"
-                    f"_c_field_{courant_str}_mpi_dim_{mpi_dim}"
+                    f"{options_str}"
+                    f"_rank_{mpi.rank()}"
+                    f"_size_{truncated_size}"
+                    f"_c_field_{courant_str}"
+                    f"_mpi_dim_{mpi_dim}"
+                    f"_n_threads={n_threads}"
                 )
             )
             shutil.rmtree(plot_path, ignore_errors=True)
@@ -169,11 +173,11 @@ def test_single_vs_multi_node(  # pylint: disable=too-many-arguments,too-many-br
                         tmp[tuple(ranges)] = dataset[
                             tuple([*ranges, slice(i, i + 1)])
                         ].squeeze()
-                        simulation.quick_look(tmp)
+                        simulation.quick_look(tmp, n_threads)
 
                         filename = f"step={i:04d}.svg"
                         pyplot.savefig(plot_path / filename)
-                        print("Saving figure")
+                        print(f"Saving figure {plot_path=} {filename=}")
                         pyplot.close()
 
     # assert