PyMPDATA_MPI.utils 

 1# pylint: disable=missing-module-docstring,missing-function-docstring
 2
 3from contextlib import contextmanager
 4
 5import numba_mpi as mpi
 6import numpy as np
 7
 8
 9@contextmanager
10def barrier_enclosed():
11    try:
12        mpi.barrier()
13        yield
14    finally:
15        mpi.barrier()
16
17
18def setup_dataset_and_sync_all_workers(storage, dataset_name):
19    dataset = storage[dataset_name]
20    dataset[:] = np.nan
21    mpi.barrier()
22    return dataset
@contextmanager
def barrier_enclosed(): 
10@contextmanager
11def barrier_enclosed():
12    try:
13        mpi.barrier()
14        yield
15    finally:
16        mpi.barrier()
def setup_dataset_and_sync_all_workers(storage, dataset_name): 
19def setup_dataset_and_sync_all_workers(storage, dataset_name):
20    dataset = storage[dataset_name]
21    dataset[:] = np.nan
22    mpi.barrier()
23    return dataset