Auto-chunking for ERA5 loading

coiled · Nov 15, 2024 · f69fdb5 · f69fdb5
1 parent 4c29fdf
commit f69fdb5
Show file tree

Hide file tree

Showing 4 changed files with 5 additions and 1 deletion.
diff --git a/tests/geospatial/workloads/atmospheric_circulation.py b/tests/geospatial/workloads/atmospheric_circulation.py
@@ -11,7 +11,7 @@ def atmospheric_circulation(
 ) -> Delayed:
     ds = xr.open_zarr(
         "gs://weatherbench2/datasets/era5/1959-2023_01_10-full_37-1h-0p25deg-chunk-1.zarr",
-        chunks={},
+        chunks={"time": "auto"},
     )
     if scale == "small":
         # 852.56 GiB (small)

diff --git a/tests/geospatial/workloads/climatology.py b/tests/geospatial/workloads/climatology.py
@@ -72,6 +72,7 @@ def rechunk_map_blocks(
     # Load dataset
     ds = xr.open_zarr(
         "gs://weatherbench2/datasets/era5/1959-2023_01_10-wb13-6h-1440x721.zarr",
+        chunks={"time": "auto"},
     )
 
     if scale == "small":
@@ -122,6 +123,7 @@ def highlevel_api(
     # Load dataset
     ds = xr.open_zarr(
         "gs://weatherbench2/datasets/era5/1959-2023_01_10-wb13-6h-1440x721.zarr",
+        chunks={"time": "auto"},
     )
 
     if scale == "small":

diff --git a/tests/geospatial/workloads/rechunking.py b/tests/geospatial/workloads/rechunking.py
@@ -11,6 +11,7 @@ def era5_rechunking(
 ) -> Delayed:
     ds = xr.open_zarr(
         "gs://weatherbench2/datasets/era5/1959-2023_01_10-full_37-1h-0p25deg-chunk-1.zarr",
+        chunks={"time": "auto"},
     ).drop_encoding()
 
     if scale == "small":

diff --git a/tests/geospatial/workloads/regridding.py b/tests/geospatial/workloads/regridding.py
@@ -13,6 +13,7 @@ def xesmf(
 ) -> Delayed:
     ds = xr.open_zarr(
         "gs://weatherbench2/datasets/era5/1959-2023_01_10-full_37-1h-0p25deg-chunk-1.zarr",
+        chunks={"time": "auto"},
     )
     # Fixed time range and variable as the interesting part of this benchmark scales with the
     # regridding matrix