anemoi-datasets 0.5.27__py3-none-any.whl → 0.5.29__py3-none-any.whl

anemoi/datasets/create/__init__.py

@@ -256,8 +256,7 @@ class Dataset:
         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 @@ class Dataset:
             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 @@ class Init(Actor, HasRegistryMixin, HasStatisticTempMixin, HasElementForDataMixi
         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 @@ class Init(Actor, HasRegistryMixin, HasStatisticTempMixin, HasElementForDataMixi
             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 @@ class Init(Actor, HasRegistryMixin, HasStatisticTempMixin, HasElementForDataMixi
         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 @@ class Init(Actor, HasRegistryMixin, HasStatisticTempMixin, HasElementForDataMixi
 
         self.dataset.check_name(
             raise_exception=self.check_name,
-            is_test=self.test,
             resolution=resolution,
             dates=dates,
             frequency=frequency,
@@ -865,7 +857,7 @@ class Load(Actor, HasRegistryMixin, HasStatisticTempMixin, HasElementForDataMixi
             # assert isinstance(group[0], datetime.datetime), type(group[0])
             LOG.debug(f"Building data for group {igroup}/{self.n_groups}")
 
-            result = self.input.select(group_of_dates=group)
+            result = self.input.select(argument=group)
             assert result.group_of_dates == group, (len(result.group_of_dates), len(group), group)
 
             # There are several groups.
@@ -1616,3 +1608,44 @@ def creator_factory(name: str, trace: str | None = None, **kwargs: Any) -> Any:
     )[name]
     LOG.debug(f"Creating {cls.__name__} with {kwargs}")
     return cls(**kwargs)
+
+
+def validate_config(config: Any) -> None:
+
+    import json
+
+    import jsonschema
+
+    def _tidy(d):
+        if isinstance(d, dict):
+            return {k: _tidy(v) for k, v in d.items()}
+
+        if isinstance(d, list):
+            return [_tidy(v) for v in d if v is not None]
+
+        # jsonschema does not support datetime.date
+        if isinstance(d, datetime.datetime):
+            return d.isoformat()
+
+        if isinstance(d, datetime.date):
+            return d.isoformat()
+
+        return d
+
+    # https://json-schema.org
+
+    with open(
+        os.path.join(
+            os.path.dirname(os.path.dirname(__file__)),
+            "schemas",
+            "recipe.json",
+        )
+    ) as f:
+        schema = json.load(f)
+
+    try:
+        jsonschema.validate(instance=_tidy(config), schema=schema)
+    except jsonschema.exceptions.ValidationError as e:
+        LOG.error("❌ Config validation failed (jsonschema):")
+        LOG.error(e.message)
+        raise

anemoi/datasets/create/config.py

@@ -18,8 +18,6 @@ from anemoi.utils.config import DotDict
 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__)
 
 
@@ -279,6 +277,8 @@ class LoadersConfig(Config):
 
         self.output.order_by = normalize_order_by(self.output.order_by)
 
+        self.setdefault("dates", Config())
+
         self.dates["group_by"] = self.build.group_by
 
         ###########
@@ -338,61 +338,13 @@ 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,
-    )
-
-    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:
-                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:
-                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
     -------
@@ -400,8 +352,6 @@ def loader_config(config: dict, is_test: bool = False) -> 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
@@ -422,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
 
 
@@ -441,3 +394,49 @@ def build_output(*args, **kwargs) -> OutputSpecs:
         The output specifications object.
     """
     return OutputSpecs(*args, **kwargs)
+
+
+def flatten_list_of_sets(list_of_sets: list[set]) -> set:
+    return {element for subset in list_of_sets for element in subset}
+
+
+def mars_str_to_set(s: str) -> set[str]:
+    """Mars strings are like 1/to/2 or 1/to/2/by/1
+
+    Returns a set of strings, e.g. {'1', '2'}
+    """
+    assert "/" in s, "mars_str_to_set expects a string with '/'"
+    lst = s.split("/")
+    assert len(lst) in (3, 5), f"mars_str_to_set expects a string like 1/to/2 or 1/to/4/by/1, got {s}"
+    if len(lst) == 3:
+        assert "to" in lst
+        start, _, end = lst
+        step = 1
+    elif len(lst) == 5:
+        assert "by" in lst and "to" in lst
+        start, _, end, _, step = lst
+    return {str(i) for i in range(int(start), int(end) + 1, int(step))}
+
+
+def get_ensembles_set(obj):
+    """Counts the number of ensembles in the configuration."""
+    if isinstance(obj, dict):
+        if "number" in obj:
+            if isinstance(obj["number"], (list, tuple)):
+                return set([str(element) for element in obj["number"]])
+            if isinstance(obj["number"], (str, int)):
+                if "/" in str(obj["number"]):
+                    return mars_str_to_set(obj["number"])
+                else:
+                    return {str(obj["number"])}
+    if isinstance(obj, (dict)):
+        return flatten_list_of_sets([get_ensembles_set(v) for v in obj.values()])
+    if isinstance(obj, (list, tuple)):
+        return flatten_list_of_sets([get_ensembles_set(v) for v in obj])
+    return {}
+
+
+def count_ensembles(config: Config) -> int:
+    """Counts the number of ensembles in the configuration."""
+    ensembles = get_ensembles_set(config.input)
+    return len(ensembles) if ensembles else 1

anemoi/datasets/create/input/__init__.py

@@ -1,4 +1,4 @@
-# (C) Copyright 2024 Anemoi contributors.
+# (C) Copyright 2024-2025 Anemoi contributors.
 #
 # This software is licensed under the terms of the Apache Licence Version 2.0
 # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
@@ -7,21 +7,15 @@
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 
-import logging
 from copy import deepcopy
+from functools import cached_property
+from typing import TYPE_CHECKING
 from typing import Any
 
-from anemoi.datasets.dates.groups import GroupOfDates
+from anemoi.datasets.create.input.context.field import FieldContext
 
-from .trace import trace_select
-
-LOG = logging.getLogger(__name__)
-
-
-class Context:
-    """Context for building input data."""
-
-    pass
+if TYPE_CHECKING:
+    from anemoi.datasets.create.input.action import Recipe
 
 
 class InputBuilder:
@@ -34,72 +28,58 @@ class InputBuilder:
         ----------
         config : dict
             Configuration dictionary.
-        data_sources : Union[dict, list]
+        data_sources : dict
             Data sources.
         **kwargs : Any
             Additional keyword arguments.
         """
         self.kwargs = kwargs
+        self.config = deepcopy(config)
+        self.data_sources = deepcopy(dict(data_sources=data_sources))
 
-        config = deepcopy(config)
-        if data_sources:
-            config = dict(
-                data_sources=dict(
-                    sources=data_sources,
-                    input=config,
-                )
-            )
-        self.config = config
-        self.action_path = ["input"]
-
-    @trace_select
-    def select(self, group_of_dates: GroupOfDates) -> Any:
-        """Select data based on the group of dates.
-
-        Parameters
-        ----------
-        group_of_dates : GroupOfDates
-            Group of dates to select data for.
-
-        Returns
-        -------
-        Any
-            Selected data.
-        """
-        from .action import ActionContext
+    @cached_property
+    def action(self) -> "Recipe":
+        """Returns the action object based on the configuration."""
+        from .action import Recipe
         from .action import action_factory
 
-        """This changes the context."""
-        context = ActionContext(**self.kwargs)
-        action = action_factory(self.config, context, self.action_path)
-        return action.select(group_of_dates)
-
-    def __repr__(self) -> str:
-        """Return a string representation of the InputBuilder.
-
-        Returns
-        -------
-        str
-            String representation.
-        """
-        from .action import ActionContext
-        from .action import action_factory
+        sources = action_factory(self.data_sources, "data_sources")
+        input = action_factory(self.config, "input")
 
-        context = ActionContext(**self.kwargs)
-        a = action_factory(self.config, context, self.action_path)
-        return repr(a)
+        return Recipe(input, sources)
 
-    def _trace_select(self, group_of_dates: GroupOfDates) -> str:
-        """Trace the select operation.
+    def select(self, argument) -> Any:
+        """Select data based on the group of dates.
 
         Parameters
         ----------
-        group_of_dates : GroupOfDates
+        argument : GroupOfDates
             Group of dates to select data for.
 
         Returns
         -------
-        str
-            Trace string.
+        Any
+            Selected data.
         """
-        return f"InputBuilder({group_of_dates})"
+        context = FieldContext(argument, **self.kwargs)
+        return context.create_result(self.action(context, argument))
+
+
+def build_input(config: dict, data_sources: dict | list, **kwargs: Any) -> InputBuilder:
+    """Build an InputBuilder instance.
+
+    Parameters
+    ----------
+    config : dict
+        Configuration dictionary.
+    data_sources : Union[dict, list]
+        Data sources.
+    **kwargs : Any
+        Additional keyword arguments.
+
+    Returns
+    -------
+    InputBuilder
+        An instance of InputBuilder.
+    """
+    return InputBuilder(config, data_sources, **kwargs)