Merge pull request #95 from EcoExtreML/fix_94

Setting chunks auto in open_mfdataset

Author: SarahAlidoost

PyStemmusScope/forcing_io.py

@@ -1,5 +1,6 @@
 """Module for forcing data input and output operations."""
 from pathlib import Path
+import dask
 import hdf5storage
 import numpy as np
 import xarray as xr
@@ -114,7 +115,7 @@ def read_forcing_data_global(  # noqa:PLR0913 (too many arguments)
     lon: float,
     start_time: np.datetime64,
     end_time: np.datetime64,
-    timestep: str = "1800S",
+    timestep: str = "1800s",
 ) -> dict:
     """Read forcing data for a certain location, based on global datasets.
 
@@ -130,12 +131,14 @@ def read_forcing_data_global(  # noqa:PLR0913 (too many arguments)
     Returns:
         Dictionary containing the forcing data.
     """
-    return global_data.collect_datasets(
-        global_data_dir=global_data_dir,
-        latlon=(lat, lon),
-        time_range=(start_time, end_time),
-        timestep=timestep,
-    )
+    # see https://docs.dask.org/en/latest/array-slicing.html#efficiency
+    with dask.config.set(**{"array.slicing.split_large_chunks": True}):  # type: ignore
+        return global_data.collect_datasets(
+            global_data_dir=global_data_dir,
+            latlon=(lat, lon),
+            time_range=(start_time, end_time),
+            timestep=timestep,
+        )
 
 
 def write_dat_files(data: dict, input_dir: Path):

PyStemmusScope/global_data/cams_co2.py

@@ -22,7 +22,7 @@ def retrieve_co2_data(
         latlon: Latitude and longitude of the site.
         time_range: Start and end time of the model run.
         timestep: Desired timestep of the model, this is derived from the forcing data.
-            In a pandas-timedelta compatible format. For example: "1800S"
+            In a pandas-timedelta compatible format. For example: "1800s"
 
     Returns:
         DataArray containing the CO2 at the specified site for the given time range.
@@ -55,12 +55,12 @@ def extract_cams_data(
         latlon: Latitude and longitude of the site.
         time_range: Start and end time of the model run.
         timestep: Desired timestep of the model, this is derived from the forcing data.
-            In a pandas-timedelta compatible format. For example: "1800S"
+            In a pandas-timedelta compatible format. For example: "1800s"
 
     Returns:
         DataArray containing the CO2 concentration.
     """
-    ds = xr.open_mfdataset(files_cams)
+    ds = xr.open_mfdataset(files_cams, chunks="auto")
 
     check_cams_dataset(cams_data=ds, latlon=latlon, time_range=time_range)
 

PyStemmusScope/global_data/cci_landcover.py

@@ -24,7 +24,7 @@ def retrieve_landcover_data(
         latlon: Latitude and longitude of the site.
         time_range: Start and end time of the model run.
         timestep: Desired timestep of the model, this is derived from the forcing data.
-            In a pandas-timedelta compatible format. For example: "1800S"
+            In a pandas-timedelta compatible format. For example: "1800s"
 
     Returns:
         Dictionary containing IGBP and LCCS land cover classes.
@@ -57,12 +57,12 @@ def extract_landcover_data(
         latlon: Latitude and longitude of the site.
         time_range: Start and end time of the model run.
         timestep: Desired timestep of the model, this is derived from the forcing data.
-            In a pandas-timedelta compatible format. For example: "1800S"
+            In a pandas-timedelta compatible format. For example: "1800s"
 
     Returns:
         Dictionary containing IGBP and LCCS land cover classes.
     """
-    cci_dataset = xr.open_mfdataset(files_cci)
+    cci_dataset = xr.open_mfdataset(files_cci, chunks="auto")
 
     check_cci_dataset(cci_dataset, latlon, time_range)  # Assert spatial/temporal bounds
 

PyStemmusScope/global_data/copernicus_lai.py

@@ -22,7 +22,7 @@ def retrieve_lai_data(
         latlon: Latitude and longitude of the site.
         time_range: Start and end time of the model run.
         timestep: Desired timestep of the model, this is derived from the forcing data.
-            In a pandas-timedelta compatible format. For example: "1800S"
+            In a pandas-timedelta compatible format. For example: "1800s"
 
     Returns:
         DataArray containing the LAI of the specified site for the given time range.
@@ -55,12 +55,12 @@ def extract_lai_data(
         latlon: Latitude and longitude of the site.
         time_range: Start and end time of the model run.
         timestep: Desired timestep of the model, this is derived from the forcing data.
-            In a pandas-timedelta compatible format. For example: "1800S"
+            In a pandas-timedelta compatible format. For example: "1800s"
 
     Returns:
         DataArray containing the LAI of the specified site for the given time range.
     """
-    ds = xr.open_mfdataset(files_lai)
+    ds = xr.open_mfdataset(files_lai, chunks="auto")
 
     check_lai_dataset(ds, latlon, time_range)
 

PyStemmusScope/global_data/era5.py

@@ -27,7 +27,7 @@ def retrieve_era5_data(
         latlon: Latitude and longitude of the site.
         time_range: Start and end time of the model run.
         timestep: Desired timestep of the model, this is derived from the forcing data.
-            In a pandas-timedelta compatible format. For example: "1800S"
+            In a pandas-timedelta compatible format. For example: "1800s"
 
     Returns:
         Dictionary containing the variables extracted from ERA5.
@@ -68,7 +68,7 @@ def load_era5_data(
         latlon: Latitude and longitude of the site.
         time_range: Start and end time of the model run.
         timestep: Desired timestep of the model, this is derived from the forcing data.
-            In a pandas-timedelta compatible format. For example: "1800S"
+            In a pandas-timedelta compatible format. For example: "1800s"
 
     Returns:
         Dictionary containing the variables extracted from ERA5.
@@ -117,14 +117,15 @@ def get_era5_dataset(
         name: Either "ERA5" or "ERA5-land".
         time_range: Start and end time of the model run.
         timestep: Desired timestep of the model, this is derived from the forcing data.
-            In a pandas-timedelta compatible format. For example: "1800S"
+            In a pandas-timedelta compatible format. For example: "1800s"
 
     Returns:
         The ERA5 or ERA5-land dataset.
     """
     tol = RESOLUTION_ERA5 if name == "ERA5" else RESOLUTION_ERA5LAND
 
-    ds = xr.open_mfdataset(files)
+    ds = xr.open_mfdataset(files, chunks="auto")
+
     check_era5_dataset(ds, name, latlon, time_range)
 
     try:

PyStemmusScope/global_data/global_data_selection.py

@@ -22,7 +22,7 @@ def collect_datasets(
         latlon: Latitude and longitude of the site.
         time_range: Start and end time of the model run.
         timestep: Desired timestep of the model, this is derived from the forcing data.
-            In a pandas-timedelta compatible format. For example: "1800S"
+            In a pandas-timedelta compatible format. For example: "1800s"
 
     Returns:
         Dictionary containing the variables extracted from the global datasets.