Fix dottest dependency and use base_comm instead of MPI.COMM_WORLD (#134)

* Replace MPI.COMM_WORLD with self.base_comm, Fix dottest

* Fix flake8 for local dev

* Minor change

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 = base_comm.Get_rank()
+        self.size = 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
 
 

setup.cfg

@@ -5,5 +5,5 @@ python_files = tests/*.py tests_nccl/*.py
 [flake8]
 ignore = E203, E501, W503, E402
 per-file-ignores =
-__init__.py: F401, F403, F405
+    __init__.py: F401, F403, F405
 max-line-length = 88

tests/test_fredholm.py

@@ -135,7 +135,7 @@ def test_Fredholm1(par):
     y_adj_dist = Fop_MPI.H @ y_dist
     y_adj = y_adj_dist.asarray()
     # Dot test
-    dottest(Fop_MPI, x, y_dist, par["nsl"] * par["nx"] * par["nz"],par["nsl"] * par["ny"] * par["nz"])
+    dottest(Fop_MPI, x, y_dist, par["nsl"] * par["nx"] * par["nz"], par["nsl"] * par["ny"] * par["nz"])
 
     if rank == 0:
         Fop = pylops.signalprocessing.Fredholm1(

tutorials/mdd.py

@@ -14,7 +14,6 @@
 """
 
 import numpy as np
-from scipy.signal import filtfilt
 from matplotlib import pyplot as plt
 from mpi4py import MPI