Deblend with pylops mpi

[sadalan97]

Seismic Deblending Multinode Processing

Description

Hello, I am developing some seismic deblending tests using the 2D Radon Transform. Currently, all processing is done in parallel on a single node, but I am now exploring how to perform the processing in a multinode setup. This multinode parallel processing is new to me, so I am facing some difficulties adjusting my code.

Questions

1. Am I on the right track for configuring multinode processing?

2. How should I adjust the variable Bop (the data blending operator) to work correctly with MPI?

3. After all operator adjustments for MPI, will it be possible to use FISTA without issues, or will additional modifications be needed in the optimization algorithm?

In any case, I truly appreciate any help you all may be able to offer.

---

Initial adjustments

Initially, I adapted the data_blended (resulting from data_blended = Bopcp * data1.ravel()) for MPI as follows:

comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()

N = data_blended.shape[0]

data_blended_dist = DistributedArray(N, dtype=np.float32)

counts = np.full(size, N // size)
counts[:N % size] += 1
starts = np.cumsum(np.concatenate([[0], counts[:-1]]))
ends = starts + counts

data_blended_dist[:] = data_blended[starts[rank]:ends[rank]]

# --- 2D Radon and Sliding Configuration ---
npx = 34
vmin = 500
vmax = 2500
winsize = 10
overlap = 5
px = np.linspace(-1/vmin, 1/vmax, npx)
taxis = np.arange(nt) * dt
haxis = np.linspace(-winsize//2, winsize//2, winsize) * dr
Rop = TR2D(
    taxis, haxis, px, kind="linear", interp=False, engine='numba',
    onthefly=False, centeredh=True, dtype="float32"
)
nop = (npx, nt)
nwins, dims, mwin_inends, dwin_inends = pylops.signalprocessing.sliding2d_design(
    (ns, nt), winsize, overlap, nop
)
Slid = pylops.signalprocessing.Sliding2D(Rop, dims, (ns, nt), winsize, overlap, tapertype=None)
logging.info(f"Sliding2D operator created, shape={Slid.shape}")

# --- Divide blocks among MPI ranks ---
blocks_per_rank = nr // size
start = rank * blocks_per_rank
end = (rank + 1) * blocks_per_rank if rank != size - 1 else nr
local_blocks = [Slid for _ in range(start, end)]

# --- Create distributed operator ---
SopDiag = MPIBlockDiag(local_blocks)
print(f"[Rank {rank}] MPIBlockDiag created with blocks {start}:{end}, shape = {SopDiag.shape}")