def test_from_regular_cube_no_chunks_and_chunks(self):
dataset = xcube.core.new.new_cube(variables=dict(rad=0.5))
gm1 = GridMapping.from_dataset(dataset)
self.assertEqual((360, 180), gm1.tile_size)
dataset = dataset.chunk(dict(lon=10, lat=20))
gm2 = GridMapping.from_dataset(dataset)
self.assertEqual((10, 20), gm2.tile_size)
def test_from_sentinel_2(self):
dataset = create_s2plus_dataset()
gm = GridMapping.from_dataset(dataset)
# Should pick the projected one which is regular
self.assertEqual("Projected CRS", gm.crs.type_name)
self.assertEqual(True, gm.is_regular)
gm = GridMapping.from_dataset(dataset, prefer_is_regular=True)
# Should pick the projected one which is regular
self.assertEqual("Projected CRS", gm.crs.type_name)
self.assertEqual(True, gm.is_regular)
gm = GridMapping.from_dataset(dataset, prefer_is_regular=False)
# Should pick the geographic one which is irregular
self.assertEqual('Geographic 2D CRS', gm.crs.type_name)
self.assertEqual(False, gm.is_regular)
gm = GridMapping.from_dataset(dataset, prefer_crs=GEO_CRS)
# Should pick the geographic one which is irregular
self.assertEqual('Geographic 2D CRS', gm.crs.type_name)
self.assertEqual(False, gm.is_regular)
gm = GridMapping.from_dataset(dataset,
prefer_crs=GEO_CRS,
prefer_is_regular=True)
# Should pick the geographic one which is irregular
self.assertEqual('Geographic 2D CRS', gm.crs.type_name)
self.assertEqual(False, gm.is_regular)
def test_select_spatial_subset_with_gm(self):
ds1 = new_cube(inverse_y=True)
ds2 = select_spatial_subset(ds1,
xy_bbox=(40., 40., 42., 42.),
grid_mapping=GridMapping.from_dataset(ds1))
self.assertEqual(((2, ), (2, )), (ds2.lon.shape, ds2.lat.shape))
ds2 = select_spatial_subset(ds1,
xy_bbox=(40., 40., 42., 42.),
geo_coding=GridMapping.from_dataset(ds1))
self.assertEqual(((2, ), (2, )), (ds2.lon.shape, ds2.lat.shape))
def step1a(input_slice):
nonlocal grid_mapping
grid_mapping = GridMapping.from_dataset(input_slice)
subset = select_spatial_subset(input_slice,
xy_bbox=output_geom.xy_bbox,
xy_border=output_geom.x_res,
ij_border=1,
grid_mapping=grid_mapping)
if subset is None:
monitor('no spatial overlap with input')
elif subset is not input_slice:
grid_mapping = GridMapping.from_dataset(subset)
return subset
def test_chunks_are_larger_than_sizes(self):
cube1 = new_cube(variables=dict(chl=0.6, tsm=0.9, flags=16))
for var in cube1.variables.values():
self.assertIsNone(var.chunks)
rc = CubeRechunker()
cube2, gm, cc = rc.transform_cube(
cube1, GridMapping.from_dataset(cube1),
CubeConfig(chunks=dict(time=64, lat=512, lon=512)))
self.assertIsInstance(cube2, xr.Dataset)
self.assertIsInstance(gm, GridMapping)
self.assertIsInstance(cc, CubeConfig)
self.assertEqual(cube1.attrs, cube2.attrs)
self.assertEqual(set(cube1.coords), set(cube2.coords))
self.assertEqual(set(cube1.data_vars), set(cube2.data_vars))
for k, v in cube2.coords.items():
if v.chunks is not None:
self.assertIsInstance(v.chunks, tuple, msg=f'{k!r}={v!r}')
self.assertNotIn('chunks', v.encoding)
for k, v in cube2.data_vars.items():
self.assertIsInstance(v.chunks, tuple, msg=f'{k!r}={v!r}')
self.assertEqual(((5, ), (180, ), (360, )), v.chunks)
self.assertNotIn('chunks', v.encoding)
def test_chunks_are_smaller_than_sizes(self):
cube1 = new_cube(variables=dict(chl=0.6, tsm=0.9, flags=16))
for var in cube1.variables.values():
self.assertIsNone(var.chunks)
rc = CubeRechunker()
cube2, gm, cc = rc.transform_cube(
cube1, GridMapping.from_dataset(cube1),
CubeConfig(chunks=dict(time=2, lat=100, lon=200)))
self.assertIsInstance(cube2, xr.Dataset)
self.assertIsInstance(gm, GridMapping)
self.assertIsInstance(cc, CubeConfig)
self.assertEqual(cube1.attrs, cube2.attrs)
self.assertEqual({'time': 2, 'lat': 100, 'lon': 200}, cc.chunks)
self.assertEqual(set(cube1.coords), set(cube2.coords))
self.assertEqual(set(cube1.data_vars), set(cube2.data_vars))
for k, v in cube2.coords.items():
if v.chunks is not None:
self.assertIsInstance(v.chunks, tuple, msg=f'{k!r}={v!r}')
self.assertNotIn('chunks', v.encoding)
# self.assertIn('chunks', v.encoding)
# self.assertEqual([v.sizes[d] for d in v.dims],
# v.encoding['chunks'])
for k, v in cube2.data_vars.items():
self.assertIsInstance(v.chunks, tuple, msg=f'{k!r}={v!r}')
self.assertEqual(((2, 2, 1), (100, 80), (200, 160)), v.chunks)
self.assertNotIn('chunks', v.encoding)
def _get_grid_mapping_lazily(self) -> GridMapping:
"""
Retrieve, i.e. read or compute, the tile grid used by the multi-level dataset.
:return: the dataset for the level at *index*.
"""
return GridMapping.from_dataset(self.get_dataset(0))
def test_from_regular_cube_with_crs(self):
dataset = xcube.core.new.new_cube(variables=dict(rad=0.5),
x_start=0,
y_start=0,
x_name='x',
y_name='y',
crs='epsg:25832')
gm1 = GridMapping.from_dataset(dataset)
self.assertEqual(pyproj.CRS.from_string('epsg:25832'), gm1.crs)
dataset = dataset.drop_vars('crs')
gm2 = GridMapping.from_dataset(dataset)
self.assertEqual(GEO_CRS, gm2.crs)
gm3 = GridMapping.from_dataset(dataset, crs=gm1.crs)
self.assertEqual(gm1.crs, gm3.crs)
self.assertEqual(('x', 'y'), gm3.xy_var_names)
self.assertEqual(('x', 'y'), gm3.xy_dim_names)
def transform_cube(self,
cube: xr.Dataset,
gm: GridMapping,
cube_config: CubeConfig) -> TransformedCube:
desired_var_names = cube_config.variable_names
if desired_var_names:
cube = select_variables_subset(cube,
var_names=desired_var_names)
cube_config = cube_config.drop_props('variable_names')
desired_bbox = cube_config.bbox
if desired_bbox is not None:
# Find out whether its possible to make a spatial subset
# without resampling. First, grid mapping must be regular.
can_do_spatial_subset = False
if gm.is_regular:
can_do_spatial_subset = True
# Current spatial resolution must be the
# desired spatial resolution, otherwise spatial resampling
# is required later, which will include the desired
# subsetting.
desired_res = cube_config.spatial_res
if desired_res is not None \
and not (math.isclose(gm.x_res, desired_res)
and math.isclose(gm.y_res, desired_res)):
can_do_spatial_subset = False
if can_do_spatial_subset:
# Finally, the desired CRS must be equal to the current
# one, or they must both be geographic.
desired_crs = cube_config.crs
if desired_crs:
desired_crs = pyproj.CRS.from_string(desired_crs)
if desired_crs != gm.crs \
and not (desired_crs.is_geographic
and gm.crs.is_geographic):
can_do_spatial_subset = False
if can_do_spatial_subset:
cube = select_spatial_subset(cube,
xy_bbox=desired_bbox)
# Now that we have a new cube subset, we must adjust
# its grid mapping.
gm = GridMapping.from_dataset(
cube,
crs=gm.crs,
xy_var_names=gm.xy_var_names,
)
# Consume spatial properties
cube_config = cube_config.drop_props(['bbox',
'spatial_res',
'crs'])
desired_time_range = cube_config.time_range
if desired_time_range:
cube = select_temporal_subset(cube,
time_range=desired_time_range)
cube_config = cube_config.drop_props('time_range')
return cube, gm, cube_config
def test_rectify_dataset(self):
source_ds = create_s2plus_dataset()
expected_data = np.array([
[nan, nan, nan, nan, nan, nan, nan, nan, nan],
[
nan, 0.019001, 0.019001, 0.008999, 0.012001, 0.012001,
0.022999, nan, nan
],
[
nan, 0.021, 0.021, 0.009998, 0.009998, 0.008999, 0.022999, nan,
nan
],
[
nan, 0.022999, 0.022999, 0.007999, 0.007999, 0.008999,
0.023998, nan, nan
],
[nan, 0.022999, 0.022999, 0.007, 0.007, 0.009998, 0.021, nan, nan],
[nan, nan, nan, nan, nan, nan, nan, nan, nan]
])
source_gm = GridMapping.from_dataset(source_ds, prefer_crs=CRS_WGS84)
target_ds = rectify_dataset(source_ds,
source_gm=source_gm,
tile_size=None)
self.assertEqual(None, target_ds.rrs_665.chunks)
print(repr(target_ds.rrs_665.values))
np.testing.assert_almost_equal(target_ds.rrs_665.values,
expected_data,
decimal=3)
target_ds = rectify_dataset(source_ds,
source_gm=source_gm,
tile_size=5)
self.assertEqual(((5, 1), (5, 4)), target_ds.rrs_665.chunks)
np.testing.assert_almost_equal(target_ds.rrs_665.values,
expected_data,
decimal=3)
target_ds = rectify_dataset(source_ds,
source_gm=source_gm,
tile_size=None,
is_j_axis_up=True)
self.assertEqual(None, target_ds.rrs_665.chunks)
np.testing.assert_almost_equal(target_ds.rrs_665.values,
expected_data[::-1],
decimal=3)
target_ds = rectify_dataset(source_ds,
source_gm=source_gm,
tile_size=5,
is_j_axis_up=True)
self.assertEqual(((5, 1), (5, 4)), target_ds.rrs_665.chunks)
np.testing.assert_almost_equal(target_ds.rrs_665.values,
expected_data[::-1],
decimal=3)
def test_non_geographical_crs(self):
cube = new_cube(x_name='x',
y_name='y',
crs='epsg:25832',
variables=dict(a=1, b=2, c=3))
gm = GridMapping.from_dataset(cube)
dataset = encode_cube(cube, gm, xr.Dataset(dict(d=True)))
self.assertIsInstance(dataset, xr.Dataset)
self.assertEqual({'a', 'b', 'c', 'd', 'crs'}, set(dataset.data_vars))
def test_affine_transform_dataset(self):
source_ds = new_cube()
source_gm = GridMapping.from_dataset(source_ds)
target_gm = GridMapping.regular(size=(8, 4),
xy_min=(0, 0),
xy_res=2,
crs=CRS_WGS84)
# TODO: assert stuff
resample_in_space(source_ds, source_gm, target_gm)
def test_geographical_crs(self):
cube = new_cube(variables=dict(a=1, b=2, c=3))
gm = GridMapping.from_dataset(cube)
dataset = encode_cube(cube, gm)
self.assertIs(cube, dataset)
dataset = encode_cube(cube, gm, xr.Dataset(dict(d=True)))
self.assertIsInstance(dataset, xr.Dataset)
self.assertEqual({'a', 'b', 'c', 'd'}, set(dataset.data_vars))
def test_easter_egg(self):
cube = new_cube()
md_adjuster = CubeMetadataAdjuster()
with self.assertRaises(ValueError) as cm:
md_adjuster.transform_cube(
cube, GridMapping.from_dataset(cube),
CubeConfig(metadata=dict(inverse_fine_structure_constant=136)))
self.assertEqual(('inverse_fine_structure_constant must be 137'
' or running in wrong universe', ),
cm.exception.args)
def test_from_regular_cube(self):
dataset = xcube.core.new.new_cube(variables=dict(rad=0.5))
gm = GridMapping.from_dataset(dataset)
self.assertEqual((360, 180), gm.size)
self.assertEqual((360, 180), gm.tile_size)
self.assertEqual(GEO_CRS, gm.crs)
self.assertEqual((1, 1), gm.xy_res)
self.assertEqual(True, gm.is_regular)
self.assertEqual(False, gm.is_lon_360)
self.assertEqual(True, gm.is_j_axis_up)
self.assertEqual((2, 180, 360), gm.xy_coords.shape)
self.assertEqual(('coord', 'lat', 'lon'), gm.xy_coords.dims)
def test_it(self):
cube = new_cube(variables=dict(a=0.5))
gm = GridMapping.from_dataset(cube)
cube_config = CubeConfig()
identity = CubeIdentity()
t_cube = identity.transform_cube(cube,
gm,
cube_config)
self.assertIsInstance(t_cube, tuple)
self.assertEqual(3, len(t_cube))
self.assertIs(cube, t_cube[0])
self.assertIs(gm, t_cube[1])
self.assertIs(cube_config, t_cube[2])
def assertCallableWorks(self, user_code_callable):
code_config = CodeConfig(_callable=user_code_callable,
callable_params=self.good_params)
executor = CubeUserCodeExecutor(code_config)
ds_input = new_cube(variables=dict(a=1))
ds_output, gm, cc = executor.transform_cube(
ds_input, GridMapping.from_dataset(ds_input), CubeConfig())
self.assertIsInstance(ds_output, xr.Dataset)
self.assertIsInstance(gm, GridMapping)
self.assertIsInstance(cc, CubeConfig)
self.assertIsNot(ds_output, ds_input)
self.assertIn('X', ds_output)
self.assertEqual(42, ds_output.X)
def test_transform_s2(self):
dataset = create_s2plus_dataset()
gm = GridMapping.from_dataset(dataset, prefer_is_regular=True)
# Assert we've picked the projected one which is regular
self.assertEqual("Projected CRS", gm.crs.type_name)
self.assertEqual(True, gm.is_regular)
gm_t = gm.transform(CRS_CRS84)
self.assertEqual(CRS_CRS84, gm_t.crs)
gm_t = gm.transform(CRS_WGS84)
self.assertEqual(CRS_WGS84, gm_t.crs)
def test_metadata_adjusted_geo_crs(self):
(x1, x2), (y1, y2) = (53, 54), (11, 12)
cube1 = new_cube(
width=1000,
height=1000,
variables=dict(chl=0.6, tsm=0.9, flags=16),
x_start=x1,
y_start=y1,
x_res=(x2 - x1) / 1000,
y_res=(x2 - x1) / 1000,
)
cube1.attrs = {}
md_adjuster = CubeMetadataAdjuster()
cube2, gm, cc = md_adjuster.transform_cube(
cube1, GridMapping.from_dataset(cube1),
CubeConfig(metadata=dict(title='S2L2A subset'),
variable_metadata=dict(
chl=dict(long_name='Chlorophyll'),
tsm=dict(long_name='Total suspended matter'),
flags=dict(long_name='Quality flags'),
)))
self.assertIsNot(cube1, cube2)
self.assertIsInstance(gm, GridMapping)
self.assertIsInstance(cc, CubeConfig)
date_created = cube2.attrs.pop('date_created', None)
self.assertIsInstance(date_created, str)
history = cube2.attrs.pop('history', None)
self.assertIsInstance(history, list)
self.assertEqual(
{
'Conventions': 'CF-1.7',
'title': 'S2L2A subset',
'geospatial_bounds_crs': 'CRS84',
'geospatial_bounds': 'POLYGON((53 11, 53 12,'
' 54 12, 54 11, 53 11))',
'geospatial_lat_max': 12,
'geospatial_lat_min': 11,
'geospatial_lat_resolution': 0.001,
'geospatial_lat_units': 'degrees_north',
'geospatial_lon_max': 54,
'geospatial_lon_min': 53,
'geospatial_lon_resolution': 0.001,
'geospatial_lon_units': 'degrees_east',
}, cube2.attrs)
self.assertEqual({'long_name': 'Chlorophyll'}, cube2.chl.attrs)
self.assertEqual({'long_name': 'Total suspended matter'},
cube2.tsm.attrs)
self.assertEqual({'long_name': 'Quality flags'}, cube2.flags.attrs)
def test_non_empty_cube(self):
cube = new_cube(variables=dict(a=0.5))
gm = GridMapping.from_dataset(cube)
cube_config = CubeConfig(tile_size=180)
t_cube = transform_cube((cube, gm, cube_config), MyTiler())
self.assertIsInstance(t_cube, tuple)
self.assertEqual(3, len(t_cube))
cube2, gm2, cc2 = t_cube
self.assertIsNot(cube, cube2)
self.assertEqual(((5,), (180,), (180, 180)), cube2.a.chunks)
self.assertIs(gm, gm2)
self.assertEqual(None, cc2.tile_size)
def test_empty_cube(self):
cube = new_cube()
gm = GridMapping.from_dataset(cube)
cube_config = CubeConfig(tile_size=180)
t_cube = transform_cube((cube, gm, cube_config), MyTiler())
self.assertIsInstance(t_cube, tuple)
self.assertEqual(3, len(t_cube))
cube2, gm2, cc2 = t_cube
self.assertIs(cube, cube2)
self.assertIs(gm, gm2)
self.assertIs(cube_config, cc2)
self.assertEqual((180, 180), cc2.tile_size)
def test_transform_no_op(self):
dataset = create_s2plus_dataset()
gm = GridMapping.from_dataset(dataset, prefer_is_regular=True)
# Assert we've picked the projected one which is regular
self.assertEqual("Projected CRS", gm.crs.type_name)
self.assertEqual(True, gm.is_regular)
gm_t = gm.transform(gm.crs)
self.assertIs(gm, gm_t)
# Almost no op
gm = GridMapping.regular(size=(3, 3), xy_min=(10, 53), xy_res=0.1, crs=CRS_CRS84)
gm_t = gm.transform(crs=gm.crs, xy_var_names=('x', 'y'))
self.assertEqual(('x', 'y'), gm_t.xy_var_names)
def test_metadata_adjusted_other_crs(self):
crs = pyproj.CRS.from_string('epsg:25832')
t = pyproj.Transformer.from_crs(crs_from=CRS_CRS84, crs_to=crs)
((x1, x2), (y1, y2)) = t.transform((53, 54), (10, 11))
cube1 = new_cube(
width=1000,
height=1000,
variables=dict(chl=0.6, tsm=0.9, flags=16),
x_start=x1,
y_start=y1,
x_res=(x2 - x1) / 1000,
y_res=(x2 - x1) / 1000,
x_name='x',
y_name='y',
crs=crs,
)
cube1.attrs = {}
md_adjuster = CubeMetadataAdjuster()
cube2, gm, cc = md_adjuster.transform_cube(
cube1, GridMapping.from_dataset(cube1),
CubeConfig(metadata=dict(title='S2L2A subset'),
variable_metadata=dict(
chl=dict(long_name='Chlorophyll'),
tsm=dict(long_name='Total suspended matter'),
flags=dict(long_name='Quality flags'),
)))
self.assertIsNot(cube1, cube2)
self.assertIsInstance(gm, GridMapping)
self.assertIsInstance(cc, CubeConfig)
self.assertAlmostEqual(53.,
cube2.attrs.get('geospatial_lon_min'),
delta=0.001)
self.assertAlmostEqual(10.,
cube2.attrs.get('geospatial_lat_min'),
delta=0.001)
self.assertAlmostEqual(0.001,
cube2.attrs.get('geospatial_lon_resolution'),
delta=0.0005)
self.assertAlmostEqual(0.001,
cube2.attrs.get('geospatial_lat_resolution'),
delta=0.0005)
def test_from_real_olci(self):
olci_l2_path = os.path.join(os.path.dirname(__file__), '..', '..',
'..', 'examples', 'notebooks', 'inputdata',
'S3-OLCI-L2A.zarr.zip')
dataset = xr.open_zarr(olci_l2_path)
gm = GridMapping.from_dataset(dataset)
self.assertEqual((1189, 1890), gm.size)
self.assertEqual((512, 512), gm.tile_size)
self.assertEqual(GEO_CRS, gm.crs)
self.assertEqual((0.0025, 0.0025), gm.xy_res)
# self.assertAlmostEqual(12.693771178309552, gm.x_min)
# self.assertAlmostEqual(20.005413821690446, gm.x_max)
# self.assertAlmostEqual(55.19965017830955, gm.y_min)
# self.assertAlmostEqual(60.63871982169044, gm.y_max)
self.assertEqual(False, gm.is_regular)
self.assertEqual(False, gm.is_lon_360)
self.assertEqual(None, gm.is_j_axis_up)
self.assertEqual((2, 1890, 1189), gm.xy_coords.shape)
self.assertEqual(('coord', 'y', 'x'), gm.xy_coords.dims)
gm = gm.to_regular()
self.assertEqual((2926, 2177), gm.size)
def test_from_non_regular_cube(self):
lon = np.array(
[[8, 9.3, 10.6, 11.9], [8, 9.2, 10.4, 11.6], [8, 9.1, 10.2, 11.3]],
dtype=np.float32)
lat = np.array([[56, 56.1, 56.2, 56.3], [55, 55.2, 55.4, 55.6],
[54, 54.3, 54.6, 54.9]],
dtype=np.float32)
rad = np.random.random(3 * 4).reshape((3, 4))
dims = ('y', 'x')
dataset = xr.Dataset(
dict(lon=xr.DataArray(lon, dims=dims),
lat=xr.DataArray(lat, dims=dims),
rad=xr.DataArray(rad, dims=dims)))
gm = GridMapping.from_dataset(dataset)
self.assertEqual((4, 3), gm.size)
self.assertEqual((4, 3), gm.tile_size)
self.assertEqual(GEO_CRS, gm.crs)
self.assertEqual(False, gm.is_regular)
self.assertEqual(False, gm.is_lon_360)
self.assertEqual(None, gm.is_j_axis_up)
self.assertEqual((2, 3, 4), gm.xy_coords.shape)
self.assertEqual(('coord', 'y', 'x'), gm.xy_coords.dims)
self.assertEqual((0.8, 0.8), gm.xy_res)
data_vars=dict(refl=xr.DataArray(np.array([
[0, 1, 0, 2, 0, 3, 0, 4],
[2, 0, 3, 0, 4, 0, 1, 0],
[0, 4, 0, 1, 0, 2, 0, 3],
[1, 0, 2, 0, 3, 0, 4, 0],
[0, 3, 0, 4, 0, 1, 0, 2],
[4, 0, 1, 0, 2, 0, 3, 0],
],
dtype=np.float64),
dims=('lat', 'lon'))),
coords=dict(lon=xr.DataArray(50.0 + res * np.arange(0, 8) + 0.5 * res,
dims='lon'),
lat=xr.DataArray(10.6 - res * np.arange(0, 6) - 0.5 * res,
dims='lat')))
source_gm = GridMapping.from_dataset(source_ds)
class AffineTransformDatasetTest(unittest.TestCase):
def test_subset(self):
target_gm = GridMapping.regular((3, 3), (50.0, 10.0), res,
source_gm.crs)
target_ds = affine_transform_dataset(source_ds, source_gm, target_gm)
self.assertIsInstance(target_ds, xr.Dataset)
self.assertEqual(set(target_ds.variables), set(source_ds.variables))
self.assertEqual((3, 3), target_ds.refl.shape)
np.testing.assert_almost_equal(
target_ds.refl.values, np.array([
[1, 0, 2],
[0, 3, 0],
[4, 0, 1],
def rectify_dataset(source_ds: xr.Dataset,
*,
var_names: Union[str, Sequence[str]] = None,
source_gm: GridMapping = None,
xy_var_names: Tuple[str, str] = None,
target_gm: GridMapping = None,
tile_size: Union[int, Tuple[int, int]] = None,
is_j_axis_up: bool = None,
output_ij_names: Tuple[str, str] = None,
compute_subset: bool = True,
uv_delta: float = 1e-3) -> Optional[xr.Dataset]:
"""
Reproject dataset *source_ds* using its per-pixel
x,y coordinates or the given *source_gm*.
The function expects *source_ds* or the given
*source_gm* to have either one- or two-dimensional
coordinate variables that provide spatial x,y coordinates
for every data variable with the same spatial dimensions.
For example, a dataset may comprise variables with
spatial dimensions ``var(..., y_dim, x_dim)``,
then one the function expects coordinates to be provided
in two forms:
1. One-dimensional ``x_var(x_dim)``
and ``y_var(y_dim)`` (coordinate) variables.
2. Two-dimensional ``x_var(y_dim, x_dim)``
and ``y_var(y_dim, x_dim)`` (coordinate) variables.
If *target_gm* is given and it defines a tile size
or *tile_size* is given, and the number of tiles is
greater than one in the output's x- or y-direction, then the
returned dataset will be composed of lazy, chunked dask
arrays. Otherwise the returned dataset will be composed
of ordinary numpy arrays.
:param source_ds: Source dataset grid mapping.
:param var_names: Optional variable name or sequence of
variable names.
:param source_gm: Target dataset grid mapping.
:param xy_var_names: Optional tuple of the x- and y-coordinate
variables in *source_ds*. Ignored if *source_gm* is given.
:param target_gm: Optional output geometry. If not given,
output geometry will be computed to spatially fit *dataset*
and to retain its spatial resolution.
:param tile_size: Optional tile size for the output.
:param is_j_axis_up: Whether y coordinates are increasing with
positive image j axis.
:param output_ij_names: If given, a tuple of variable names in
which to store the computed source pixel coordinates in
the returned output.
:param compute_subset: Whether to compute a spatial subset
from *dataset* using *output_geom*. If set, the function
may return ``None`` in case there is no overlap.
:param uv_delta: A normalized value that is used to determine
whether x,y coordinates in the output are contained
in the triangles defined by the input x,y coordinates.
The higher this value, the more inaccurate the rectification
will be.
:return: a reprojected dataset, or None if the requested output
does not intersect with *dataset*.
"""
if source_gm is None:
source_gm = GridMapping.from_dataset(source_ds,
xy_var_names=xy_var_names)
src_attrs = dict(source_ds.attrs)
if target_gm is None:
target_gm = source_gm.to_regular(tile_size=tile_size)
elif compute_subset:
source_ds_subset = select_spatial_subset(
source_ds,
xy_bbox=target_gm.xy_bbox,
ij_border=1,
xy_border=0.5 * (target_gm.x_res + target_gm.y_res),
grid_mapping=source_gm)
if source_ds_subset is None:
return None
if source_ds_subset is not source_ds:
# TODO: GridMapping.from_dataset() may be expensive.
# Find a more effective way.
source_gm = GridMapping.from_dataset(source_ds_subset)
source_ds = source_ds_subset
# if src_geo_coding.xy_var_names != output_geom.xy_var_names:
# output_geom = output_geom.derive(
# xy_var_names=src_geo_coding.xy_var_names
# )
# if src_geo_coding.xy_dim_names != output_geom.xy_dim_names:
# output_geom = output_geom.derive(
# xy_dim_names=src_geo_coding.xy_dim_names
# )
if tile_size is not None or is_j_axis_up is not None:
target_gm = target_gm.derive(tile_size=tile_size,
is_j_axis_up=is_j_axis_up)
src_vars = _select_variables(source_ds, source_gm, var_names)
if target_gm.is_tiled:
compute_dst_src_ij_images = _compute_ij_images_xarray_dask
compute_dst_var_image = _compute_var_image_xarray_dask
else:
compute_dst_src_ij_images = _compute_ij_images_xarray_numpy
compute_dst_var_image = _compute_var_image_xarray_numpy
dst_src_ij_array = compute_dst_src_ij_images(source_gm, target_gm,
uv_delta)
dst_x_dim, dst_y_dim = target_gm.xy_dim_names
dst_dims = dst_y_dim, dst_x_dim
dst_ds_coords = target_gm.to_coords()
dst_vars = dict()
for src_var_name, src_var in src_vars.items():
dst_var_dims = src_var.dims[0:-2] + dst_dims
dst_var_coords = {
d: src_var.coords[d]
for d in dst_var_dims if d in src_var.coords
}
dst_var_coords.update(
{d: dst_ds_coords[d]
for d in dst_var_dims if d in dst_ds_coords})
dst_var_array = compute_dst_var_image(src_var,
dst_src_ij_array,
fill_value=np.nan)
dst_var = xr.DataArray(dst_var_array,
dims=dst_var_dims,
coords=dst_var_coords,
attrs=src_var.attrs)
dst_vars[src_var_name] = dst_var
if output_ij_names:
output_i_name, output_j_name = output_ij_names
dst_ij_coords = {
d: dst_ds_coords[d]
for d in dst_dims if d in dst_ds_coords
}
dst_vars[output_i_name] = xr.DataArray(dst_src_ij_array[0],
dims=dst_dims,
coords=dst_ij_coords)
dst_vars[output_j_name] = xr.DataArray(dst_src_ij_array[1],
dims=dst_dims,
coords=dst_ij_coords)
return xr.Dataset(dst_vars, coords=dst_ds_coords, attrs=src_attrs)
def select_spatial_subset(dataset: xr.Dataset,
ij_bbox: Tuple[int, int, int, int] = None,
ij_border: int = 0,
xy_bbox: Tuple[float, float, float, float] = None,
xy_border: float = 0.,
grid_mapping: GridMapping = None,
geo_coding: GridMapping = None,
xy_names: Tuple[str, str] = None) \
-> Optional[xr.Dataset]:
"""
Select a spatial subset of *dataset* for the
bounding box *ij_bbox* or *xy_bbox*.
*ij_bbox* or *xy_bbox* must not be given both.
:param xy_bbox: Bounding box in coordinates of the dataset's CRS.
:param xy_border: Extra border added to *xy_bbox*.
:param dataset: Source dataset.
:param ij_bbox: Bounding box (i_min, i_min, j_max, j_max)
in pixel coordinates.
:param ij_border: Extra border added to *ij_bbox*
in number of pixels
:param xy_bbox: The bounding box in x,y coordinates.
:param xy_border: Border in units of the x,y coordinates.
:param grid_mapping: Optional dataset grid mapping.
:param geo_coding: Deprecated. Use *grid_mapping* instead.
:param xy_names: Optional tuple of the x- and y-coordinate
variables in *dataset*. Ignored if *geo_coding* is given.
:return: Spatial dataset subset
"""
if ij_bbox is None and xy_bbox is None:
raise ValueError('One of ij_bbox and xy_bbox must be given')
if ij_bbox and xy_bbox:
raise ValueError('Only one of ij_bbox and xy_bbox can be given')
if geo_coding:
warnings.warn(
'keyword "geo_coding" has been deprecated,'
' use "grid_mapping" instead', DeprecationWarning)
grid_mapping = grid_mapping or geo_coding
if grid_mapping is None:
grid_mapping = GridMapping.from_dataset(dataset, xy_var_names=xy_names)
x_name, y_name = grid_mapping.xy_var_names
x = dataset[x_name]
y = dataset[y_name]
if x.ndim == 1 and y.ndim == 1:
# Hotfix für #981 and #985
if xy_bbox:
if y.values[0] < y.values[-1]:
ds = dataset.sel(
**{
x_name:
slice(xy_bbox[0] - xy_border, xy_bbox[2] + xy_border),
y_name:
slice(xy_bbox[1] - xy_border, xy_bbox[3] + xy_border)
})
else:
ds = dataset.sel(
**{
x_name:
slice(xy_bbox[0] - xy_border, xy_bbox[2] + xy_border),
y_name:
slice(xy_bbox[3] + xy_border, xy_bbox[1] - xy_border)
})
return ds
else:
return dataset.isel(
**{
x_name: slice(ij_bbox[0] - ij_border, ij_bbox[2] +
ij_border),
y_name: slice(ij_bbox[1] - ij_border, ij_bbox[3] +
ij_border)
})
else:
if xy_bbox:
ij_bbox = grid_mapping.ij_bbox_from_xy_bbox(xy_bbox,
ij_border=ij_border,
xy_border=xy_border)
if ij_bbox[0] == -1:
return None
width, height = grid_mapping.size
i_min, j_min, i_max, j_max = ij_bbox
if i_min > 0 or j_min > 0 or i_max < width - 1 or j_max < height - 1:
x_dim, y_dim = grid_mapping.xy_dim_names
i_slice = slice(i_min, i_max + 1)
j_slice = slice(j_min, j_max + 1)
return dataset.isel({x_dim: i_slice, y_dim: j_slice})
return dataset
def resample_in_space(dataset: xr.Dataset,
source_gm: GridMapping = None,
target_gm: GridMapping = None,
var_configs: Mapping[Hashable, Mapping[str,
Any]] = None):
"""
Resample a dataset in the spatial dimensions.
If the source grid mapping *source_gm* is not given,
it is derived from *dataset*:
``source_gm = GridMapping.from_dataset(dataset)``.
If the target grid mapping *target_gm* is not given,
it is derived from *source_gm*:
``target_gm = source_gm.to_regular()``.
If *source_gm* is almost equal to *target_gm*, this
function is a no-op and *dataset* is returned unchanged.
Otherwise the function computes a spatially
resampled version of *dataset* and returns it.
Using *var_configs*, the resampling of individual
variables can be configured. If given, *var_configs*
must be a mapping from variable names to configuration
dictionaries which can have the following properties:
* ``spline_order`` (int) - The order of spline polynomials
used for interpolating. It is used for upsampling only.
Possible values are 0 to 5.
Default is 1 (bi-linear) for floating point variables,
and 0 (= nearest neighbor) for integer and bool variables.
* ``aggregator`` (str) - An optional aggregating
function. It is used for downsampling only.
Examples are numpy.nanmean, numpy.nanmin, numpy.nanmax.
Default is numpy.nanmean for floating point variables,
and None (= nearest neighbor) for integer and bool variables.
* ``recover_nan`` (bool) - whether a special algorithm
shall be used that is able to recover values that would
otherwise yield NaN during resampling.
Default is True for floating point variables,
and False for integer and bool variables.
Note that *var_configs* is only used if the resampling involves
an affine transformation. This is true if the CRS of
*source_gm* and CRS of *target_gm* are equal and one of two
cases is given:
1. *source_gm* is regular.
In this case the resampling is the affine transformation.
and the result is returned directly.
2. *source_gm* is not regular and has a lower resolution
than *target_cm*.
In this case *dataset* is downsampled first using an affine
transformation. Then the result is rectified.
In all other cases, no affine transformation is applied and
the resampling is a direct rectification.
:param dataset: The source dataset.
:param source_gm: The source grid mapping.
:param target_gm: The target grid mapping. Must be regular.
:param var_configs: Optional resampling configurations
for individual variables.
:return: The spatially resampled dataset.
"""
if source_gm is None:
# No source grid mapping given, so do derive it from dataset
source_gm = GridMapping.from_dataset(dataset)
if target_gm is None:
# No target grid mapping given, so do derive it from source
target_gm = source_gm.to_regular()
if source_gm.is_close(target_gm):
# If source and target grid mappings are almost equal
return dataset
# target_gm must be regular
GridMapping.assert_regular(target_gm, name='target_gm')
# Are source and target both geographic grid mappings?
both_geographic = source_gm.crs.is_geographic \
and target_gm.crs.is_geographic
if both_geographic or source_gm.crs == target_gm.crs:
# If CRSes are both geographic or their CRSes are equal:
if source_gm.is_regular:
# If also the source is regular, then resampling reduces
# to an affine transformation.
return affine_transform_dataset(
dataset,
source_gm=source_gm,
target_gm=target_gm,
var_configs=var_configs,
)
# If the source is not regular, we need to rectify it,
# so the target is regular. Our rectification implementation
# works only correctly if source pixel size >= target pixel
# size. Therefore check if we must downscale source first.
x_scale = source_gm.x_res / target_gm.x_res
y_scale = source_gm.y_res / target_gm.y_res
if x_scale > _SCALE_LIMIT and y_scale > _SCALE_LIMIT:
# Source pixel size >= target pixel size.
# We can rectify.
return rectify_dataset(dataset,
source_gm=source_gm,
target_gm=target_gm)
# Source has higher resolution than target.
# Downscale first, then rectify
if source_gm.is_regular:
# If source is regular
downscaled_gm = source_gm.scale((x_scale, y_scale))
downscaled_dataset = resample_dataset(
dataset,
((x_scale, 1, 0), (1, y_scale, 0)),
size=downscaled_gm.size,
tile_size=source_gm.tile_size,
xy_dim_names=source_gm.xy_dim_names,
var_configs=var_configs,
)
else:
_, downscaled_size = scale_xy_res_and_size(source_gm.xy_res,
source_gm.size,
(x_scale, y_scale))
downscaled_dataset = resample_dataset(
dataset,
((x_scale, 1, 0), (1, y_scale, 0)),
size=downscaled_size,
tile_size=source_gm.tile_size,
xy_dim_names=source_gm.xy_dim_names,
var_configs=var_configs,
)
downscaled_gm = GridMapping.from_dataset(
downscaled_dataset,
tile_size=source_gm.tile_size,
prefer_crs=source_gm.crs)
return rectify_dataset(downscaled_dataset,
source_gm=downscaled_gm,
target_gm=target_gm)
# If CRSes are not both geographic and their CRSes are different
# transform the source_gm so its CRS matches the target CRS:
transformed_source_gm = source_gm.transform(target_gm.crs)
transformed_x, transformed_y = transformed_source_gm.xy_coords
reprojected_dataset = resample_in_space(dataset.assign(
transformed_x=transformed_x, transformed_y=transformed_y),
source_gm=transformed_source_gm,
target_gm=target_gm)
if not target_gm.crs.is_geographic:
# Add 'crs' variable according to CF conventions
reprojected_dataset = reprojected_dataset.assign(
crs=xr.DataArray(0, attrs=target_gm.crs.to_cf()))
return reprojected_dataset
def write_data(self,
data: Union[xr.Dataset, MultiLevelDataset],
data_id: str,
replace: bool = False,
**write_params) -> str:
assert_instance(data, (xr.Dataset, MultiLevelDataset), name='data')
assert_instance(data_id, str, name='data_id')
tile_size = write_params.pop('tile_size', None)
if isinstance(data, MultiLevelDataset):
ml_dataset = data
if tile_size:
warnings.warn('tile_size is ignored for multi-level datasets')
else:
base_dataset: xr.Dataset = data
if tile_size:
assert_instance(tile_size, int, name='tile_size')
gm = GridMapping.from_dataset(base_dataset)
x_name, y_name = gm.xy_dim_names
base_dataset = base_dataset.chunk({
x_name: tile_size,
y_name: tile_size
})
ml_dataset = BaseMultiLevelDataset(base_dataset)
fs, root, write_params = self.load_fs(write_params)
consolidated = write_params.pop('consolidated', True)
use_saved_levels = write_params.pop('use_saved_levels', False)
base_dataset_id = write_params.pop('base_dataset_id', None)
if use_saved_levels:
ml_dataset = BaseMultiLevelDataset(ml_dataset.get_dataset(0),
tile_grid=ml_dataset.tile_grid)
path_class = get_fs_path_class(fs)
data_path = path_class(data_id)
fs.mkdirs(str(data_path), exist_ok=replace)
for index in range(ml_dataset.num_levels):
level_dataset = ml_dataset.get_dataset(index)
if base_dataset_id and index == 0:
# Write file "0.link" instead of copying
# level zero dataset to "0.zarr".
# Compute a relative base dataset path first
base_dataset_path = path_class(root, base_dataset_id)
data_parent_path = data_path.parent
try:
base_dataset_path = base_dataset_path.relative_to(
data_parent_path)
except ValueError as e:
raise DataStoreError(
f'invalid base_dataset_id: {base_dataset_id}') from e
base_dataset_path = '..' / base_dataset_path
# Then write relative base dataset path into link file
link_path = data_path / f'{index}.link'
with fs.open(str(link_path), mode='w') as fp:
fp.write(f'{base_dataset_path}')
else:
# Write level "{index}.zarr"
level_path = data_path / f'{index}.zarr'
zarr_store = fs.get_mapper(str(level_path), create=True)
try:
level_dataset.to_zarr(zarr_store,
mode='w' if replace else None,
consolidated=consolidated,
**write_params)
except ValueError as e:
# TODO: remove already written data!
raise DataStoreError(f'Failed to write'
f' dataset {data_id}: {e}') from e
if use_saved_levels:
level_dataset = xr.open_zarr(zarr_store,
consolidated=consolidated)
ml_dataset.set_dataset(index, level_dataset)
return data_id