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)