Check for smaller values

Author: rohanbabbar04

.github/workflows/build.yml

@@ -43,4 +43,4 @@ jobs:
       - name: Install pylops-mpi
         run: pip install .
       - name: Testing using pytest-mpi
-        run: mpiexec -n ${{ matrix.rank }} pytest tests/ --with-mpi -v -x
+        run: mpiexec -n ${{ matrix.rank }} pytest tests/test_blockdiag.py --with-mpi -v -x

tests/test_blockdiag.py

@@ -11,10 +11,10 @@
 import pylops_mpi
 from pylops_mpi.utils.dottest import dottest
 
-par1 = {'ny': 101, 'nx': 101, 'dtype': np.float64}
-par1j = {'ny': 101, 'nx': 101, 'dtype': np.complex128}
-par2 = {'ny': 301, 'nx': 101, 'dtype': np.float64}
-par2j = {'ny': 301, 'nx': 101, 'dtype': np.complex128}
+par1 = {'ny': 31, 'nx': 31, 'dtype': np.float64}
+par1j = {'ny': 41, 'nx': 41, 'dtype': np.complex128}
+par2 = {'ny': 31, 'nx': 41, 'dtype': np.float64}
+par2j = {'ny': 31, 'nx': 41, 'dtype': np.complex128}
 
 np.random.seed(42)
 
@@ -57,49 +57,49 @@ def test_blockdiag(par):
         assert_allclose(y_rmat_mpi, y_rmat_np, rtol=1e-14)
 
 
