anemoi-utils 0.1.8__py3-none-any.whl → 0.1.9__py3-none-any.whl

anemoi/utils/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.1.8'
-__version_tuple__ = version_tuple = (0, 1, 8)
+__version__ = version = '0.1.9'
+__version_tuple__ = version_tuple = (0, 1, 9)

anemoi/utils/dates.py

@@ -10,6 +10,16 @@ import calendar
 import datetime
 
 
+def normalise_frequency(frequency):
+    if isinstance(frequency, int):
+        return frequency
+    assert isinstance(frequency, str), (type(frequency), frequency)
+
+    unit = frequency[-1].lower()
+    v = int(frequency[:-1])
+    return {"h": v, "d": v * 24}[unit]
+
+
 def no_time_zone(date):
     """Remove time zone information from a date.
 
@@ -27,6 +37,7 @@ def no_time_zone(date):
     return date.replace(tzinfo=None)
 
 
+# this function is use in anemoi-datasets
 def as_datetime(date):
     """Convert a date to a datetime object, removing any time zone information.
 
@@ -162,11 +173,15 @@ class HindcastDatesTimes:
         """
 
         self.reference_dates = reference_dates
-        self.years = (1, years + 1)
+
+        if isinstance(years, list):
+            self.years = years
+        else:
+            self.years = range(1, years + 1)
 
     def __iter__(self):
         for reference_date in self.reference_dates:
-            for year in range(*self.years):
+            for year in self.years:
                 if reference_date.month == 2 and reference_date.day == 29:
                     date = datetime.datetime(reference_date.year - year, 2, 28)
                 else:
@@ -246,3 +261,61 @@ class Autumn(DateTimes):
             _description_
         """
         super().__init__(datetime.datetime(year, 9, 1), datetime.datetime(year, 11, 30), **kwargs)
+
+
+class ConcatDateTimes:
+    def __init__(self, *dates):
+        if len(dates) == 1 and isinstance(dates[0], list):
+            dates = dates[0]
+
+        self.dates = dates
+
+    def __iter__(self):
+        for date in self.dates:
+            yield from date
+
+
+class EnumDateTimes:
+    def __init__(self, dates):
+        self.dates = dates
+
+    def __iter__(self):
+        for date in self.dates:
+            yield as_datetime(date)
+
+
+def datetimes_factory(*args, **kwargs):
+    if args and kwargs:
+        raise ValueError("Cannot provide both args and kwargs for a list of dates")
+
+    if not args and not kwargs:
+        raise ValueError("No dates provided")
+
+    if kwargs:
+        name = kwargs.get("name")
+
+        if name == "hindcast":
+            reference_dates = kwargs["reference_dates"]
+            reference_dates = datetimes_factory(reference_dates)
+            years = kwargs["years"]
+            return HindcastDatesTimes(reference_dates=reference_dates, years=years)
+
+        kwargs = kwargs.copy()
+        if "frequency" in kwargs:
+            freq = kwargs.pop("frequency")
+            kwargs["increment"] = normalise_frequency(freq)
+        return DateTimes(**kwargs)
+
+    if not any((isinstance(x, dict) or isinstance(x, list)) for x in args):
+        return EnumDateTimes(args)
+
+    if len(args) == 1:
+        a = args[0]
+
+        if isinstance(a, dict):
+            return datetimes_factory(**a)
+
+        if isinstance(a, list):
+            return datetimes_factory(*a)
+
+    return ConcatDateTimes(*[datetimes_factory(a) for a in args])

anemoi/utils/mars/__init__.py

@@ -0,0 +1,76 @@
+# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# 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.
+
+
+"""Utilities for working with Mars requests.
+
+Has some konwledge of how certain streams are organised in Mars.
+
+"""
+
+import datetime
+import logging
+import os
+
+import yaml
+
+LOG = logging.getLogger(__name__)
+
+DEFAULT_MARS_LABELLING = {
+    "class": "od",
+    "type": "an",
+    "stream": "oper",
+    "expver": "0001",
+}
+
+
+def _expand_mars_labelling(request):
+    """Expand the request with the default Mars labelling.
+
+    The default Mars labelling is:
+
+    {'class': 'od',
+     'type': 'an',
+     'stream': 'oper',
+     'expver': '0001'}
+
+    """
+    result = DEFAULT_MARS_LABELLING.copy()
+    result.update(request)
+    return result
+
+
+STREAMS = None
+
+
+def _lookup_mars_stream(request):
+    global STREAMS
+
+    if STREAMS is None:
+
+        with open(os.path.join(os.path.dirname(__file__), "mars.yaml")) as f:
+            STREAMS = yaml.safe_load(f)
+
+    request = _expand_mars_labelling(request)
+    for s in STREAMS:
+        match = s["match"]
+        if all(request.get(k) == v for k, v in match.items()):
+            return s["info"]
+
+
+def recenter(date, center, members):
+
+    center = _lookup_mars_stream(center)
+    members = _lookup_mars_stream(members)
+
+    return (center, members)
+
+
+if __name__ == "__main__":
+    date = datetime.datetime(2024, 5, 9, 0)
+
+    print(recenter(date, {"type": "an"}, {"stream": "elda"}))

