feat: added masked option to asarray

Author: mrava87

examples/plot_distributed_array.py

@@ -140,6 +140,11 @@
                                 mask=mask)
 x[:] = (MPI.COMM_WORLD.Get_rank() % subsize + 1.) * np.ones(local_shape)
 xloc = x.asarray()
+xloc1 = x.asarray(masked=True)
+
+if rank == 0:
+    print('xloc (with repeated portions)', xloc)
+    print('xloc (only effective portions):', xloc1)
 
 # Dot product
 dot = x.dot(x)

pylops_mpi/DistributedArray.py

@@ -359,10 +359,15 @@ def sub_comm(self):
         """
         return self._sub_comm
 
-    def asarray(self):
+    def asarray(self, masked: bool = False):
         """Global view of the array
 
-        Gather all the local arrays
+        Gather all the local arrays from base communicator or subcommunicator
+
+        Parameters
+        ----------
+        masked : :obj:`bool`
+            Return local arrays of the subcommunicator (`True`) or base communicator (`False`).
 
         Returns
         -------
@@ -375,10 +380,14 @@ def asarray(self):
             return self.local_array
 
         if deps.nccl_enabled and getattr(self, "base_comm_nccl"):
-            return nccl_asarray(self.base_comm_nccl, self.local_array, self.local_shapes, self.axis)
+            return nccl_asarray(self.sub_comm if masked else self.base_comm,
+                                self.local_array, self.local_shapes, self.axis)
         else:
             # Gather all the local arrays and apply concatenation.
-            final_array = self._allgather(self.local_array)
+            if masked:
+                final_array = self._allgather_subcomm(self.local_array)
+            else:
+                final_array = self._allgather(self.local_array)
             return np.concatenate(final_array, axis=self.axis)
 
     @classmethod
@@ -503,6 +512,17 @@ def _allgather(self, send_buf, recv_buf=None):
             self.base_comm.Allgather(send_buf, recv_buf)
             return recv_buf
 
+    def _allgather_subcomm(self, send_buf, recv_buf=None):
+        """Allgather operation with subcommunicator
+        """
+        if deps.nccl_enabled and getattr(self, "base_comm_nccl"):
+            return nccl_allgather(self.sub_comm, send_buf, recv_buf)
+        else:
+            if recv_buf is None:
+                return self.sub_comm.allgather(send_buf)
+            self.sub_comm.Allgather(send_buf, recv_buf)
+            return recv_buf
+
     def __neg__(self):
         arr = DistributedArray(global_shape=self.global_shape,
                                base_comm=self.base_comm,

tests/test_distributedarray.py

@@ -194,12 +194,48 @@ def test_distributed_norm(par):
     assert_allclose(arr.norm(), np.linalg.norm(par['x'].flatten()), rtol=1e-13)
 
 
+@pytest.mark.mpi(min_size=2)
+@pytest.mark.parametrize("par", [(par6), (par8)])
+def test_distributed_masked(par):
+    """Test Asarray with masked array"""
+    # Number of subcommunicators
+    if MPI.COMM_WORLD.Get_size() % 2 == 0:
+        nsub = 2
+    elif MPI.COMM_WORLD.Get_size() % 3 == 0:
+        nsub = 3
+    else:
+        pass
+    subsize = max(1, MPI.COMM_WORLD.Get_size() // nsub)
+    mask = np.repeat(np.arange(nsub), subsize)
+
+    # Replicate x as required in masked arrays
+    x = par['x']
+    if par['axis'] != 0:
+        x = np.swapaxes(x, par['axis'], 0)
+    for isub in range(1, nsub):
+        x[(x.shape[0] // nsub) * isub:(x.shape[0] // nsub) * (isub + 1)] = x[:x.shape[0] // nsub]
+    if par['axis'] != 0:
+        x = np.swapaxes(x, 0, par['axis'])
+
+    arr = DistributedArray.to_dist(x=x, partition=par['partition'], mask=mask, axis=par['axis'])
+
+    # Global view
+    xloc = arr.asarray()
+    assert xloc.shape == x.shape
+
+    # Global masked view
+    xmaskedloc = arr.asarray(masked=True)
+    xmasked_shape = list(x.shape)
+    xmasked_shape[par['axis']] = int(xmasked_shape[par['axis']] // nsub)
+    assert xmaskedloc.shape == tuple(xmasked_shape)
+
+
 @pytest.mark.mpi(min_size=2)
 @pytest.mark.parametrize("par1, par2", [(par6, par7), (par6b, par7b),
                                         (par8, par9), (par8b, par9b)])
 def test_distributed_maskeddot(par1, par2):
     """Test Distributed Dot product with masked array"""
-    # number of subcommunicators
+    # Number of subcommunicators
     if MPI.COMM_WORLD.Get_size() % 2 == 0:
         nsub = 2
     elif MPI.COMM_WORLD.Get_size() % 3 == 0:
@@ -208,7 +244,7 @@ def test_distributed_maskeddot(par1, par2):
         pass
     subsize = max(1, MPI.COMM_WORLD.Get_size() // nsub)
     mask = np.repeat(np.arange(nsub), subsize)
-    print('subsize, mask', subsize, mask)
+
     # Replicate x1 and x2 as required in masked arrays
     x1, x2 = par1['x'], par2['x']
     if par1['axis'] != 0:
@@ -234,7 +270,7 @@ def test_distributed_maskeddot(par1, par2):
                                  (par8), (par8b), (par9), (par9b)])
 def test_distributed_maskednorm(par):
     """Test Distributed numpy.linalg.norm method with masked array"""
-    # number of subcommunicators
+    # Number of subcommunicators
     if MPI.COMM_WORLD.Get_size() % 2 == 0:
         nsub = 2
     elif MPI.COMM_WORLD.Get_size() % 3 == 0: