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

anemoi/datasets/create/sources/recentre.py

@@ -12,7 +12,8 @@ from typing import Any
 
 from anemoi.datasets.compute.recentre import recentre as _recentre
 
-from .legacy import legacy_source
+from . import source_registry
+from .legacy import LegacySource
 from .mars import mars
 
 
@@ -105,43 +106,43 @@ def load_if_needed(context: Any, dates: Any, dict_or_dataset: dict | Any) -> Any
     return dict_or_dataset
 
 
-@legacy_source(__file__)
-def recentre(
-    context: Any,
-    dates: Any,
-    members: dict | Any,
-    centre: dict | Any,
-    alpha: float = 1.0,
-    remapping: dict = {},
-    patches: dict = {},
-) -> Any:
-    """Recentres the members dataset using the centre dataset.
-
-    Parameters
-    ----------
-    context : Any
-        The context for recentering.
-    dates : Any
-        The dates for recentering.
-    members : Union[dict, Any]
-        The members dataset or request dictionary.
-    centre : Union[dict, Any]
-        The centre dataset or request dictionary.
-    alpha : float, optional
-        The alpha value for recentering. Defaults to 1.0.
-    remapping : dict, optional
-        The remapping dictionary. Defaults to {}.
-    patches : dict, optional
-        The patches dictionary. Defaults to {}.
-
-    Returns
-    -------
-    Any
-        The recentred dataset.
-    """
-    members = load_if_needed(context, dates, members)
-    centre = load_if_needed(context, dates, centre)
-    return _recentre(members=members, centre=centre, alpha=alpha)
-
-
-execute = recentre
+@source_registry.register("recentre")
+class RecentreSource(LegacySource):
+
+    @staticmethod
+    def _execute(
+        context: Any,
+        dates: Any,
+        members: dict | Any,
+        centre: dict | Any,
+        alpha: float = 1.0,
+        remapping: dict = {},
+        patches: dict = {},
+    ) -> Any:
+        """Recentres the members dataset using the centre dataset.
+
+        Parameters
+        ----------
+        context : Any
+            The context for recentering.
+        dates : Any
+            The dates for recentering.
+        members : Union[dict, Any]
+            The members dataset or request dictionary.
+        centre : Union[dict, Any]
+            The centre dataset or request dictionary.
+        alpha : float, optional
+            The alpha value for recentering. Defaults to 1.0.
+        remapping : dict, optional
+            The remapping dictionary. Defaults to {}.
+        patches : dict, optional
+            The patches dictionary. Defaults to {}.
+
+        Returns
+        -------
+        Any
+            The recentred dataset.
+        """
+        members = load_if_needed(context, dates, members)
+        centre = load_if_needed(context, dates, centre)
+        return _recentre(members=members, centre=centre, alpha=alpha)

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/tendencies.py

@@ -14,9 +14,9 @@ from typing import Any
 from earthkit.data.core.temporary import temp_file
 from earthkit.data.readers.grib.output import new_grib_output
 
-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
 
 
 def _date_to_datetime(d: Any) -> Any:
@@ -83,116 +83,89 @@ def group_by_field(ds: Any) -> dict[tuple, list[Any]]:
     return d
 
 
-@legacy_source(__file__)
-def tendencies(dates: list[datetime.datetime], time_increment: Any, **kwargs: Any) -> Any:
-    """Computes tendencies for the given dates and time increment.
+@source_registry.register("tendencies")
+class TendenciesSource(LegacySource):
 
-    Parameters
-    ----------
-    dates : List[datetime.datetime]
-        A list of datetime objects.
-    time_increment : Any
-        A time increment string ending with 'h' or a datetime.timedelta object.
-    **kwargs : Any
-        Additional keyword arguments.
-
-    Returns
-    -------
-    Any
-        A dataset object with computed tendencies.
-    """
-    print("✅", kwargs)
-    time_increment = normalise_time_delta(time_increment)
-
-    shifted_dates = [d - time_increment for d in dates]
-    all_dates = sorted(list(set(dates + shifted_dates)))
+    @staticmethod
+    def _execute(dates: list[datetime.datetime], time_increment: Any, **kwargs: Any) -> Any:
+        """Computes tendencies for the given dates and time increment.
 
-    # from .mars import execute as mars
-    from anemoi.datasets.create.mars import execute as mars
+        Parameters
+        ----------
+        dates : List[datetime.datetime]
+            A list of datetime objects.
+        time_increment : Any
+            A time increment string ending with 'h' or a datetime.timedelta object.
+        **kwargs : Any
+            Additional keyword arguments.
 
-    ds = mars(dates=all_dates, **kwargs)
-
-    dates_in_data = ds.unique_values("valid_datetime", progress_bar=False)["valid_datetime"]
-    for d in all_dates:
-        assert d.isoformat() in dates_in_data, d
-
-    ds1 = ds.sel(valid_datetime=[d.isoformat() for d in dates])
-    ds2 = ds.sel(valid_datetime=[d.isoformat() for d in shifted_dates])
-
-    assert len(ds1) == len(ds2), (len(ds1), len(ds2))
-
-    group1 = group_by_field(ds1)
-    group2 = group_by_field(ds2)
+        Returns
+        -------
+        Any
+            A dataset object with computed tendencies.
+        """
+        print("✅", kwargs)
+        time_increment = normalise_time_delta(time_increment)
 