anemoi/utils/mars/mars.yaml

@@ -0,0 +1,5 @@
+- match:
+    class: od
+    stream: elda
+  info:
+    runs: [6, 18]

{anemoi_utils-0.1.8.dist-info → anemoi_utils-0.1.9.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: anemoi-utils
-Version: 0.1.8
+Version: 0.1.9
 Summary: A package to hold various functions to support training of ML models on ECMWF data.
 Author-email: "European Centre for Medium-Range Weather Forecasts (ECMWF)" <software.support@ecmwf.int>
 License: Apache License
@@ -223,6 +223,7 @@ Classifier: Operating System :: OS Independent
 Requires-Python: >=3.9
 License-File: LICENSE
 Requires-Dist: tomli
+Requires-Dist: pyyaml
 Provides-Extra: all
 Requires-Dist: tomli ; extra == 'all'
 Requires-Dist: GitPython ; extra == 'all'

{anemoi_utils-0.1.8.dist-info → anemoi_utils-0.1.9.dist-info}/RECORD

@@ -1,15 +1,17 @@
 anemoi/utils/__init__.py,sha256=zZZpbKIoGWwdCOuo6YSruLR7C0GzvzI1Wzhyqaa0K7M,456
-anemoi/utils/_version.py,sha256=PdJ7dZoz_SyEgX0MdrMfQYBFlGcwpemv6ibF8NKALBY,411
+anemoi/utils/_version.py,sha256=NWmu2cvzOcqY9v-ee-qFLmtXRczssdN-cFGZ9qMNSmY,411
 anemoi/utils/caching.py,sha256=HrC9aFHlcCTaM2Z5u0ivGIXz7eFu35UQQhUuwwuG2pk,1743
 anemoi/utils/checkpoints.py,sha256=IR86FFNh5JR_uQVlgybnZG74PyU0CNLhyocqARwZIrs,2069
 anemoi/utils/config.py,sha256=XEesqODvkuE3ZA7dnEnZ-ooBRtU6ecPmkfP65FtialA,2147
-anemoi/utils/dates.py,sha256=ZabXLecHMio9ecIJbUxNcfOFoGfkcotSwy_xw1kNQ9w,6474
+anemoi/utils/dates.py,sha256=Ot9OTY1uFvHxW1EU4DPv3oUqmzvkXTwKuwhlfVlY788,8426
 anemoi/utils/grib.py,sha256=gVfo4KYQv31iRyoqRDwk5tiqZDUgOIvhag_kO0qjYD0,3067
 anemoi/utils/humanize.py,sha256=LD6dGnqChxA5j3tMhSybsAGRQzi33d_qS9pUoUHubkc,10330
 anemoi/utils/provenance.py,sha256=v54L9jF1JgYcclOhg3iojRl1v3ajbiWz_oc289xTgO4,9574
 anemoi/utils/text.py,sha256=pGWtDvRFoDxAnSuZJiA-GOGJOJLHsw2dAm0tfVvPKno,8599
-anemoi_utils-0.1.8.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-anemoi_utils-0.1.8.dist-info/METADATA,sha256=bEA_bMbWJubncfBmc6tMi6cJLtFgZrgth899Z60i5Yg,15079
-anemoi_utils-0.1.8.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-anemoi_utils-0.1.8.dist-info/top_level.txt,sha256=DYn8VPs-fNwr7fNH9XIBqeXIwiYYd2E2k5-dUFFqUz0,7
-anemoi_utils-0.1.8.dist-info/RECORD,,
+anemoi/utils/mars/__init__.py,sha256=RAeY8gJ7ZvsPlcIvrQ4fy9xVHs3SphTAPw_XJDtNIKo,1750
+anemoi/utils/mars/mars.yaml,sha256=R0dujp75lLA4wCWhPeOQnzJ45WZAYLT8gpx509cBFlc,66
+anemoi_utils-0.1.9.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+anemoi_utils-0.1.9.dist-info/METADATA,sha256=wVXW4E6hpkTjCEHUCWDiQXm-pZfJl81Mrzdy6kwKbfA,15101
+anemoi_utils-0.1.9.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+anemoi_utils-0.1.9.dist-info/top_level.txt,sha256=DYn8VPs-fNwr7fNH9XIBqeXIwiYYd2E2k5-dUFFqUz0,7
+anemoi_utils-0.1.9.dist-info/RECORD,,