Concatenating domain decomposition shot gathers for disk write

[dimitri-voytan]

I modified the seismic acoustic example to use mpi for domain decomposition and output the pressure at the receiver locations by

from devito import configuration
configuration['mpi'] = 'full'
.
.
.
_, _, _, out = run(shape=shape, spacing=spacing, nbl=args.nbl, tn=tn, fs=args.fs ...)
rec = out[0].data #receiver data
print(rec.data.shape)

If I run e.g. mpirun -n 4 python3 acoustic_example.py, the rec.data.shape will be smaller by a factor of ~4 along the last axis (spatial axis) relative to a run with -n 1. This is an expected behavior. However, I would like to organize this data into a single array (undoing the domain decomposition) to save it to disk.

Is there a way to get the portion of the domain allocated to each MPI process so that the receiver data can be concatenated (e.g. on the rank 0 process)?

[mloubout]

check the gather_data(rank) utility.

[mloubout]

See here , my bad for the typo

devito/devito/types/dense.py

Line 426 in e0ec639

def data_gather(self, start=None, stop=None, step=1, rank=0):

[dimitri-voytan]

Thanks, it looks like rec.data._gather() fails with:

../devito/data/data.py in __getitem__(self, glb_idx, comm_type, gather_rank)
    202 
    203             owners, send, global_si, local_si = \
--> 204                 mpi_index_maps(loc_idx, local_val.shape, self._distributor.topology,
    205                                self._distributor.all_coords, comm)
    206 
AttributeError: 'SparseDistributor' object has no attribute 'topology'

I can see the how the domain was decomposed with rec.data._distributor.decomposition , but is there another way to gather with SparseDistributor?

[mloubout]

Oh ok interesting will look into it

[dimitri-voytan]

Thanks for looking into it.

Just FYI, mpi_index_maps(loc_idx, local_val.shape, self._distributor.topology, self._distributor.all_coords, comm) will also fail because SparseDistributor doesn't have an all_coords attribute.

[dimitri-voytan]

If anyone else is running into this issue, you could try this. No warranty that this is correct.

def gather_data(rec, domain, comm, root=0):
    '''
    param rec   :: reciever object output from the forward modeling run 
    param domain :: domain assigned to each process from dist._decomposition
    param root  :: rank of process that will hold result of gather
    '''
    #send both the data and indices associated with the domain
    # in the same mpi gather to make sure they can be arranged in the correct order

    recv = comm.gather((np.asarray(rec.data), domain[0][rank]), root=root) 

    #Collect and process on rank 0
    if rank == root:

        #Separate the data and indices in the gather
        data = []
        indices = []

        for item in recv:
            data.append(item[0])
            indices.append(item[1])
        
        indices = np.concatenate(indices, axis=-1)
        data = np.concatenate(data, axis=-1)
        
        #Sort by the provided indices.
        data = data[:,indices]
        return data
    else:
        return None

#Then use
dist = rec.data._distributor
rank = dist.myrank
domain = dist._decomposition
comm = dist.comm
root = 0

#Gather recievers from subdomains
data = gather_data(rec, domain, comm, root=root)