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 get_communicator(comm, layout):
partitioner = fv3util.CubedSpherePartitioner(
fv3util.TilePartitioner(layout))
communicator = fv3util.CubedSphereCommunicator(comm, partitioner)
return communicator
# 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"],
bk=state["atmosphere_hybrid_b_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