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

anemoi/datasets/create/sources/repeated_dates.py

@@ -0,0 +1,44 @@
+# (C) Copyright 2024 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.
+#
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+
+
+import logging
+from typing import Any
+
+from anemoi.transform.fields import new_field_with_valid_datetime
+from anemoi.transform.fields import new_fieldlist_from_list
+
+from anemoi.datasets.create.input.repeated_dates import DateMapper
+from anemoi.datasets.create.source import Source
+from anemoi.datasets.create.sources import source_registry
+
+LOG = logging.getLogger(__name__)
+
+
+@source_registry.register("repeated_dates")
+class RepeatedDatesSource(Source):
+
+    def __init__(self, context, source: Any, mode: str, **kwargs) -> None:
+        # assert False, (context, source, mode, kwargs)
+        super().__init__(context, **kwargs)
+        self.mapper = DateMapper.from_mode(mode, source, kwargs)
+        self.source = source
+
+    def execute(self, group_of_dates):
+        source = self.context.create_source(self.source, "data_sources", str(id(self)))
+
+        result = []
+        for one_date_group, many_dates_group in self.mapper.transform(group_of_dates):
+            print(f"one_date_group: {one_date_group}, many_dates_group: {many_dates_group}")
+            source_results = source(self.context, one_date_group)
+            for field in source_results:
+                for date in many_dates_group:
+                    result.append(new_field_with_valid_datetime(field, date))
+
+        return new_fieldlist_from_list(result)

anemoi/datasets/create/sources/source.py

@@ -12,58 +12,36 @@ from typing import Any
 
 from earthkit.data import from_source
 
-from anemoi.datasets.create.utils import to_datetime_list
+from anemoi.datasets.create.sources import source_registry
 
-from .legacy import legacy_source
+from .legacy import LegacySource
 
 
-@legacy_source(__file__)
-def source(context: Any | None, dates: list[datetime], **kwargs: Any) -> Any:
-    """Generates a source based on the provided context, dates, and additional keyword arguments.
+@source_registry.register("source")
+class GenericSource(LegacySource):
 
-    Parameters
-    ----------
-    context : Optional[Any]
-        The context in which the source is generated.
-    dates : List[datetime]
-        A list of datetime objects representing the dates.
-    **kwargs : Any
-        Additional keyword arguments for the source generation.
+    @staticmethod
+    def _execute(context: Any | None, dates: list[datetime], **kwargs: Any) -> Any:
+        """Generates a source based on the provided context, dates, and additional keyword arguments.
 
-    Returns
-    -------
-    Any
-        The generated source.
-    """
-    name = kwargs.pop("name")
-    context.trace("✅", f"from_source({name}, {dates}, {kwargs}")
-    if kwargs["date"] == "$from_dates":
-        kwargs["date"] = list({d.strftime("%Y%m%d") for d in dates})
-    if kwargs["time"] == "$from_dates":
-        kwargs["time"] = list({d.strftime("%H%M") for d in dates})
-    return from_source(name, **kwargs)
+        Parameters
+        ----------
+        context : Optional[Any]
+            The context in which the source is generated.
+        dates : List[datetime]
+            A list of datetime objects representing the dates.
+        **kwargs : Any
+            Additional keyword arguments for the source generation.
 
-
-execute = source
-
-if __name__ == "__main__":
-    import yaml
-
-    config: dict[str, Any] = yaml.safe_load(
+        Returns
+        -------
+        Any
+            The generated source.
         """
