Replace MPI.COMM_WORLD with self.base_comm, Fix dottest

Author: rohanbabbar04

pylops_mpi/DistributedArray.py

@@ -256,7 +256,7 @@ def mask(self):
 
         Returns
         -------
-        engine : :obj:`list`
+        mask : :obj:`list`
         """
         return self._mask
 
@@ -288,11 +288,7 @@ def rank(self):
         -------
         rank : :obj:`int`
         """
-        # cp.cuda.Device().id will give local rank
-        # It works ok in the single-node multi-gpu environment.
-        # But in multi-node environment, the function will break.
-        # So we have to use MPI.COMM_WORLD() in both cases of base_comm (MPI and NCCL)
-        return MPI.COMM_WORLD.Get_rank()
+        return self.base_comm.Get_rank()
 
     @property
     def size(self):
@@ -303,7 +299,7 @@ def size(self):
         -------
         size : :obj:`int`
         """
-        return MPI.COMM_WORLD.Get_size()
+        return self.base_comm.Get_size()
 
     @property
     def axis(self):
@@ -812,8 +808,8 @@ def __init__(self, distarrays: List, base_comm: MPI.Comm = MPI.COMM_WORLD):
         self.distarrays = distarrays
         self.narrays = len(distarrays)
         self.base_comm = base_comm
-        self.rank = MPI.COMM_WORLD.Get_rank()
-        self.size = MPI.COMM_WORLD.Get_size()
+        self.rank = self.base_comm.Get_rank()
+        self.size = self.base_comm.Get_size()
 
     def __getitem__(self, index):
         return self.distarrays[index]

pylops_mpi/utils/__init__.py

@@ -1,5 +1,4 @@
 # isort: skip_file
 
-# currently dottest create circular dependency with DistributedArray.py
-# from .dottest import *
+from .dottest import *
 from .deps import *

pylops_mpi/utils/dottest.py

@@ -4,7 +4,7 @@
 
 import numpy as np
 
-from pylops_mpi.DistributedArray import DistributedArray
+from pylops_mpi import DistributedArray
 from pylops.utils.backend import to_numpy