feat: abstracting accumulation

docs/howtos/create/05-create-accumulations.rst

@@ -0,0 +1,129 @@
+.. _create-accumulations:
+
+##########################################
+ Create a dataset with accumulated fields
+##########################################
+
+Many fields come as accumulations over time, e.g `tp` (total
+precipitations), `sd` (now depth) or `ssr` (surface shortwave
+radiation). Given one dataset, one may want to accumulate some of its
+fields on specific periods of time.
+
+This depends on the data's native format. For an accumulated field (say
+`tp` for simplicity), one needs to know:
+
+-  the `accumulation_period` over which to accumulate (e.g 6h).
+-  the desired `validity_time` at which accumulation stops and for which
+   the value is valid.
+-  the `data_accumulation_period`, that is the duration over which the
+   data is already accumulated.
+
+The resulting field is then "`tp` accumulated over `accumulation_period`
+hours up to `validity_time`". In a common case, dataset features, e.g.,
+1h-accumulated `tp` at a 1 hour frequency, and each raw file features
+`tp` as accumulated over the *last* hour. So having 6h-accumulated `tp`
+consists in taking all 6 files before (and including) `validity_time`
+and summing fields in them.
+
+The resulting accumulated field can be treated as a normal anemoi
+`source` in recipes (e.g, filters can be applied to the source).
+
+Note that depending on how your native dataset is built (e.g, your
+native files feature the accumulation on the *next* hour), the
+calculation can be very different. See $Subtleties below with the
+associated recipes.
+
+***************************************
+ Using accumulations in recipes : mars
+***************************************
+
+In the example below we see recipes to create accumulations from MARS
+data. To keep older recipes working, there are two equivalent ways to do
+so. The first one is a generic way working for MARS and grib-index
+sources.
+
+.. literalinclude:: yaml/recipe-accumulate-era.yaml
+
+That recipe will generate the following dataset:
+
+.. code:: bash
+
+   📦 Path       : recipe-accumulate.zarr
+   🔢 Format version: 0.30.0
+
+   📅 Start      : 2021-01-10 18:00
+   📅 End        : 2021-01-12 12:00
+   ⏰ Frequency  : 6h
+   🚫 Missing    : 0
+   🌎 Resolution : 20.0
+   🌎 Field shape: [162]
+
+   📐 Shape      : 8 × 2 × 1 × 162 (10.1 KiB)
+   💽 Size       : 23.2 KiB (23.2 KiB)
+   📁 Files      : 52
+
+      Index │ Variable │ Min │       Max │        Mean │      Stdev
+      ──────┼──────────┼─────┼───────────┼─────────────┼───────────
+         0  │ cp       │   0 │ 0.0110734 │ 0.000244731 │ 0.00103593
+         1  │ tp       │   0 │ 0.0333021 │  0.00058075 │ 0.00210331
+      ──────┴──────────┴─────┴───────────┴─────────────┴───────────
+   🔋 Dataset ready, last update 26 seconds ago.
+   📊 Statistics ready.
+
+The "legacy" way to do is the following (syntax is only slightly
+different)
+
+.. literalinclude:: yaml/recipe-accumulation-era.yaml
+
+The resulting dataset is:
+
+.. code:: bash
+
+   📦 Path       : recipe-accumulation.zarr
+   🔢 Format version: 0.30.0
+
+   📅 Start      : 2021-01-10 18:00
+   📅 End        : 2021-01-12 12:00
+   ⏰ Frequency  : 6h
+   🚫 Missing    : 0
+   🌎 Resolution : 20.0
+   🌎 Field shape: [9, 18]
+
+   📐 Shape      : 8 × 2 × 1 × 162 (10.1 KiB)
+   💽 Size       : 22.2 KiB (22.2 KiB)
+   📁 Files      : 52
+
+      Index │ Variable │ Min │       Max │        Mean │      Stdev
+      ──────┼──────────┼─────┼───────────┼─────────────┼───────────
+         0  │ cp       │   0 │ 0.0110734 │ 0.000244739 │ 0.00103593
+         1  │ tp       │   0 │ 0.0333023 │ 0.000580769 │ 0.00210332
+      ──────┴──────────┴─────┴───────────┴─────────────┴───────────
+   🔋 Dataset ready, last update 3 minutes ago.
+   📊 Statistics ready.
+
+Note that statitics for the two datasets are equal up to `1e-6`, this is
+due to rounding errors that can accumulate. Larger discrepancies are a
+sign something might be wrong.
+
+*********************************************
+ Using accumulations in recipes : grib files
+*********************************************
+
+If your data source is grib files, you can use a grib-index as a source.
+First create a `grib-index
+<https://anemoi.readthedocs.io/projects/datasets/en/latest/howtos/create/01-grib-data.html#using-an-index-file>`_
+
+that creates a database to query fields. Then, say we want to accumulate
+3h-data over 6h.
+
+.. literalinclude:: yaml/recipe-accumulate-gribindex.yaml
+
+Note that we also added a filter at the end of the recipe to rename `tp`
+to `tp_6h`. The frequency of the dataset is `1h`, so the accumulation is
+a moving window.
+
+************
+ Subtleties
+************
+
+Some datasets (such as ERA5) feature