-      name: mars
-      class: ea
-      expver: '0001'
-      grid: 20.0/20.0
-      levtype: sfc
-      param: [2t]
-      number: [0, 1]
-      date: $from_dates
-      time: $from_dates
-    """
-    )
-    dates: list[str] = yaml.safe_load("[2022-12-30 18:00, 2022-12-31 00:00, 2022-12-31 06:00, 2022-12-31 12:00]")
-    dates = to_datetime_list(dates)
-
-    for f in source(None, dates, **config):
-        print(f, f.to_numpy().mean())
+        name = kwargs.pop("name")
+        context.trace("✅", f"from_source({name}, {dates}, {kwargs}")
+        if kwargs["date"] == "$from_dates":
+            kwargs["date"] = list({d.strftime("%Y%m%d") for d in dates})
+        if kwargs["time"] == "$from_dates":
+            kwargs["time"] = list({d.strftime("%H%M") for d in dates})
+        return from_source(name, **kwargs)

anemoi/datasets/create/sources/xarray_support/__init__.py

@@ -17,7 +17,8 @@ from earthkit.data.core.fieldlist import MultiFieldList
 
 from anemoi.datasets.create.sources.patterns import iterate_patterns
 
-from ..legacy import legacy_source
+from .. import source_registry
+from ..legacy import LegacySource
 from .fieldlist import XarrayFieldList
 
 LOG = logging.getLogger(__name__)
@@ -96,6 +97,7 @@ def load_one(
     if isinstance(dataset, xr.Dataset):
         data = dataset
     else:
+        print(f"Opening dataset {dataset} with options {options}")
         data = xr.open_dataset(dataset, **options)
 
     fs = XarrayFieldList.from_xarray(data, flavour=flavour, patch=patch)
@@ -152,26 +154,30 @@ def load_many(emoji: str, context: Any, dates: list[datetime.datetime], pattern:
     return MultiFieldList(result)
 
 
-@legacy_source("xarray")
-def execute(context: Any, dates: list[str], url: str, *args: Any, **kwargs: Any) -> ekd.FieldList:
-    """Executes the loading of datasets.
-
-    Parameters
-    ----------
-    context : Any
-        Context object.
-    dates : List[str]
-        List of dates.
-    url : str
-        URL pattern for loading datasets.
-    *args : Any
-        Additional arguments.
-    **kwargs : Any
-        Additional keyword arguments.
-
-    Returns
-    -------
-    ekd.FieldList
-        The loaded datasets.
-    """
-    return load_many("🌐", context, dates, url, *args, **kwargs)
+@source_registry.register("xarray")
+class LegacyXarraySource(LegacySource):
+    name = "xarray"
+
+    @staticmethod
+    def _execute(context: Any, dates: list[str], url: str, *args: Any, **kwargs: Any) -> ekd.FieldList:
+        """Executes the loading of datasets.
+
+        Parameters
+        ----------
+        context : Any
+            Context object.
+        dates : List[str]
+            List of dates.
+        url : str
+            URL pattern for loading datasets.
+        *args : Any
+            Additional arguments.
+        **kwargs : Any
+            Additional keyword arguments.
+
+        Returns
+        -------
+        ekd.FieldList
+            The loaded datasets.
+        """
+        return load_many("🌐", context, dates, url, *args, **kwargs)

anemoi/datasets/create/sources/xarray_support/coordinates.py

@@ -223,13 +223,10 @@ class Coordinate:
         # Assume the array is sorted
 
         index = np.searchsorted(values, value)
-        index = index[index < len(values)]
-
-        if np.all(values[index] == value):
+        if np.all(index < len(values)) and np.all(values[index] == value):
             return index
 
         # If not found, we need to check if the value is in the array
-
         index = np.where(np.isin(values, value))[0]
 
         # We could also return incomplete matches

anemoi/datasets/create/sources/xarray_support/field.py

@@ -121,16 +121,16 @@ class XArrayField(Field):
             Index to select a specific element, by default None.
         """
         if index is not None:
-            values = self.selection[index]
+            values = self.selection[index].values
         else:
-            values = self.selection
+            values = self.selection.values
 
         assert dtype is None
 
         if flatten:
-            return values.values.flatten()
+            return values.flatten()
 
-        return values  # .reshape(self.shape)
+        return values
 
     @cached_property
     def _metadata(self) -> XArrayMetadata:

anemoi/datasets/create/sources/xarray_support/flavour.py

@@ -557,10 +557,10 @@ class DefaultCoordinateGuesser(CoordinateGuesser):
         super().__init__(ds)
 
     def _is_point(self, c: xr.DataArray, attributes: CoordinateAttributes) -> PointCoordinate | None:
-        if attributes.standard_name in ["cell", "station", "poi", "point"]:
+        if attributes.standard_name in ["location", "cell", "id", "station", "poi", "point"]:
             return PointCoordinate(c)
 
-        if attributes.name in ["cell", "station", "poi", "point"]:  # WeatherBench
+        if attributes.name in ["location", "cell", "id", "station", "poi", "point"]:  # WeatherBench
             return PointCoordinate(c)
 
         return None

anemoi/datasets/create/sources/xarray_support/patch.py

@@ -10,13 +10,14 @@
 
 import logging
 from typing import Any
+from typing import Literal
 
 import xarray as xr
 
 LOG = logging.getLogger(__name__)
 
 
