Fix method plot_2d_decomposition of TensorFemSpace

.github/workflows/testing.yml

@@ -195,7 +195,7 @@ jobs:
       - name: Run coverage tests on macOS
         if: matrix.os == 'macos-14'
         working-directory: ./pytest
-        run: >- 
+        run: >-
           python -m pytest -n auto
           --cov psydac
           --cov-config $GITHUB_WORKSPACE/pyproject.toml

examples/poisson_2d_mapping.py

@@ -781,7 +781,8 @@ def main(*, test_case, ncells, degree, nquads,
     ##########
 
     # Plot domain decomposition (master only)
-    V.plot_2d_decomposition(model.mapping, refine=N)
+    fig = V.plot_2d_decomposition(mapping, refine=N)
+    fig.show()
 
     # Perform other visualization using master or all processes
     if not distribute_viz:

psydac/fem/tensor.py

@@ -1184,24 +1184,77 @@ def set_refined_space(self, ncells, new_space):
         self._refined_space[tuple(ncells)] = new_space
 
     # ...
-    def plot_2d_decomposition(self, mapping=None, refine=10):
+    def plot_2d_decomposition(self, mapping=None, *, refine=10, fig=None, ax=None, mpi_root=0):
+        """
+        Plot decomposition of 2D TensorFemSpace w/ mapping to 2D physical space
+
+        Plot the domain decomposition across MPI processes of a 2D
+        TensorFemSpace with a mapping between 2D logical and 2D physical spaces.
+        This function must be called collectively, and only the root process will make
+        the plot. On non-root processes the arguments `fig` and `ax` must be None.
+
+        Parameters
+        ----------
+        mapping : BasicCallableMapping
+            Mapping from (eta1, eta2) to (x1, x2).
+
+        refine : int, default=10
+            Cell refinement along the logical dimensions eta1 and eta2.
+
+        fig : plt.Figure, optional
+            Figure where the plot should be made. Must be None on non-root processes.
 
+        ax : plt.Axes, optional
+            Axes where the plot should be made. Must be None on non-root processes.
+
+        mpi_root: int, default=0
+            The rank of the MPI root process which should create the plot.
+
+        Returns
+        -------
+        plt.Figure
+            Figure where the plot was made. Coincides with `fig` if provided.
+        """
         import matplotlib.pyplot as plt
         from matplotlib.patches  import Polygon, Patch
+        from sympde.topology.mapping import BasicCallableMapping
         from psydac.utilities.utils import refine_array_1d
 
+        # Sanity check
+        assert self.ldim == 2, "Function only works in 2D"
+
+        # Check mapping
         if mapping is None:
             mapping = lambda eta: eta
         else:
-            assert mapping.ldim == self.ldim == 2
-            assert mapping.pdim == self.ldim == 2
+            assert isinstance(mapping, BasicCallableMapping)
+            assert mapping.ldim == 2, "Domain of argument `mapping` must be 2D"
+            assert mapping.pdim == 2, "Codomain of argument `mapping` must be 2D"
 
-        assert refine >= 1
-        N = int(refine)
-        V1, V2 = self.spaces
+        # Check refine argument
+        assert isinstance(refine, int), f"Argument `refine` must be int, got {type(refine)} instead"
+        assert refine >= 1, f"Argument `refine` must be >= 1, got {refine} instead"
 
+        # Extract information about MPI communicator
         mpi_comm = self.coeff_space.cart.comm
         mpi_rank = mpi_comm.rank
+        mpi_size = mpi_comm.size
+
+        # Check mpi_root argument
+        assert isinstance(mpi_root, int), f"Argument `mpi_root` must be int, got {type(mpi_root)} instead"
+        assert mpi_root >= 0, f"Argument `mpi_root` must be >= 0, got {mpi_root} instead"
+        assert mpi_root < mpi_size, f"Argument `mpi_root` must be smaller than communicator size ({mpi_size}), got {mpi_root} instead"
+
+        # Check fig and ax arguments
+        if mpi_rank == mpi_root:
+            assert isinstance(fig, plt.Figure) or fig is None, f"Argument `fig` must be matplotlib Figure, got {type(fig)} instead"
+            assert isinstance(ax, plt.Axes) or ax is None, f"Argument `ax` must be matplotlib Axes, got {type(ax)} instead"
+        else:
+            assert fig is None, f"Argument `fig` must be None on non-root process with rank {mpi_rank}"
+            assert ax is None, f"Argument `ax` must be None on non-root process with rank {mpi_rank}"
+
+        N = refine
+        V1, V2 = self.spaces
 
         # Local grid, refined
         [sk1, sk2], [ek1, ek2] = self.local_domain
@@ -1218,23 +1271,62 @@ def plot_2d_decomposition(self, mapping=None, refine=10):
         poly = Polygon(xy, edgecolor='None')
 
         # Gather polygons on master process
-        polys = mpi_comm.gather(poly)
+        polys = mpi_comm.gather(poly, root=mpi_root)
+
+        # Gather (s1, s2, e1, e2) on root
+        if mpi_rank == mpi_root:
+            s1_all = np.empty(mpi_size, dtype=int)
+            s2_all = np.empty(mpi_size, dtype=int)
+            e1_all = np.empty(mpi_size, dtype=int)
+            e2_all = np.empty(mpi_size, dtype=int)
+        else:
+            s1_all = None
+            s2_all = None
+            e1_all = None
+            e2_all = None
+
+        mpi_comm.Gather(sk1 * N, s1_all, root=mpi_root)
+        mpi_comm.Gather(sk2 * N, s2_all, root=mpi_root)
+        mpi_comm.Gather((ek1 + 1) * N, e1_all, root=mpi_root)
+        mpi_comm.Gather((ek2 + 1) * N, e2_all, root=mpi_root)
+
+        # Gather pcoords on root
+        # TODO: use Gatherv, and NumPy arrays as buffers
+        gathered_pcoords = mpi_comm.gather(pcoords, root=mpi_root)
 
         #-------------------------------
         # Non-master processes stop here
-        if mpi_rank != 0:
+        if mpi_rank != mpi_root:
             return
         #-------------------------------
 
-        # Global grid, refined
-        eta1    = refine_array_1d(V1.breaks, N)
-        eta2    = refine_array_1d(V2.breaks, N)
-        pcoords = np.array([[mapping(e1, e2) for e2 in eta2] for e1 in eta1])
-        xx      = pcoords[:, :, 0]
-        yy      = pcoords[:, :, 1]
+        # Reconstruct global grid (refined) on root process
+        global_shape   = ((V1.breaks.size - 1) * N + 1,
+                          (V2.breaks.size - 1) * N + 1,
+                          2)
+        pcoords_global = np.empty(global_shape)
+
+        for rank in range(mpi_comm.size):
+            s1 = s1_all[rank]
+            e1 = e1_all[rank]
+            s2 = s2_all[rank]
+            e2 = e2_all[rank]
+            pcoords_global[s1:e1+1, s2:e2+1, :] = gathered_pcoords[rank]
+
+        xx = pcoords_global[:, :, 0]
+        yy = pcoords_global[:, :, 1]
+
+        # If fig or ax are given, get one from the other. Otherwise create new ones
+        if fig and ax:
+            assert ax in fig.axes, "Argument `ax` must be in `fig.axes`"
+        elif fig:
+            ax = fig.gca()
+        elif ax:
+            fig = ax.figure
+        else:
+            fig, ax = plt.subplots(1, 1)
 
         # Plot decomposed domain
-        fig, ax = plt.subplots(1, 1)
         colors  = itertools.cycle(plt.rcParams['axes.prop_cycle'].by_key()['color'])
         handles = []
         for i, (poly, color) in enumerate(zip(polys, colors)):
@@ -1253,7 +1345,8 @@ def plot_2d_decomposition(self, mapping=None, refine=10):
         ax.set_aspect('equal')
         ax.legend(handles=handles, bbox_to_anchor=(1.05, 1), loc=2)
         fig.tight_layout()
-        fig.show()
+
+        return fig
 
     # ...
     def __str__(self):