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

anemoi/datasets/create/sources/anemoi_dataset.py

@@ -9,65 +9,69 @@
 
 import numpy as np
 
-from .legacy import legacy_source
+from . import source_registry
+from .legacy import LegacySource
 
 
-@legacy_source(__file__)
-def execute(context, dates, params=None, **kwargs):
-    import earthkit.data as ekd
+@source_registry.register("anemoi_dataset")
+class AnemoiDatasetSource(LegacySource):
 
-    from anemoi.datasets import open_dataset
+    @staticmethod
+    def _execute(context, dates, params=None, **kwargs):
+        import earthkit.data as ekd
 
-    ds = open_dataset(**kwargs)
-    # dates_to_index = {date: i for i, date in enumerate(ds.dates)}
+        from anemoi.datasets import open_dataset
 
-    indices = []
-    for date in dates:
-        idx = np.where(ds.dates == date)[0]
-        if len(idx) == 0:
-            continue
-        indices.append((int(idx[0]), date))
+        ds = open_dataset(**kwargs)
+        # dates_to_index = {date: i for i, date in enumerate(ds.dates)}
 
-    vars = ds.variables
-    if params is None:
-        params = vars
+        indices = []
+        for date in dates:
+            idx = np.where(ds.dates == date)[0]
+            if len(idx) == 0:
+                continue
+            indices.append((int(idx[0]), date))
 
-    if not isinstance(params, (list, tuple, set)):
-        params = [params]
+        vars = ds.variables
+        if params is None:
+            params = vars
 
-    params = set(params)
-    results = []
+        if not isinstance(params, (list, tuple, set)):
+            params = [params]
 
-    ensemble = ds.shape[2] > 1
-    latitudes = ds.latitudes
-    longitudes = ds.longitudes
+        params = set(params)
+        results = []
 
-    for idx, date in indices:
+        ensemble = ds.shape[2] > 1
+        latitudes = ds.latitudes
+        longitudes = ds.longitudes
 
-        metadata = dict(valid_datetime=date, latitudes=latitudes, longitudes=longitudes)
+        for idx, date in indices:
 
-        for j, y in enumerate(ds[idx]):
+            metadata = dict(valid_datetime=date, latitudes=latitudes, longitudes=longitudes)
 
-            param = vars[j]
-            if param not in params:
-                continue
+            for j, y in enumerate(ds[idx]):
+
+                param = vars[j]
+                if param not in params:
+                    continue
 
-            # metadata['name'] = param
-            # metadata['param_level'] = param
-            metadata["param"] = param
+                # metadata['name'] = param
+                # metadata['param_level'] = param
+                metadata["param"] = param
 
-            for k, e in enumerate(y):
-                if ensemble:
-                    metadata["number"] = k + 1
+                for k, e in enumerate(y):
+                    if ensemble:
+                        metadata["number"] = k + 1
 
-                metadata["values"] = e
+                    metadata["values"] = e
 
-                results.append(metadata.copy())
+                    results.append(metadata.copy())
 
-    print(results[0].keys())
+        print(results[0].keys())
 
-    # "list-of-dicts" does support resolution
-    results = ekd.from_source("list-of-dicts", results)
+        # "list-of-dicts" does support resolution
+        results = ekd.from_source("list-of-dicts", results)
 
-    # return new_fieldlist_from_list([new_field_from_latitudes_longitudes(x, latitudes, longitudes) for x in results])
-    return results
+        # return new_fieldlist_from_list([new_field_from_latitudes_longitudes(x, latitudes, longitudes) for x in results])
+        return results

anemoi/datasets/create/sources/constants.py

@@ -11,41 +11,42 @@ from typing import Any
 
 from earthkit.data import from_source
 
