def extract_local_sparse_matrix(self, target_rank):
logger.debug("Extract local sparse matrix for rank{}".format(target_rank))
t_rank = target_rank
dsts = self.dsts
srcs = self.srcs
wgts = self.wgts
rank_dsts = self.rank_dsts
rank_srcs = self.rank_srcs
t_rank_dsts = rank_dsts == t_rank # bool type array
t_rank_srcs = rank_srcs == t_rank
local_idxs = np.where(t_rank_dsts * t_rank_srcs)[0]
send_idxs = np.where(np.invert(t_rank_dsts) * t_rank_srcs)[0]
recv_idxs = np.where(t_rank_dsts * np.invert(t_rank_srcs))[0]
arr_dict = dict()
arr_dict["spmat_size"] = self.spmat_size
arr_dict["local_dsts"] = dsts[local_idxs]
arr_dict["local_srcs"] = srcs[local_idxs]
arr_dict["local_wgts"] = wgts[local_idxs]
arr_dict["send_ranks"] = rank_dsts[send_idxs]
arr_dict["send_dsts"] = dsts[send_idxs]
arr_dict["send_srcs"] = srcs[send_idxs]
arr_dict["send_wgts"] = wgts[send_idxs]
arr_dict["recv_ranks"] = rank_srcs[recv_idxs]
arr_dict["recv_dsts"] = dsts[recv_idxs]
return arr_dict
def distribute_local_sparse_matrix(self, comm):
logger.debug('Distribute local sparse matrixes')
if self.myrank == 0:
req_list = list()
for target_rank in range(1,self.nproc):
arr_dict = self.extract_local_sparse_matrix(target_rank)
req = comm.isend(arr_dict, dest=target_rank, tag=10)
req_list.append(req)
for req in req_list: req.wait()
else:
self.arr_dict = comm.recv(source=0, tag=10)
def read_sparse_matrix(self):
logger.debug("Read a NetCDF file as sparse matrix")
ne = self.ne
ngq = self.ngq
method = self.method
ranks = self.ranks
fname = "spmat_{}_ne{:03d}np{}.nc".format(method.lower(), ne, ngq)
spmat_fpath = os.path.join(self.spmat_dpath, fname)
assert os.path.exists(spmat_fpath), "{} is not found.".format(spmat_fpath)
spmat_ncf = nc.Dataset(spmat_fpath, "r")
self.spmat_size = len(spmat_ncf.dimensions["spmat_size"])
self.dsts = spmat_ncf.variables["dsts"][:]
self.srcs = spmat_ncf.variables["srcs"][:]
self.wgts = spmat_ncf.variables["weights"][:]
self.rank_dsts = ranks[self.dsts] # rank number of destinations
self.rank_srcs = ranks[self.srcs] # rank number of sources
def save_netcdf(self, base_dir, target_method, nc_format):
logger.debug("Save the mpi tables as NetCDF")
ncf = nc.Dataset(
base_dir + "/nproc{}_rank{}.nc".format(self.nproc, self.myrank), "w", format=nc_format
) # 'NETCDF4', 'NETCDF3_CLASSIC'
ncf.description = "MPI index tables with the SFC partitioning on the cubed-sphere"
ncf.target_method = target_method
ncf.ne = self.ne
ncf.ngq = self.ngq
ncf.nproc = self.nproc
ncf.myrank = self.myrank
ncf.ep_size = self.cubegrid.ep_size
ncf.up_size = self.cubegrid.up_size
ncf.local_ep_size = self.cubegrid.local_ep_size
ncf.local_up_size = self.cubegrid.local_up_size
ncf.spmat_size = self.spmat_size
ncf.createDimension("local_ep_size", self.cubegrid.local_ep_size)
ncf.createDimension("send_sche_size", len(self.send_schedule))
ncf.createDimension("recv_sche_size", len(self.recv_schedule))
ncf.createDimension("send_size", len(self.send_dsts))
ncf.createDimension("recv_size", len(self.recv_dsts))
ncf.createDimension("3", 3)
vlocal_gids = ncf.createVariable("local_gids", "i4", ("local_ep_size",))
vlocal_gids.long_name = "Global index of local points"
vbuf_sizes = ncf.createVariable("buf_sizes", "i4", ("3",))
vbuf_sizes.long_name = "[local_src_size, send_buf_size, recv_buf_size]"
vsend_schedule = ncf.createVariable("send_schedule", "i4", ("send_sche_size", "3"))
vsend_schedule.long_name = "[rank, start, size]"
vrecv_schedule = ncf.createVariable("recv_schedule", "i4", ("recv_sche_size", "3"))
vrecv_schedule.long_name = "[rank, start, size]"
vsend_dsts = ncf.createVariable("send_dsts", "i4", ("send_size",))
vsend_dsts.long_name = "Destination index for local and send buffer"
vsend_srcs = ncf.createVariable("send_srcs", "i4", ("send_size",))
vsend_srcs.long_name = "Source index for local and send buffer"
vsend_wgts = ncf.createVariable("send_wgts", "f8", ("send_size",))
vsend_wgts.long_name = "Weight value for local and send buffer"
vrecv_dsts = ncf.createVariable("recv_dsts", "i4", ("recv_size",))
vrecv_dsts.long_name = "Destination index for recv buffer"
vrecv_srcs = ncf.createVariable("recv_srcs", "i4", ("recv_size",))
vrecv_srcs.long_name = "Source index for recv buffer"
vlocal_gids[:] = self.cubegrid.local_gids[:]
vbuf_sizes[:] = (self.local_src_size, self.send_buf_size, self.recv_buf_size)
vsend_schedule[:] = self.send_schedule[:]
vrecv_schedule[:] = self.recv_schedule[:]
vsend_dsts[:] = self.send_dsts[:]
vsend_srcs[:] = self.send_srcs[:]
vsend_wgts[:] = self.send_wgts[:]
vrecv_dsts[:] = self.recv_dsts[:]
vrecv_srcs[:] = self.recv_srcs[:]
ncf.close()
def __init__(self, ne, ngq, nproc, myrank, cs_grid_dpath, is_rotate=False, homme_style=False):
self.ne = ne
self.ngq = ngq
self.nproc = nproc
self.myrank = myrank
# -----------------------------------------------------
# Read the grid indices
# -----------------------------------------------------
fname_tag = "_rotated" if is_rotate else ""
fname = "cs_grid_ne{:03d}np{}{}.nc".format(ne, ngq, fname_tag)
cs_fpath = os.path.join(cs_grid_dpath, fname)
assert os.path.exists(cs_fpath), "{} is not found.".format(cs_fpath)
cs_ncf = nc.Dataset(cs_fpath, "r")
ep_size = len(cs_ncf.dimensions["ep_size"])
up_size = len(cs_ncf.dimensions["up_size"])
gq_indices = cs_ncf.variables["gq_indices"][:] # (ep_size,5)
mvps = cs_ncf.variables["mvps"][:] # (ep_size,4)
is_uvps = cs_ncf.variables["is_uvps"][:] # (ep_size)
uids = cs_ncf.variables["uids"][:] # (ep_size)
gids = cs_ncf.variables["gids"][:] # (up_size)
xyzs = cs_ncf.variables["xyzs"][:] # (up_size,3)
latlons = cs_ncf.variables["latlons"][:] # (up_size,2)
alpha_betas = cs_ncf.variables["alpha_betas"][:] # (ep_size,2)
# mvps = cs_ncf.variables['mvps'][:] # (ep_size,4)
# nbrs = cs_ncf.variables['nbrs'][:] # (up_size,8)
# -----------------------------------------------------
# Set the rank and local indices
# -----------------------------------------------------
logger.debug("Set the rank and local indices")
partition = CubePartition(ne, nproc, homme_style)
local_nelem = partition.nelems[myrank]
local_ep_size = local_nelem * ngq * ngq
local_gids = np.zeros(local_ep_size, "i4")
local_uids = np.zeros(local_ep_size, "i4")
local_is_uvps = np.zeros(local_ep_size, "bool")
local_mvps = np.zeros((local_ep_size, 4), "i4")
local_gq_indices = np.zeros((local_ep_size, 5), "i4")
local_alpha_betas = np.zeros((local_ep_size, 2), "f8")
local_latlons = np.zeros((local_ep_size, 2), "f8")
local_xyzs = np.zeros((local_ep_size, 3), "f8")
# MPI rank at each grid point
gq_eijs = gq_indices[:, :3] - 1
idxs = gq_eijs[:, 0] * ne * ne + gq_eijs[:, 1] * ne + gq_eijs[:, 2]
ranks = partition.elem_proc.ravel()[idxs]
local_gids[:] = np.where(ranks == myrank)[0]
local_is_uvps[:] = is_uvps[local_gids]
local_mvps[:] = mvps[local_gids, :]
local_gq_indices[:] = gq_indices[local_gids, :]
local_alpha_betas[:] = alpha_betas[local_gids, :]
local_uids[:] = uids[local_gids]
local_xyzs[:] = xyzs[local_uids, :]
local_latlons[:] = latlons[local_uids, :]
# Local index at each grid point
lids = np.zeros(ep_size, "i4")
for proc in range(nproc):
local_ep_size_tmp = partition.nelems[proc] * ngq * ngq
idxs = np.where(ranks == proc)[0]
lids[idxs] = np.arange(local_ep_size_tmp, dtype="i4")
# -----------------------------------------------------
# Public variables
# -----------------------------------------------------
self.partition = partition
self.ranks = ranks
self.lids = lids
self.ep_size = ep_size
self.up_size = up_size
self.local_ep_size = local_ep_size
self.local_up_size = local_is_uvps.sum()
self.local_nelem = local_nelem
self.local_gids = local_gids
self.local_uids = local_uids
self.local_is_uvps = local_is_uvps
self.local_mvps = local_mvps
self.local_gq_indices = local_gq_indices
self.local_alpha_betas = local_alpha_betas
self.local_latlons = local_latlons
self.local_xyzs = local_xyzs
def make_mpi_tables(self):
"""
Destination, source, weight from the sparse matrix
Make Generate the meta index grouped by rank
local_group: {dst:[(src,wgt),...]}
send_group: {rank:{dst:[(src,wgt),...]),...}
recv_group: {rank:[dst,...],...}
All dictionaries are OrderedDicts.
"""
logger.debug("Make MPI tables")
lids = self.lids
arr_dict = self.arr_dict
self.spmat_size = arr_dict["spmat_size"]
# ---------------------------------------
# local_group
# ---------------------------------------
local_dsts = arr_dict["local_dsts"]
local_srcs = arr_dict["local_srcs"]
local_wgts = arr_dict["local_wgts"]
dsw_list = [(d, s, w) for d, s, w in zip(local_dsts, local_srcs, local_wgts)]
local_group = OrderedDict(
[(dst, [(s, w) for (d, s, w) in val]) for (dst, val) in groupby(dsw_list, lambda x: x[0])]
)
local_src_size = len(dsw_list)
local_buf_size = len(local_group)
# ---------------------------------------
# send_group
# ---------------------------------------
send_ranks = arr_dict["send_ranks"]
send_dsts = arr_dict["send_dsts"]
send_srcs = arr_dict["send_srcs"]
send_wgts = arr_dict["send_wgts"]
rdsw_list = [(r, d, s, w) for r, d, s, w in zip(send_ranks, send_dsts, send_srcs, send_wgts)]
sorted_rdsw_list = sorted(rdsw_list, key=lambda x: x[0])
send_group_tmp = OrderedDict(
[(rank, [(d, s, w) for (r, d, s, w) in val]) for (rank, val) in groupby(sorted_rdsw_list, lambda x: x[0])]
)
send_group = OrderedDict()
for rank, dsw_list in send_group_tmp.items():
send_group[rank] = OrderedDict(
[(dst, [(s, w) for (d, s, w) in val]) for (dst, val) in groupby(dsw_list, lambda x: x[0])]
)
# ---------------------------------------
# recv_group
# ---------------------------------------
recv_ranks = arr_dict["recv_ranks"]
recv_dsts = arr_dict["recv_dsts"]
rd_list = [(r, d) for r, d in zip(recv_ranks, recv_dsts)]
sorted_rd_list = sorted(rd_list, key=lambda x: x[0])
recv_group = OrderedDict(
[(rank, np.unique([d for (r, d) in val])) for (rank, val) in groupby(sorted_rd_list, lambda x: x[0])]
)
# -----------------------------------------------------
# Make the send_schedule, send_dsts, send_srcs, send_wgts
# -----------------------------------------------------
logger.debug("Make the send_schedule, send_dsts, send_srcs, send_wgts")
# ---------------------------------------
# size and allocation
# ---------------------------------------
send_sche_size = len(send_group)
send_buf_size = np.unique(send_dsts).size
send_map_size = local_dsts.size + send_dsts.size
send_schedule = np.zeros((send_sche_size, 3), "i4") # (rank,start,size)
send_dsts = np.zeros(send_map_size, "i4")
send_srcs = np.zeros(send_map_size, "i4")
send_wgts = np.zeros(send_map_size, "f8")
send_buf = np.zeros(send_buf_size, "i4") # global dst index
# ---------------------------------------
# send_schedule
# ---------------------------------------
send_buf_seq = 0
for seq, rank in enumerate(send_group.keys()):
start = send_buf_seq
size = len(send_group[rank])
send_schedule[seq][:] = (rank, start, size)
send_buf_seq += size
if send_buf_size != send_buf_seq:
logger.error("Error: send_buf_size({}) != send_buf_seq({})".format(send_buf_size, send_buf_seq))
raise SystemError
# ---------------------------------------
# send local indices in myrank
# directly go to the recv_buf, not to the send_buf
# ---------------------------------------
seq = 0
recv_buf_seq = 0
for dst, sw_list in local_group.items():
for src, wgt in sw_list:
send_dsts[seq] = recv_buf_seq
send_srcs[seq] = lids[src]
send_wgts[seq] = wgt
seq += 1
recv_buf_seq += 1
# ---------------------------------------
# send indices for the other ranks
# ---------------------------------------
send_buf_seq = 0
for rank, dst_dict in send_group.items():
for dst, sw_list in dst_dict.items():
for src, wgt in sw_list:
send_dsts[seq] = send_buf_seq
send_srcs[seq] = lids[src]
send_wgts[seq] = wgt
seq += 1
send_buf[send_buf_seq] = dst # for diagnostics
send_buf_seq += 1
if seq != send_map_size:
logger.error("Error: seq({}) != send_map_size({})".format(seq, send_map_size))
raise SystemError
if send_buf_seq != send_buf_size:
logger.error("Error: send_buf_seq({}) != send_buf_size({})".format(send_buf_seq, send_buf_size))
raise SystemError
# -----------------------------------------------------
# Make the recv_schedule, recv_dsts, recv_srcs
# -----------------------------------------------------
logger.debug("Make the recv_schedule, recv_dsts, recv_srcs")
# ---------------------------------------
# size and allocation
# ---------------------------------------
recv_sche_size = len(recv_group)
recv_buf_size = local_buf_size + int(np.sum([d_unique.size for d_unique in recv_group.values()]))
recv_map_size = recv_buf_size
recv_schedule = np.zeros((recv_sche_size, 3), "i4") # (rank,start,size)
recv_dsts = np.zeros(recv_map_size, "i4")
recv_srcs = np.zeros(recv_map_size, "i4")
# ---------------------------------------
# recv_schedule
# ---------------------------------------
recv_buf_seq = local_buf_size
for seq, (rank, d_unique) in enumerate(recv_group.items()):
start = recv_buf_seq
size = d_unique.size
recv_schedule[seq][:] = (rank, start, size)
recv_buf_seq += size
# ---------------------------------------
# recv indices
# ---------------------------------------
recv_buf_list = list(local_group.keys()) # destinations
for rank, d_unique in recv_group.items():
recv_buf_list.extend(d_unique)
recv_buf = np.array(recv_buf_list, "i4")
unique_dsts = np.unique(recv_buf)
seq = 0
for dst in unique_dsts:
for bsrc in np.where(recv_buf == dst)[0]:
recv_dsts[seq] = lids[dst] # local index
recv_srcs[seq] = bsrc # buffer index
seq += 1
# -----------------------------------------------------
# Public variables for diagnostic
# -----------------------------------------------------
self.local_group = local_group
self.send_group = send_group
self.recv_group = recv_group
self.send_buf = send_buf # global dst index
self.recv_buf = recv_buf # global dst index
# -----------------------------------------------------
# Public variables
# -----------------------------------------------------
self.local_src_size = local_src_size
self.send_buf_size = send_buf_size
self.recv_buf_size = recv_buf_size
self.send_schedule = send_schedule # (rank,start,size)
self.send_dsts = np.array(send_dsts, "i4") # to buffer
self.send_srcs = np.array(send_srcs, "i4") # from local
self.send_wgts = np.array(send_wgts, "f8")
self.recv_schedule = recv_schedule # (rank,start,size)
self.recv_dsts = np.array(recv_dsts, "i4") # to local
self.recv_srcs = np.array(recv_srcs, "i4") # from buffer
def save_netcdf(self, base_dir, target_method, nc_format):
logger.debug('Save the mpi tables as NetCDF')
ncf = nc.Dataset(base_dir + "/nproc{}_rank{}.nc".format(self.nproc,self.myrank), 'w', format=nc_format) # 'NETCDF4', 'NETCDF3_CLASSIC'
ncf.description = 'MPI index tables with the SFC partitioning on the cubed-sphere'
ncf.target_method = target_method
ncf.ne = self.ne
ncf.ngq = self.ngq
ncf.nproc = self.nproc
ncf.myrank = self.myrank
ncf.ep_size = self.cubegrid.ep_size
ncf.up_size = self.cubegrid.up_size
ncf.local_ep_size = self.cubegrid.local_ep_size
ncf.local_up_size = self.cubegrid.local_up_size
ncf.spmat_size = self.spmat_size
ncf.createDimension('local_ep_size', self.cubegrid.local_ep_size)
ncf.createDimension('send_sche_size', len(self.send_schedule))
ncf.createDimension('recv_sche_size', len(self.recv_schedule))
ncf.createDimension('send_size', len(self.send_dsts))
ncf.createDimension('recv_size', len(self.recv_dsts))
ncf.createDimension('3', 3)
vlocal_gids = ncf.createVariable('local_gids', 'i4', ('local_ep_size',))
vlocal_gids.long_name = 'Global index of local points'
vbuf_sizes = ncf.createVariable('buf_sizes', 'i4', ('3',))
vbuf_sizes.long_name = '[local_src_size, send_buf_size, recv_buf_size]'
vsend_schedule = ncf.createVariable('send_schedule', 'i4', ('send_sche_size','3',))
vsend_schedule.long_name = '[rank, start, size]'
vrecv_schedule = ncf.createVariable('recv_schedule', 'i4', ('recv_sche_size','3',))
vrecv_schedule.long_name = '[rank, start, size]'
vsend_dsts = ncf.createVariable('send_dsts', 'i4', ('send_size',))
vsend_dsts.long_name = 'Destination index for local and send buffer'
vsend_srcs = ncf.createVariable('send_srcs', 'i4', ('send_size',))
vsend_srcs.long_name = 'Source index for local and send buffer'
vsend_wgts = ncf.createVariable('send_wgts', 'f8', ('send_size',))
vsend_wgts.long_name = 'Weight value for local and send buffer'
vrecv_dsts = ncf.createVariable('recv_dsts', 'i4', ('recv_size',))
vrecv_dsts.long_name = 'Destination index for recv buffer'
vrecv_srcs = ncf.createVariable('recv_srcs', 'i4', ('recv_size',))
vrecv_srcs.long_name = 'Source index for recv buffer'
vlocal_gids[:] = self.cubegrid.local_gids[:]
vbuf_sizes[:] = (self.local_src_size, \
self.send_buf_size, \
self.recv_buf_size)
vsend_schedule[:] = self.send_schedule[:]
vrecv_schedule[:] = self.recv_schedule[:]
vsend_dsts[:] = self.send_dsts[:]
vsend_srcs[:] = self.send_srcs[:]
vsend_wgts[:] = self.send_wgts[:]
vrecv_dsts[:] = self.recv_dsts[:]
vrecv_srcs[:] = self.recv_srcs[:]
ncf.close()