-# @pytest.mark.mpi(min_size=2)
-# @pytest.mark.parametrize("par", [(par1), (par1j), (par2), (par2j)])
-# def test_stacked_blockdiag(par):
-#     """Tests for MPIStackedBlogDiag"""
-#     size = MPI.COMM_WORLD.Get_size()
-#     rank = MPI.COMM_WORLD.Get_rank()
-#     Op = pylops.MatrixMult(A=((rank + 1) * np.ones(shape=(par['ny'], par['nx']))).astype(par['dtype']))
-#     BDiag_MPI = pylops_mpi.MPIBlockDiag(ops=[Op, ])
-#     FirstDeriv_MPI = pylops_mpi.MPIFirstDerivative(dims=(par['ny'], par['nx']), dtype=par['dtype'])
-#     StackedBDiag_MPI = pylops_mpi.MPIStackedBlockDiag(ops=[BDiag_MPI, FirstDeriv_MPI])
-#
-#     dist1 = pylops_mpi.DistributedArray(global_shape=size * par['nx'], dtype=par['dtype'])
-#     dist1[:] = np.ones(dist1.local_shape, dtype=par['dtype'])
-#     dist2 = pylops_mpi.DistributedArray(global_shape=par['nx'] * par['ny'], dtype=par['dtype'])
-#     dist2[:] = np.ones(dist2.local_shape, dtype=par['dtype'])
-#     x = pylops_mpi.StackedDistributedArray(distarrays=[dist1, dist2])
-#     x_global = x.asarray()
-#
-#     dist1 = pylops_mpi.DistributedArray(global_shape=size * par['ny'], dtype=par['dtype'])
-#     dist1[:] = np.ones(dist1.local_shape, dtype=par['dtype'])
-#     dist2 = pylops_mpi.DistributedArray(global_shape=par['nx'] * par['ny'], dtype=par['dtype'])
-#     dist2[:] = np.ones(dist2.local_shape, dtype=par['dtype'])
-#     y = pylops_mpi.StackedDistributedArray(distarrays=[dist1, dist2])
-#     y_global = y.asarray()
-#
-#     # Forward
-#     x_mat = StackedBDiag_MPI @ x
-#     # Adjoint
-#     y_rmat = StackedBDiag_MPI.H @ y
-#     assert isinstance(x_mat, pylops_mpi.StackedDistributedArray)
-#     assert isinstance(y_rmat, pylops_mpi.StackedDistributedArray)
-#     # Dot test
-#     dottest(StackedBDiag_MPI, x, y, size * par['ny'] + par['nx'] * par['ny'], size * par['nx'] + par['nx'] * par['ny'])
-#
-#     x_mat_mpi = x_mat.asarray()
-#     y_rmat_mpi = y_rmat.asarray()
-#
-#     if rank == 0:
-#         ops = [pylops.MatrixMult((i + 1) * np.ones(shape=(par['ny'], par['nx'])).astype(par['dtype'])) for i in range(size)]
-#         BDiag = pylops.BlockDiag(ops=ops)
-#         FirstDeriv = pylops.FirstDerivative(dims=(par['ny'], par['nx']), axis=0, dtype=par['dtype'])
-#         BDiag_final = pylops.BlockDiag([BDiag, FirstDeriv])
-#         x_mat_np = BDiag_final @ x_global
-#         y_rmat_np = BDiag_final.H @ y_global
-#         assert_allclose(x_mat_mpi, x_mat_np, rtol=1e-14)
-#         assert_allclose(y_rmat_mpi, y_rmat_np, rtol=1e-14)
+@pytest.mark.mpi(min_size=2)
+@pytest.mark.parametrize("par", [(par1), (par1j), (par2), (par2j)])
+def test_stacked_blockdiag(par):
+    """Tests for MPIStackedBlogDiag"""
+    size = MPI.COMM_WORLD.Get_size()
+    rank = MPI.COMM_WORLD.Get_rank()
+    Op = pylops.MatrixMult(A=((rank + 1) * np.ones(shape=(par['ny'], par['nx']))).astype(par['dtype']))
+    BDiag_MPI = pylops_mpi.MPIBlockDiag(ops=[Op, ])
+    FirstDeriv_MPI = pylops_mpi.MPIFirstDerivative(dims=(par['ny'], par['nx']), dtype=par['dtype'])
+    StackedBDiag_MPI = pylops_mpi.MPIStackedBlockDiag(ops=[BDiag_MPI, FirstDeriv_MPI])
+
+    dist1 = pylops_mpi.DistributedArray(global_shape=size * par['nx'], dtype=par['dtype'])
+    dist1[:] = np.ones(dist1.local_shape, dtype=par['dtype'])
+    dist2 = pylops_mpi.DistributedArray(global_shape=par['nx'] * par['ny'], dtype=par['dtype'])
+    dist2[:] = np.ones(dist2.local_shape, dtype=par['dtype'])
+    x = pylops_mpi.StackedDistributedArray(distarrays=[dist1, dist2])
+    x_global = x.asarray()
+
+    dist1 = pylops_mpi.DistributedArray(global_shape=size * par['ny'], dtype=par['dtype'])
+    dist1[:] = np.ones(dist1.local_shape, dtype=par['dtype'])
+    dist2 = pylops_mpi.DistributedArray(global_shape=par['nx'] * par['ny'], dtype=par['dtype'])
+    dist2[:] = np.ones(dist2.local_shape, dtype=par['dtype'])
+    y = pylops_mpi.StackedDistributedArray(distarrays=[dist1, dist2])
+    y_global = y.asarray()
+
+    # Forward
+    x_mat = StackedBDiag_MPI @ x
+    # Adjoint
+    y_rmat = StackedBDiag_MPI.H @ y
+    assert isinstance(x_mat, pylops_mpi.StackedDistributedArray)
+    assert isinstance(y_rmat, pylops_mpi.StackedDistributedArray)
+    # Dot test
+    dottest(StackedBDiag_MPI, x, y, size * par['ny'] + par['nx'] * par['ny'], size * par['nx'] + par['nx'] * par['ny'])
+
+    x_mat_mpi = x_mat.asarray()
+    y_rmat_mpi = y_rmat.asarray()
+
+    if rank == 0:
+        ops = [pylops.MatrixMult((i + 1) * np.ones(shape=(par['ny'], par['nx'])).astype(par['dtype'])) for i in range(size)]
+        BDiag = pylops.BlockDiag(ops=ops)
+        FirstDeriv = pylops.FirstDerivative(dims=(par['ny'], par['nx']), axis=0, dtype=par['dtype'])
+        BDiag_final = pylops.BlockDiag([BDiag, FirstDeriv])
+        x_mat_np = BDiag_final @ x_global
+        y_rmat_np = BDiag_final.H @ y_global
+        assert_allclose(x_mat_mpi, x_mat_np, rtol=1e-14)
+        assert_allclose(y_rmat_mpi, y_rmat_np, rtol=1e-14)