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}")

[mrava87]

Hello @sadalan97,
thanks for your interest in PyLops.

Short answer, you are the first to explore deblending with pylops-mpi, so what I say may turn out to be incorrect as I have never tried myself 😉

Having said that:

1. I am not sure from your code what is TR2D? Also the way you setup SopDiag would take the overall data and split it across the first dimension so you will have a sliding operator with overlap within each rank but across rank the windows will not overlap. Whilst this is not ideal, I think it is currently the only option you have. We have not implement sliding/patching operators in pylops-mpi yet and this PR Implementation of distributed Halo Operator pylops-mpi#166 may provide a way to add halos to any operator without re-implementing all operators by hand (but is not ready yet)
2. For Bop it is actually easier as you will have a gather of ns x nt that will turn into a single long trace with nttot. So you could possibly just split shots across ranks and then use a MPIHStack to stack all partial results. But, so far we do not expose nttot to the user and this is computed internally based on the max times, so different rank may end up with different nttot. In this PR (will be merged soon), you will get exactly that feat: added multithreading support to stacking ops #708. We are doing it for multiprocessor within a single node/rank but you can also exploit it to make multi-node Bop.
3. No, you will need to implement a distributed FISTA following what we did for cg/cgls. Should be pretty easy.... We have plans to extend the solver suite in pylops-mpi but we have not been able to do it yet... if you are interested, feel free to start and open a PR.