def get_variable_value(variable, dimensions):
if dimensions is not None and len(dimensions) > 0:
to_format = [None] * len(dimensions)
for idx in range(len(dimensions)):
current_dimension = dimensions[idx]
si_type = current_dimension._src_idx_type
if si_type is None:
if current_dimension.bounds_local is None:
to_insert = slice(0, len(current_dimension))
else:
to_insert = slice(*current_dimension.bounds_local)
elif si_type == SourceIndexType.FANCY:
to_insert = current_dimension._src_idx
elif si_type == SourceIndexType.BOUNDS:
to_insert = slice(*current_dimension._src_idx)
else:
raise NotImplementedError(si_type)
to_format[idx] = to_insert
slc = get_formatted_slice(to_format, len(dimensions))
else:
slc = slice(None)
try:
ret = variable.__getitem__(slc)
except IndexError:
# TODO: Hack! Slicing the top-level MFTime variable does not work with multifiles.
ret = super(MFTime, variable).__getitem__(slc)
return ret
def __getitem__(self,slc):
slc = get_formatted_slice(slc,2)
uid = self.uid[slc]
if self._value is not None:
value = self._value[:,slc[0],slc[1]]
else:
value = None
if self.row is not None:
row = self.row[slc[0]]
col = self.col[slc[1]]
else:
row = None
col = None
ret = copy(self)
if self._row_src_idx is not None:
ret._row_src_idx = self._row_src_idx[slc[0]]
ret._col_src_idx = self._col_src_idx[slc[1]]
ret.uid = uid
ret._value = value
ret.row = row
ret.col = col
# ret = SpatialGridDimension(value=value,uid=uid,row=row,col=col,name_value=self.name_value,
# units=self.units,meta=self.meta,name=self.name,name_uid=self.name_uid)
return(ret)
def __getitem__(self, slc):
"""
Dimensions may be sliced like other sliceable Python objects. A shallow copy of the dimension is created before
slicing. Use :meth:`~ocgis.Dimension.get_distributed_slice` for parallel slicing.
:param slc: A :class:`slice`-like object.
:rtype: :class:`~ocgis.Dimension`
:raises: IndexError
>>> dim = Dimension('five', 5)
>>> sub = dim[2:4]
>>> assert len(sub) == 2
>>> assert id(dim) != id(sub)
"""
# We cannot slice zero length dimensions.
if len(self) == 0:
raise IndexError('Zero-length dimensions are not slicable.')
slc = get_formatted_slice(slc, 1)[0]
ret = self.copy()
# Slicing work is done here.
self.__getitem_main__(ret, slc)
return ret
def get_variable_value(variable, dimensions):
if dimensions is not None and len(dimensions) > 0:
to_format = [None] * len(dimensions)
for idx in range(len(dimensions)):
current_dimension = dimensions[idx]
si_type = current_dimension._src_idx_type
if si_type is None:
if current_dimension.bounds_local is None:
to_insert = slice(0, len(current_dimension))
else:
to_insert = slice(*current_dimension.bounds_local)
elif si_type == SourceIndexType.FANCY:
to_insert = current_dimension._src_idx
elif si_type == SourceIndexType.BOUNDS:
to_insert = slice(*current_dimension._src_idx)
else:
raise NotImplementedError(si_type)
to_format[idx] = to_insert
slc = get_formatted_slice(to_format, len(dimensions))
else:
slc = slice(None)
try:
ret = variable.__getitem__(slc)
except IndexError:
# TODO: Hack! Slicing the top-level MFTime variable does not work with multifiles.
ret = super(MFTime, variable).__getitem__(slc)
return ret
def get_distributed_slice(self, slc, **kwargs):
"""
Slice a distributed object.
:param slc: A slice-like object.
:type slc: <varying>
:param dict kwargs: Optional arguments to :meth:`~ocgis.Variable.get_distributed_slice`.
:rtype: :class:`ocgis.spatial.geomc.AbstractGeometryCoordinates`
"""
slc = get_formatted_slice(slc, self.ndim)
if self.cindex is None:
target = self.x
else:
target = self.cindex
if len(slc) != target.ndim:
dslc = get_dslice(self.dimensions, slc)
for dim in target.dimensions:
if dim.name not in dslc:
dslc[dim.name] = slice(None)
slc = [dslc[dim.name] for dim in target.dimensions]
new_parent = target.get_distributed_slice(slc, **kwargs).parent
ret = self.copy()
ret.parent = new_parent
return ret
def __getitem__(self,slc):
slc = get_formatted_slice(slc,2)
uid = self.uid[slc]
if self._value is not None:
value = self._value[:,slc[0],slc[1]]
else:
value = None
if self.row is not None:
row = self.row[slc[0]]
col = self.col[slc[1]]
else:
row = None
col = None
ret = copy(self)
if self._row_src_idx is not None:
ret._row_src_idx = self._row_src_idx[slc[0]]
ret._col_src_idx = self._col_src_idx[slc[1]]
ret.uid = uid
ret._value = value
ret.row = row
ret.col = col
return(ret)
def __getitem__(self, slc):
"""
Dimensions may be sliced like other sliceable Python objects. A shallow copy of the dimension is created before
slicing. Use :meth:`~ocgis.Dimension.get_distributed_slice` for parallel slicing.
:param slc: A :class:`slice`-like object.
:rtype: :class:`~ocgis.Dimension`
:raises: IndexError
>>> dim = Dimension('five', 5)
>>> sub = dim[2:4]
>>> assert len(sub) == 2
>>> assert id(dim) != id(sub)
"""
# We cannot slice zero length dimensions.
if len(self) == 0:
raise IndexError('Zero-length dimensions are not slicable.')
slc = get_formatted_slice(slc, 1)[0]
ret = self.copy()
# Slicing work is done here.
self.__getitem_main__(ret, slc)
return ret
def __getitem__(self, slc):
slc = get_formatted_slice(slc, self.ndim)
if not isinstance(slc, dict):
slc = get_dslice(self.dimensions, slc)
ret = self.copy()
new_parent = ret.parent[slc]
ret.parent = new_parent
return ret
def __getitem__(self, slc):
slc = get_formatted_slice(slc, self._ndims)
ret = copy(self)
for attr in self._attrs_slice:
ref_set = get_none_or_slice(getattr(ret, attr), slc)
setattr(ret, attr, ref_set)
ret.properties = self._get_sliced_properties_(slc)
ret = self._format_slice_state_(ret, slc)
return ret
def __getitem__(self, item):
ret = SpatialGridDimension.__getitem__(self, item)
if ret._src_idx is not None:
slice_row, slice_col = get_formatted_slice(item, 2)
src_idx = {}
for key, slc in izip(['row', 'col'], [slice_row, slice_col]):
src_idx[key] = np.atleast_1d(ret._src_idx[key][slc])
ret._src_idx = src_idx
return ret
def get_variable(self, entry_key, parent=None, nullable=False):
"""
Get the coordinate variable name for the dimension map entry ``entry_key``.
:param str entry_key: See :class:`ocgis.constants.DimensionMapKey` for valid entry keys.
:param parent: If present, use the returned variable name to return the variable object form ``parent``.
:type parent: :class:`~ocgis.VariableCollection`
:param bool nullable: If ``True`` and ``parent`` is not ``None``, return ``None`` if the variable is not found
in ``parent``.
:rtype: str
"""
ret = self._get_element_(entry_key, DMK.VARIABLE, None)
# If there is an entry and a parent is provided, get the variable from the parent.
if ret is not None and parent is not None:
to_remove = []
base_variable = parent[ret]
base_variable_ndim = base_variable.ndim
has_sections = False
for ii in [DMK.X, DMK.Y]:
entry = self._get_entry_(ii)
section = entry.get(DMK.SECTION)
# Sections allow use to use a single variable as a source for multiple variables. In general, variables
# are atomic in the dimension map (x-coordinate is one variable). However, some metadata formats put
# both coordinates in a single variable (x/y-coordinate is one variable with the dimension name
# determining what the values represent.
if section is not None:
has_sections = True
section = get_formatted_slice(section, base_variable_ndim)
new_variable = base_variable[section]
new_dimensions = [
d for d in new_variable.dimensions if d.size > 1
]
new_variable.reshape(new_dimensions)
new_variable.set_name(ii)
entry[DMK.VARIABLE] = new_variable.name
entry.pop(DMK.SECTION)
new_variable = new_variable.extract()
parent.add_variable(new_variable)
if base_variable.name not in to_remove:
to_remove.append(base_variable.name)
if has_sections:
for tt in to_remove:
parent.remove_variable(tt)
ret = self._get_element_(entry_key, DMK.VARIABLE, None)
if ret is not None and parent is not None:
# Check if the variable has bounds.
bnds = self.get_bounds(entry_key)
ret = get_variable_from_field(ret, parent, nullable)
# Set the bounds on the outgoing variable if they are not already set by the object.
if bnds is not None and not ret.has_bounds:
ret.set_bounds(get_variable_from_field(bnds, parent, False),
force=True)
return ret
def __getitem__(self, slc):
slc = get_formatted_slice(slc, 5)
ret = copy(self)
ret.realization = get_none_or_slice(self.realization, slc[0])
ret.temporal = get_none_or_slice(self.temporal, slc[1])
ret.level = get_none_or_slice(self.level, slc[2])
ret.spatial = get_none_or_slice(self.spatial, (slc[3], slc[4]))
ret.variables = self.variables.get_sliced_variables(slc)
return ret
def get_variable(self, entry_key, parent=None, nullable=False):
"""
Get the coordinate variable name for the dimension map entry ``entry_key``.
:param str entry_key: See :class:`ocgis.constants.DimensionMapKey` for valid entry keys.
:param parent: If present, use the returned variable name to return the variable object from ``parent``.
:type parent: :class:`~ocgis.VariableCollection`
:param bool nullable: If ``True`` and ``parent`` is not ``None``, return ``None`` if the variable is not found
in ``parent``.
:rtype: str | None
"""
ret = self._get_element_(entry_key, DMK.VARIABLE, None)
# If there is an entry and a parent is provided, get the variable from the parent.
if ret is not None and parent is not None:
to_remove = []
base_variable = parent[ret]
base_variable_ndim = base_variable.ndim
has_sections = False
for ii in [DMK.X, DMK.Y]:
entry = self._get_entry_(ii)
section = entry.get(DMK.SECTION)
# Sections allow use to use a single variable as a source for multiple variables. In general, variables
# are atomic in the dimension map (x-coordinate is one variable). However, some metadata formats put
# both coordinates in a single variable (x/y-coordinate is one variable with the dimension name
# determining what the values represent.
if section is not None:
has_sections = True
section = get_formatted_slice(section, base_variable_ndim)
new_variable = base_variable[section]
new_dimensions = [d for d in new_variable.dimensions if d.size > 1]
new_variable.reshape(new_dimensions)
new_variable.set_name(base_variable.name + '_' + ii)
entry[DMK.VARIABLE] = new_variable.name
entry.pop(DMK.SECTION)
new_variable = new_variable.extract()
parent.add_variable(new_variable)
if base_variable.name not in to_remove:
to_remove.append(base_variable.name)
self.set_variable(ii, new_variable)
if has_sections:
for tt in to_remove:
parent.remove_variable(tt)
ret = self._get_element_(entry_key, DMK.VARIABLE, None)
if ret is not None and parent is not None:
# Check if the variable has bounds.
bnds = self.get_bounds(entry_key)
ret = get_variable_from_field(ret, parent, nullable)
# Set the bounds on the outgoing variable if they are not already set by the object.
if bnds is not None and not ret.has_bounds:
ret.set_bounds(get_variable_from_field(bnds, parent, False), force=True)
return ret
def test_get_formatted_slc(self):
ret = get_formatted_slice(slice(None,None,None),10)
self.assertEqual(ret,[slice(None,None,None)]*10)
ret = get_formatted_slice(0,1)
self.assertEqual(slice(0,1),ret)
with self.assertRaises(IndexError):
get_formatted_slice(slice(0,1),2)
ret = get_formatted_slice((slice(0,1),0),2)
self.assertEqual(ret,[slice(0,1,None),slice(0,1,None)])
ret = get_formatted_slice([(1,2,3),slice(None)],2)
self.assertNumpyAll(ret[0],np.arange(1,4))
self.assertEqual(ret[1],slice(None))
self.assertEqual(len(ret),2)
ret = get_formatted_slice((1,2),1)
self.assertNumpyAll(ret,np.array([1,2]))
ret = get_formatted_slice((1,),1)
self.assertEqual(ret,slice(1))
def get_variable_value(variable, dimensions):
if dimensions is not None and len(dimensions) > 0:
to_format = [None] * len(dimensions)
for idx in range(len(dimensions)):
current_dimension = dimensions[idx]
si_type = current_dimension._src_idx_type
if si_type is None:
if current_dimension.bounds_local is None:
to_insert = slice(0, len(current_dimension))
else:
to_insert = slice(*current_dimension.bounds_local)
elif si_type == SourceIndexType.FANCY:
to_insert = current_dimension._src_idx
elif si_type == SourceIndexType.BOUNDS:
to_insert = slice(*current_dimension._src_idx)
else:
raise NotImplementedError(si_type)
to_format[idx] = to_insert
slc = get_formatted_slice(to_format, len(dimensions))
else:
slc = slice(None)
ret = variable.__getitem__(slc)
return ret
def get_index_slice_for_iteration(slc):
slc = get_formatted_slice(slc, 1)[0]
return slc
def get_index_slice_for_iteration(slc):
slc = get_formatted_slice(slc, 1)[0]
if isinstance(slc, slice):
slc = np.arange(slc.start, slc.stop)
return slc
def get_distributed_slice(self, slc):
"""
Slice the dimension in parallel. The sliced dimension object is a shallow copy. The returned dimension may be
empty.
:param slc: A :class:`slice`-like object or a fancy slice. If this is a fancy slice, ``slc`` must be
processor-local. If the fancy slice uses integer indices, the indices must be local. In other words, a fancy
``slc`` is not manipulated or redistributed prior to slicing.
:rtype: :class:`~ocgis.Dimension`
:raises: :class:`~ocgis.exc.EmptyObjectError`
"""
raise_if_empty(self)
slc = get_formatted_slice(slc, 1)[0]
is_fancy = not isinstance(slc, slice)
if not is_fancy and slc == slice(None):
ret = self.copy()
# Use standard slicing for non-distributed dimensions.
elif not self.dist:
ret = self[slc]
else:
if is_fancy:
local_slc = slc
else:
local_slc = get_global_to_local_slice((slc.start, slc.stop),
self.bounds_local)
if local_slc is not None:
local_slc = slice(*local_slc)
# Slice does not overlap local bounds. The dimension is now empty with size 0.
if local_slc is None:
ret = self.copy()
ret.convert_to_empty()
dimension_size = 0
# Slice overlaps so do a slice on the dimension using the local slice.
else:
ret = self[local_slc]
dimension_size = len(ret)
assert dimension_size >= 0
dimension_sizes = vm.gather(dimension_size)
if vm.rank == 0:
sum_dimension_size = 0
for ds in dimension_sizes:
try:
sum_dimension_size += ds
except TypeError:
pass
bounds_global = (0, sum_dimension_size)
else:
bounds_global = None
bounds_global = vm.bcast(bounds_global)
if not ret.is_empty:
ret.bounds_global = bounds_global
# Normalize the local bounds on live ranks.
inner_live_ranks = get_nonempty_ranks(ret, vm)
with vm.scoped('bounds normalization', inner_live_ranks):
if not vm.is_null:
if vm.rank == 0:
adjust = len(ret)
else:
adjust = None
adjust = vm.bcast(adjust)
for current_rank in vm.ranks:
if vm.rank == current_rank:
if vm.rank != 0:
ret.bounds_local = [
b + adjust for b in ret.bounds_local
]
adjust += len(ret)
vm.barrier()
adjust = vm.bcast(adjust, root=current_rank)
return ret
def get_distributed_slice(self, slc):
"""
Slice the dimension in parallel. The sliced dimension object is a shallow copy. The returned dimension may be
empty.
:param slc: A :class:`slice`-like object or a fancy slice. If this is a fancy slice, ``slc`` must be
processor-local. If the fancy slice uses integer indices, the indices must be local. In other words, a fancy
``slc`` is not manipulated or redistributed prior to slicing.
:rtype: :class:`~ocgis.Dimension`
:raises: :class:`~ocgis.exc.EmptyObjectError`
"""
raise_if_empty(self)
slc = get_formatted_slice(slc, 1)[0]
is_fancy = not isinstance(slc, slice)
if not is_fancy and slc == slice(None):
ret = self.copy()
# Use standard slicing for non-distributed dimensions.
elif not self.dist:
ret = self[slc]
else:
if is_fancy:
local_slc = slc
else:
local_slc = get_global_to_local_slice((slc.start, slc.stop), self.bounds_local)
if local_slc is not None:
local_slc = slice(*local_slc)
# Slice does not overlap local bounds. The dimension is now empty with size 0.
if local_slc is None:
ret = self.copy()
ret.convert_to_empty()
dimension_size = 0
# Slice overlaps so do a slice on the dimension using the local slice.
else:
ret = self[local_slc]
dimension_size = len(ret)
assert dimension_size >= 0
dimension_sizes = vm.gather(dimension_size)
if vm.rank == 0:
sum_dimension_size = 0
for ds in dimension_sizes:
try:
sum_dimension_size += ds
except TypeError:
pass
bounds_global = (0, sum_dimension_size)
else:
bounds_global = None
bounds_global = vm.bcast(bounds_global)
if not ret.is_empty:
ret.bounds_global = bounds_global
# Normalize the local bounds on live ranks.
inner_live_ranks = get_nonempty_ranks(ret, vm)
with vm.scoped('bounds normalization', inner_live_ranks):
if not vm.is_null:
if vm.rank == 0:
adjust = len(ret)
else:
adjust = None
adjust = vm.bcast(adjust)
for current_rank in vm.ranks:
if vm.rank == current_rank:
if vm.rank != 0:
ret.bounds_local = [b + adjust for b in ret.bounds_local]
adjust += len(ret)
vm.barrier()
adjust = vm.bcast(adjust, root=current_rank)
return ret
def _get_value_from_source_(self):
try:
ret = SpatialGridDimension._get_value_(self)
except AttributeError:
if self.row is None or self.col is None:
ds = self._request_dataset.driver.open()
try:
slices = {k: get_formatted_slice(self._src_idx[k], 1) for k in self._src_idx.keys()}
slice_row = slices['row']
slice_col = slices['col']
variable_row = ds.variables[self.name_row]
variable_col = ds.variables[self.name_col]
# Load values ######################################################################################
value_row = np.atleast_2d(variable_row[slice_row, slice_col])
value_col = np.atleast_2d(variable_col[slice_row, slice_col])
fill = np.zeros([2] + list(value_row.shape), dtype=value_row.dtype)
try:
fill_value = value_row.fill_value
except AttributeError:
fill_value = None
fill = np.ma.array(fill, fill_value=fill_value, mask=False)
fill[0, :, :] = value_row
fill[1, :, :] = value_col
ret = fill
# Load corners #####################################################################################
try:
name_row_corners = variable_row.corners
except AttributeError:
# Likely no corners.
pass
else:
name_col_corners = variable_col.corners
value_row_corners = ds.variables[name_row_corners][slice_row, slice_col, :]
value_col_corners = ds.variables[name_col_corners][slice_row, slice_col, :]
# A reshape may be required if this is a singleton slice operation.
def _reshape_corners_(arr):
if arr.ndim < 3:
assert arr.shape == (1, 4)
arr = arr.reshape(1, 1, 4)
return arr
value_row_corners = _reshape_corners_(value_row_corners)
value_col_corners = _reshape_corners_(value_col_corners)
fill = np.zeros([2] + list(value_row_corners.shape), dtype=value_row_corners.dtype)
try:
fill_value = value_row_corners.fill_value
except AttributeError:
fill_value = None
fill = np.ma.array(fill, fill_value=fill_value, mask=False)
fill[0, :, :, :] = value_row_corners
fill[1, :, :, :] = value_col_corners
self.corners = fill
finally:
self._request_dataset.driver.close(ds)
else:
raise
return ret