-    assert group1.keys() == group2.keys(), (group1.keys(), group2.keys())
+        shifted_dates = [d - time_increment for d in dates]
+        all_dates = sorted(list(set(dates + shifted_dates)))
 
-    # prepare output tmp file so we can read it back
-    tmp = temp_file()
-    path = tmp.path
-    out = new_grib_output(path)
+        from .mars import mars
 
-    for k in group1:
-        assert len(group1[k]) == len(group2[k]), k
-        print()
-        print("❌", k)
+        ds = mars(dates=all_dates, **kwargs)
 
-        for field, b_field in zip(group1[k], group2[k]):
-            for k in ["param", "level", "number", "grid", "shape"]:
-                assert field.metadata(k) == b_field.metadata(k), (
-                    k,
-                    field.metadata(k),
-                    b_field.metadata(k),
-                )
+        dates_in_data = ds.unique_values("valid_datetime", progress_bar=False)["valid_datetime"]
+        for d in all_dates:
+            assert d.isoformat() in dates_in_data, d
 
-            c = field.to_numpy()
-            b = b_field.to_numpy()
-            assert c.shape == b.shape, (c.shape, b.shape)
+        ds1 = ds.sel(valid_datetime=[d.isoformat() for d in dates])
+        ds2 = ds.sel(valid_datetime=[d.isoformat() for d in shifted_dates])
 
-            ################
-            # Actual computation happens here
-            x = c - b
-            ################
+        assert len(ds1) == len(ds2), (len(ds1), len(ds2))
 
-            assert x.shape == c.shape, c.shape
-            print(f"Computing data for {field.metadata('valid_datetime')}={field}-{b_field}")
-            out.write(x, template=field)
+        group1 = group_by_field(ds1)
+        group2 = group_by_field(ds2)
 
-    out.close()
+        assert group1.keys() == group2.keys(), (group1.keys(), group2.keys())
 
-    from earthkit.data import from_source
+        # prepare output tmp file so we can read it back
+        tmp = temp_file()
+        path = tmp.path
+        out = new_grib_output(path)
 
-    ds = from_source("file", path)
-    # save a reference to the tmp file so it is deleted
-    # only when the dataset is not used anymore
-    ds._tmp = tmp
+        for k in group1:
+            assert len(group1[k]) == len(group2[k]), k
+            print()
+            print("❌", k)
 
-    return ds
+            for field, b_field in zip(group1[k], group2[k]):
+                for k in ["param", "level", "number", "grid", "shape"]:
+                    assert field.metadata(k) == b_field.metadata(k), (
+                        k,
+                        field.metadata(k),
+                        b_field.metadata(k),
+                    )
 
+                c = field.to_numpy()
+                b = b_field.to_numpy()
+                assert c.shape == b.shape, (c.shape, b.shape)
 
-execute = tendencies
+                ################
+                # Actual computation happens here
+                x = c - b
+                ################
 
-if __name__ == "__main__":
-    import yaml
+                assert x.shape == c.shape, c.shape
+                print(f"Computing data for {field.metadata('valid_datetime')}={field}-{b_field}")
+                out.write(x, template=field)
 
-    config = yaml.safe_load(
-        """
+        out.close()
 
-    config:
-      time_increment: 12h
-      database: marser
-      class: ea
-      # date: computed automatically
-      # time: computed automatically
-      expver: "0001"
-      grid: 20.0/20.0
-      levtype: sfc
-      param: [2t]
-    """
-    )["config"]
+        from earthkit.data import from_source
 
-    dates = 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)
+        ds = from_source("file", path)
+        # save a reference to the tmp file so it is deleted
+        # only when the dataset is not used anymore
+        ds._tmp = tmp
 
-    DEBUG = True
-    for f in tendencies(dates, **config):
-        print(f, f.to_numpy().mean())
+        return ds

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__)
@@ -152,26 +153,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/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_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)