def _gs_initialize_(self, regridding_role):
if regridding_role == RegriddingRole.SOURCE:
name = GridSplitterConstants.IndexFile.NAME_SRCIDX_GUID
elif regridding_role == RegriddingRole.DESTINATION:
name = GridSplitterConstants.IndexFile.NAME_DSTIDX_GUID
else:
raise NotImplementedError(regridding_role)
element_dim = self.element_dim
src_index = arange_from_dimension(element_dim, start=1, dtype=env.NP_INT)
src_index_var = Variable(name=name, value=src_index, dimensions=element_dim)
self.parent.add_variable(src_index_var)
def reduce_reindex_coordinate_index(cindex, start_index=0):
"""
Reindex a subset of global coordinate indices contained in the ``cindex`` variable.
The starting index value (``0`` or ``1``) is set by ``start_index`` for the re-indexing procedure.
Function will not respect masks.
The function returns a two-element tuple:
* First element --> A :class:`numpy.ndarray` with the same dimension as ``cindex`` containing the new indexing.
* Second element --> A :class:`numpy.ndarray` containing the unique indices that may be used to reduce an external
coordinate storage variable or array.
:param cindex: A variable containing coordinate index integer values. This variable may be distributed. This may
also be a NumPy array.
:type cindex: :class:`~ocgis.Variable` | :class:`~numpy.ndarray`
:param int start_index: The first index to use for the re-indexing of ``cindex``. This may be ``0`` or ``1``.
:rtype: tuple
"""
# Get the coordinate index values as a NumPy array.
try:
cindex = cindex.get_value()
except AttributeError:
# Assume this is already a NumPy array.
pass
# Only work with 1D arrays.
cindex = np.atleast_1d(cindex)
# Used to return the coordinate index to the original shape of the incoming coordinate index.
original_shape = cindex.shape
cindex = cindex.flatten()
# Create the unique coordinate index array.
# barrier_print('before create_unique_global_array')
u = np.array(create_unique_global_array(cindex))
# barrier_print('after create_unique_global_array')
# Synchronize the data type for the new coordinate index.
lrank = vm.rank
if lrank == 0:
dtype = u.dtype
else:
dtype = None
dtype = vm.bcast(dtype)
# Flag to indicate if the current rank has any unique values.
has_u = len(u) > 0
# Create the new coordinate index.
new_u_dimension = create_distributed_dimension(len(u), name='__new_u_dimension__')
new_u = arange_from_dimension(new_u_dimension, start=start_index, dtype=dtype)
# Create a hash for the new index. This is used to remap the old coordinate index.
if has_u:
uidx = {ii: jj for ii, jj in zip(u, new_u)}
else:
uidx = None
vm.barrier()
# Construct local bounds for the rank's unique value. This is used as a cheap index when ranks are looking for
# index overlaps.
if has_u:
local_bounds = min(u), max(u)
else:
local_bounds = None
# Put a copy for the bounds indexing on each rank.
lb_global = vm.gather(local_bounds)
lb_global = vm.bcast(lb_global)
# Find the vm ranks the local rank cares about. It cares if unique values have overlapping unique bounds.
overlaps = []
for rank, lb in enumerate(lb_global):
if rank == lrank:
continue
if lb is not None:
contains = lb[0] <= cindex
contains = np.logical_and(lb[1] >= cindex, contains)
if np.any(contains):
overlaps.append(rank)
# Ranks must be able to identify which ranks will be asking them for data.
global_overlaps = vm.gather(overlaps)
global_overlaps = vm.bcast(global_overlaps)
destinations = [ii for ii, jj in enumerate(global_overlaps) if vm.rank in jj]
# MPI communication tags used in the algorithm.
tag_search = MPITag.REDUCE_REINDEX_SEARCH
tag_success = MPITag.REDUCE_REINDEX_SUCCESS
tag_child_finished = MPITag.REDUCE_REINDEX_CHILD_FINISHED
tag_found = MPITag.REDUCE_REINDEX_FOUND
# Fill array for the new coordinate index.
new_cindex = np.empty_like(cindex)
# vm.barrier_print('starting run_rr')
# Fill the new coordinate indexing.
if lrank == 0:
run_rr_root(new_cindex, cindex, uidx, destinations, tag_child_finished, tag_found, tag_search, tag_success)
else:
run_rr_nonroot(new_cindex, cindex, uidx, destinations, has_u, overlaps, tag_child_finished, tag_found,
tag_search,
tag_success)
# vm.barrier_print('finished run_rr')
# Return array to its original shape.
new_cindex = new_cindex.reshape(*original_shape)
vm.barrier()
return new_cindex, u
def get_subset(bbox, subset_filename, buffer_width, rhs_tol=10.):
rd = ocgis.RequestDataset(uri=IN_PATH)
rd.metadata['dimensions']['nlandmesh_face']['dist'] = True
vc = rd.get_raw_field()
# ------------------------------------------------------------------------------------------------------------------
# Subset the face centers and accumulate the indices of face centers occurring inside the bounding box selection.
start_index = vc[MESHVAR].attrs.get('start_index', 0)
# Stores indices of faces contained in the bounding box.
px = vc[FACE_CENTER_X].extract().get_value()
py = vc[FACE_CENTER_Y].extract().get_value()
# Handle bounding box wrapping. This requires creating two bounding boxes to capture the left and right side of the
# sphere.
buffered_bbox = box(*bbox).buffer(buffer_width).envelope.bounds
if buffered_bbox[0] < 0:
bbox_rhs = list(deepcopy(buffered_bbox))
bbox_rhs[0] = buffered_bbox[0] + 360.
bbox_rhs[2] = 360. + rhs_tol
bboxes = [buffered_bbox, bbox_rhs]
else:
bboxes = [buffered_bbox]
initial = None
for ctr, curr_bbox in enumerate(bboxes):
select = create_boolean_select_array(curr_bbox, px, py, initial=initial)
initial = select
# ------------------------------------------------------------------------------------------------------------------
# Use the selection criteria to extract associated nodes and reindex the new coordinate arrays.
from ocgis.vmachine.mpi import rank_print
# Retrieve the live ranks following the subset.
has_select = ocgis.vm.gather(select.any())
if ocgis.vm.rank == 0:
live_ranks = np.array(ocgis.vm.ranks)[has_select]
else:
live_ranks = None
live_ranks = ocgis.vm.bcast(live_ranks)
with ocgis.vm.scoped('live ranks', live_ranks):
if not ocgis.vm.is_null:
has_subset = True
rank_print('live ranks:', ocgis.vm.ranks)
sub = vc[FACE_NODE].get_distributed_slice([select, slice(None)]).parent
cindex = sub[FACE_NODE]
cindex_original_shape = cindex.shape
cindex_value = cindex.get_value().flatten()
if start_index > 0:
cindex_value -= start_index
vc_coords = vc[XVAR][cindex_value].parent
archetype_dim = vc_coords[XVAR].dimensions[0]
arange_dimension = create_distributed_dimension(cindex_value.shape[0], name='arange_dim')
new_cindex_value = arange_from_dimension(arange_dimension, start=start_index)
cindex.set_value(new_cindex_value.reshape(*cindex_original_shape))
new_vc_coords_dimension = create_distributed_dimension(vc_coords[XVAR].shape[0], name=archetype_dim.name,
src_idx=archetype_dim._src_idx)
vc_coords.dimensions[archetype_dim.name] = new_vc_coords_dimension
# ------------------------------------------------------------------------------------------------------------------
# Format the new variable collection and write out the new data file.
# Remove old coordinate variables.
for to_modify in [XVAR, YVAR]:
sub[to_modify].extract(clean_break=True)
for to_add in [XVAR, YVAR]:
var_to_add = vc_coords[to_add].extract()
sub.add_variable(var_to_add)
rank_print('start sub.write')
sub.write(subset_filename)
rank_print('finished sub.write')
if ocgis.vm.rank == 0:
print 'subset x extent:', sub[FACE_CENTER_X].extent
print 'subset y extent:', sub[FACE_CENTER_Y].extent
ocgis.RequestDataset(subset_filename).inspect()
else:
has_subset = False
return has_subset