diff --git a/README.md b/README.md
index 18f9708..5af6b2b 100644
--- 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)
- 
+ 
-### 2 workers
+### 2 workers (MPI_DIM = 0, n_threads = 1)
- 
- 
+ 
+ 
### 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)
- 
+ 
-### 2 workers (MPI_DIM = 0)
+### 2 workers (MPI_DIM = 0, n_threads = 3)
- 
- 
+ 
+ 
-### 2 workers (MPI_DIM = -1)
+### 2 workers (MPI_DIM = -1, n_threads = 3)
- 
- 
+ 
+ 
-### 3 workers (MPI_DIM = 0)
+### 3 workers (MPI_DIM = 0, n_threads = 3)
- 
- 
- 
+ 
+ 
+ 
-### 3 workers (MPI_DIM = -1)
+### 3 workers (MPI_DIM = -1, n_threads = 3)
- 
- 
- 
+ 
+ 
+ 
## Package architecture
diff --git a/scenarios/cartesian.py b/scenarios/cartesian.py
index cdeab96..7aaa578 100644
--- 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
index 7264cd3..87de896 100644
--- 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
index 337656e..77410b3 100644
--- 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