-from .legacy import legacy_source
-
-
-@legacy_source(__file__)
-def constants(context: Any, dates: list[str], template: dict[str, Any], param: str) -> Any:
-    """Deprecated function to retrieve constants data.
-
-    Parameters
-    ----------
-    context : Any
-        The context object for tracing.
-    dates : list of str
-        List of dates for which data is required.
-    template : dict of str to Any
-        Template dictionary for the data source.
-    param : str
-        Parameter to retrieve.
-
-    Returns
-    -------
-    Any
-        Data retrieved from the source.
-    """
-    from warnings import warn
-
-    warn(
-        "The source `constants` is deprecated, use `forcings` instead.",
-        DeprecationWarning,
-        stacklevel=2,
-    )
-    context.trace("✅", f"from_source(constants, {template}, {param}")
-    if len(template) == 0:
-        raise ValueError("Forcings template is empty.")
-
-    return from_source("forcings", source_or_dataset=template, date=dates, param=param)
-
-
-execute: Any = constants
+from . import source_registry
+from .legacy import LegacySource
+
+
+@source_registry.register("constants")
+class ConstantsSource(LegacySource):
+
+    @staticmethod
+    def _execute(context: Any, dates: list[str], template: dict[str, Any], param: str) -> Any:
+        """Deprecated function to retrieve constants data.
+
+        Parameters
+        ----------
+        context : Any
+            The context object for tracing.
+        dates : list of str
+            List of dates for which data is required.
+        template : dict of str to Any
+            Template dictionary for the data source.
+        param : str
+            Parameter to retrieve.
+
+        Returns
+        -------
+        Any
+            Data retrieved from the source.
+        """
+        from warnings import warn
+
+        warn(
+            "The source `constants` is deprecated, use `forcings` instead.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        context.trace("✅", f"from_source(constants, {template}, {param}")
+        if len(template) == 0:
+            raise ValueError("Forcings template is empty.")
+
+        return from_source("forcings", source_or_dataset=template, date=list(dates), param=param)

anemoi/datasets/create/sources/empty.py

@@ -12,25 +12,29 @@ from typing import Any
 
 import earthkit.data as ekd
 
-from .legacy import legacy_source
-
-
-@legacy_source(__file__)
-def execute(context: Any, dates: list[str], **kwargs: Any) -> ekd.FieldList:
-    """Executes the loading of an empty data source.
-
-    Parameters
-    ----------
-    context : object
-        The context in which the function is executed.
-    dates : list
-        List of dates for which data is to be loaded.
-    **kwargs : dict
-        Additional keyword arguments.
-
-    Returns
-    -------
-    ekd.FieldList
-        Loaded empty data source.
-    """
-    return ekd.from_source("empty")
+from . import source_registry
+from .legacy import LegacySource
+
+
+@source_registry.register("empty")
+class EmptySource(LegacySource):
+
+    @staticmethod
+    def _execute(context: Any, dates: list[str], **kwargs: Any) -> ekd.FieldList:
+        """Executes the loading of an empty data source.
+
+        Parameters
+        ----------
+        context : object
+            The context in which the function is executed.
+        dates : list
+            List of dates for which data is to be loaded.
+        **kwargs : dict
+            Additional keyword arguments.
+
+        Returns
+        -------
+        ekd.FieldList
+            Loaded empty data source.
+        """
+        return ekd.from_source("empty")

anemoi/datasets/create/sources/forcings.py

@@ -11,31 +11,32 @@ from typing import Any
 
 from earthkit.data import from_source
 
-from .legacy import legacy_source
-
-
-@legacy_source(__file__)
-def forcings(context: Any, dates: list[str], template: str, param: str) -> Any:
-    """Loads forcing data from a specified source.
-
-    Parameters
-    ----------
-    context : object
-        The context in which the function is executed.
-    dates : list
-        List of dates for which data is to be loaded.
-    template : FieldList
-        Template for the data source.
-    param : str
-        Parameter for the data source.
-
-    Returns
-    -------
-    object
-        Loaded forcing data.
-    """
-    context.trace("✅", f"from_source(forcings, {template}, {param}")
-    return from_source("forcings", source_or_dataset=template, date=dates, param=param)
-
-
-execute = forcings
+from . import source_registry
+from .legacy import LegacySource
+
+
+@source_registry.register("forcings")
+class ForcingsSource(LegacySource):
+
+    @staticmethod
+    def _execute(context: Any, dates: list[str], template: str, param: str) -> Any:
+        """Loads forcing data from a specified source.
+
+        Parameters
+        ----------
+        context : object
+            The context in which the function is executed.
+        dates : list
+            List of dates for which data is to be loaded.
+        template : FieldList
+            Template for the data source.
+        param : str
+            Parameter for the data source.
+
+        Returns
+        -------
+        object
+            Loaded forcing data.
+        """
+        context.trace("✅", f"from_source(forcings, {template}, {param}")
+        return from_source("forcings", source_or_dataset=template, date=list(dates), param=param)

anemoi/datasets/create/sources/grib.py

@@ -20,7 +20,8 @@ from anemoi.transform.grids import grid_registry
 from earthkit.data import from_source
 from earthkit.data.utils.patterns import Pattern
 
-from .legacy import legacy_source
+from . import source_registry
+from .legacy import LegacySource
 
 LOG = logging.getLogger(__name__)
 
@@ -47,6 +48,14 @@ def check(ds: Any, paths: list[str], **kwargs: Any) -> None:
         if isinstance(v, (tuple, list)):
             count *= len(v)
 
+    # in the case of static data (e.g repeated dates) dates might be empty
+    if len(ds) != count and kwargs.get("dates", []) == []:
+        LOG.warning(
+            f"Expected {count} fields, got {len(ds)} (kwargs={kwargs}, paths={paths})"
+            f" Received empty dates - assuming this is static data."
+        )
+        return
+
     if len(ds) != count:
         raise ValueError(f"Expected {count} fields, got {len(ds)} (kwargs={kwargs}, paths={paths})")
 
@@ -73,74 +82,85 @@ def _expand(paths: list[str]) -> Any:
             yield path
 
 
-@legacy_source(__file__)
-def execute(
-    context: Any,
-    dates: list[Any],
-    path: str | list[str],
-    flavour: str | dict[str, Any] | None = None,
-    grid_definition: dict[str, Any] | None = None,
-    *args: Any,
-    **kwargs: Any,
-) -> ekd.FieldList:
-    """Executes the function to load data from GRIB files.
-
-    Parameters
-    ----------
-    context : Any
-        The context in which the function is executed.
-    dates : list of Any
-        List of dates.
-    path : str or list of str
-        Path or list of paths to the GRIB files.
-    flavour : str or dict of str to Any, optional
-        Flavour information, by default None.
-    grid_definition : dict of str to Any, optional
-        Grid definition configuration to create a Grid object, by default None.
-    *args : Any
-        Additional positional arguments.
-    **kwargs : Any
-        Additional keyword arguments.
-
-    Returns
-    -------
-    Any
-        The loaded dataset.
-    """
-    given_paths = path if isinstance(path, list) else [path]
-    if flavour is not None:
-        flavour = RuleBasedFlavour(flavour)
-
-    if grid_definition is not None:
-        grid = grid_registry.from_config(grid_definition)
-    else:
-        grid = None
-
-    ds = from_source("empty")
-    dates = [d.isoformat() for d in dates]
-
-    for path in given_paths:
-        paths = Pattern(path).substitute(*args, date=dates, allow_extra=True, **kwargs)
-
-        for name in ("grid", "area", "rotation", "frame", "resol", "bitmap"):
-            if name in kwargs:
-                raise ValueError(f"MARS interpolation parameter '{name}' not supported")
-
-        for path in _expand(paths):
-            context.trace("📁", "PATH", path)
-            s = from_source("file", path)
-            if flavour is not None:
-                s = flavour.map(s)
-            s = s.sel(valid_datetime=dates, **kwargs)
-            ds = ds + s
-
-    if kwargs and not context.partial_ok:
-        check(ds, given_paths, valid_datetime=dates, **kwargs)
-
-    if grid is not None:
-        ds = new_fieldlist_from_list([new_field_from_grid(f, grid) for f in ds])
-
-    if len(ds) == 0:
-        LOG.warning(f"No fields found for {dates} in {given_paths} (kwargs={kwargs})")
-
-    return ds
+@source_registry.register("grib")
+class GribSource(LegacySource):
+
+    @staticmethod
+    def _execute(
+        context: Any,
+        dates: list[Any],
+        path: str | list[str],
+        flavour: str | dict[str, Any] | None = None,
+        grid_definition: dict[str, Any] | None = None,
+        *args: Any,
+        **kwargs: Any,
+    ) -> ekd.FieldList:
+        """Executes the function to load data from GRIB files.
+
+        Parameters
+        ----------
+        context : Any
+            The context in which the function is executed.
+        dates : list of Any
+            List of dates.
+        path : str or list of str
+            Path or list of paths to the GRIB files.
+        flavour : str or dict of str to Any, optional
+            Flavour information, by default None.
+        grid_definition : dict of str to Any, optional
+            Grid definition configuration to create a Grid object, by default None.
+        *args : Any
+            Additional positional arguments.
+        **kwargs : Any
+            Additional keyword arguments.
+
+        Returns
+        -------
+        Any
+            The loaded dataset.
+        """
+        given_paths = path if isinstance(path, list) else [path]
+        if flavour is not None:
+            flavour = RuleBasedFlavour(flavour)
+
+        if grid_definition is not None:
+            grid = grid_registry.from_config(grid_definition)
+        else:
+            grid = None
+
+        ds = from_source("empty")
+        dates = [d.isoformat() for d in dates]
+
+        for path in given_paths:
+
+            # do not substitute if not needed
+            if "{" not in path:
+                paths = [path]
+            else:
+                paths = Pattern(path).substitute(*args, date=dates, allow_extra=True, **kwargs)
+
+            for name in ("grid", "area", "rotation", "frame", "resol", "bitmap"):
+                if name in kwargs:
+                    raise ValueError(f"MARS interpolation parameter '{name}' not supported")
+
+            for path in _expand(paths):
+                context.trace("📁", "PATH", path)
+                s = from_source("file", path)
+                if flavour is not None:
+                    s = flavour.map(s)
+                sel_kwargs = kwargs.copy()
+                if dates != []:
+                    sel_kwargs["valid_datetime"] = dates
+                s = s.sel(**sel_kwargs)
+                ds = ds + s
+
+        if kwargs and not context.partial_ok:
+            check(ds, given_paths, valid_datetime=dates, **kwargs)
+
+        if grid is not None:
+            ds = new_fieldlist_from_list([new_field_from_grid(f, grid) for f in ds])
+
+        if len(ds) == 0:
+            LOG.warning(f"No fields found for {dates} in {given_paths} (kwargs={kwargs})")
+
+        return ds

anemoi/datasets/create/sources/grib_index.py

@@ -19,7 +19,8 @@ from anemoi.transform.flavour import RuleBasedFlavour
 from cachetools import LRUCache
 from earthkit.data.indexing.fieldlist import FieldArray
 
-from .legacy import legacy_source
+from . import source_registry
+from .legacy import LegacySource
 
 LOG = logging.getLogger(__name__)
 
@@ -569,44 +570,47 @@ class GribIndex:
             yield data
 
 
-@legacy_source(__file__)
-def execute(
-    context: Any,
-    dates: list[Any],
-    indexdb: str,
-    flavour: str | None = None,
-    **kwargs: Any,
-) -> FieldArray:
-    """Execute the GRIB data retrieval process.
-
-    Parameters
-    ----------
-    context : Any
-        The execution context.
-    dates : List[Any]
-        List of dates to retrieve data for.
-    indexdb : str
-        Path to the GRIB index database.
-    flavour : Optional[str], optional
-        Flavour configuration for mapping fields, by default None.
-    **kwargs : Any
-        Additional filtering criteria.
-
-    Returns
-    -------
-    FieldArray
-        An array of retrieved GRIB fields.
-    """
-    index = GribIndex(indexdb)
-    result = []
-
-    if flavour is not None:
-        flavour = RuleBasedFlavour(flavour)
-
-    for grib in index.retrieve(dates, **kwargs):
-        field = ekd.from_source("memory", grib)[0]
-        if flavour:
-            field = flavour.apply(field)
-        result.append(field)
-
-    return FieldArray(result)
+@source_registry.register("grib_index")
+class GribIndexSource(LegacySource):
+
+    @staticmethod
+    def _execute(
+        context: Any,
+        dates: list[Any],
+        indexdb: str,
+        flavour: str | None = None,
+        **kwargs: Any,
+    ) -> FieldArray:
+        """Execute the GRIB data retrieval process.
+
+        Parameters
+        ----------
+        context : Any
+            The execution context.
+        dates : List[Any]
+            List of dates to retrieve data for.
+        indexdb : str
+            Path to the GRIB index database.
+        flavour : Optional[str], optional
+            Flavour configuration for mapping fields, by default None.
+        **kwargs : Any
+            Additional filtering criteria.
+
+        Returns
+        -------
+        FieldArray
+            An array of retrieved GRIB fields.
+        """
+        index = GribIndex(indexdb)
+        result = []
+
+        if flavour is not None:
+            flavour = RuleBasedFlavour(flavour)
+
+        for grib in index.retrieve(dates, **kwargs):
+            field = ekd.from_source("memory", grib)[0]
+            if flavour:
+                field = flavour.apply(field)
+            result.append(field)
+
+        return FieldArray(result)