docs/howtos/create/yaml/recipe-accumulate-era.yaml

@@ -0,0 +1,17 @@
+dates:
+  start: 2021-01-10 18:00:00
+  end: 2021-01-12 12:00:00
+  frequency: 6h
+
+input:
+    accumulate:
+      source:
+        mars:
+          expver: "0001"
+          class: ea
+          stream: oper
+          grid: 20./20.
+          levtype: sfc
+          param: [ tp, cp]
+          accumulation_period: 6
+          data_accumulation_period: 6 # this argument will be ignored because of ERA-like accumulations

docs/howtos/create/yaml/recipe-accumulate-gribindex.yaml

@@ -0,0 +1,18 @@
+dates:
+  end: 2020-01-06 12:00:00+00:00
+  start: 2020-01-03 06:00:00+00:00
+  frequency: 1h
+
+input:
+  pipe:
+    - accumulate:
+        source:
+          grib-index:
+            indexdb: /path/to/gribindex.db
+            levtype: sfc
+            param:
+            - tp
+            accumulation_period: 6
+            data_accumulation_period: 3
+    - rename:
+        tp: tp_6h

docs/howtos/create/yaml/recipe-accumulation-era.yaml

@@ -0,0 +1,14 @@
+dates:
+  start: 2021-01-10 18:00:00
+  end: 2021-01-12 12:00:00
+  frequency: 6h
+
+input:
+    accumulations:
+      expver: "0001"
+      class: ea
+      stream: oper
+      grid: 20./20.
+      levtype: sfc
+      param: [ tp, cp]
+      accumulation_period: 6

pyproject.toml

