def test_iter_regridded_fields_problem_bounds(self):
"""Test a dataset with crap bounds will work when with_corners is False."""
dst = self.test_data.get_rd("cancm4_tas").get()[:, :, :, 20:25, 30:35]
dst.spatial.crs = Spherical()
src = deepcopy(dst[0, 0, 0, :, :])
egrid_dst = get_esmf_grid_from_sdim(dst.spatial)
egrid_src = get_esmf_grid_from_sdim(src.spatial)
self.assertEqual(egrid_dst.mask[0].sum(), 25)
self.assertEqual(egrid_src.mask[0].sum(), 25)
self.assertNumpyAll(egrid_dst.coords[0][0], egrid_src.coords[0][0])
self.assertNumpyAll(egrid_dst.coords[0][1], egrid_src.coords[0][1])
ret = list(iter_esmf_fields(src))
self.assertEqual(len(ret), 1)
variable_alias, efield, tidx = ret[0]
self.assertNumpyAll(efield.grid.coords[0][0], dst.spatial.grid.value.data[1])
self.assertNumpyAll(efield.grid.coords[0][1], dst.spatial.grid.value.data[0])
ret = list(iter_regridded_fields([src], dst, with_corners=False))[0]
actual = dst[0, 0, 0, :, :].variables.first().value
to_test = ret.variables.first().value
self.assertFalse(to_test.mask.any())
self.assertNumpyAllClose(to_test.data, actual.data)
self.assertNumpyAll(to_test.mask, actual.mask)
def test_get_esmf_grid_from_sdim_with_mask(self):
"""Test with masked data."""
rd = self.test_data.get_rd("cancm4_tas")
ops = ocgis.OcgOperations(
dataset=rd, geom="state_boundaries", select_ugid=[23], snippet=True, vector_wrap=False
)
ret = ops.execute()
field = ret[23]["tas"]
egrid = get_esmf_grid_from_sdim(field.spatial)
actual = np.array([[0, 0, 1], [1, 1, 1], [1, 1, 1], [1, 1, 1]], dtype=np.int32)
self.assertNumpyAll(egrid.mask[0], actual)
sdim = get_sdim_from_esmf_grid(egrid)
self.assertNumpyAll(sdim.get_mask(), field.spatial.get_mask())
actual = np.array(
[
[
[[True, True, True, True], [True, True, True, True], [False, False, False, False]],
[[False, False, False, False], [False, False, False, False], [False, False, False, False]],
[[False, False, False, False], [False, False, False, False], [False, False, False, False]],
[[False, False, False, False], [False, False, False, False], [False, False, False, False]],
],
[
[[True, True, True, True], [True, True, True, True], [False, False, False, False]],
[[False, False, False, False], [False, False, False, False], [False, False, False, False]],
[[False, False, False, False], [False, False, False, False], [False, False, False, False]],
[[False, False, False, False], [False, False, False, False], [False, False, False, False]],
],
]
)
self.assertNumpyAll(actual, sdim.grid.corners.mask)
def test_get_sdim_from_esmf_grid(self):
rd = ocgis.RequestDataset(**self.get_dataset())
keywords = dict(
has_corners=[True, False], has_mask=[True, False], crs=[None, CoordinateReferenceSystem(epsg=4326)]
)
for k in itr_products_keywords(keywords, as_namedtuple=True):
field = rd.get()
sdim = field.spatial
if not k.has_corners:
sdim.grid.row.remove_bounds()
sdim.grid.col.remove_bounds()
self.assertIsNone(sdim.grid.corners)
egrid = get_esmf_grid_from_sdim(sdim)
if not k.has_mask:
egrid.mask[0][2, 2] = 0
sdim.grid.value.mask[:, 2, 2] = True
if k.has_corners:
sdim.grid.corners.mask[:, 2, 2] = True
nsdim = get_sdim_from_esmf_grid(egrid, crs=k.crs)
self.assertEqual(nsdim.crs, k.crs)
self.assertNumpyAll(sdim.grid.value, nsdim.grid.value)
if k.has_corners:
self.assertNumpyAll(sdim.grid.corners, nsdim.grid.corners)
else:
self.assertIsNone(nsdim.grid.corners)
def test_get_sdim_from_esmf_grid(self):
rd = ocgis.RequestDataset(**self.get_dataset())
keywords = dict(
has_corners=[True, False], has_mask=[True, False], crs=[None, CoordinateReferenceSystem(epsg=4326)]
)
for k in itr_products_keywords(keywords, as_namedtuple=True):
field = rd.get()
sdim = field.spatial
egrid = get_esmf_grid_from_sdim(sdim)
if not k.has_mask:
# set the grid flag to indicate no mask is present
egrid.item_done[ESMF.StaggerLoc.CENTER][0] = False
# remove the mask from the grid
egrid.mask[0] = [1]
else:
egrid.mask[0][2, 2] = 0
sdim.grid.value.mask[:, 2, 2] = True
sdim.grid.corners.mask[:, 2, 2] = True
if not k.has_corners:
egrid.coords[ESMF.StaggerLoc.CORNER] = [np.array(0.0), np.array(0.0)]
egrid.coords_done[ESMF.StaggerLoc.CORNER] = [False, False]
nsdim = get_sdim_from_esmf_grid(egrid, crs=k.crs)
self.assertEqual(nsdim.crs, k.crs)
self.assertNumpyAll(sdim.grid.value, nsdim.grid.value)
if k.has_corners:
self.assertNumpyAll(sdim.grid.corners, nsdim.grid.corners)
else:
self.assertIsNone(nsdim.grid.corners)
def test_get_esmf_grid_from_sdim_with_corners(self):
"""Test with the with_corners option set to False."""
rd = ocgis.RequestDataset(**self.get_dataset())
field = rd.get()
self.assertIsNotNone(field.spatial.grid.corners)
egrid = get_esmf_grid_from_sdim(field.spatial, with_corners=False)
corner = egrid.coords[ESMF.StaggerLoc.CORNER]
for idx in [0, 1]:
self.assertIsNone(corner[idx])
def test_get_esmf_grid_from_sdim_value_mask(self):
"""Test with an additional mask."""
rd = ocgis.RequestDataset(**self.get_dataset())
field = rd.get()
np.random.seed(1)
self.assertFalse(np.any(field.spatial.get_mask()))
value_mask = np.random.randint(0, 2, field.spatial.get_mask().shape)
egrid = get_esmf_grid_from_sdim(field.spatial, value_mask=value_mask)
self.assertNumpyAll(egrid.mask[0], np.invert(value_mask.astype(bool)).astype(egrid.mask[0].dtype))
def test_get_esmf_grid_from_sdim_real_data(self):
"""Test creating ESMF field from real data using an OCGIS spatial dimension."""
rd = self.test_data.get_rd("cancm4_tas")
field = rd.get()
egrid = get_esmf_grid_from_sdim(field.spatial)
for idx_esmf, idx_ocgis in zip([0, 1], [1, 0]):
coords = egrid.coords[ESMF.StaggerLoc.CENTER][idx_esmf]
self.assertNumpyAll(coords, field.spatial.grid.value[idx_ocgis, ...].data)
def test_get_esmf_grid_from_sdim(self):
rd = ocgis.RequestDataset(**self.get_dataset())
have_ocgis_bounds = [True, False]
for h in have_ocgis_bounds:
field = rd.get()
sdim = field.spatial
if not h:
sdim.grid.row.bounds
sdim.grid.col.bounds
sdim.grid.row.bounds = None
sdim.grid.col.bounds = None
self.assertIsNone(sdim.grid.row.bounds)
self.assertIsNone(sdim.grid.col.bounds)
egrid = get_esmf_grid_from_sdim(sdim)
# ocgis is row major with esmf being column major (i.e. in ocgis rows are stored in the zero index)
for idx_esmf, idx_ocgis in zip([0, 1], [1, 0]):
coords = egrid.coords[ESMF.StaggerLoc.CENTER][idx_esmf]
self.assertNumpyAll(coords, field.spatial.grid.value[idx_ocgis, ...].data)
corner = egrid.coords[ESMF.StaggerLoc.CORNER]
if h:
corner_row = corner[1]
corner_row_actual = np.array(
[
[36.5, 36.5, 36.5, 36.5, 36.5],
[37.5, 37.5, 37.5, 37.5, 37.5],
[38.5, 38.5, 38.5, 38.5, 38.5],
[39.5, 39.5, 39.5, 39.5, 39.5],
[40.5, 40.5, 40.5, 40.5, 40.5],
],
dtype=field.spatial.grid.value.dtype,
)
self.assertNumpyAll(corner_row, corner_row_actual)
corner = egrid.coords[ESMF.StaggerLoc.CORNER]
corner_col = corner[0]
corner_col_actual = np.array(
[
[-105.5, -104.5, -103.5, -102.5, -101.5],
[-105.5, -104.5, -103.5, -102.5, -101.5],
[-105.5, -104.5, -103.5, -102.5, -101.5],
[-105.5, -104.5, -103.5, -102.5, -101.5],
[-105.5, -104.5, -103.5, -102.5, -101.5],
],
dtype=field.spatial.grid.value.dtype,
)
self.assertNumpyAll(corner_col, corner_col_actual)
else:
for idx in [0, 1]:
self.assertNumpyAll(corner[idx], np.array(0.0))
def write(self):
#todo: doc
for coll in self.colls:
""":type coll: :class:`ocgis.api.collection.SpatialCollection`"""
for row in coll.get_iter_melted():
field = row['field']
variable = row['variable']
egrid = get_esmf_grid_from_sdim(field.spatial, with_corners=self.with_corners,
value_mask=self.value_mask)
esmf_field_name = self.esmf_field_name or variable.alias
efield = ESMF.Field(egrid, esmf_field_name, ndbounds=field.shape[0:-2], mask_values=[0])
efield[:] = variable.value
return efield
def test_get_esmf_grid_from_sdim_change_origin_row_and_col(self):
"""Test with different row and column grid origin."""
rd = ocgis.RequestDataset(**self.get_dataset())
field = rd.get()
field.spatial.grid.row._value = np.flipud(field.spatial.grid.row.value)
field.spatial.grid.row.bounds = np.fliplr(np.flipud(field.spatial.grid.row.bounds))
field.spatial.grid.col._value = np.flipud(field.spatial.grid.col.value)
field.spatial.grid.col.bounds = np.fliplr(np.flipud(field.spatial.grid.col.bounds))
egrid = get_esmf_grid_from_sdim(field.spatial)
for idx_esmf, idx_ocgis in zip([0, 1], [1, 0]):
coords = egrid.coords[ESMF.StaggerLoc.CENTER][idx_esmf]
self.assertNumpyAll(coords, field.spatial.grid.value[idx_ocgis, ...].data)
corner = egrid.coords[ESMF.StaggerLoc.CORNER]
corner_row = corner[1]
corner_row_actual = np.array(
[
[40.5, 40.5, 40.5, 40.5, 40.5],
[39.5, 39.5, 39.5, 39.5, 39.5],
[38.5, 38.5, 38.5, 38.5, 38.5],
[37.5, 37.5, 37.5, 37.5, 37.5],
[36.5, 36.5, 36.5, 36.5, 36.5],
],
dtype=field.spatial.grid.value.dtype,
)
self.assertNumpyAll(corner_row, corner_row_actual)
corner = egrid.coords[ESMF.StaggerLoc.CORNER]
corner_col = corner[0]
corner_col_actual = np.array(
[
[-101.5, -102.5, -103.5, -104.5, -105.5],
[-101.5, -102.5, -103.5, -104.5, -105.5],
[-101.5, -102.5, -103.5, -104.5, -105.5],
[-101.5, -102.5, -103.5, -104.5, -105.5],
[-101.5, -102.5, -103.5, -104.5, -105.5],
],
dtype=field.spatial.grid.value.dtype,
)
self.assertNumpyAll(corner_col, corner_col_actual)
