def test_get_esmf_grid_with_mask(self):
"""Test with masked data."""
from ocgis.regrid.base import get_esmf_grid
x = Variable(name='x', value=[1, 2, 3], dimensions='x')
y = Variable(name='y', value=[4, 5, 6], dimensions='y')
grid = Grid(x, y, crs=Spherical())
gmask = grid.get_mask(create=True)
gmask[1, 1] = True
grid.set_mask(gmask)
self.assertEqual(grid.get_mask().sum(), 1)
egrid = get_esmf_grid(grid)
egrid_mask_inverted = np.invert(np.array(egrid.mask[0], dtype=bool))
self.assertNumpyAll(grid.get_mask(), egrid_mask_inverted)
# Test with a value mask.
value_mask = np.zeros(grid.shape, dtype=bool)
value_mask[-1, -1] = True
egrid = get_esmf_grid(grid, value_mask=value_mask)
egrid_mask_inverted = np.invert(np.array(egrid.mask[0], dtype=bool))
self.assertNumpyAll(egrid_mask_inverted,
np.logical_or(grid.get_mask(), value_mask))
def test_system(self):
from ocgis.regrid.base import create_esmf_grid, iter_esmf_fields, RegridOperation, destroy_esmf_objects
import ESMF
yc = Variable(name='yc', value=np.arange(-90 + (45 / 2.), 90, 45), dimensions='ydim', dtype=float)
xc = Variable(name='xc', value=np.arange(15, 360, 30), dimensions='xdim', dtype=float)
ogrid = Grid(y=yc, x=xc, crs=Spherical())
ogrid.set_extrapolated_bounds('xc_bounds', 'yc_bounds', 'bounds')
np.random.seed(1)
mask = np.random.rand(*ogrid.shape)
mask = mask > 0.5
self.assertTrue(mask.sum() > 3)
ogrid.set_mask(mask)
egrid = create_esmf_grid(ogrid)
actual_shape = egrid.size[0].tolist()
desired_shape = np.flipud(ogrid.shape).tolist()
self.assertEqual(actual_shape, desired_shape)
desired = ogrid.get_value_stacked()
desired = np.ma.array(desired, mask=False)
desired.mask[0, :, :] = ogrid.get_mask()
desired.mask[1, :, :] = ogrid.get_mask()
desired = desired.sum()
actual_col = egrid.get_coords(0)
actual_row = egrid.get_coords(1)
actual_mask = np.invert(egrid.mask[0].astype(bool))
actual = np.ma.array(actual_row, mask=actual_mask).sum() + np.ma.array(actual_col, mask=actual_mask).sum()
self.assertEqual(actual, desired)
desired = 9900.0
corners = egrid.coords[ESMF.StaggerLoc.CORNER]
actual = corners[0].sum() + corners[1].sum()
self.assertEqual(actual, desired)
ofield = create_exact_field(ogrid, 'data', ntime=3, crs=Spherical())
variable_name, efield, tidx = list(iter_esmf_fields(ofield, split=False))[0]
desired_value = ofield['data'].get_value()
self.assertAlmostEqual(efield.data.sum(), desired_value.sum(), places=3)
destroy_esmf_objects([egrid, efield])
ofield.grid.set_mask(ofield.grid.get_mask(), cascade=True)
desired_value = ofield['data'].get_masked_value()
keywords = dict(split=[False, True])
for k in self.iter_product_keywords(keywords):
opts = {'split': k.split}
dofield = ofield.deepcopy()
dofield['data'].get_value().fill(0)
ro = RegridOperation(ofield, dofield, regrid_options=opts)
actual_field = ro.execute()
actual_value = actual_field['data'].get_masked_value()
self.assertAlmostEqual(0.0, np.abs(desired_value - actual_value).max())
def test_get_intersects_bounds_sequence(self):
keywords = dict(bounds=[True, False], use_bounds=[True, False])
for ctr, k in enumerate(self.iter_product_keywords(keywords)):
# self.log.debug([ctr, k])
# if ctr != 1: continue
# Bounds may not be used if they are not present.
if k.use_bounds and not k.bounds:
continue
y = self.get_variable_y(bounds=k.bounds)
x = self.get_variable_x(bounds=k.bounds)
grid = Grid(x, y)
bounds_sequence = box(-99, 39, -98, 39)
bg = grid.get_intersects(bounds_sequence)
self.assertNotEqual(grid.shape, bg.shape)
self.assertTrue(bg.is_vectorized)
with self.assertRaises(EmptySubsetError):
bounds_sequence = box(1000, 1000, 1001, 10001)
grid.get_intersects(bounds_sequence, use_bounds=k.use_bounds)
bounds_sequence = box(-99999, 1, 1, 1000)
bg2 = grid.get_intersects(bounds_sequence, use_bounds=k.use_bounds)
for target in ['x', 'y']:
original = getattr(grid, target).get_value()
sub = getattr(bg2, target).get_value()
self.assertNumpyAll(original, sub)
# Test mask is not shared with subsetted grid.
grid = self.get_gridxy()
new_mask = grid.get_mask(create=True)
new_mask[:, 1] = True
grid.set_mask(new_mask)
self.assertFalse(grid.has_bounds)
bounds_sequence = box(101.5, 40.5, 103.5, 42.5)
sub = grid.get_intersects(bounds_sequence, use_bounds=False)
new_mask = sub.get_mask()
new_mask.fill(True)
sub.set_mask(new_mask)
self.assertEqual(grid.get_mask().sum(), 4)
def test_system_masking(self):
"""Test behavior of the grid mask. This is an independently managed variable."""
x = Variable('xc', value=[1, 2, 3], dimensions='dimx')
y = Variable('yc', value=[10, 20, 30, 40], dimensions='dimy')
grid = Grid(x, y)
data = Variable('data',
value=np.zeros(grid.shape),
dimensions=['dimy', 'dimx'])
grid.parent.add_variable(data)
gmask = grid.get_mask()
self.assertIsNone(gmask)
self.assertIsNone(grid.mask_variable)
new_mask = np.zeros(grid.shape, dtype=bool)
new_mask[1, 1] = True
grid.set_mask(new_mask, cascade=True)
self.assertIsInstance(grid.mask_variable, Variable)
actual = grid.get_mask()
self.assertNumpyAll(actual, new_mask)
actual = get_variable_names(grid.get_member_variables())
desired = [x.name, y.name, grid._mask_name]
self.assertAsSetEqual(actual, desired)
self.assertNumpyAll(grid.get_mask(), data.get_mask())
path = self.get_temporary_file_path('foo.nc')
grid.parent.write(path)
with self.nc_scope(path) as ds:
actual = ds.variables[grid.mask_variable.name]
self.assertNumpyAll(grid.get_mask(), actual[:].mask)
# Test mask is used when read from file.
actual_field = RequestDataset(path).get()
self.assertNumpyAll(grid.get_mask(), actual_field.grid.get_mask())
self.assertEqual(actual_field.grid.get_mask().sum(), 1)
self.assertTrue(actual_field.grid.is_vectorized)
self.assertEqual(actual_field.grid.get_mask().dtype, bool)
actual_field.set_abstraction_geom()
self.assertNumpyAll(actual_field.geom.get_mask(), grid.get_mask())
def test_get_esmf_grid_with_mask(self):
"""Test with masked data."""
from ocgis.regrid.base import create_esmf_grid
x = Variable(name='x', value=[1, 2, 3], dimensions='x')
y = Variable(name='y', value=[4, 5, 6], dimensions='y')
grid = Grid(x, y, crs=Spherical())
gmask = grid.get_mask(create=True)
gmask[1, 1] = True
grid.set_mask(gmask)
self.assertEqual(grid.get_mask().sum(), 1)
egrid = create_esmf_grid(grid)
egrid_mask_inverted = np.invert(np.array(egrid.mask[0], dtype=bool))
self.assertNumpyAll(grid.get_mask(), egrid_mask_inverted)
# Test with a value mask.
value_mask = np.zeros(grid.shape, dtype=bool)
value_mask[-1, -1] = True
egrid = create_esmf_grid(grid, value_mask=value_mask)
egrid_mask_inverted = np.invert(np.array(egrid.mask[0], dtype=bool))
self.assertNumpyAll(egrid_mask_inverted, np.logical_or(grid.get_mask(), value_mask))
def get_ocgis_grid_from_esmf_grid(egrid, crs=None, dimension_map=None):
"""
Create an OCGIS :class:`~ocgis.interface.base.dimension.spatial.SpatialDimension` object from an ESMF
:class:`~ESMF.driver.grid.Grid`.
:type egrid: :class:`ESMF.driver.grid.Grid`
:param crs: The coordinate system to attach to the output spatial dimension.
:type crs: :class:`ocgis.interface.base.crs.CoordinateReferenceSystem`
:param dimension_map: Dimension map for the outgoing OCGIS field/grid.
:type dimension_map: :class:`ocgis.DimensionMap`
:rtype: :class:`~ocgis.Grid`
"""
if dimension_map is None:
dimension_map = {
DimensionMapKey.X: {
'variable': 'x',
'bounds': 'x_bounds',
DimensionMapKey.DIMENSION: ['x']
},
DimensionMapKey.Y: {
'variable': 'y',
'bounds': 'y_bounds',
DimensionMapKey.DIMENSION: ['y']
}
}
dimension_map = DimensionMap.from_dict(dimension_map)
else:
assert isinstance(dimension_map, DimensionMap)
# OCGIS grid values are built on centers.
coords = egrid.coords[ESMF.StaggerLoc.CENTER]
shape_coords_list = list(coords[0].shape)
dtype_coords = coords[0].dtype
# construct the ocgis grid array and fill
grid_value = np.zeros([2] + shape_coords_list, dtype=dtype_coords)
grid_value[0, ...] = coords[1]
grid_value[1, ...] = coords[0]
# Build OCGIS corners array if corners are present on the ESMF grid object.
has_corners = get_esmf_grid_has_corners(egrid)
if has_corners:
corner = egrid.coords[ESMF.StaggerLoc.CORNER]
grid_corners = np.zeros([2] + shape_coords_list + [4],
dtype=dtype_coords)
slices = [(0, 0), (0, 1), (1, 1), (1, 0)]
for ii, jj in iter_array(coords[0], use_mask=False):
row_slice = slice(ii, ii + 2)
col_slice = slice(jj, jj + 2)
row_corners = corner[1][row_slice, col_slice]
col_corners = corner[0][row_slice, col_slice]
for kk, slc in enumerate(slices):
grid_corners[:, ii, jj,
kk] = row_corners[slc], col_corners[slc]
else:
grid_corners = None
# Does the grid have a mask?
has_mask = False
if egrid.mask is not None:
if egrid.mask[ESMF.StaggerLoc.CENTER] is not None:
has_mask = True
if has_mask:
# if there is a mask, update the grid values
egrid_mask = egrid.mask[ESMF.StaggerLoc.CENTER]
egrid_mask = np.invert(egrid_mask.astype(bool))
# actually construct the masked arrays
grid_value = np.ma.array(grid_value)
if grid_corners is not None:
grid_corners = np.ma.array(grid_corners)
grid_dimensions = [
dimension_map.get_dimension(DimensionMapKey.Y)[0],
dimension_map.get_dimension(DimensionMapKey.X)[0]
]
if grid_corners is not None:
grid_bounds_dimensions = deepcopy(grid_dimensions)
grid_bounds_dimensions.append(constants.DEFAULT_NAME_CORNERS_DIMENSION)
name = dimension_map.get_bounds(DimensionMapKey.X)
x_bounds = Variable(
name=name,
value=grid_corners[1, ...],
dimensions=['y', 'x', constants.DEFAULT_NAME_CORNERS_DIMENSION])
name = dimension_map.get_bounds(DimensionMapKey.Y)
y_bounds = Variable(
name=name,
value=grid_corners[0, ...],
dimensions=['y', 'x', constants.DEFAULT_NAME_CORNERS_DIMENSION])
else:
x_bounds, y_bounds = [None] * 2
name = dimension_map.get_variable(DimensionMapKey.X)
x = Variable(name=name,
dimensions=grid_dimensions,
value=grid_value[1, ...],
bounds=x_bounds)
name = dimension_map.get_variable(DimensionMapKey.Y)
y = Variable(name=name,
dimensions=grid_dimensions,
value=grid_value[0, ...],
bounds=y_bounds)
ogrid = Grid(x, y, crs=crs)
if has_mask:
ogrid.set_mask(egrid_mask)
return ogrid
def test_system(self):
from ocgis.regrid.base import get_esmf_grid, iter_esmf_fields, RegridOperation, destroy_esmf_objects
import ESMF
yc = Variable(name='yc',
value=np.arange(-90 + (45 / 2.), 90, 45),
dimensions='ydim',
dtype=float)
xc = Variable(name='xc',
value=np.arange(15, 360, 30),
dimensions='xdim',
dtype=float)
ogrid = Grid(y=yc, x=xc, crs=Spherical())
ogrid.set_extrapolated_bounds('xc_bounds', 'yc_bounds', 'bounds')
np.random.seed(1)
mask = np.random.rand(*ogrid.shape)
mask = mask > 0.5
self.assertTrue(mask.sum() > 3)
ogrid.set_mask(mask)
egrid = get_esmf_grid(ogrid)
actual_shape = egrid.size[0].tolist()
desired_shape = np.flipud(ogrid.shape).tolist()
self.assertEqual(actual_shape, desired_shape)
desired = ogrid.get_value_stacked()
desired = np.ma.array(desired, mask=False)
desired.mask[0, :, :] = ogrid.get_mask()
desired.mask[1, :, :] = ogrid.get_mask()
desired = desired.sum()
actual_col = egrid.get_coords(0)
actual_row = egrid.get_coords(1)
actual_mask = np.invert(egrid.mask[0].astype(bool))
actual = np.ma.array(actual_row, mask=actual_mask).sum() + np.ma.array(
actual_col, mask=actual_mask).sum()
self.assertEqual(actual, desired)
desired = 9900.0
corners = egrid.coords[ESMF.StaggerLoc.CORNER]
actual = corners[0].sum() + corners[1].sum()
self.assertEqual(actual, desired)
ofield = create_exact_field(ogrid, 'data', ntime=3, crs=Spherical())
variable_name, efield, tidx = list(
iter_esmf_fields(ofield, split=False))[0]
desired_value = ofield['data'].get_value()
self.assertAlmostEqual(efield.data.sum(),
desired_value.sum(),
places=3)
destroy_esmf_objects([egrid, efield])
ofield.grid.set_mask(ofield.grid.get_mask(), cascade=True)
desired_value = ofield['data'].get_masked_value()
keywords = dict(split=[False, True])
for k in self.iter_product_keywords(keywords):
opts = {'split': k.split}
dofield = ofield.deepcopy()
dofield['data'].get_value().fill(0)
ro = RegridOperation(ofield, dofield, regrid_options=opts)
actual_field = ro.execute()
actual_value = actual_field['data'].get_masked_value()
self.assertAlmostEqual(0.0,
np.abs(desired_value - actual_value).max())
def get_ocgis_grid_from_esmf_grid(egrid):
"""
Create an OCGIS grid from an ESMF grid.
:param egrid: The input ESMF grid to convert to an OCGIS grid.
:type egrid: :class:`ESMF.Grid`
:return: :class:`~ocgis.Grid`
"""
dmap = egrid._ocgis['dimension_map']
edims = list(egrid._ocgis['dimnames'])
odims = egrid._ocgis['dimnames_backref']
coords = egrid.coords[ESMF.StaggerLoc.CENTER]
var_x = Variable(name=dmap.get_variable(DMK.X),
value=coords[0],
dimensions=edims)
var_y = Variable(name=dmap.get_variable(DMK.Y),
value=coords[1],
dimensions=edims)
# Build OCGIS corners array if corners are present on the ESMF grid object.
has_corners = esmf_grid_has_corners(egrid)
if has_corners:
corner = egrid.coords[ESMF.StaggerLoc.CORNER]
if egrid.periodic_dim == 0:
xcorner = np.zeros([corner[0].shape[0] + 1, corner[0].shape[1]],
dtype=corner[0].dtype)
xcorner[0:corner[0].shape[0], :] = corner[0]
xcorner[-1, :] = corner[0][0, :]
ycorner = np.zeros([corner[1].shape[0] + 1, corner[1].shape[1]],
dtype=corner[1].dtype)
ycorner[0:corner[1].shape[0], :] = corner[1]
ycorner[-1, :] = corner[1][0, :]
else:
xcorner = corner[0]
ycorner = corner[1]
ocorner_x = create_ocgis_corners_from_esmf_corners(xcorner)
ocorner_y = create_ocgis_corners_from_esmf_corners(ycorner)
cdims = deepcopy(edims)
cdims.append(constants.DEFAULT_NAME_CORNERS_DIMENSION)
vocorner_x = Variable(name=dmap.get_bounds(DMK.X),
value=ocorner_x,
dimensions=cdims)
vocorner_y = Variable(name=dmap.get_bounds(DMK.Y),
value=ocorner_y,
dimensions=cdims)
crs = get_crs_from_esmf(egrid)
ogrid = Grid(x=var_x, y=var_y, crs=crs)
# Does the grid have a mask?
has_mask = False
if egrid.mask is not None:
if egrid.mask[ESMF.StaggerLoc.CENTER] is not None:
has_mask = True
if has_mask:
# if there is a mask, update the grid values
egrid_mask = egrid.mask[ESMF.StaggerLoc.CENTER]
egrid_mask = np.invert(egrid_mask.astype(bool))
ogrid.set_mask(egrid_mask)
ogrid.parent.dimension_map = dmap
if tuple(odims) != tuple(edims):
broadcast_variable(var_x, odims)
broadcast_variable(var_y, odims)
if has_corners:
broadcast_variable(
vocorner_x,
list(odims) + [constants.DEFAULT_NAME_CORNERS_DIMENSION])
broadcast_variable(
vocorner_y,
list(odims) + [constants.DEFAULT_NAME_CORNERS_DIMENSION])
if has_corners:
var_x.set_bounds(vocorner_x)
var_y.set_bounds(vocorner_y)
return ogrid
def get_ocgis_grid_from_esmf_grid(egrid):
"""
Create an OCGIS grid from an ESMF grid.
:param egrid: The input ESMF grid to convert to an OCGIS grid.
:type egrid: :class:`ESMF.Grid`
:return: :class:`~ocgis.Grid`
"""
dmap = egrid._ocgis['dimension_map']
edims = list(egrid._ocgis['dimnames'])
odims = egrid._ocgis['dimnames_backref']
coords = egrid.coords[ESMF.StaggerLoc.CENTER]
var_x = Variable(name=dmap.get_variable(DMK.X), value=coords[0], dimensions=edims)
var_y = Variable(name=dmap.get_variable(DMK.Y), value=coords[1], dimensions=edims)
# Build OCGIS corners array if corners are present on the ESMF grid object.
has_corners = esmf_grid_has_corners(egrid)
if has_corners:
corner = egrid.coords[ESMF.StaggerLoc.CORNER]
if egrid.periodic_dim == 0:
xcorner = np.zeros([corner[0].shape[0] + 1, corner[0].shape[1]], dtype=corner[0].dtype)
xcorner[0:corner[0].shape[0], :] = corner[0]
xcorner[-1, :] = corner[0][0, :]
ycorner = np.zeros([corner[1].shape[0] + 1, corner[1].shape[1]], dtype=corner[1].dtype)
ycorner[0:corner[1].shape[0], :] = corner[1]
ycorner[-1, :] = corner[1][0, :]
else:
xcorner = corner[0]
ycorner = corner[1]
ocorner_x = create_ocgis_corners_from_esmf_corners(xcorner)
ocorner_y = create_ocgis_corners_from_esmf_corners(ycorner)
cdims = deepcopy(edims)
cdims.append(constants.DEFAULT_NAME_CORNERS_DIMENSION)
vocorner_x = Variable(name=dmap.get_bounds(DMK.X), value=ocorner_x, dimensions=cdims)
vocorner_y = Variable(name=dmap.get_bounds(DMK.Y), value=ocorner_y, dimensions=cdims)
crs = get_crs_from_esmf(egrid)
ogrid = Grid(x=var_x, y=var_y, crs=crs)
# Does the grid have a mask?
has_mask = False
if egrid.mask is not None:
if egrid.mask[ESMF.StaggerLoc.CENTER] is not None:
has_mask = True
if has_mask:
# if there is a mask, update the grid values
egrid_mask = egrid.mask[ESMF.StaggerLoc.CENTER]
egrid_mask = np.invert(egrid_mask.astype(bool))
ogrid.set_mask(egrid_mask)
ogrid.parent.dimension_map = dmap
if tuple(odims) != tuple(edims):
broadcast_variable(var_x, odims)
broadcast_variable(var_y, odims)
if has_corners:
broadcast_variable(vocorner_x, list(odims) + [constants.DEFAULT_NAME_CORNERS_DIMENSION])
broadcast_variable(vocorner_y, list(odims) + [constants.DEFAULT_NAME_CORNERS_DIMENSION])
if has_corners:
var_x.set_bounds(vocorner_x)
var_y.set_bounds(vocorner_y)
return ogrid