removed normalization restored normal behaviour

Author: astroC86

examples/plot_halo.py

@@ -209,8 +209,7 @@ def local_extent_from_slice(local_shape, local_slice, halo):
         dims=dims,
         halo=halo,
         proc_grid_shape=proc_grid_shape,
-        comm=comm,
-        normalize=True, # why?
+        comm=comm
     )
 
     # Global array
@@ -261,8 +260,7 @@ def local_extent_from_slice(local_shape, local_slice, halo):
         dims=dims,
         halo=halo,
         proc_grid_shape=proc_grid_shape,
-        comm=comm,
-        normalize=False, # why?
+        comm=comm
     )
 
     for axis in [0, 1]:
@@ -307,8 +305,7 @@ def local_extent_from_slice(local_shape, local_slice, halo):
         dims=dims,
         halo=halo,
         proc_grid_shape=proc_grid_shape,
-        comm=comm,
-        normalize=False, # why?
+        comm=comm
     )
 
     # Global array

examples/plot_halo_nsconv.py

@@ -86,8 +86,7 @@ def local_extent_from_slice(local_shape, local_slice, halo):
     dims=(n, ),
     halo=halo,
     proc_grid_shape=proc_grid_shape,
-    comm=comm,
-    normalize=True,
+    comm=comm
 )
 
 # Distributed array

pylops_mpi/basicoperators/Halo.py

@@ -4,6 +4,7 @@
 
 from pylops.utils.backend import get_module
 from pylops_mpi import MPILinearOperator, DistributedArray, Partition
+from pylops_mpi.Distributed import DistributedMixIn
 
 
 def halo_block_split(global_shape: tuple, comm, grid_shape: tuple = None) -> tuple:
@@ -31,20 +32,18 @@ def halo_block_split(global_shape: tuple, comm, grid_shape: tuple = None) -> tup
     return tuple(slices)
 
 
-class MPIHalo(MPILinearOperator):
+class MPIHalo(DistributedMixIn, MPILinearOperator):
     def __init__(
         self,
         dims: tuple,
         halo,
         proc_grid_shape: tuple = None,
         comm: MPI.Comm = MPI.COMM_WORLD,
         dtype=np.float64,
-        normalize: bool = False,
     ):
         self.global_dims = tuple(dims)
         self.ndim = len(dims)
 
-        self.normalize = normalize
         self.comm = comm
         self.dtype = dtype
 
