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

anemoi/datasets/create/sources/tendencies.py

@@ -1,171 +0,0 @@
-# (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 datetime
-from collections import defaultdict
-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.sources import source_registry
-
-from .legacy import LegacySource
-
-
-def _date_to_datetime(d: Any) -> Any:
-    """Converts a date string or a list/tuple of date strings to datetime objects.
-
-    Parameters
-    ----------
-    d : Any
-        A date string or a list/tuple of date strings.
-
-    Returns
-    -------
-    Any
-        A datetime object or a list/tuple of datetime objects.
-    """
-    if isinstance(d, (list, tuple)):
-        return [_date_to_datetime(x) for x in d]
-    return datetime.datetime.fromisoformat(d)
-
-
-def normalise_time_delta(t: Any) -> datetime.timedelta:
-    """Normalizes a time delta string to a datetime.timedelta object.
-
-    Parameters
-    ----------
-    t : Any
-        A time delta string ending with 'h' or a datetime.timedelta object.
-
-    Returns
-    -------
-    datetime.timedelta
-        A normalized datetime.timedelta object.
-    """
-    if isinstance(t, datetime.timedelta):
-        assert t == datetime.timedelta(hours=t.hours), t
-
-    assert t.endswith("h"), t
-
-    t = int(t[:-1])
-    t = datetime.timedelta(hours=t)
-    return t
-
-
-def group_by_field(ds: Any) -> dict[tuple, list[Any]]:
-    """Groups fields by their metadata excluding 'date', 'time', and 'step'.
-
-    Parameters
-    ----------
-    ds : Any
-        A dataset object.
-
-    Returns
-    -------
-    Dict[Tuple, List[Any]]
-        A dictionary where keys are tuples of metadata items and values are lists of fields.
-    """
-    d = defaultdict(list)
-    for field in ds.order_by("valid_datetime"):
-        m = field.metadata(namespace="mars")
-        for k in ("date", "time", "step"):
-            m.pop(k, None)
-        keys = tuple(m.items())
-        d[keys].append(field)
-    return d
-
-
-@source_registry.register("tendencies")
-class TendenciesSource(LegacySource):
-
-    @staticmethod
-    def _execute(dates: list[datetime.datetime], time_increment: Any, **kwargs: Any) -> Any:
-        """Computes tendencies for the given dates and time increment.
-
-        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)))
-
-        from .mars import mars
-
-        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)
-
-        assert group1.keys() == group2.keys(), (group1.keys(), group2.keys())
-
-        # prepare output tmp file so we can read it back
-        tmp = temp_file()
-        path = tmp.path
-        out = new_grib_output(path)
-
-        for k in group1:
-            assert len(group1[k]) == len(group2[k]), k
-            print()
-            print("❌", k)
-
-            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)
-
-                ################
-                # Actual computation happens here
-                x = c - b
-                ################
-
-                assert x.shape == c.shape, c.shape
-                print(f"Computing data for {field.metadata('valid_datetime')}={field}-{b_field}")
-                out.write(x, template=field)
-
-        out.close()
-
-        from earthkit.data import from_source
-
-        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
-
-        return ds