def get_mock_communicator_list(layout):
total_ranks = 6 * fv3util.TilePartitioner(layout).total_ranks
shared_buffer = {}
communicators = []
for rank in range(total_ranks):
comm = fv3util.testing.DummyComm(rank,
total_ranks,
buffer_dict=shared_buffer)
communicator = get_communicator(comm, layout)
communicators.append(communicator)
return communicators
def set_up_communicator(
disable_halo_exchange: bool,
) -> Tuple[Optional[MPI.Comm], Optional[util.CubedSphereCommunicator]]:
layout = spec.namelist.layout
partitioner = util.CubedSpherePartitioner(util.TilePartitioner(layout))
if MPI is not None:
comm = MPI.COMM_WORLD
else:
comm = None
if not disable_halo_exchange:
assert comm is not None
cube_comm = util.CubedSphereCommunicator(comm, partitioner)
else:
cube_comm = util.CubedSphereCommunicator(NullComm(0, 0), partitioner)
return comm, cube_comm
def __init__(self, indices, shape_params, rank, layout, data_fields={}):
self.rank = rank
self.partitioner = fv3util.TilePartitioner(layout)
self.subtile_index = self.partitioner.subtile_index(self.rank)
self.layout = layout
for s in self.shape_params:
setattr(self, s, int(shape_params[s]))
self.subtile_width_x = int((self.npx - 1) / self.layout[0])
self.subtile_width_y = int((self.npy - 1) / self.layout[1])
for ivar, jvar in self.index_pairs:
local_i, local_j = self.global_to_local_indices(
int(indices[ivar]), int(indices[jvar]))
setattr(self, ivar, local_i)
setattr(self, jvar, local_j)
self.nid = int(self.ied - self.isd + 1)
self.njd = int(self.jed - self.jsd + 1)
self.nic = int(self.ie - self.is_ + 1)
self.njc = int(self.je - self.js + 1)
self.halo = utils.halo
# TODO: do we want to set face indices this way?
self.isf = 0
self.ief = self.npx - 1
self.jsf = 0
self.jef = self.npy - 1
self.global_is, self.global_js = self.local_to_global_indices(
self.is_, self.js)
self.global_ie, self.global_je = self.local_to_global_indices(
self.ie, self.je)
self.global_isd, self.global_jsd = self.local_to_global_indices(
self.isd, self.jsd)
self.global_ied, self.global_jed = self.local_to_global_indices(
self.ied, self.jed)
self.west_edge = self.global_is == self.halo
self.east_edge = self.global_ie == self.npx + self.halo - 2
self.south_edge = self.global_js == self.halo
self.north_edge = self.global_je == self.npy + self.halo - 2
self.j_offset = self.js - self.jsd - 1
self.i_offset = self.is_ - self.isd - 1
self.sw_corner = self.west_edge and self.south_edge
self.se_corner = self.east_edge and self.south_edge
self.nw_corner = self.west_edge and self.north_edge
self.ne_corner = self.east_edge and self.north_edge
self.data_fields = {}
self.add_data(data_fields)
self._sizer = None
self._quantity_factory = None
def get_communicator(comm, layout):
partitioner = fv3util.CubedSpherePartitioner(
fv3util.TilePartitioner(layout))
communicator = fv3util.CubedSphereCommunicator(comm, partitioner)
return communicator
mpi_comm = MPI.COMM_WORLD
# get grid from serialized data
grid_savepoint = serializer.get_savepoint("Grid-Info")[0]
grid_data = {}
grid_fields = serializer.fields_at_savepoint(grid_savepoint)
for field in grid_fields:
grid_data[field] = serializer.read(field, grid_savepoint)
if len(grid_data[field].flatten()) == 1:
grid_data[field] = grid_data[field][0]
grid = fv3core.testing.TranslateGrid(grid_data, rank).python_grid()
spec.set_grid(grid)
# set up grid-dependent helper structures
layout = spec.namelist.layout
partitioner = util.CubedSpherePartitioner(util.TilePartitioner(layout))
communicator = util.CubedSphereCommunicator(mpi_comm, partitioner)
# create a state from serialized data
savepoint_in = serializer.get_savepoint("FVDynamics-In")[0]
driver_object = fv3core.testing.TranslateFVDynamics([grid])
input_data = driver_object.collect_input_data(serializer, savepoint_in)
input_data["comm"] = communicator
state = driver_object.state_from_inputs(input_data)
dycore = fv3core.DynamicalCore(
comm=communicator,
grid_data=spec.grid.grid_data,
stencil_factory=spec.grid.stencil_factory,
damping_coefficients=spec.grid.damping_coefficients,
config=spec.namelist.dynamical_core,
ak=state["atmosphere_hybrid_a_coordinate"],
experiment_name = yaml.safe_load(
open(
args.data_dir + "/input.yml",
"r",
))["experiment_name"]
# set up of helper structures
serializer = serialbox.Serializer(
serialbox.OpenModeKind.Read,
args.data_dir,
"Generator_rank" + str(rank),
)
cube_comm = util.CubedSphereCommunicator(
comm,
util.CubedSpherePartitioner(
util.TilePartitioner(spec.namelist.layout)),
)
# get grid from serialized data
grid_savepoint = serializer.get_savepoint("Grid-Info")[0]
grid_data = {}
grid_fields = serializer.fields_at_savepoint(grid_savepoint)
for field in grid_fields:
grid_data[field] = serializer.read(field, grid_savepoint)
if len(grid_data[field].flatten()) == 1:
grid_data[field] = grid_data[field][0]
grid = fv3core.testing.TranslateGrid(grid_data, rank).python_grid()
spec.set_grid(grid)
# set up grid-dependent helper structures
layout = spec.namelist.layout