2D-distribution

.github/workflows/build.yml

@@ -14,8 +14,9 @@ jobs:
     strategy:
       matrix:
         os: [ubuntu-latest, macos-latest]
-        python-version: ['3.8', '3.9', '3.10', '3.11', '3.12']
+        python-version: ['3.9', '3.10', '3.11', '3.12']
         mpi: ['mpich', 'openmpi', 'intelmpi']
+        rank: ['2', '3', '4']
         exclude:
           - os: macos-latest
             mpi: 'intelmpi'
@@ -52,4 +53,4 @@ jobs:
       - name: Install pylops-mpi
         run: pip install .
       - name: Testing using pytest-mpi
-        run: mpiexec -n 2 pytest --with-mpi
+        run: mpiexec -n ${{ matrix.rank }} pytest --with-mpi

examples/plot_distributed_array.py

@@ -16,13 +16,18 @@
 plt.close("all")
 np.random.seed(42)
 
+# MPI parameters
+size = MPI.COMM_WORLD.Get_size()  # number of nodes
+rank = MPI.COMM_WORLD.Get_rank()  # rank of current node
+
+
 # Defining the global shape of the distributed array
 global_shape = (10, 5)
 
 ###############################################################################
-# Let's start by defining the
-# class with the input parameters ``global_shape``,
-# ``partition``, and ``axis``. Here's an example implementation of the class with ``axis=0``.
+# Let's start by defining the class with the input parameters ``global_shape``,
+# ``partition``, and ``axis``. Here's an example implementation of the class
+# with ``axis=0``.
 arr = pylops_mpi.DistributedArray(global_shape=global_shape,
                                   partition=pylops_mpi.Partition.SCATTER,
                                   axis=0)
@@ -72,6 +77,9 @@
 pylops_mpi.plot_local_arrays(arr2, "Distributed Array - 2", vmin=0, vmax=1)
 
 ###############################################################################
+# Let's move now to consider various operations that one can perform on
+# :py:class:`pylops_mpi.DistributedArray` objects.
+#
 # **Scaling** - Each process operates on its local portion of
 # the array and scales the corresponding elements by a given scalar.
 scale_arr = .5 * arr1
@@ -101,3 +109,95 @@
 # of the array and multiplies the corresponding elements together.
 mult_arr = arr1 * arr2
 pylops_mpi.plot_local_arrays(mult_arr, "Multiplication", vmin=0, vmax=1)
+
+###############################################################################
+# Finally, let's look at the case where parallelism could be applied over
+# multiple axes - and more specifically one belonging to the model/data and one
+# to the operator. This kind of "2D"-parallelism requires repeating parts of
+# the model/data over groups of ranks. However, when global operations such as
+# ``dot`` or ``norm`` are applied on a ``pylops_mpi.DistributedArray`` of
+# this kind, we need to ensure that the repeated portions to do all contribute
+# to the computation. This can be achieved via the ``mask`` input parameter:
+# a list of size equal to the number of ranks, whose elements contain the index
+# of the subgroup/subcommunicator (with partial arrays in different groups
+# are identical to each other).
+
+# Defining the local and global shape of the distributed array
+local_shape = 5
+global_shape = local_shape * size
+
+# Create mask
+nsub = 2
+subsize = max(1, size // nsub)
+mask = np.repeat(np.arange(size // subsize), subsize)
+if rank == 0:
+    print("1D masked arrays")
+    print(f"Mask: {mask}")
+
+# Create and fill the distributed array
+x = pylops_mpi.DistributedArray(global_shape=global_shape,
+                                partition=Partition.SCATTER,
+                                mask=mask)
+x[:] = (MPI.COMM_WORLD.Get_rank() % subsize + 1.) * np.ones(local_shape)
+xloc = x.asarray()
+
+# Dot product
+dot = x.dot(x)
+dotloc = np.dot(xloc[local_shape * subsize * (rank // subsize):local_shape * subsize * (rank // subsize + 1)],
+                xloc[local_shape * subsize * (rank // subsize):local_shape * subsize * (rank // subsize + 1)])
+print(f"Dot check (Rank {rank}): {np.allclose(dot, dotloc)}")
+
+# Norm
+norm = x.norm(ord=2)
+normloc = np.linalg.norm(xloc[local_shape * subsize * (rank // subsize):local_shape * subsize * (rank // subsize + 1)],
+                         ord=2)
+print(f"Norm check (Rank {rank}): {np.allclose(norm, normloc)}")
+
+###############################################################################
+# And with 2d-arrays distributed over axis=1
+extra_dim_shape = 2
+if rank == 0:
+    print("2D masked arrays (over axis=1)")
+
+# Create and fill the distributed array
+x = pylops_mpi.DistributedArray(global_shape=(extra_dim_shape, global_shape),
+                                partition=Partition.SCATTER,
+                                axis=1, mask=mask)
+x[:] = (MPI.COMM_WORLD.Get_rank() % subsize + 1.) * np.ones((extra_dim_shape, local_shape))
+xloc = x.asarray()
+
+# Dot product
+dot = x.dot(x)
+dotloc = np.dot(xloc[:, local_shape * subsize * (rank // subsize):local_shape * subsize * (rank // subsize + 1)].ravel(),
+                xloc[:, local_shape * subsize * (rank // subsize):local_shape * subsize * (rank // subsize + 1)].ravel())
+print(f"Dot check (Rank {rank}): {np.allclose(dot, dotloc)}")
+
+# Norm
+norm = x.norm(ord=2, axis=1)
+normloc = np.linalg.norm(xloc[:, local_shape * subsize * (rank // subsize):local_shape * subsize * (rank // subsize + 1)],
+                         ord=2, axis=1)
+print(f"Norm check (Rank {rank}): {np.allclose(norm, normloc)}")
+
+###############################################################################
+# And finally with 2d-arrays distributed over axis=0
+if rank == 0:
+    print("2D masked arrays (over axis=0)")
+
+# Create and fill the distributed array
+x = pylops_mpi.DistributedArray(global_shape=(global_shape, extra_dim_shape),
+                                partition=Partition.SCATTER,
+                                axis=0, mask=mask)
+x[:] = (MPI.COMM_WORLD.Get_rank() % subsize + 1.) * np.ones((local_shape, extra_dim_shape))
+xloc = x.asarray()
+
+# Dot product
+dot = x.dot(x)
+dotloc = np.dot(xloc[local_shape * subsize * (rank // subsize):local_shape * subsize * (rank // subsize + 1)].ravel(),
+                xloc[local_shape * subsize * (rank // subsize):local_shape * subsize * (rank // subsize + 1)].ravel())
+print(f"Dot check (Rank {rank}): {np.allclose(dot, dotloc)}")
+
+# Norm
+norm = x.norm(ord=2, axis=0)
+normloc = np.linalg.norm(xloc[local_shape * subsize * (rank // subsize):local_shape * subsize * (rank // subsize + 1)],
+                         ord=2, axis=0)
+print(f"Norm check (Rank {rank}): {np.allclose(norm, normloc)}")