@@ -55,13 +54,16 @@ def __init__(
 
         self.local_dims = self._calc_local_dims()
         self.local_extent = self._calc_local_extent()
-        self._local_extent_sizes = self.comm.allgather(int(np.prod(self.local_extent)))
+        self._local_extent_sizes = self._allgather(
+            self.comm,
+            None,
+            int(np.prod(self.local_extent)),
+        )
 
         self.shape = (
             int(np.sum(self._local_extent_sizes)),
             int(np.prod(self.global_dims)),
         )
-        self._norm_factors = self._compute_norm_factors() if self.normalize else None
         super().__init__(shape=self.shape, dtype=np.dtype(dtype), base_comm=comm)
 
     def _parse_halo(self, h):
@@ -117,26 +119,6 @@ def _calc_local_extent(self):
             ext.append(self.local_dims[ax] + minus_halo + plus_halo)
         return tuple(ext)
 
-    def _compute_norm_factors(self):
-        weights = np.ones(self.local_dims, dtype=np.float64)
-        for ax in range(self.ndim):
-            before, after = self.halo[2 * ax], self.halo[2 * ax + 1]
-            if before == 0 and after == 0:
-                continue
-            n_local = self.local_dims[ax]
-            factors = np.ones(n_local, dtype=np.float64)
-            minus_nbr, plus_nbr = self.neigh[("-", ax)], self.neigh[("+", ax)]
-            minus_copy = minus_nbr != MPI.PROC_NULL
-            plus_copy = plus_nbr != MPI.PROC_NULL
-            if before and minus_copy:
-                factors[:before] += 1.0
-            if after and plus_copy:
-                factors[n_local - after:] += 1.0
-            shape = [1] * self.ndim
-            shape[ax] = n_local
-            weights *= factors.reshape(shape)
-        return (1.0 / np.sqrt(weights)).astype(self.dtype, copy=False)
-
     def _exchange_along_axis(self, ncp, arr, axis, before, after):
         minus_nbr, plus_nbr = self.neigh[("-", axis)], self.neigh[("+", axis)]
         # slice definitions
@@ -162,46 +144,6 @@ def _exchange_along_axis(self, ncp, arr, axis, before, after):
             self.cart_comm.Sendrecv(snd, dest=plus_nbr, recvbuf=rcv, source=plus_nbr)
             arr[tuple(rcv_s)] = rcv
 
-    def _exchange_adjoint_along_axis(self, ncp, arr, axis, before, after):
-        minus_nbr, plus_nbr = self.neigh[("-", axis)], self.neigh[("+", axis)]
-        out = arr.copy()
-        if before == 0 and after == 0:
-            return out
-
-        def axis_slice(start, end):
-            s = [slice(None)] * self.ndim
-            s[axis] = slice(start, end)
-            return tuple(s)
-
-        empty = axis_slice(0, 0)
-        halo_minus = axis_slice(0, before) if before else None
-        halo_plus = axis_slice(-after, None) if after else None
-        core_minus = axis_slice(before, 2 * before) if before else None
-        core_plus = axis_slice(-2 * after, -after) if after else None
-
-        # remove contributions from halo slabs along this axis (forward overwrites them)
-        if halo_minus:
-            out[halo_minus] = 0
-        if halo_plus:
-            out[halo_plus] = 0
-
-        # send right halo to plus neighbor, receive right halo from minus neighbor
-        snd_slice = halo_plus or empty
-        rcv_slice = halo_minus or empty
-        rcv = ncp.zeros_like(arr[rcv_slice])
-        self.cart_comm.Sendrecv(arr[snd_slice].copy(), dest=plus_nbr, recvbuf=rcv, source=minus_nbr)
-        if core_minus:
-            out[core_minus] += rcv
-
-        # send left halo to minus neighbor, receive left halo from plus neighbor
-        snd_slice = halo_minus or empty
-        rcv_slice = halo_plus or empty
-        rcv = ncp.zeros_like(arr[rcv_slice])
-        self.cart_comm.Sendrecv(arr[snd_slice].copy(), dest=minus_nbr, recvbuf=rcv, source=plus_nbr)
-        if core_plus:
-            out[core_plus] += rcv
-        return out
-
     def _matvec(self, x):
         ncp = get_module(x.engine)
         if x.partition != Partition.SCATTER:
@@ -211,12 +153,13 @@ def _matvec(self, x):
             global_shape=self.shape[0],
             partition=Partition.SCATTER,
             local_shapes=self._local_extent_sizes,
+            base_comm=x.base_comm,
+            base_comm_nccl=x.base_comm_nccl,
+            engine=x.engine,
+            dtype=self.dtype,
         )
 
         core = x.local_array.reshape(self.local_dims)