-def patch_attributes(ds: xr.Dataset, attributes: dict[str, dict[str, Any]]) -> Any:
+def patch_attributes(ds: xr.Dataset, attributes: dict[str, dict[str, Any]]) -> xr.Dataset:
     """Patch the attributes of the dataset.
 
     Parameters
@@ -38,7 +39,7 @@ def patch_attributes(ds: xr.Dataset, attributes: dict[str, dict[str, Any]]) -> A
     return ds
 
 
-def patch_coordinates(ds: xr.Dataset, coordinates: list[str]) -> Any:
+def patch_coordinates(ds: xr.Dataset, coordinates: list[str]) -> xr.Dataset:
     """Patch the coordinates of the dataset.
 
     Parameters
@@ -59,7 +60,7 @@ def patch_coordinates(ds: xr.Dataset, coordinates: list[str]) -> Any:
     return ds
 
 
-def patch_rename(ds: xr.Dataset, renames: dict[str, str]) -> Any:
+def patch_rename(ds: xr.Dataset, renames: dict[str, str]) -> xr.Dataset:
     """Rename variables in the dataset.
 
     Parameters
@@ -77,7 +78,7 @@ def patch_rename(ds: xr.Dataset, renames: dict[str, str]) -> Any:
     return ds.rename(renames)
 
 
-def patch_sort_coordinate(ds: xr.Dataset, sort_coordinates: list[str]) -> Any:
+def patch_sort_coordinate(ds: xr.Dataset, sort_coordinates: list[str]) -> xr.Dataset:
     """Sort the coordinates of the dataset.
 
     Parameters
@@ -98,11 +99,175 @@ def patch_sort_coordinate(ds: xr.Dataset, sort_coordinates: list[str]) -> Any:
     return ds
 
 
