Add regridding benchmark

AB_environments/AB_baseline.conda.yaml

@@ -19,7 +19,7 @@ dependencies:
   - fsspec ==2024.6.1
   - s3fs ==2024.6.1
   - gcsfs ==2024.6.1
-  - pyarrow ==17.0.0
+  - pyarrow ==15.0.0
   - jupyterlab ==4.2.4
   - lz4 ==4.3.3
   - ipywidgets ==8.1.3
@@ -45,6 +45,9 @@ dependencies:
   - bokeh ==3.5.1
   - gilknocker ==0.4.1
   - openssl >1.1.0g
+  - rioxarray ==0.17.0
+  - h5netcdf ==1.3.0
+  - xesmf ==0.8.7
   # End copy-paste
 
   - pip:

AB_environments/AB_sample.conda.yaml

@@ -25,7 +25,7 @@ dependencies:
   - fsspec ==2024.6.1
   - s3fs ==2024.6.1
   - gcsfs ==2024.6.1
-  - pyarrow ==17.0.0
+  - pyarrow ==15.0.0
   - jupyterlab ==4.2.4
   - lz4 ==4.3.3
   - ipywidgets ==8.1.3
@@ -51,6 +51,9 @@ dependencies:
   - bokeh ==3.5.1
   - gilknocker ==0.4.1
   - openssl >1.1.0g
+  - rioxarray ==0.17.0
+  - h5netcdf ==1.3.0
+  - xesmf ==0.8.7
   # End copy-paste
 
   - pip:

ci/environment.yml

@@ -49,6 +49,7 @@ dependencies:
   - openssl >1.1.0g
   - rioxarray ==0.17.0
   - h5netcdf ==1.3.0
+  - xesmf ==0.8.7
 
 ########################################################
 # PLEASE READ:

tests/geospatial/test_regrid.py

@@ -0,0 +1,67 @@
+import numpy as np
+import pytest
+import xarray as xr
+import xesmf as xe
+from coiled.credentials.google import CoiledShippedCredentials
+
+
+@pytest.mark.parametrize("output_resolution", [1.5, 0.1])
+def test_xesmf(
+    gcs_url,
+    scale,
+    client_factory,
+    output_resolution,
+    cluster_kwargs={
+        "workspace": "dask-engineering-gcp",
+        "region": "us-central1",
+        "wait_for_workers": True,
+    },
+    scale_kwargs={
+        "small": {"n_workers": 10},
+        "medium": {"n_workers": 100},
+        "large": {"n_workers": 100},
+    },
+):
+    with client_factory(
+        **scale_kwargs[scale], **cluster_kwargs
+    ) as client:  # noqa: F841
+        # Load dataset
+        ds = xr.open_zarr(
+            "gs://weatherbench2/datasets/era5/1959-2023_01_10-wb13-6h-1440x721.zarr",
+        )
+
+        if scale == "small":
+            # 101.83 GiB (small)
+            time_range = slice("2020-01-01", "2022-12-31")
+            variables = ["sea_surface_temperature"]
+        elif scale == "medium":
+            # 2.12 TiB (medium)
+            time_range = slice("1959-01-01", "2022-12-31")
+            variables = ["sea_surface_temperature"]
+        else:
+            # 4.24 TiB (large)
+            # This currently doesn't complete successfully.
+            time_range = slice("1959-01-01", "2022-12-31")
+            variables = ["sea_surface_temperature", "snow_depth"]
+        ds = ds[variables].sel(time=time_range)
+
+        # 240x121
+        out_grid = xr.Dataset(
+            {
+                "latitude": (
+                    ["latitude"],
+                    np.arange(90, -90 - output_resolution, -output_resolution),
+                    {"units": "degrees_north"},
+                ),
+                "longitude": (
+                    ["longitude"],
+                    np.arange(0, 360, output_resolution),
+                    {"units": "degrees_east"},
+                ),
+            }
+        )
+        regridder = xe.Regridder(ds, out_grid, "bilinear", periodic=True)
+        regridded = regridder(ds, keep_attrs=True)
+
+        result = regridded.chunk(time="auto")
+        result.to_zarr(gcs_url, storage_options={"token": CoiledShippedCredentials()})