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_get_esmf_grid_bilinear_regrid_method(self): """Test with a regrid method that does not require corners.""" from ocgis.regrid.base import get_esmf_grid from ESMF import RegridMethod import ESMF rd = RequestDataset(**self.get_dataset()) field = rd.get() self.assertTrue(field.grid.has_bounds) egrid = get_esmf_grid(field.grid, regrid_method=RegridMethod.BILINEAR) corner = egrid.coords[ESMF.StaggerLoc.CORNER] for idx in [0, 1]: self.assertIsNone(corner[idx])
def test_get_esmf_grid_periodicity(self): """Test periodicity parameters generate reasonable output.""" from ocgis.regrid.base import get_esmf_grid lon_in = np.arange(-180, 180, 10) lat_in = np.arange(-90, 90.1, 4) lon = Variable('lon', lon_in, 'dlon') lat = Variable('lat', lat_in, 'dlat') ogrid = Grid(x=lon, y=lat, crs=Spherical()) egrid = get_esmf_grid(ogrid) self.assertEqual(egrid.periodic_dim, 0) self.assertEqual(egrid.num_peri_dims, 1) self.assertEqual(egrid.pole_dim, 1)
def test_get_esmf_grid(self): import ESMF rd = RequestDataset(**self.get_dataset()) have_ocgis_bounds = [True, False] for h in have_ocgis_bounds: field = rd.get() ogrid = field.grid if not h: ogrid.remove_bounds() self.assertFalse(ogrid.has_bounds) from ocgis.regrid.base import get_esmf_grid egrid = get_esmf_grid(ogrid) # 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] desired = ogrid.get_value_stacked()[idx_ocgis, :] self.assertNumpyAll(coords, desired) 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=ogrid.archetype.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=ogrid.archetype.dtype) self.assertNumpyAll(corner_col, corner_col_actual) else: # No corners should be on the ESMF grid. for idx in [0, 1]: self.assertIsNone(corner[idx])
def test_get_ocgis_grid_from_esmf_grid(self): from ocgis.regrid.base import get_esmf_grid from ocgis.regrid.base import get_ocgis_grid_from_esmf_grid rd = 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() grid = field.grid if k.has_mask: gmask = grid.get_mask(create=True) gmask[1, :] = True grid.set_mask(gmask) if not k.has_corners: grid.x.set_bounds(None) grid.y.set_bounds(None) self.assertFalse(grid.has_bounds) egrid = get_esmf_grid(grid) ogrid = get_ocgis_grid_from_esmf_grid(egrid, crs=k.crs) if k.has_mask: actual_mask = ogrid.get_mask() self.assertEqual(actual_mask.sum(), 4) self.assertTrue(actual_mask[1, :].all()) self.assertEqual(ogrid.crs, k.crs) self.assertNumpyAll(grid.get_value_stacked(), ogrid.get_value_stacked()) if k.has_corners: desired = grid.x.bounds.get_value() actual = ogrid.x.bounds.get_value() self.assertNumpyAll(actual, desired) desired = grid.y.bounds.get_value() actual = ogrid.y.bounds.get_value() self.assertNumpyAll(actual, desired) else: self.assertFalse(ogrid.has_bounds)
def test_get_esmf_grid_change_origin_row_and_col(self): """Test with different row and column grid origin.""" from ocgis.regrid.base import get_esmf_grid import ESMF rd = RequestDataset(**self.get_dataset()) field = rd.get() field.grid.y.set_value(np.flipud(field.grid.y.get_value())) field.grid.y.bounds.set_value( np.fliplr(np.flipud(field.grid.y.bounds.get_value()))) field.grid.x.set_value(np.flipud(field.grid.x.get_value())) field.grid.x.bounds.set_value( np.fliplr(np.flipud(field.grid.x.bounds.get_value()))) egrid = get_esmf_grid(field.grid) for idx_esmf, idx_ocgis in zip([0, 1], [1, 0]): coords = egrid.coords[ESMF.StaggerLoc.CENTER][idx_esmf] self.assertNumpyAll(coords, field.grid.get_value_stacked()[idx_ocgis, ...]) corner = egrid.coords[ESMF.StaggerLoc.CORNER] corner_row = corner[1] dtype = field.grid.dtype 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=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=dtype) self.assertNumpyAll(corner_col, corner_col_actual)
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())