def test_expand_grid(self):
x = [101, 102, 103]
y = [40, 41, 42, 43]
vx = Variable('x', value=x, dtype=float, dimensions='xdim')
vx.set_extrapolated_bounds('x_bnds', 'bounds')
vy = Variable('y', value=y, dtype=float, dimensions='ydim')
vy.set_extrapolated_bounds('y_bnds', 'bounds')
grid = Grid(vx, vy)
for variable in [vx, vy]:
self.assertEqual(grid.parent[variable.name].ndim, 1)
expand_grid(grid)
for variable in [vx, vy]:
self.assertEqual(grid.parent[variable.name].ndim, 2)
def get_esmf_grid(ogrid, regrid_method='auto', value_mask=None):
"""
Create an ESMF :class:`~ESMF.driver.grid.Grid` object from an OCGIS :class:`ocgis.Grid` object.
:param ogrid: The target OCGIS grid to convert to an ESMF grid.
:type ogrid: :class:`~ocgis.Grid`
:param regrid_method: If ``'auto'`` or :attr:`ESMF.api.constants.RegridMethod.CONSERVE`, use corners/bounds from the
grid object. If :attr:`ESMF.api.constants.RegridMethod.CONSERVE`, the corners/bounds must be present on the grid.
:param value_mask: If an ``bool`` :class:`numpy.array` with same shape as ``ogrid``, use a logical *or* operation
with the OCGIS field mask when creating the input grid's mask. Values of ``True`` in ``value_mask`` are assumed to
be masked. If ``None`` is provided (the default) then the mask will be set using the spatial mask on ``ogrid``.
:type value_mask: :class:`numpy.array`
:rtype: :class:`ESMF.driver.grid.Grid`
"""
pkwds = get_periodicity_parameters(ogrid)
egrid = ESMF.Grid(max_index=np.array(ogrid.shape),
staggerloc=ESMF.StaggerLoc.CENTER,
coord_sys=ESMF.CoordSys.SPH_DEG,
num_peri_dims=pkwds['num_peri_dims'],
pole_dim=pkwds['pole_dim'],
periodic_dim=pkwds['periodic_dim'])
ovalue_stacked = ogrid.get_value_stacked()
row = egrid.get_coords(1, staggerloc=ESMF.StaggerLoc.CENTER)
row[:] = ovalue_stacked[0, ...]
col = egrid.get_coords(0, staggerloc=ESMF.StaggerLoc.CENTER)
col[:] = ovalue_stacked[1, ...]
# Use a logical or operation to merge with value_mask if present
if value_mask is not None:
# convert to boolean to make sure
value_mask = value_mask.astype(bool)
# do the logical or operation selecting values
value_mask = np.logical_or(value_mask, ogrid.get_mask(create=True))
else:
value_mask = ogrid.get_mask()
# Follows SCRIP convention where 1 is unmasked and 0 is masked.
if value_mask is not None:
esmf_mask = np.invert(value_mask).astype(np.int32)
egrid.add_item(ESMF.GridItem.MASK,
staggerloc=ESMF.StaggerLoc.CENTER,
from_file=False)
egrid.mask[0][:] = esmf_mask
# Attempt to access corners if requested.
if regrid_method == 'auto' and ogrid.has_bounds:
regrid_method = RegridMethod.CONSERVE
if regrid_method == RegridMethod.CONSERVE:
# Conversion to ESMF objects requires an expanded grid (non-vectorized).
# TODO: Create ESMF grids from vectorized OCGIS grids.
expand_grid(ogrid)
# Convert to ESMF corners from OCGIS corners.
corners_esmf = np.zeros(
[2] + [element + 1 for element in ogrid.x.bounds.shape[0:2]],
dtype=ogrid.archetype.bounds.dtype)
get_esmf_corners_from_ocgis_corners(ogrid.y.bounds.get_value(),
fill=corners_esmf[0, :, :])
get_esmf_corners_from_ocgis_corners(ogrid.x.bounds.get_value(),
fill=corners_esmf[1, :, :])
# adding corners. first tell the grid object to allocate corners
egrid.add_coords(staggerloc=[ESMF.StaggerLoc.CORNER])
# get the coordinate pointers and set the coordinates
grid_corner = egrid.coords[ESMF.StaggerLoc.CORNER]
# Note indexing is reversed for ESMF v. OCGIS. In ESMF, the x-coordinate (longitude) is the first
# coordinate. If this is periodic, the last column corner wraps/connect to the first. This coordinate must
# be removed.
if pkwds['num_peri_dims'] is not None:
grid_corner[0][:] = corners_esmf[1][:, 0:-1]
grid_corner[1][:] = corners_esmf[0][:, 0:-1]
else:
grid_corner[0][:] = corners_esmf[1]
grid_corner[1][:] = corners_esmf[0]
return egrid
def test_get_intersects_ordering(self):
"""Test grid ordering/origins do not influence grid subsetting."""
keywords = {
KeywordArgument.OPTIMIZED_BBOX_SUBSET: [False, True],
'should_wrap': [False, True],
'reverse_x': [False, True],
'reverse_y': [False, True],
'should_expand': [False, True],
}
x_value = np.array(
[155., 160., 165., 170., 175., 180., 185., 190., 195., 200., 205.])
y_value = np.array([-20., -15., -10., -5., 0., 5., 10., 15., 20.])
bbox = [168., -12., 191., 5.3]
for k in self.iter_product_keywords(keywords, as_namedtuple=False):
reverse_x = k.pop('reverse_x')
reverse_y = k.pop('reverse_y')
should_expand = k.pop('should_expand')
should_wrap = k.pop('should_wrap')
ompi = OcgDist()
ompi.create_dimension('dx', len(x_value), dist=True)
ompi.create_dimension('dy', len(y_value))
ompi.update_dimension_bounds()
if reverse_x:
new_x_value = x_value.copy()
new_x_value = np.flipud(new_x_value)
else:
new_x_value = x_value
if reverse_y:
new_y_value = y_value.copy()
new_y_value = np.flipud(new_y_value)
else:
new_y_value = y_value
if MPI_RANK == 0:
x = Variable('x', new_x_value, 'dx')
y = Variable('y', new_y_value, 'dy')
else:
x, y = [None, None]
x = variable_scatter(x, ompi)
y = variable_scatter(y, ompi)
grid = Grid(x, y, crs=Spherical())
with vm.scoped_by_emptyable('scattered', grid):
if not vm.is_null:
if should_expand:
expand_grid(grid)
if should_wrap:
grid = deepcopy(grid)
grid.wrap()
actual_bbox = MultiPolygon([
box(-180, -12, -169, 5.3),
box(168, -12, 180, 5.3)
])
else:
actual_bbox = box(*bbox)
live_ranks = vm.get_live_ranks_from_object(grid)
with vm.scoped('grid.get_intersects', live_ranks):
if not vm.is_null:
sub = grid.get_intersects(actual_bbox, **k)
with vm.scoped_by_emptyable('sub grid', sub):
if not vm.is_null:
if should_wrap:
current_x_value = sub.x.get_value()
current_x_value[
sub.x.get_value() < 0] += 360
self.assertEqual(
sub.extent_global,
(170.0, -10.0, 190.0, 5.0))
if should_expand:
desired = False
else:
desired = True
self.assertEqual(grid.is_vectorized,
desired)
self.assertEqual(sub.is_vectorized,
desired)
self.assertFalse(grid.has_allocated_point)
self.assertFalse(
grid.has_allocated_polygon)
def create_esmf_grid(ogrid, regrid_method='auto', value_mask=None):
"""
Create an ESMF :class:`~ESMF.driver.grid.Grid` object from an OCGIS :class:`ocgis.Grid` object.
:param ogrid: The target OCGIS grid to convert to an ESMF grid.
:type ogrid: :class:`~ocgis.Grid`
:param regrid_method: If ``'auto'`` or :attr:`ESMF.api.constants.RegridMethod.CONSERVE`, use corners/bounds from the
grid object. If :attr:`ESMF.api.constants.RegridMethod.CONSERVE`, the corners/bounds must be present on the grid.
:param value_mask: If an ``bool`` :class:`numpy.array` with same shape as ``ogrid``, use a logical *or* operation
with the OCGIS field mask when creating the input grid's mask. Values of ``True`` in ``value_mask`` are assumed to
be masked. If ``None`` is provided (the default) then the mask will be set using the spatial mask on ``ogrid``.
:type value_mask: :class:`numpy.array`
:rtype: :class:`ESMF.driver.grid.Grid`
"""
assert isinstance(ogrid, Grid)
assert create_crs(ogrid.crs) == Spherical()
pkwds = get_periodicity_parameters(ogrid)
# Fill ESMPy grid coordinate values. ###############################################################################
dimensions = ogrid.parent.dimensions
get_dimension = ogrid.dimension_map.get_dimension
y_dimension = get_dimension(DMK.Y, dimensions=dimensions)
x_dimension = get_dimension(DMK.X, dimensions=dimensions)
# ESMPy has index 0 = x-coordinate and index 1 = y-coordinate.
if vm.size == 1:
max_index = np.array([x_dimension.size, y_dimension.size],
dtype=np.int32)
else:
max_index = np.array(
[x_dimension.size_global, y_dimension.size_global], dtype=np.int32)
pkwds['coord_sys'] = ESMF.CoordSys.SPH_DEG
pkwds['staggerloc'] = ESMF.StaggerLoc.CENTER
pkwds['max_index'] = max_index
egrid = ESMF.Grid(**pkwds)
# msg = (egrid.lower_bounds, egrid.upper_bounds)
# ocgis_lh(msg, level=logging.WARN, logger='ocli.chunked_rwg', force=True)
egrid._ocgis['dimnames'] = (x_dimension.name, y_dimension.name)
egrid._ocgis['dimnames_backref'] = get_dimension_names(ogrid.dimensions)
egrid._ocgis['dimension_map'] = deepcopy(ogrid.dimension_map)
ogrid_dimnames = (y_dimension.name, x_dimension.name)
is_yx_order = get_dimension_names(ogrid.dimensions) == ogrid_dimnames
ovalue_stacked = ogrid.get_value_stacked()
row = egrid.get_coords(1, staggerloc=ESMF.StaggerLoc.CENTER)
orow = ovalue_stacked[0, ...]
if is_yx_order:
orow = np.swapaxes(orow, 0, 1)
row[:] = orow
col = egrid.get_coords(0, staggerloc=ESMF.StaggerLoc.CENTER)
ocol = ovalue_stacked[1, ...]
if is_yx_order:
ocol = np.swapaxes(ocol, 0, 1)
col[:] = ocol
####################################################################################################################
# Use a logical or operation to merge with value_mask if present
if value_mask is not None:
# convert to boolean to make sure
value_mask = value_mask.astype(bool)
# do the logical or operation selecting values
value_mask = np.logical_or(value_mask, ogrid.get_mask(create=True))
else:
value_mask = ogrid.get_mask()
# Follows SCRIP convention where 1 is unmasked and 0 is masked.
if value_mask is not None:
esmf_mask = np.invert(value_mask).astype(np.int32)
esmf_mask = np.swapaxes(esmf_mask, 0, 1)
egrid.add_item(ESMF.GridItem.MASK,
staggerloc=ESMF.StaggerLoc.CENTER,
from_file=False)
egrid.mask[0][:] = esmf_mask
# Attempt to access corners if requested.
if regrid_method == 'auto' and ogrid.has_bounds:
regrid_method = ESMF.RegridMethod.CONSERVE
if regrid_method == ESMF.RegridMethod.CONSERVE:
# Conversion to ESMF objects requires an expanded grid (non-vectorized).
# TODO: Create ESMF grids from vectorized OCGIS grids.
expand_grid(ogrid)
# Convert to ESMF corners from OCGIS corners.
corners_esmf = np.zeros([2] + [element + 1 for element in max_index],
dtype=col.dtype)
get_esmf_corners_from_ocgis_corners(ogrid.y.bounds.get_value(),
fill=corners_esmf[0, :, :])
get_esmf_corners_from_ocgis_corners(ogrid.x.bounds.get_value(),
fill=corners_esmf[1, :, :])
# adding corners. first tell the grid object to allocate corners
egrid.add_coords(staggerloc=[ESMF.StaggerLoc.CORNER])
# get the coordinate pointers and set the coordinates
grid_corner = egrid.coords[ESMF.StaggerLoc.CORNER]
# Note indexing is reversed for ESMF v. OCGIS. In ESMF, the x-coordinate (longitude) is the first
# coordinate. If this is periodic, the last column corner wraps/connect to the first. This coordinate must
# be removed.
if pkwds['num_peri_dims'] is not None:
grid_corner[0][:] = corners_esmf[1][0:-1, :]
grid_corner[1][:] = corners_esmf[0][0:-1, :]
else:
grid_corner[0][:] = corners_esmf[1]
grid_corner[1][:] = corners_esmf[0]
return egrid
def create_esmf_grid(ogrid, regrid_method='auto', value_mask=None):
"""
Create an ESMF :class:`~ESMF.driver.grid.Grid` object from an OCGIS :class:`ocgis.Grid` object.
:param ogrid: The target OCGIS grid to convert to an ESMF grid.
:type ogrid: :class:`~ocgis.Grid`
:param regrid_method: If ``'auto'`` or :attr:`ESMF.api.constants.RegridMethod.CONSERVE`, use corners/bounds from the
grid object. If :attr:`ESMF.api.constants.RegridMethod.CONSERVE`, the corners/bounds must be present on the grid.
:param value_mask: If an ``bool`` :class:`numpy.array` with same shape as ``ogrid``, use a logical *or* operation
with the OCGIS field mask when creating the input grid's mask. Values of ``True`` in ``value_mask`` are assumed to
be masked. If ``None`` is provided (the default) then the mask will be set using the spatial mask on ``ogrid``.
:type value_mask: :class:`numpy.array`
:rtype: :class:`ESMF.driver.grid.Grid`
"""
assert isinstance(ogrid, Grid)
assert create_crs(ogrid.crs) == Spherical()
pkwds = get_periodicity_parameters(ogrid)
# Fill ESMPy grid coordinate values. ###############################################################################
dimensions = ogrid.parent.dimensions
get_dimension = ogrid.dimension_map.get_dimension
y_dimension = get_dimension(DMK.Y, dimensions=dimensions)
x_dimension = get_dimension(DMK.X, dimensions=dimensions)
# ESMPy has index 0 = x-coordinate and index 1 = y-coordinate.
if vm.size == 1:
max_index = np.array([x_dimension.size, y_dimension.size], dtype=np.int32)
else:
max_index = np.array([x_dimension.size_global, y_dimension.size_global], dtype=np.int32)
pkwds['coord_sys'] = ESMF.CoordSys.SPH_DEG
pkwds['staggerloc'] = ESMF.StaggerLoc.CENTER
pkwds['max_index'] = max_index
egrid = ESMF.Grid(**pkwds)
# msg = (egrid.lower_bounds, egrid.upper_bounds)
# ocgis_lh(msg, level=logging.WARN, logger='ocli.chunked_rwg', force=True)
egrid._ocgis['dimnames'] = (x_dimension.name, y_dimension.name)
egrid._ocgis['dimnames_backref'] = get_dimension_names(ogrid.dimensions)
egrid._ocgis['dimension_map'] = deepcopy(ogrid.dimension_map)
ogrid_dimnames = (y_dimension.name, x_dimension.name)
is_yx_order = get_dimension_names(ogrid.dimensions) == ogrid_dimnames
ovalue_stacked = ogrid.get_value_stacked()
row = egrid.get_coords(1, staggerloc=ESMF.StaggerLoc.CENTER)
orow = ovalue_stacked[0, ...]
if is_yx_order:
orow = np.swapaxes(orow, 0, 1)
row[:] = orow
col = egrid.get_coords(0, staggerloc=ESMF.StaggerLoc.CENTER)
ocol = ovalue_stacked[1, ...]
if is_yx_order:
ocol = np.swapaxes(ocol, 0, 1)
col[:] = ocol
####################################################################################################################
# Use a logical or operation to merge with value_mask if present
if value_mask is not None:
# convert to boolean to make sure
value_mask = value_mask.astype(bool)
# do the logical or operation selecting values
value_mask = np.logical_or(value_mask, ogrid.get_mask(create=True))
else:
value_mask = ogrid.get_mask()
# Follows SCRIP convention where 1 is unmasked and 0 is masked.
if value_mask is not None:
esmf_mask = np.invert(value_mask).astype(np.int32)
esmf_mask = np.swapaxes(esmf_mask, 0, 1)
egrid.add_item(ESMF.GridItem.MASK, staggerloc=ESMF.StaggerLoc.CENTER, from_file=False)
egrid.mask[0][:] = esmf_mask
# Attempt to access corners if requested.
if regrid_method == 'auto' and ogrid.has_bounds:
regrid_method = ESMF.RegridMethod.CONSERVE
if regrid_method == ESMF.RegridMethod.CONSERVE:
# Conversion to ESMF objects requires an expanded grid (non-vectorized).
# TODO: Create ESMF grids from vectorized OCGIS grids.
expand_grid(ogrid)
# Convert to ESMF corners from OCGIS corners.
corners_esmf = np.zeros([2] + [element + 1 for element in max_index], dtype=col.dtype)
get_esmf_corners_from_ocgis_corners(ogrid.y.bounds.get_value(), fill=corners_esmf[0, :, :])
get_esmf_corners_from_ocgis_corners(ogrid.x.bounds.get_value(), fill=corners_esmf[1, :, :])
# adding corners. first tell the grid object to allocate corners
egrid.add_coords(staggerloc=[ESMF.StaggerLoc.CORNER])
# get the coordinate pointers and set the coordinates
grid_corner = egrid.coords[ESMF.StaggerLoc.CORNER]
# Note indexing is reversed for ESMF v. OCGIS. In ESMF, the x-coordinate (longitude) is the first
# coordinate. If this is periodic, the last column corner wraps/connect to the first. This coordinate must
# be removed.
if pkwds['num_peri_dims'] is not None:
grid_corner[0][:] = corners_esmf[1][0:-1, :]
grid_corner[1][:] = corners_esmf[0][0:-1, :]
else:
grid_corner[0][:] = corners_esmf[1]
grid_corner[1][:] = corners_esmf[0]
return egrid