How to correctly use MPI with Devito?

[AileniMahesh785]

When I use MPI to parallelize the wavefield simulation with the method AcousticWaveSolver imported from devito, the shapes of the simulated wavefields are not as expected.
The following code illustrates the way in which the complete wavefield is not stored along the spatial dimensions.
u.shape is expected to be (246186, 281, 281), where the second and third axes represent spatial dimensions.
But the code produces a different output.

`from mpi4py import MPI
from devito import configuration

from examples.seismic import Model, RickerSource, Receiver, TimeAxis, AcquisitionGeometry
from examples.seismic.acoustic import AcousticWaveSolver
import numpy as np

------------------------------------------

MPI setup

------------------------------------------

comm = MPI.COMM_WORLD
rank = comm.rank
size = comm.size

------------------------------------------

Define source locations (one per rank)

------------------------------------------

num_shots = 10
source_locations = [(200. * i, 1000.) for i in range(num_shots)]

if size != num_shots:
raise ValueError(f"Run with {num_shots} MPI ranks, got {size}.")

my_source_coords = source_locations[0]

------------------------------------------

Create velocity model

------------------------------------------

shape = (201, 201) # grid points
extent = (2000., 2000.) # domain size (m)
spacing = tuple(e/s for e, s in zip(extent, shape))

vp = 1500. * np.ones(shape, dtype=np.float32)

model = Model(vp=vp, origin=(0., 0.), shape=shape,
spacing=spacing, space_order=2, nbl=40, bcs="damp")

------------------------------------------

Time axis

------------------------------------------

t0, tn, dt = 0., 1000., 1.0
time_axis = TimeAxis(start=t0, stop=tn, step=dt)

------------------------------------------

Define source and receiver geometry

------------------------------------------

src = RickerSource(name='src', grid=model.grid, f0=0.01, time_range=time_axis)

src.coordinates.data[0, :] = np.array(my_source_coords)

nrec = 101
rec = Receiver(name='rec', grid=model.grid, npoint=nrec, time_range=time_axis)
rec.coordinates.data[:, 0] = np.linspace(0, extent[0], nrec)
rec.coordinates.data[:, 1] = extent[1] / 2.
geometry = AcquisitionGeometry(model, rec.coordinates.data,
src.coordinates.data, t0, tn)

------------------------------------------

Set up and run the forward solver

------------------------------------------

solver = AcousticWaveSolver(model, geometry, time_order=2, space_order=4)

rec_data, u, _ = solver.forward(save=True)

Save wavefield and receiver data for this shot

#np.save(f"wavefield_shot{rank}.npy", u.data)
#np.save(f"receivers_shot{rank}.npy", rec_data.data)

print(f"Rank {rank}: u.shape = {u.shape}, rec.shape = {rec_data.shape}")
`

`Rank 1: u.shape = (246186, 57, 140)

Rank 4: u.shape = (246186, 56, 141)

Rank 5: u.shape = (246186, 56, 140)

Rank 8: u.shape = (246186, 56, 141)

Rank 9: u.shape = (246186, 56, 140)

Rank 3: u.shape = (246186, 56, 140)

Rank 6: u.shape = (246186, 56, 141)

Rank 7: u.shape = (246186, 56, 140)

Rank 2: u.shape = (246186, 56, 141)

Rank 0: u.shape = (246186, 57, 141)`

What am I missing here? How to get the complete wavefield corresponding to each source location at once using mpi4py, rather than performing serial simulations?

Many Thanks!
Aileni Mahesh

[mloubout]

MPI is used fror domain decomposition,i.e the domain is slit into smaller domain, one for each rank. So you can either gather the wavefield at the end (which defeat the purpose so why use mpi in the first place and just use openmp) or work with each domain individually.

[AileniMahesh785]

Thanks!