def test_create_index_variable_global(self): raise SkipTest('not implemented') dsrc_size = 5 ddst_size = 7 dsrc = Dimension('dsrc', dsrc_size, dist=True) src_dist = OcgDist() src_dist.add_dimension(dsrc) src_dist.update_dimension_bounds() ddst = Dimension('ddst', ddst_size, dist=True) dst_dist = OcgDist() dst_dist.add_dimension(ddst) dst_dist.update_dimension_bounds() if vm.rank == 0: np.random.seed(1) dst = np.random.rand(ddst_size) src = np.random.choice(dst, size=dsrc_size, replace=False) src = Variable(name='src', value=src, dimensions=dsrc.name) # TODO: move create_ugid_global to create_global_index on a standard variable object dst = GeometryVariable(name='dst', value=dst, dimensions=ddst.name) else: src, dst = [None] * 2 src = variable_scatter(src, src_dist) dst = variable_scatter(dst, dst_dist) actual = create_index_variable_global('index_array', src, dst) self.assertNumpyAll(dst.get_value()[actual.get_value()], src.get_value())
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_create_unique_global_array(self): dist = OcgDist() dist.create_dimension('dim', 9, dist=True) dist.update_dimension_bounds() values = [ [4, 2, 1, 2, 1, 4, 1, 4, 2], [44, 25, 16, 27, 18, 49, 10, 41, 22], [44, 25, 16, 27, 44, 49, 10, 41, 44], [1, 1, 1, 1, 1, 1, 1, 1, 1] ] for v in values: if vm.rank == 0: index = Variable(name='cindex', value=v, dimensions='dim') desired = np.unique(index.get_value()) desired_length = len(desired) else: index = None index = variable_scatter(index, dist) with vm.scoped_by_emptyable('not empty', index): if not vm.is_null: uvar = create_unique_global_array(index.get_value()) uvar_gathered = vm.gather(uvar) if vm.rank == 0: uvar_gathered = hgather(uvar_gathered) self.assertEqual(len(uvar_gathered), desired_length) self.assertEqual(set(uvar_gathered), set(desired))
def test_create_unique_global_array(self): dist = OcgDist() dist.create_dimension('dim', 9, dist=True) dist.update_dimension_bounds() values = [[4, 2, 1, 2, 1, 4, 1, 4, 2], [44, 25, 16, 27, 18, 49, 10, 41, 22], [44, 25, 16, 27, 44, 49, 10, 41, 44], [1, 1, 1, 1, 1, 1, 1, 1, 1]] for v in values: if vm.rank == 0: index = Variable(name='cindex', value=v, dimensions='dim') desired = np.unique(index.get_value()) desired_length = len(desired) else: index = None index = variable_scatter(index, dist) with vm.scoped_by_emptyable('not empty', index): if not vm.is_null: uvar = create_unique_global_array(index.get_value()) uvar_gathered = vm.gather(uvar) if vm.rank == 0: uvar_gathered = hgather(uvar_gathered) self.assertEqual(len(uvar_gathered), desired_length) self.assertEqual(set(uvar_gathered), set(desired))
def test_system_process_geometries(self): """Test multiple geometries with coordinate system update.""" a = 'POLYGON((-105.21347987288135073 40.21514830508475313,-104.39928495762711691 40.21514830508475313,-104.3192002118643984 39.5677966101694949,-102.37047139830508513 39.61451271186440692,-102.12354343220337682 37.51896186440677639,-105.16009004237288593 37.51896186440677639,-105.21347987288135073 40.21514830508475313))' b = 'POLYGON((-104.15235699152542281 39.02722457627118757,-103.71189088983049942 39.44099576271186436,-102.71750529661017026 39.28082627118644155,-102.35712394067796538 37.63908898305084705,-104.13900953389830306 37.63241525423728717,-104.15235699152542281 39.02722457627118757))' geom = [{'geom': wkt.loads(xx), 'properties': {'UGID': ugid}} for ugid, xx in enumerate([a, b])] grid_value = [ [[37.0, 37.0, 37.0, 37.0], [38.0, 38.0, 38.0, 38.0], [39.0, 39.0, 39.0, 39.0], [40.0, 40.0, 40.0, 40.0]], [[-105.0, -104.0, -103.0, -102.0], [-105.0, -104.0, -103.0, -102.0], [-105.0, -104.0, -103.0, -102.0], [-105.0, -104.0, -103.0, -102.0]]] output_crs = CoordinateReferenceSystem( value={'a': 6370997, 'lon_0': -100, 'y_0': 0, 'no_defs': True, 'proj': 'laea', 'x_0': 0, 'units': 'm', 'b': 6370997, 'lat_0': 45}) x = Variable('x', grid_value[1], dimensions=['lat', 'lon']) y = Variable('y', grid_value[0], dimensions=['lat', 'lon']) grid = Grid(x, y) field = Field(grid=grid, crs=Spherical()) ops = OcgOperations(dataset=field, geom=geom, output_crs=output_crs) ret = ops.execute() expected = {0: -502052.79407259845, 1: -510391.37909706926} for field, container in ret.iter_fields(yield_container=True): self.assertAlmostEqual(field.grid.get_value_stacked().mean(), expected[container.geom.ugid.get_value()[0]])
def create_mftime_nc_files(test_obj, with_all_cf=False, units_on_time_bounds=False, calendar_on_second=True): value = [0, 1, 2] units = ['days since 2000-1-1', 'days since 2001-1-1'] names = ['time_2000.nc', 'time_2001.nc'] paths = [] if units_on_time_bounds: tv = Variable(name='time_bounds_maker', value=value, dimensions='tbmdim') tv.set_extrapolated_bounds('time_bounds', 'bounds') for ctr, (unit, name) in enumerate(zip(units, names)): path = test_obj.get_temporary_file_path(name) paths.append(path) with test_obj.nc_scope(path, 'w', format='NETCDF4_CLASSIC') as f: f.createDimension('time') vtime = f.createVariable('time', np.float32, dimensions=('time', )) vtime[:] = value vtime.units = unit if ctr == 0 or (calendar_on_second and ctr == 1): vtime.calendar = 'standard' vtime.axis = 'T' if units_on_time_bounds: vtime.calendar = 'noleap' vtime.bounds = 'time_bounds' if units_on_time_bounds: f.createDimension('bounds', 2) vbtime = f.createVariable('time_bounds', np.float32, dimensions=('time', 'bounds')) vbtime[:] = tv.bounds.get_value() # Note no calendar on bounds variable. vbtime.units = units if with_all_cf: f.createDimension('x', 4) f.createDimension('y', 5) x = f.createVariable('x', float, dimensions=('x', )) x[:] = [100., 101., 102., 103.] x.axis = 'X' y = f.createVariable('y', float, dimensions=('y', )) y[:] = [40., 41., 42., 43., 44.] y.axis = 'Y' data = f.createVariable('data', float, dimensions=('time', 'y', 'x')) data[:] = 1 return paths
def test_array_resolution_called(self): """Test the driver's array resolution method is called appropriately.""" m_DriverNetcdfSCRIP = mock.create_autospec(DriverNetcdfSCRIP) with mock.patch('ocgis.driver.registry.get_driver_class', return_value=m_DriverNetcdfSCRIP): x = Variable(name='x', value=[1, 2, 3], dimensions='dimx') y = Variable(name='y', value=[4, 5, 6], dimensions='dimy') pgc = PointGC(x=x, y=y) _ = pgc.resolution_x _ = pgc.resolution_y self.assertEqual(m_DriverNetcdfSCRIP.array_resolution.call_count, 2)
def test_init_dimension_map(self): """Test initializing with a dimension map only.""" dmap = DimensionMap() x = Variable(value=[1, 2, 3], dimensions='elements', name='x') y = Variable(value=[4, 5, 6], dimensions='elements', name='y') topo = dmap.get_topology(Topology.POINT, create=True) topo.set_variable(DMK.X, x) topo.set_variable(DMK.Y, y) f = Field(variables=[x, y], dimension_map=dmap) p = PointGC(parent=f) self.assertNumpyAll(x.get_value(), p.x.get_value()) self.assertNumpyAll(y.get_value(), p.y.get_value())
def test_reduce_reindex_coordinate_index(self): dist = OcgDist() dist.create_dimension('dim', 12, dist=True) dist.update_dimension_bounds() global_cindex_arr = np.array([4, 2, 1, 2, 1, 4, 1, 4, 2, 5, 6, 7]) if vm.rank == 0: var_cindex = Variable('cindex', value=global_cindex_arr, dimensions='dim') else: var_cindex = None var_cindex = variable_scatter(var_cindex, dist) vm.create_subcomm_by_emptyable('test', var_cindex, is_current=True) if vm.is_null: return raise_if_empty(var_cindex) coords = np.array([ 0, 11, 22, 33, 44, 55, 66, 77, 88, 99, 100, 110, 120, 130, 140, 150 ]) coords = Variable(name='coords', value=coords, dimensions='coord_dim') new_cindex, u_indices = reduce_reindex_coordinate_index(var_cindex) desired = coords[global_cindex_arr].get_value() if len(u_indices) > 0: new_coords = coords[u_indices].get_value() else: new_coords = np.array([]) gathered_new_coords = vm.gather(new_coords) gathered_new_cindex = vm.gather(new_cindex) if vm.rank == 0: gathered_new_coords = hgather(gathered_new_coords) gathered_new_cindex = hgather(gathered_new_cindex) actual = gathered_new_coords[gathered_new_cindex] self.assertAsSetEqual(gathered_new_cindex.tolist(), [2, 1, 0, 3, 4, 5]) desired_new_coords = [11, 22, 44, 55, 66, 77] self.assertAsSetEqual(gathered_new_coords.tolist(), desired_new_coords) self.assertEqual(len(gathered_new_coords), len(desired_new_coords)) self.assertNumpyAll(actual, desired)
def test_chunked_rwg_spatial_subset(self): env.CLOBBER_UNITS_ON_BOUNDS = False src_grid = create_gridxy_global(crs=Spherical()) src_field = create_exact_field(src_grid, 'foo') xvar = Variable(name='x', value=[-90., -80.], dimensions='xdim') yvar = Variable(name='y', value=[40., 50.], dimensions='ydim') dst_grid = Grid(x=xvar, y=yvar, crs=Spherical()) if ocgis.vm.rank == 0: source = self.get_temporary_file_path('source.nc') else: source = None source = ocgis.vm.bcast(source) src_field.write(source) if ocgis.vm.rank == 0: destination = self.get_temporary_file_path('destination.nc') else: destination = None destination = ocgis.vm.bcast(destination) dst_grid.parent.write(destination) wd = os.path.join(self.current_dir_output, 'chunks') weight = os.path.join(self.current_dir_output, 'weights.nc') spatial_subset = os.path.join(self.current_dir_output, 'spatial_subset.nc') runner = CliRunner() cli_args = [ 'chunked-rwg', '--source', source, '--destination', destination, '--wd', wd, '--spatial_subset', '--spatial_subset_path', spatial_subset, '--weight', weight, '--esmf_regrid_method', 'BILINEAR', '--persist' ] result = runner.invoke(ocli, args=cli_args, catch_exceptions=False) self.assertEqual(result.exit_code, 0) actual = RequestDataset(uri=spatial_subset).create_field() actual_ymean = actual.grid.get_value_stacked()[0].mean() actual_xmean = actual.grid.get_value_stacked()[1].mean() self.assertEqual(actual_ymean, 45.) self.assertEqual(actual_xmean, -85.) self.assertEqual(actual.grid.shape, (14, 14)) self.assertTrue(os.path.exists(weight)) actual = RequestDataset(weight, driver='netcdf').create_field() self.assertIn('history', actual.attrs)
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 fixture_cindex(start_index): assert start_index == 0 or start_index == 1 value = np.arange(0 + start_index, 6 + start_index).reshape(-1, 1) return Variable(name='cindex', value=value, dimensions=['elements', 'misc'], attrs={AttributeName.START_INDEX: start_index})
def test_system_regridding_crs(self): """Test with coordinate systems.""" dest_crs = WGS84() grid_spherical = self.get_gridxy_global(resolution=10.0, wrapped=False, crs=Spherical()) self.assertEqual(grid_spherical.crs, Spherical()) coords = grid_spherical.get_value_stacked() data_value = self.get_exact_field_value(coords[1], coords[0]) desired = data_value.copy() data_var = Variable(name='data_src', value=data_value, dimensions=grid_spherical.dimensions) source = Field(grid=grid_spherical, is_data=data_var, crs=grid_spherical.crs) self.assertEqual(source.crs, Spherical()) destination = deepcopy(source) destination.update_crs(dest_crs) source_expanded = deepcopy(source.grid) source_expanded.expand() diff = np.abs(destination.y.get_value() - source_expanded.y.get_value()) self.assertAlmostEqual(diff.max(), 0.19231511439) for output_crs in [None, WGS84()]: ops = OcgOperations(dataset=source, regrid_destination=destination, output_crs=output_crs) ret = ops.execute() actual = ret.get_element(variable_name=data_var.name) if output_crs is None: self.assertEqual(actual.parent.crs, Spherical()) else: self.assertEqual(actual.parent.crs, WGS84()) actual = actual.get_value() diff = np.abs(actual - desired) self.assertTrue(diff.max() < 1e-5)
def test_iter_geometries_with_cindex(self): cindex = Variable(name='cindex', value=[2, 2], dimensions='elements') p = self.fixture(cindex=cindex) self.assertIsNotNone(p.cindex) actual = list(p.iter_geometries()) actual = np.array([ii[1].xy for ii in actual]).tolist() desired = [[[2.0], [8.0]], [[2.0], [8.0]]] self.assertEqual(actual, desired)
def test_get_ocgis_field_from_esmf_spatial_only(self): """Test with spatial information only.""" from ocgis.regrid.base import get_esmf_field_from_ocgis_field from ocgis.regrid.base import get_ocgis_field_from_esmf_field row = Variable(name='row', value=[5, 6], dimensions='row') col = Variable(name='col', value=[7, 8], dimensions='col') grid = Grid(col, row) ofield = Field(grid=grid, crs=Spherical()) efield = get_esmf_field_from_ocgis_field(ofield) ofield_actual = get_ocgis_field_from_esmf_field(efield) self.assertEqual(len(ofield_actual.data_variables), 0) self.assertNumpyAll(grid.get_value_stacked(), ofield_actual.grid.get_value_stacked())
def create_mftime_nc_files(test_obj, with_all_cf=False, units_on_time_bounds=False, calendar_on_second=True): value = [0, 1, 2] units = ['days since 2000-1-1', 'days since 2001-1-1'] names = ['time_2000.nc', 'time_2001.nc'] paths = [] if units_on_time_bounds: tv = Variable(name='time_bounds_maker', value=value, dimensions='tbmdim') tv.set_extrapolated_bounds('time_bounds', 'bounds') for ctr, (unit, name) in enumerate(zip(units, names)): path = test_obj.get_temporary_file_path(name) paths.append(path) with test_obj.nc_scope(path, 'w', format='NETCDF4_CLASSIC') as f: f.createDimension('time') vtime = f.createVariable('time', np.float32, dimensions=('time',)) vtime[:] = value vtime.units = unit if ctr == 0 or (calendar_on_second and ctr == 1): vtime.calendar = 'standard' vtime.axis = 'T' if units_on_time_bounds: vtime.calendar = 'noleap' vtime.bounds = 'time_bounds' if units_on_time_bounds: f.createDimension('bounds', 2) vbtime = f.createVariable('time_bounds', np.float32, dimensions=('time', 'bounds')) vbtime[:] = tv.bounds.get_value() # Note no calendar on bounds variable. vbtime.units = unit if with_all_cf: f.createDimension('x', 4) f.createDimension('y', 5) x = f.createVariable('x', float, dimensions=('x',)) x[:] = [100., 101., 102., 103.] x.axis = 'X' y = f.createVariable('y', float, dimensions=('y',)) y[:] = [40., 41., 42., 43., 44.] y.axis = 'Y' data = f.createVariable('data', float, dimensions=('time', 'y', 'x')) data[:] = 1 return paths
def create_regridding_field(grid, name_variable): col_shape = grid.shape[1] row_shape = grid.shape[0] value = np.ones(col_shape * row_shape, dtype=float).reshape( grid.shape) * 15. variable = Variable(name=name_variable, value=value, dimensions=grid.dimensions) field = Field(is_data=variable, grid=grid, crs=Spherical()) return field
def fixture(self, **kwargs): kwargs = kwargs.copy() node_dimension_name = kwargs.pop('node_dimension_name', 'dnodes') x = Variable(name='xc', value=[0, 1, 2, 3, 4, 5], dimensions=node_dimension_name, dtype=float) y = Variable(name='yc', value=[6, 7, 8, 9, 10, 11], dimensions=node_dimension_name, dtype=float) z = Variable(name='zc', value=[12, 13, 14, 15, 16, 17], dimensions=node_dimension_name, dtype=float) p = PointGC(x, y, z=z, **kwargs) return p
def create_host_attribute_variable(dimension_map, name=VariableName.UGRID_HOST_VARIABLE): point = dimension_map.get_topology(Topology.POINT) dkeys = [DMK.X, DMK.Y, DMK.LEVEL] face_coordinates = None if point is not None and point != Topology.AUTO: x_repr, y_repr, z_repr = [point.get_variable(k) for k in dkeys] if x_repr is not None: face_coordinates = [x_repr, y_repr] if z_repr is not None: face_coordinates.append(z_repr) face_coordinates = ' '.join(face_coordinates) poly = dimension_map.get_topology(Topology.POLYGON) z = None if poly is not None: x, y, z = [poly.get_variable(k) for k in dkeys] node_coordinates = [x, y] if z is not None: node_coordinates.append(z) node_coordinates = ' '.join(node_coordinates) face_node_connectivity = poly.get_variable( DMK.ELEMENT_NODE_CONNECTIVITY) else: node_coordinates = None if z is None: dimension = 2 else: dimension = 3 locations = [] if point is not None: locations.append('face') if poly is not None: locations.append('node') locations = ' '.join(locations) attrs = { 'standard_name': 'mesh_topology', 'cf_role': 'mesh_topology', 'dimension': dimension, 'locations': locations, 'node_coordinates': node_coordinates, 'face_coordinates': face_coordinates } if poly is not None: attrs['face_node_connectivity'] = face_node_connectivity return Variable(name=name, attrs=attrs)
def test_set_variable_spatial_mask(self): value = np.random.rand(10, 3, 4) value = np.ma.array(value, mask=False) mask_spatial = np.zeros((3, 4), dtype=bool) mask_spatial[2, 3] = True value.mask[:, 1, 1] = True var = Variable(name='var', value=value, dimensions=['time', 'y', 'x']) slice_row = slice(None) slice_col = slice(None) self.assertFalse(var.get_mask()[:, 2, 3].any()) set_variable_spatial_mask(var, mask_spatial, slice_row, slice_col) self.assertTrue(var.get_mask().any()) self.assertTrue(var.get_mask()[:, 1, 1].all()) vmask = var.get_mask() vmask[:, 1, 1] = False vmask[:, 2, 3] = False var.set_mask(vmask) self.assertFalse(var.get_mask().any())
def test_get_intersects_wrapped(self): """Test using a subset geometry that needs to be wrapped.""" coords = (-10.0, 40.0, 50.0, 50.0) bbox = box(*coords) gvar = GeometryVariable(value=bbox, dimensions='geom', crs=Spherical(), is_bbox=True, wrapped_state=WrappedState.UNWRAPPED) x = [60.0, 355.0] x = Variable(name='x', value=x, dimensions='n_nodes') y = [45.0, 45.0] y = Variable(name='y', value=y, dimensions='n_nodes') pgc = PointGC(x, y) ret = pgc.get_intersects(gvar) self.assertEqual(ret.archetype.size, 1) geom = list(ret.iter_geometries())[0][1] actual = np.array(geom).tolist() self.assertEqual(actual, [355., 45.])
def create_spatial_mask_variable(name, mask_value, dimensions): """ Create an OCGIS spatial mask variable with standard attributes. By default, the value of the returned variable is allocated with zeroes. :param str name: Variable name :param mask_value: Boolean array with dimension matching ``dimensions`` :type mask_value: :class:`numpy.ndarray` :param dimensions: Dimension sequence for the new variable :type dimensions: tuple(:class:`ocgis.Dimension`, ...) :rtype: :class:`ocgis.Variable` """ mask_variable = Variable( name, mask=mask_value, dtype=np.dtype('i1'), dimensions=dimensions, attrs={ 'ocgis_role': 'spatial_mask', 'description': 'values matching fill value are spatially masked' }) mask_variable.allocate_value(fill=0) return mask_variable
def __init__(self, **kwargs): self._name_ugid = None self._geom_type = kwargs.pop(KeywordArgument.GEOM_TYPE, 'auto') if kwargs.get(KeywordArgument.NAME) is None: kwargs[KeywordArgument.NAME] = 'geom' ugid = kwargs.pop(KeywordArgument.UGID, None) super(GeometryVariable, self).__init__(**kwargs) if ugid is not None: ugid_var = Variable(name=HeaderName.ID_SELECTION_GEOMETRY, value=[ugid], dimensions=self.dimensions) self.set_ugid(ugid_var)
def fixture_driver_scrip_netcdf_field(self): xvalue = np.arange(10., 35., step=5) yvalue = np.arange(45., 85., step=10) grid_size = xvalue.shape[0] * yvalue.shape[0] dim_grid_size = Dimension(name='grid_size', size=grid_size) x = Variable(name='grid_center_lon', dimensions=dim_grid_size) y = Variable(name='grid_center_lat', dimensions=dim_grid_size) for idx, (xv, yv) in enumerate(itertools.product(xvalue, yvalue)): x.get_value()[idx] = xv y.get_value()[idx] = yv gc = PointGC(x=x, y=y, crs=Spherical(), driver=DriverNetcdfSCRIP) grid = GridUnstruct(geoms=[gc]) ret = Field(grid=grid, driver=DriverNetcdfSCRIP) grid_dims = Variable(name='grid_dims', value=[yvalue.shape[0], xvalue.shape[0]], dimensions='grid_rank') ret.add_variable(grid_dims) return ret
def test_system_predicate(self): """Test creating a request dataset with a predicate.""" path = self.get_temporary_file_path('foo.nc') field = self.get_field() to_exclude = Variable(name='exclude') field.add_variable(to_exclude) field.write(path) rd = RequestDataset(uri=path, predicate=lambda x: not x.startswith('exclude')) self.assertNotIn('exclude', rd.metadata['variables']) actual = rd.get() self.assertNotIn('exclude', actual) # Test predicate affects data variable identification. path = self.get_temporary_file_path('foo.nc') rd = RequestDataset(uri=path, predicate=lambda x: x != 'foo') with self.assertRaises(NoDataVariablesFound): assert rd.variable
def create_ugid(self, name, start=1): """ Create a unique identifier variable for the geometry variable. The returned variable will have the same dimensions. :param str name: The name for the new geometry variable. :param int start: Starting value for the unique identifier. :rtype: :class:`~ocgis.Variable` """ if self.is_empty: value = None else: value = np.arange(start, start + self.size).reshape(self.shape) ret = Variable(name=name, value=value, dimensions=self.dimensions, is_empty=self.is_empty) self.set_ugid(ret) return ret
def test_broadcast_variable(self): value = np.random.rand(3, 4, 5) desired_value = deepcopy(value) mask = desired_value > 0.5 desired_mask = deepcopy(mask) original_dimensions = ['time', 'lat', 'lon'] src = Variable(name='src', value=value, mask=mask, dimensions=original_dimensions) dst_names = ['lon', 'lat', 'time'] broadcast_variable(src, dst_names) self.assertEqual(src.shape, (5, 4, 3)) self.assertEqual(src.get_value().shape, (5, 4, 3)) self.assertEqual(desired_value.sum(), src.get_value().sum()) broadcast_variable(src, original_dimensions) self.assertNumpyAll(desired_value, src.get_value()) self.assertNumpyMayShareMemory(value, src.get_value()) self.assertNumpyAll(desired_mask, src.get_mask())
def test_set_variable(self): var = Variable(name='test', value=[1, 2], dimensions='two') dmap = DimensionMap() dmap.set_variable(DMK.X, var, dimension='not_two') actual = dmap.get_dimension(DMK.X) desired = ['not_two'] self.assertEqual(actual, desired) # Test setting with a variable and bounds. var = Variable(name='center', value=[1, 2, 3], dtype=float, dimensions='one') var.set_extrapolated_bounds('bounds_data', 'bounds') dmap = DimensionMap() dmap.set_variable(DMK.Y, var) self.assertEqual(dmap.get_bounds(DMK.Y), var.bounds.name) new_var = Variable(name='new_center', value=[1, 2, 3], dtype=float, dimensions='one') dmap.set_variable(DMK.Y, new_var) self.assertIsNone(dmap.get_bounds(DMK.Y)) # Test a dimension position argument is needed if the variable has more than one dimension. var = Variable(name='two_dims', value=np.zeros((4, 5)), dimensions=['one', 'two']) dmap = DimensionMap() with self.assertRaises(DimensionMapError): dmap.set_variable(DMK.X, var) dmap.set_variable(DMK.X, var, pos=1) self.assertEqual(dmap.get_dimension(DMK.X), ['two']) # Test a scalar dimension. var = Variable(name='scalar_dimension', dimensions=[]) dmap = DimensionMap() dmap.set_variable(DMK.LEVEL, var) self.assertEqual(dmap.get_variable(DMK.LEVEL), 'scalar_dimension') self.assertIsNone(dmap.get_bounds(DMK.LEVEL)) self.assertEqual(dmap.get_dimension(DMK.LEVEL), []) # Test dimensionless variable. var = Variable(name='two_dims', value=np.zeros((4, 5)), dimensions=['one', 'two']) dmap = DimensionMap() dmap.set_variable(DMK.X, var, dimensionless=True) self.assertEqual(dmap.get_dimension(DMK.X), [])
def get_masking_slice(intersects_mask_value, target, apply_slice=True): """ Collective! :param intersects_mask_value: The mask to use for creating the slice indices. :type intersects_mask_value: :class:`numpy.ndarray`, dtype=bool :param target: The target slicable object to slice. :param bool apply_slice: If ``True``, apply the slice. """ raise_if_empty(target) if intersects_mask_value is None: local_slice = None else: if intersects_mask_value.all(): local_slice = None elif not intersects_mask_value.any(): shp = intersects_mask_value.shape local_slice = [(0, shp[0]), (0, shp[1])] else: _, local_slice = get_trimmed_array_by_mask(intersects_mask_value, return_adjustments=True) local_slice = [(l.start, l.stop) for l in local_slice] if local_slice is not None: offset_local_slice = [None] * len(local_slice) for idx in range(len(local_slice)): offset = target.dimensions[idx].bounds_local[0] offset_local_slice[idx] = (local_slice[idx][0] + offset, local_slice[idx][1] + offset) else: offset_local_slice = None gathered_offset_local_slices = vm.gather(offset_local_slice) if vm.rank == 0: gathered_offset_local_slices = [ g for g in gathered_offset_local_slices if g is not None ] if len(gathered_offset_local_slices) == 0: raise_empty_subset = True else: raise_empty_subset = False offset_array = np.array(gathered_offset_local_slices) global_slice = [None] * offset_array.shape[1] for idx in range(len(global_slice)): global_slice[idx] = (np.min(offset_array[:, idx, :]), np.max(offset_array[:, idx, :])) else: global_slice = None raise_empty_subset = None raise_empty_subset = vm.bcast(raise_empty_subset) if raise_empty_subset: raise EmptySubsetError global_slice = vm.bcast(global_slice) global_slice = tuple([slice(g[0], g[1]) for g in global_slice]) intersects_mask = Variable(name='mask_gather', value=intersects_mask_value, dimensions=target.dimensions, dtype=bool) if apply_slice: if vm.size_global > 1: ret = target.get_distributed_slice(global_slice) ret_mask = intersects_mask.get_distributed_slice(global_slice) else: ret = target.__getitem__(global_slice) ret_mask = intersects_mask.__getitem__(global_slice) else: ret = target ret_mask = intersects_mask return ret, ret_mask, global_slice
# Create an exact field on the grid. field = create_exact_field(grid, 'foo') # Create a subset geometry. subset_geom = box(30., 20., 40., 25.) # Subset the field using the geometry. Note we subset the grid and return its parent (the field). sub = field.grid.get_intersects(subset_geom).parent # Subset geometries are themselves treated as fields. Convert the geometry to an OCGIS geometry variable. gvar = GeometryVariable.from_shapely(subset_geom, ugid=11, crs=crs.Spherical()) # Add some descriptive variables. info = Variable(name='desc', value=['random bbox'], dtype=str, dimensions=gvar.dimensions[0]) gvar.parent.add_variable(info, is_data=True) height = Variable(name='height', value=[30], dimensions=gvar.dimensions[0], attrs={'made_up': 'yes'}) gvar.parent.add_variable(height, is_data=True) # Create an empty spatial collection. sc = SpatialCollection() # Add the subsetted field with its containing geometry. sc.add_field(sub, gvar.parent) # These are the container geometries for the spatial collection.
def write_subsets(self, src_template, dst_template, wgt_template, index_path): """ Write grid subsets to netCDF files using the provided filename templates. The template must contain the full file path with a single curly-bracer pair to insert the combination counter. ``wgt_template`` should not be a full path. This name is used when generating weight files. >>> template_example = '/path/to/data_{}.nc' :param str src_template: The template for the source subset file. :param str dst_template: The template for the destination subset file. :param str wgt_template: The template for the weight filename. >>> wgt_template = 'esmf_weights_{}.nc' :param index_path: Path to the output indexing netCDF. """ src_filenames = [] dst_filenames = [] wgt_filenames = [] dst_slices = [] # nzeros = len(str(reduce(lambda x, y: x * y, self.nsplits_dst))) for ctr, (sub_src, sub_dst, dst_slc) in enumerate(self.iter_src_grid_subsets(yield_dst=True), start=1): # padded = create_zero_padded_integer(ctr, nzeros) src_path = src_template.format(ctr) dst_path = dst_template.format(ctr) wgt_filename = wgt_template.format(ctr) src_filenames.append(os.path.split(src_path)[1]) dst_filenames.append(os.path.split(dst_path)[1]) wgt_filenames.append(wgt_filename) dst_slices.append(dst_slc) for target, path in zip([sub_src, sub_dst], [src_path, dst_path]): if target.is_empty: is_empty = True target = None else: is_empty = False field = Field(grid=target, is_empty=is_empty) ocgis_lh(msg='writing: {}'.format(path), level=logging.DEBUG) with vm.scoped_by_emptyable('field.write', field): if not vm.is_null: field.write(path) ocgis_lh(msg='finished writing: {}'.format(path), level=logging.DEBUG) with vm.scoped('index write', [0]): if not vm.is_null: dim = Dimension('nfiles', len(src_filenames)) vname = ['source_filename', 'destination_filename', 'weights_filename'] values = [src_filenames, dst_filenames, wgt_filenames] grid_splitter_destination = GridSplitterConstants.IndexFile.NAME_DESTINATION_VARIABLE attrs = [{'esmf_role': 'grid_splitter_source'}, {'esmf_role': grid_splitter_destination}, {'esmf_role': 'grid_splitter_weights'}] vc = VariableCollection() grid_splitter_index = GridSplitterConstants.IndexFile.NAME_INDEX_VARIABLE vidx = Variable(name=grid_splitter_index) vidx.attrs['esmf_role'] = grid_splitter_index vidx.attrs['grid_splitter_source'] = 'source_filename' vidx.attrs[GridSplitterConstants.IndexFile.NAME_DESTINATION_VARIABLE] = 'destination_filename' vidx.attrs['grid_splitter_weights'] = 'weights_filename' x_bounds = GridSplitterConstants.IndexFile.NAME_X_BOUNDS_VARIABLE vidx.attrs[x_bounds] = x_bounds y_bounds = GridSplitterConstants.IndexFile.NAME_Y_BOUNDS_VARIABLE vidx.attrs[y_bounds] = y_bounds vc.add_variable(vidx) for idx in range(len(vname)): v = Variable(name=vname[idx], dimensions=dim, dtype=str, value=values[idx], attrs=attrs[idx]) vc.add_variable(v) bounds_dimension = Dimension(name='bounds', size=2) xb = Variable(name=x_bounds, dimensions=[dim, bounds_dimension], attrs={'esmf_role': 'x_split_bounds'}, dtype=int) yb = Variable(name=y_bounds, dimensions=[dim, bounds_dimension], attrs={'esmf_role': 'y_split_bounds'}, dtype=int) x_name = self.dst_grid.x.dimensions[0].name y_name = self.dst_grid.y.dimensions[0].name for idx, slc in enumerate(dst_slices): xb.get_value()[idx, :] = slc[x_name].start, slc[x_name].stop yb.get_value()[idx, :] = slc[y_name].start, slc[y_name].stop vc.add_variable(xb) vc.add_variable(yb) vc.write(index_path) vm.barrier()
def iter_src_grid_subsets(self, yield_dst=False): """ Yield source grid subsets using the extent of its associated destination grid subset. :param bool yield_dst: If ``True``, yield the destination subset as well as the source grid subset. :return: The source grid if ``yield_dst`` is ``False``, otherwise a three-element tuple in the form ``(<source grid subset>, <destination grid subset>, <destination grid slice>)``. :rtype: :class:`ocgis.Grid` or (:class:`ocgis.Grid`, :class:`ocgis.Grid`, dict) """ if yield_dst: yield_slice = True else: yield_slice = False dst_grid_resolution = self.dst_grid.resolution src_grid_resolution = self.src_grid.resolution if dst_grid_resolution <= src_grid_resolution: target_resolution = dst_grid_resolution else: target_resolution = src_grid_resolution buffer_value = 2 * target_resolution for yld in self.iter_dst_grid_subsets(yield_slice=yield_slice): if yield_slice: dst_grid_subset, dst_slice = yld else: dst_grid_subset = yld dst_box = None with vm.scoped_by_emptyable('extent_global', dst_grid_subset): if not vm.is_null: if self.check_contains: dst_box = box(*dst_grid_subset.extent_global) # Use the envelope! A buffer returns "fancy" borders. We just want to expand the bounding box. sub_box = box(*dst_grid_subset.extent_global).buffer(buffer_value).envelope ocgis_lh(msg=str(sub_box.bounds), level=logging.DEBUG) else: sub_box, dst_box = [None, None] live_ranks = vm.get_live_ranks_from_object(dst_grid_subset) sub_box = vm.bcast(sub_box, root=live_ranks[0]) if self.check_contains: dst_box = vm.bcast(dst_box, root=live_ranks[0]) src_grid_subset = self.src_grid.get_intersects(sub_box, keep_touches=False, cascade=False, optimized_bbox_subset=True) if not self.allow_masked: gmask = self.src_grid.get_mask() if not self.allow_masked: if gmask is not None and gmask.any(): raise ValueError('Masked values in source grid subset.') with vm.scoped_by_emptyable('src_grid_subset', src_grid_subset): if not vm.is_null: if self.check_contains: src_box = box(*src_grid_subset.extent_global) if not does_contain(src_box, dst_box): raise ValueError('Contains check failed.') else: src_grid_subset = Grid(Variable('x', is_empty=True), Variable('y', is_empty=True)) if yield_dst: yld = (src_grid_subset, dst_grid_subset, dst_slice) else: yld = src_grid_subset yield yld