-        if self.normalize:
-            norm = self._norm_factors if ncp is np else ncp.asarray(self._norm_factors)
-            core = core * norm
         halo_arr = ncp.zeros(self.local_extent, dtype=self.dtype)
         # insert core
         core_slices = [
@@ -234,25 +177,16 @@ def _matvec(self, x):
         return y
 
     def _rmatvec(self, x):
-        ncp = get_module(x.engine)
         if x.partition != Partition.SCATTER:
             raise ValueError(f"x should have partition={Partition.SCATTER} Got {x.partition} instead...")
-
-        y = DistributedArray(global_shape=self.shape[1], partition=Partition.SCATTER)
-
-        arr = x.local_array.reshape(self.local_extent).copy()
-        # adjoint of halo exchange (reverse axis order because of the Transpose)
-        for ax in reversed(range(self.ndim)):
-            before, after = self.halo[2 * ax], self.halo[2 * ax + 1]
-            arr = self._exchange_adjoint_along_axis(ncp, arr, axis=ax, before=before, after=after)
-        core_slices = [
-            slice(left, left + ldim)
-            for left, ldim in zip(self.halo[::2], self.local_dims)
-        ]
+        res = DistributedArray(global_shape=self.shape[1],
+                               partition=Partition.SCATTER,
+                               base_comm=x.base_comm,
+                               base_comm_nccl=x.base_comm_nccl,
+                               engine=x.engine,
+                               dtype=self.dtype)
+        arr = x.local_array.reshape(self.local_extent)
+        core_slices = [slice(left, left + ldim) for left, ldim in zip(self.halo[::2], self.local_dims)]
         core = arr[tuple(core_slices)]
-        if self.normalize:
-            norm = self._norm_factors if ncp is np else ncp.asarray(self._norm_factors)
-            core = core * norm
-
-        y[:] = core.ravel()
-        return y
+        res[:] = core.ravel()
+        return res

tests/test_halo.py

@@ -6,9 +6,10 @@
 else:
     import numpy as np
     backend = "numpy"
-import numpy as npp
 from mpi4py import MPI
 import pytest
+import pylops
+from numpy.testing import assert_allclose
 
 import pylops_mpi
 from pylops_mpi.basicoperators.Halo import MPIHalo
@@ -21,28 +22,24 @@
     np.cuda.Device(device_id).use()
 
 
-@pytest.mark.mpi(min_size=8)
-def test_halo_dottest_plot_config():
+@pytest.mark.mpi(min_size=2)
+def test_halo_composed_dottest_matches_serial():
     comm = MPI.COMM_WORLD
     size = comm.Get_size()
-    p_prime = int(round(size ** (1 / 3)))
-    if p_prime ** 3 != size:
-        pytest.skip("MPI size must be a perfect cube for 3D halo grid")
-
-    gdim = (4 * p_prime, 4 * p_prime, 4 * p_prime)
-    g_shape = (p_prime, p_prime, p_prime)
+    nlocal = 16
+    n = nlocal * size
     halo = 1
 
     halo_op = MPIHalo(
-        dims=gdim,
+        dims=(n,),
         halo=halo,
-        proc_grid_shape=g_shape,
+        proc_grid_shape=(size,),
         comm=comm,
         dtype=np.float64,
     )
 
     x_dist = pylops_mpi.DistributedArray(
-        global_shape=npp.prod(gdim),
+        global_shape=n,
         base_comm=comm,
         partition=pylops_mpi.Partition.SCATTER,
         engine=backend,
@@ -51,12 +48,40 @@ def test_halo_dottest_plot_config():
     x_dist[:] = np.random.normal(0.0, 1.0, x_dist.local_array.shape)
 
     y_dist = pylops_mpi.DistributedArray(
-        global_shape=halo_op.shape[0],
+        global_shape=n,
         base_comm=comm,
         partition=pylops_mpi.Partition.SCATTER,
         engine=backend,
         dtype=np.float64,
     )
     y_dist[:] = np.random.normal(0.0, 1.0, y_dist.local_array.shape)
 
-    dottest(halo_op, x_dist, y_dist)
+    local_extent = halo_op.local_extent[0]
+    DOp = pylops.FirstDerivative(
+        dims=local_extent,
+        axis=0,
+        kind="forward",
+        dtype=np.float64,
+    )
+    DOp_dist = pylops_mpi.MPIBlockDiag([DOp], base_comm=comm, dtype=np.float64)
+    Op_dist = halo_op.H @ DOp_dist @ halo_op
+
+    dottest(Op_dist, x_dist, y_dist, n, n)
+
+    y_dist = Op_dist @ x_dist
+    y_adj_dist = Op_dist.H @ x_dist
+    y = y_dist.asarray()
+    y_adj = y_adj_dist.asarray()
+
+    x_global = x_dist.asarray()
+    if rank == 0:
+        DOp_serial = pylops.FirstDerivative(
+            dims=n,
+            axis=0,
+            kind="forward",
+            dtype=np.float64,
+        )
+        y_serial = DOp_serial @ x_global
+        y_adj_serial = DOp_serial.H @ x_global
+        assert_allclose(y, y_serial, rtol=1e-14)
+        assert_allclose(y_adj, y_adj_serial, rtol=1e-14)