Ensure single chunk for all datasets

Author: treigerm

src/climatebenchpress/data_loader/datasets/cams.py

@@ -37,7 +37,7 @@ def open(download_path: Path) -> xr.Dataset:
 
         # Restrict data to a single day.
         # The specific day is arbitrary.
-        ds = ds.sel(valid_time=slice("2023-06-15", "2023-06-15"))
+        ds = ds.sel(valid_time=slice("2023-06-15", "2023-06-15")).chunk(-1)
         # Needed to make the dataset CF-compliant.
         ds.longitude.attrs["axis"] = "X"
         ds.latitude.attrs["axis"] = "Y"

src/climatebenchpress/data_loader/datasets/era5.py

@@ -34,7 +34,7 @@ def download(download_path: Path, progress: bool = True):
                 "10m_u_component_of_wind",
                 "10m_v_component_of_wind",
             ]
-        ]
+        ].chunk(-1)
         # Needed to make the dataset CF-compliant.
         ds.time.attrs["standard_name"] = "time"
         ds.longitude.attrs["axis"] = "X"

src/climatebenchpress/data_loader/datasets/esa_biomass_cci.py

@@ -6,7 +6,6 @@
 import xarray as xr
 
 from .. import (
-    _rechunk_dataset,
     open_downloaded_canonicalized_dataset,
     open_downloaded_tiny_canonicalized_dataset,
 )
@@ -51,8 +50,7 @@ def open(download_path: Path) -> xr.Dataset:
         ds = ds.sel(
             lon=slice(FRANCE_BBOX[0], FRANCE_BBOX[2]),
             lat=slice(FRANCE_BBOX[3], FRANCE_BBOX[1]),
-        )
-        ds = _rechunk_dataset(ds)
+        ).chunk(-1)
         return ds[["agb"]]
 
 

src/climatebenchpress/data_loader/datasets/nextgems.py

@@ -56,7 +56,7 @@ def download(download_path: Path, progress: bool = True):
         idx = _get_nn_lon_lat_index(
             2**ZOOM, np.linspace(-180, 180, NUM_LON), np.linspace(-90, 90, NUM_LAT)
         )
-        ds = ds.isel(cell=idx).chunk({"time": 1, "lat": NUM_LAT, "lon": NUM_LON})
+        ds = ds.isel(cell=idx).chunk(-1)
         ds.lon.attrs["axis"] = "X"
         ds.lat.attrs["axis"] = "Y"