@@ -54,6 +54,8 @@ dependencies = [
   "jsonschema",
   "numcodecs<0.16",                   # Until we move to zarr3
   "numpy",
+  "pytest>=8.4.1",
+  "pytest-xdist>=3.7",
   "pyyaml",
   "ruamel-yaml",
   "semantic-version",

src/anemoi/datasets/create/__init__.py

@@ -256,8 +256,7 @@ def check_name(
         resolution: str,
         dates: list[datetime.datetime],
         frequency: datetime.timedelta,
-        raise_exception: bool = True,
-        is_test: bool = False,
+        raise_exception: bool = False,
     ) -> None:
         """Check the name of the dataset.
 
@@ -271,15 +270,13 @@ def check_name(
             The frequency of the dataset.
         raise_exception : bool, optional
             Whether to raise an exception if the name is invalid.
-        is_test : bool, optional
-            Whether this is a test.
         """
         basename, _ = os.path.splitext(os.path.basename(self.path))
         try:
             DatasetName(basename, resolution, dates[0], dates[-1], frequency).raise_if_not_valid()
         except Exception as e:
-            if raise_exception and not is_test:
-                raise e
+            if raise_exception:
+                raise
             else:
                 LOG.warning(f"Dataset name error: {e}")
 
@@ -577,7 +574,6 @@ def __init__(
         use_threads: bool = False,
         statistics_temp_dir: str | None = None,
         progress: Any = None,
-        test: bool = False,
         cache: str | None = None,
         **kwargs: Any,
     ):
@@ -599,8 +595,6 @@ def __init__(
             The directory for temporary statistics.
         progress : Any, optional
             The progress indicator.
-        test : bool, optional
-            Whether this is a test.
         cache : Optional[str], optional
             The cache directory.
         """
@@ -613,9 +607,8 @@ def __init__(
         self.use_threads = use_threads
         self.statistics_temp_dir = statistics_temp_dir
         self.progress = progress
-        self.test = test
 
-        self.main_config = loader_config(config, is_test=test)
+        self.main_config = loader_config(config)
 
         # self.registry.delete() ??
         self.tmp_statistics.delete()
@@ -748,7 +741,6 @@ def _run(self) -> int:
 
         self.dataset.check_name(
             raise_exception=self.check_name,
-            is_test=self.test,
             resolution=resolution,
             dates=dates,
             frequency=frequency,

src/anemoi/datasets/create/config.py

@@ -18,8 +18,6 @@
 from anemoi.utils.config import load_any_dict_format
 from earthkit.data.core.order import normalize_order_by
 
-from anemoi.datasets.dates.groups import Groups
-
 LOG = logging.getLogger(__name__)
 
 
@@ -340,72 +338,20 @@ def _prepare_serialisation(o: Any) -> Any:
     return str(o)
 
 
-def set_to_test_mode(cfg: dict) -> None:
-    """Modifies the configuration to run in test mode.
-
-    Parameters
-    ----------
-    cfg : dict
-        The configuration dictionary.
-    """
-    NUMBER_OF_DATES = 4
-
-    LOG.warning(f"Running in test mode. Changing the list of dates to use only {NUMBER_OF_DATES}.")
-    groups = Groups(**LoadersConfig(cfg).dates)
-
-    dates = groups.provider.values
-    cfg["dates"] = dict(
-        start=dates[0],
-        end=dates[NUMBER_OF_DATES - 1],
-        frequency=groups.provider.frequency,
-        group_by=NUMBER_OF_DATES,
-    )
-
-    num_ensembles = count_ensembles(cfg)
-
-    def set_element_to_test(obj):
-        if isinstance(obj, (list, tuple)):
-            for v in obj:
-                set_element_to_test(v)
-            return
-        if isinstance(obj, (dict, DotDict)):
-            if "grid" in obj and num_ensembles > 1:
-                previous = obj["grid"]
-                obj["grid"] = "20./20."
-                LOG.warning(f"Running in test mode. Setting grid to {obj['grid']} instead of {previous}")
-            if "number" in obj and num_ensembles > 1:
-                if isinstance(obj["number"], (list, tuple)):
-                    previous = obj["number"]
-                    obj["number"] = previous[0:3]
-                    LOG.warning(f"Running in test mode. Setting number to {obj['number']} instead of {previous}")
-            for k, v in obj.items():
-                set_element_to_test(v)
-            if "constants" in obj:
-                constants = obj["constants"]
-                if "param" in constants and isinstance(constants["param"], list):
-                    constants["param"] = ["cos_latitude"]
-
-    set_element_to_test(cfg)
-
-
-def loader_config(config: dict, is_test: bool = False) -> LoadersConfig:
+def loader_config(config: dict) -> LoadersConfig:
     """Loads and validates the configuration for dataset loaders.
 
     Parameters
     ----------
     config : dict
         The configuration dictionary.
-    is_test : bool, optional
-        Whether to run in test mode. Defaults to False.
 
     Returns
     -------
     LoadersConfig
         The validated configuration object.
     """
     config = Config(config)
-    if is_test:
-        set_to_test_mode(config)
     obj = LoadersConfig(config)
 
     # yaml round trip to check that serialisation works as expected
@@ -426,6 +372,9 @@ def loader_config(config: dict, is_test: bool = False) -> LoadersConfig:
             LOG.info(f"Setting env variable {k}={v}")
             os.environ[k] = str(v)
 
+    # Used by pytest only
+    # copy.pop('checks', None)
+
     return copy
 
 

src/anemoi/datasets/create/input/result/field.py

@@ -329,8 +329,7 @@ def get_cube(self) -> Any:
             LOG.debug(f"Sorting done in {seconds_to_human(time.time()-start)}.")
         except ValueError:
             self.explain(ds, order_by, remapping=remapping, patches=patches)
-            # raise ValueError(f"Error in {self}")
-            exit(1)
+            raise ValueError(f"Error in {self}")
 
         if LOG.isEnabledFor(logging.DEBUG):
             LOG.debug("Cube shape: %s", cube)
@@ -497,7 +496,6 @@ def explain(self, ds: Any, *args: Any, remapping: Any, patches: Any) -> None:
         print()
         print("❌" * 40)
         print()
-        exit(1)
 
     def build_coords(self) -> None:
         """Build the coordinates for the result."""