Support flexible timestamp recipe (#38)

* flexible timestamp recipe * fix tests * update docs * add note for timestamp format in docs * check time format and give nice error message * Update src/zampy/recipe.py Co-authored-by: Bart Schilperoort <[email protected]> * add test for time check function --------- Co-authored-by: Bart Schilperoort <[email protected]>
EcoExtreML · Oct 30, 2023 · 26beebb · 26beebb
1 parent fa72728
commit 26beebb
Show file tree

Hide file tree

Showing 4 changed files with 32 additions and 8 deletions.
diff --git a/docs/using_zampy.md b/docs/using_zampy.md
@@ -33,7 +33,7 @@ A "recipe" is a file with `yml` extension and has the following structure:
 name: "test_recipe"
 
 download:
-  years: [2020, 2020]
+  time: [2020-01-01, 2020-01-31] # must follow the numpy.datetime64 format.
   bbox: [54, 6, 50, 3] # NESW
 
   datasets:

diff --git a/src/zampy/recipe.py b/src/zampy/recipe.py
@@ -65,10 +65,10 @@ def __init__(self, recipe_path: Path) -> None:
         # Load & parse recipe
         recipe = recipe_loader(recipe_path)
 
-        self.start_year, self.end_year = recipe["download"]["years"]
+        self.start_time, self.end_time = recipe["download"]["time"]
         self.timebounds = TimeBounds(
-            np.datetime64(f"{self.start_year}-01-01T00:00"),
-            np.datetime64(f"{self.end_year}-12-31T23:59"),
+            convert_time(f"{self.start_time}"),
+            convert_time(f"{self.end_time}"),
         )
         self.spatialbounds = SpatialBounds(*recipe["download"]["bbox"])
 
@@ -123,12 +123,28 @@ def run(self) -> None:
 
             comp = dict(zlib=True, complevel=5)
             encoding = {var: comp for var in ds.data_vars}
-            fname = (  # e.g. "era5_2010-2020.nc"
-                f"{dataset_name.lower()}_{self.start_year}-{self.end_year}.nc"
-            )
+            time_start = str(self.timebounds.start.astype("datetime64[Y]"))
+            time_end = str(self.timebounds.end.astype("datetime64[Y]"))
+            # e.g. "era5_2010-2020.nc"
+            fname = f"{dataset_name.lower()}_{time_start}-{time_end}.nc"
             ds.to_netcdf(path=self.data_dir / fname, encoding=encoding)
 
         print(
             "Finished running the recipe. Output data can be found at:\n"
             f"    {self.data_dir}"
         )
+
+
+def convert_time(time: str) -> np.datetime64:
+    """Check input time and convert to np.datetime64."""
+    try:
+        timestamp = np.datetime64(time)
+    except ValueError as err:
+        msg = (
+            "The input format of timestamp in the recipe is incorrect. \n Please"
+            " follow the format of `numpy.datetime64` and update the input time,"
+            " e.g. 'YYYY-MM-DD'."
+        )
+        raise ValueError(msg) from err
+
+    return timestamp
diff --git a/tests/test_recipes/recipes/era5_recipe.yml b/tests/test_recipes/recipes/era5_recipe.yml
@@ -1,7 +1,7 @@
 name: "era5_recipe"
 
 download:
-  years: [2020, 2020]
+  time: [2020-01-01, 2020-12-31]
   bbox: [51, 4, 50, 3] # NESW
 
   datasets:

diff --git a/tests/test_recipes/test_simple_recipe.py b/tests/test_recipes/test_simple_recipe.py
@@ -3,12 +3,14 @@
 from unittest.mock import patch
 import generate_test_data
 import numpy as np
+import pytest
 import xarray as xr
 from zampy.datasets import DATASETS
 from zampy.datasets.dataset_protocol import SpatialBounds
 from zampy.datasets.dataset_protocol import TimeBounds
 from zampy.datasets.dataset_protocol import write_properties_file
 from zampy.recipe import RecipeManager
+from zampy.recipe import convert_time
 
 
 RECIPE_FILE = Path(__file__).parent / "recipes" / "era5_recipe.yml"
@@ -42,3 +44,9 @@ def test_recipe(tmp_path: Path, mocker):
 
         ds = xr.open_mfdataset(str(tmp_path / "output" / "era5_recipe" / "*.nc"))
         assert all(var in ds.data_vars for var in ["Psurf", "Wind_N"])
+
+
+def test_invalid_time_format():
+    time_from_recipe = "2020-1-01"
+    with pytest.raises(ValueError, match="The input format of timestamp"):
+        convert_time(time_from_recipe)