+def patch_subset_dataset(ds: xr.Dataset, selection: dict[str, Any]) -> xr.Dataset:
+    """Select a subset of the dataset using xarray's sel method.
+
+    Parameters
+    ----------
+    ds : xr.Dataset
+        The dataset to patch.
+    selection : dict[str, Any]
+        Dictionary mapping dimension names to selection criteria.
+        Keys must be existing dimension names in the dataset.
+        Values can be any type accepted by xarray's sel method, including:
+        - Single values (int, float, str, datetime)
+        - Lists or arrays of values
+        - Slices (using slice() objects)
+        - Boolean arrays
+
+    Returns
+    -------
+    xr.Dataset
+        The patched dataset containing only the selected subset.
+
+    Examples
+    --------
+    >>> # Select specific time and pressure level
+    >>> patch_subset_dataset(ds, {
+    ...     'time': '2020-01-01',
+    ...     'pressure': 500
+    ... })
+
+    >>> # Select a range using slice
+    >>> patch_subset_dataset(ds, {
+    ...     'lat': slice(-90, 90),
+    ...     'lon': slice(0, 180)
+    ... })
+    """
+
+    ds = ds.sel(selection)
+
+    return ds
+
+
+def patch_analysis_lead_to_valid_time(
+    ds: xr.Dataset,
+    time_coord_names: dict[Literal["analysis_time_coordinate", "lead_time_coordinate", "valid_time_coordinate"], str],
+) -> xr.Dataset:
+    """Convert analysis time and lead time coordinates to valid time.
+
+    This function creates a new valid time coordinate by adding the analysis time
+    and lead time coordinates, then stacks and reorganizes the dataset to use
+    valid time as the primary time dimension.
+
+    Parameters
+    ----------
+    ds : xr.Dataset
+        The dataset to patch.
+    time_coord_names : dict[str, str]
+        Dictionary mapping required keys to coordinate names in the dataset:
+
+        - 'analysis_time_coordinate' : str
+            Name of the analysis/initialization time coordinate.
+        - 'lead_time_coordinate' : str
+            Name of the forecast lead time coordinate.
+        - 'valid_time_coordinate' : str
+            Name for the new valid time coordinate to create.
+
+    Returns
+    -------
+    xr.Dataset
+        The patched dataset with valid time as the primary time coordinate.
+        The analysis and lead time coordinates are removed.
+
+    Examples
+    --------
+    >>> patch_analysis_lead_to_valid_time(ds, {
+    ...     'analysis_time_coordinate': 'forecast_reference_time',
+    ...     'lead_time_coordinate': 'step',
+    ...     'valid_time_coordinate': 'time'
+    ... })
+    """
+
+    assert time_coord_names.keys() == {
+        "analysis_time_coordinate",
+        "lead_time_coordinate",
+        "valid_time_coordinate",
+    }, "time_coord_names must contain exactly keys 'analysis_time_coordinate', 'lead_time_coordinate', and 'valid_time_coordinate'"
+
+    analysis_time_coordinate = time_coord_names["analysis_time_coordinate"]
+    lead_time_coordinate = time_coord_names["lead_time_coordinate"]
+    valid_time_coordinate = time_coord_names["valid_time_coordinate"]
+
+    valid_time = ds[analysis_time_coordinate] + ds[lead_time_coordinate]
+
+    ds = (
+        ds.assign_coords({valid_time_coordinate: valid_time})
+        .stack(time_index=[analysis_time_coordinate, lead_time_coordinate])
+        .set_index(time_index=valid_time_coordinate)
+        .rename(time_index=valid_time_coordinate)
+        .drop_vars([analysis_time_coordinate, lead_time_coordinate])
+    )
+
+    return ds
+
+
+def patch_rolling_operation(
+    ds: xr.Dataset, vars_operation_config: dict[Literal["dim", "steps", "vars", "operation"], str | int | list[str]]
+) -> xr.Dataset:
+    """Apply a rolling operation to specified variables in the dataset.
+
+    This function calculates a rolling operation over a specified dimension for selected
+    variables. The rolling window requires all periods to be present (min_periods=steps).
+
+    Parameters
+    ----------
+    ds : xr.Dataset
+        The dataset to patch.
+    vars_operation_config: dict
+        Configuration for the rolling operation with the following keys:
+
+        - 'dim' : str
+            The dimension along which to apply the rolling operation (e.g., 'time').
+        - 'steps' : int
+            The number of steps in the rolling window.
+        - 'vars' : list[str]
+            List of variable names to apply the rolling operation to.
+        - 'operation' : str
+            The operation to apply ('sum', 'mean', 'min', 'max', 'std', etc.).
+
+    Returns
+    -------
+    xr.Dataset
+        The patched dataset with rolling operations applied to the specified variables.
+
+    Examples
+    --------
+    >>> patch_rolling_operation(ds, {
+    ...     'dim': 'time',
+    ...     'steps': 3,
+    ...     'vars': ['precipitation', 'radiation'],
+    ...     'operation': 'sum'
+    ... })
+    """
+
+    assert vars_operation_config.keys() == {
+        "dim",
+        "steps",
+        "vars",
+        "operation",
+    }, "vars_operation_config must contain exactly keys 'dim', 'steps', 'vars', and 'operation'"
+
+    dim = vars_operation_config["dim"]
+    steps = vars_operation_config["steps"]
+    vars = vars_operation_config["vars"]
+    operation = vars_operation_config["operation"]
+
+    for var in vars:
+        rolling = ds[var].rolling(dim={dim: steps}, min_periods=steps)
+        ds[var] = getattr(rolling, operation)()
+
+    return ds
+
+
 PATCHES = {
     "attributes": patch_attributes,
     "coordinates": patch_coordinates,
     "rename": patch_rename,
     "sort_coordinates": patch_sort_coordinate,
+    "analysis_lead_to_valid_time": patch_analysis_lead_to_valid_time,
+    "rolling_operation": patch_rolling_operation,
+    "subset_dataset": patch_subset_dataset,
 }
 
 
