Merge branch 'main' into new-demo

ci/environment.yml

@@ -4,6 +4,7 @@ channels:
   - nodefaults
 dependencies:
   - dask
+  - ipykernel
   - jinja2
   - esmpy>=8.2.0
   - mercantile

ndpyramid/common.py

@@ -16,14 +16,15 @@ def __init__(self, **data) -> None:
 
         super().__init__(**data)
         epsg_codes = {'web-mercator': 'EPSG:3857', 'equidistant-cylindrical': 'EPSG:4326'}
+        # Area extent for pyresample's `create_area_def` is (lower_left_x, lower_left_y, upper_right_x, upper_right_y)
         area_extents = {
             'web-mercator': (
                 -20037508.342789244,
                 -20037508.342789248,
-                20037508.342789248,
                 20037508.342789244,
+                20037508.342789248,
             ),
-            'equidistant-cylindrical': (-180, 180, 90, -90),
+            'equidistant-cylindrical': (-180, 90, 180, -90),
         }
         self._crs = epsg_codes[self.name]
         self._proj = pyproj.Proj(self._crs)

ndpyramid/testing.py

@@ -8,16 +8,18 @@ def _bounds(ds):
     right = ds.x[-1] + (ds.x[-1] - ds.x[-2]) / 2
     top = ds.y[0] - (ds.y[1] - ds.y[0]) / 2
     bottom = ds.y[-1] + (ds.y[-1] - ds.y[-2]) / 2
-    return (left.data, bottom.data, right.data, top.data)
+    return np.array([left.data, bottom.data, right.data, top.data])
 
 
 def verify_xy_bounds(ds, zoom):
     """
     Verifies that the bounds of a chunk conforms to expectations for a WebMercatorQuad.
     """
     tile = mercantile.tile(ds.x[0], ds.y[0], zoom)
-    expected = mercantile.xy_bounds(tile)
-    np.testing.assert_allclose(expected, _bounds(ds))
+    bbox = mercantile.xy_bounds(tile)
+    expected = np.array([bbox.left, bbox.bottom, bbox.right, bbox.top])
+    actual = _bounds(ds)
+    np.testing.assert_allclose(actual, expected)
     return ds
 
 

pyproject.toml

@@ -30,12 +30,20 @@ dependencies = [
     "zarr",
     "pydantic>=1.10",
     "pyproj",
-    "rasterio"
+    "rasterio",
 ]
 
 [project.optional-dependencies]