@@ -122,7 +287,15 @@ def patch_dataset(ds: xr.Dataset, patch: dict[str, dict[str, Any]]) -> Any:
         The patched dataset.
     """
 
-    ORDER = ["coordinates", "attributes", "rename", "sort_coordinates"]
+    ORDER = [
+        "coordinates",
+        "attributes",
+        "rename",
+        "sort_coordinates",
+        "subset_dataset",
+        "analysis_lead_to_valid_time",
+        "rolling_operation",
+    ]
     for what, values in sorted(patch.items(), key=lambda x: ORDER.index(x[0])):
         if what not in PATCHES:
             raise ValueError(f"Unknown patch type {what!r}")

anemoi/datasets/create/sources/xarray_zarr.py

@@ -11,30 +11,34 @@ from typing import Any
 
 import earthkit.data as ekd
 
-from .legacy import legacy_source
+from . import source_registry
+from .legacy import LegacySource
 from .xarray import load_many
 
 
-@legacy_source(__file__)
-def execute(context: Any, dates: list[str], url: str, *args: Any, **kwargs: Any) -> ekd.FieldList:
-    """Execute the data loading process.
-
-    Parameters
-    ----------
-    context : Any
-        The context in which the execution occurs.
-    dates : List[str]
-        List of dates for which data is to be loaded.
-    url : str
-        The URL from which data is to be loaded.
-    *args : tuple
-        Additional positional arguments.
-    **kwargs : dict
-        Additional keyword arguments.
-
-    Returns
-    -------
-    ekd.FieldList
-        The loaded data.
-    """
-    return load_many("🇿", context, dates, url, *args, **kwargs)
+@source_registry.register("xarray_zarr")
+class XarrayZarrSource(LegacySource):
+
+    @staticmethod
+    def _execute(context: Any, dates: list[str], url: str, *args: Any, **kwargs: Any) -> ekd.FieldList:
+        """Execute the data loading process.
+
+        Parameters
+        ----------
+        context : Any
+            The context in which the execution occurs.
+        dates : List[str]
+            List of dates for which data is to be loaded.
+        url : str
+            The URL from which data is to be loaded.
+        *args : tuple
+            Additional positional arguments.
+        **kwargs : dict
+            Additional keyword arguments.
+
+        Returns
+        -------
+        ekd.FieldList
+            The loaded data.
+        """
+        return load_many("🇿", context, dates, url, *args, **kwargs)

anemoi/datasets/create/sources/zenodo.py

@@ -14,54 +14,58 @@ import earthkit.data as ekd
 from earthkit.data.core.fieldlist import MultiFieldList
 from earthkit.data.sources.url import download_and_cache
 
-from .legacy import legacy_source
+from . import source_registry
+from .legacy import LegacySource
 from .patterns import iterate_patterns
 from .xarray import load_one
 
 
-@legacy_source(__file__)
-def execute(context: Any, dates: Any, record_id: str, file_key: str, *args: Any, **kwargs: Any) -> ekd.FieldList:
-    """Executes the download and processing of files from Zenodo.
+@source_registry.register("zenodo")
+class ZenodoSource(LegacySource):
 
-    Parameters
-    ----------
-    context : Any
-        The context in which the function is executed.
-    dates : Any
-        The dates for which the data is required.
-    record_id : str
-        The Zenodo record ID.
-    file_key : str
-        The key to identify the file.
-    *args : Any
-        Additional arguments.
-    **kwargs : Any
-        Additional keyword arguments.
+    @staticmethod
+    def _execute(context: Any, dates: Any, record_id: str, file_key: str, *args: Any, **kwargs: Any) -> ekd.FieldList:
+        """Executes the download and processing of files from Zenodo.
 
-    Returns
-    -------
-    MultiFieldList
-        A list of fields loaded from the downloaded files.
-    """
-    import requests
+        Parameters
+        ----------
+        context : Any
+            The context in which the function is executed.
+        dates : Any
+            The dates for which the data is required.
+        record_id : str
+            The Zenodo record ID.
+        file_key : str
+            The key to identify the file.
+        *args : Any
+            Additional arguments.
+        **kwargs : Any
+            Additional keyword arguments.
 
-    result: list[Any] = []
+        Returns
+        -------
+        MultiFieldList
+            A list of fields loaded from the downloaded files.
+        """
+        import requests
 
-    URLPATTERN = "https://zenodo.org/api/records/{record_id}"
-    url = URLPATTERN.format(record_id=record_id)
-    r = requests.get(url)
-    r.raise_for_status()
-    record: dict[str, Any] = r.json()
+        result: list[Any] = []
 
-    urls: dict[str, str] = {}
-    for file in record["files"]:
-        urls[file["key"]] = file["links"]["self"]
+        URLPATTERN = "https://zenodo.org/api/records/{record_id}"
+        url = URLPATTERN.format(record_id=record_id)
+        r = requests.get(url)
+        r.raise_for_status()
+        record: dict[str, Any] = r.json()
 
-    for url, dates in iterate_patterns(file_key, dates, **kwargs):
-        if url not in urls:
-            continue
+        urls: dict[str, str] = {}
+        for file in record["files"]:
+            urls[file["key"]] = file["links"]["self"]
 
-        path = download_and_cache(urls[url])
-        result.append(load_one("?", context, dates, path, options={}, flavour=None, **kwargs))
+        for url, dates in iterate_patterns(file_key, dates, **kwargs):
+            if url not in urls:
+                continue
 
-    return MultiFieldList(result)
+            path = download_and_cache(urls[url])
+            result.append(load_one("?", context, dates, path, options={}, flavour=None, **kwargs))
+
+        return MultiFieldList(result)