+complete = [
+    "ndpyramid[dask,jupyter,xesmf]",
+    "xarray[complete]",
+    "mercantile",
+    "cftime",
+    "scipy",
+    "rioxarray"
+]
 dask = [
-    "dask",
+    "dask[complete]",
     "pyresample",
 ]
 jupyter = [
@@ -47,17 +55,13 @@ jupyter = [
 xesmf = ["xesmf"]
 
 test = [
-    "dask",
-    "mercantile",
-    "pooch",
+    "ndpyramid[complete]",
     "pre-commit",
     "pytest-benchmark",
     "pytest-codspeed",
     "pytest-cov",
     "pytest-mypy",
     "pytest",
-    "rioxarray",
-    "scipy",
 ]
 
 

tests/conftest.py

@@ -1,13 +1,33 @@
 import numpy as np
 import pandas as pd
 import pytest
+import rioxarray  # noqa
 import xarray as xr
 
 
 @pytest.fixture
 def temperature():
     ds = xr.tutorial.open_dataset('air_temperature')
-    ds['air'].encoding = {}
+    # Update time coordinate to avoid errors when casting to int32 for zarr-js
+    time = ds.time.astype('datetime64[s]')
+    ds = ds.assign_coords(time=time)
+    # Update lon coordinates to -180 to 180 for compatibility with pyresample
+    # Logic copied from carbonplan/cmip6-downscaling
+    lon = ds['lon'].where(ds['lon'] < 180, ds['lon'] - 360)
+    ds = ds.assign_coords(lon=lon)
+
+    if not (ds['lon'].diff(dim='lon') > 0).all():
+        ds = ds.reindex(lon=np.sort(ds['lon'].data))
+
+    if 'lon_bounds' in ds.variables:
+        lon_b = ds['lon_bounds'].where(ds['lon_bounds'] < 180, ds['lon_bounds'] - 360)
+        ds = ds.assign_coords(lon_bounds=lon_b)
+    # Transpose coordinates for compatibility with pyresample
+    ds = ds.transpose('time', 'lat', 'lon')
+    # Write crs for pyramid_resampled and pyramid_reproject
+    ds = ds.rio.write_crs('EPSG:4326')
+    # Chunk the dataset for testing `pyramid_resample`
+    ds = ds.chunk({'time': 1000, 'lat': 20, 'lon': 20})
     return ds
 
 

tests/test_pyramid_regrid.py

@@ -12,6 +12,8 @@
 @pytest.mark.parametrize('regridder_apply_kws', [None, {'keep_attrs': False}])
 def test_regridded_pyramid(temperature, regridder_apply_kws, benchmark):
     pytest.importorskip('xesmf')
+    # Select subset to speed up tests
+    temperature = temperature.isel(time=slice(0, 5))
     pyramid = benchmark(
         lambda: pyramid_regrid(
             temperature, levels=2, regridder_apply_kws=regridder_apply_kws, other_chunks={'time': 2}
@@ -34,6 +36,8 @@ def test_regridded_pyramid(temperature, regridder_apply_kws, benchmark):
 def test_regridded_pyramid_with_weights(temperature, benchmark):
     pytest.importorskip('xesmf')
     levels = 2
+    # Select subset to speed up tests
+    temperature = temperature.isel(time=slice(0, 5))
     weights_pyramid = generate_weights_pyramid(temperature.isel(time=0), levels)
     pyramid = benchmark(
         lambda: pyramid_regrid(

tests/test_pyramid_resample.py

@@ -15,10 +15,6 @@ def test_resampled_pyramid(temperature, benchmark, resampling):
     pytest.importorskip('pyresample')
     pytest.importorskip('rioxarray')
     levels = 2
-    temperature = temperature.rio.write_crs('EPSG:4326')
-    temperature = temperature.transpose('time', 'lat', 'lon')
-    # import pdb; pdb.set_trace()
-
     pyramid = benchmark(
         lambda: pyramid_resample(
             temperature, levels=levels, x='lon', y='lat', resampling=resampling
@@ -32,12 +28,12 @@ def test_resampled_pyramid(temperature, benchmark, resampling):
     pyramid.to_zarr(MemoryStore())
 
 
+@pytest.mark.xfail(reseason='Need to fix resampling of 2D data (tied to other_chunks issue)')
 @pytest.mark.parametrize('method', ['bilinear', 'nearest', {'air': 'nearest'}])
 def test_resampled_pyramid_2D(temperature, method, benchmark):
     pytest.importorskip('pyresample')
     pytest.importorskip('rioxarray')
     levels = 2
-    temperature = temperature.rio.write_crs('EPSG:4326')
     temperature = temperature.isel(time=0).drop_vars('time')
     pyramid = benchmark(
         lambda: pyramid_resample(temperature, levels=levels, x='lon', y='lat', resampling=method)
@@ -92,24 +88,24 @@ def test_resampled_pyramid_fill(temperature, benchmark):
     """
     pytest.importorskip('pyresample')
     pytest.importorskip('rioxarray')
-    temperature = temperature.rio.write_crs('EPSG:4326')
     pyramid = benchmark(lambda: pyramid_resample(temperature, levels=1, x='lon', y='lat'))
     assert np.isnan(pyramid['0'].air.isel(time=0, x=0, y=0).values)
 
 
-@pytest.mark.xfail(reseason='Differences between rasterio and pyresample to be investigated')
-def test_reprojected_resample_pyramid_values(temperature, benchmark):
+@pytest.mark.parametrize(
+    'method',
+    [
+        pytest.param(
+            'bilinear',
+            marks=pytest.mark.xfail(reason='Need to investigate differences for bilinear'),
+        ),
+        'nearest',
+    ],
+)
+def test_reprojected_resample_pyramid_values(dataset_3d, method, benchmark):
     pytest.importorskip('rioxarray')
     levels = 2
-    temperature = temperature.rio.write_crs('EPSG:4326')
-    temperature = temperature.chunk({'time': 10, 'lat': 10, 'lon': 10})
-    reprojected = benchmark(
-        lambda: pyramid_reproject(temperature, levels=levels, resampling='nearest')
-    )
-    resampled = benchmark(
-        lambda: pyramid_resample(
-            temperature, levels=levels, x='lon', y='lat', resampling='nearest_neighbour'
-        )
-    )
+    reprojected = pyramid_reproject(dataset_3d, levels=levels, resampling=method)
+    resampled = pyramid_resample(dataset_3d, levels=levels, x='x', y='y', resampling=method)
     xr.testing.assert_allclose(reprojected['0'].ds, resampled['0'].ds)
     xr.testing.assert_allclose(reprojected['1'].ds, resampled['1'].ds)

tests/test_pyramids.py

@@ -49,7 +49,6 @@ def sel_coarsen(ds, factor, dims, **kwargs):
 def test_reprojected_pyramid(temperature, benchmark):
     pytest.importorskip('rioxarray')
     levels = 2
-    temperature = temperature.rio.write_crs('EPSG:4326')
     pyramid = benchmark(lambda: pyramid_reproject(temperature, levels=levels))
     verify_bounds(pyramid)
     assert pyramid.ds.attrs['multiscales']
@@ -89,7 +88,7 @@ def test_reprojected_pyramid_4d(dataset_4d, benchmark):
     pytest.importorskip('rioxarray')
     levels = 2
     with pytest.raises(Exception):
-        pyramid = benchmark(lambda: pyramid_reproject(dataset_4d, levels=levels))
+        pyramid = pyramid_reproject(dataset_4d, levels=levels)
     pyramid = benchmark(lambda: pyramid_reproject(dataset_4d, levels=levels, extra_dim='band'))
     verify_bounds(pyramid)
     assert pyramid.ds.attrs['multiscales']
@@ -104,6 +103,5 @@ def test_reprojected_pyramid_fill(temperature, benchmark):
     Test for https://github.com/carbonplan/ndpyramid/issues/93.
     """
     pytest.importorskip('rioxarray')
-    temperature = temperature.rio.write_crs('EPSG:4326')
     pyramid = benchmark(lambda: pyramid_reproject(temperature, levels=1))
     assert np.isnan(pyramid['0'].air.isel(time=0, x=0, y=0).values)