anemoi-utils 0.1.6__py3-none-any.whl

anemoi/utils/__init__.py

@@ -0,0 +1,9 @@
+# (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.
+
+
+from ._version import __version__

anemoi/utils/_version.py

@@ -0,0 +1,16 @@
+# file generated by setuptools_scm
+# don't change, don't track in version control
+TYPE_CHECKING = False
+if TYPE_CHECKING:
+    from typing import Tuple, Union
+    VERSION_TUPLE = Tuple[Union[int, str], ...]
+else:
+    VERSION_TUPLE = object
+
+version: str
+__version__: str
+__version_tuple__: VERSION_TUPLE
+version_tuple: VERSION_TUPLE
+
+__version__ = version = '0.1.6'
+__version_tuple__ = version_tuple = (0, 1, 6)

anemoi/utils/checkpoints.py

@@ -0,0 +1,76 @@
+# (C) Copyright 2024 ECMWF.
+#
+# 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.
+
+"""
+Read and write extra metadata in PyTorch checkpoints files. These files
+are zip archives containing the model weights.
+"""
+
+import json
+import logging
+import os
+import zipfile
+
+LOG = logging.getLogger(__name__)
+
+DEFAULT_NAME = "anemoi-metadata.json"
+
+
+def load_metadata(path: str, name: str = DEFAULT_NAME):
+    """Load metadata from a checkpoint file
+
+    Parameters
+    ----------
+    path : str
+        The path to the checkpoint file
+    name : str, optional
+        The name of the metadata file in the zip archive
+
+    Returns
+    -------
+    JSON
+        The content of the metadata file
+
+    Raises
+    ------
+    ValueError
+        If the metadata file is not found
+    """
+    with zipfile.ZipFile(path, "r") as f:
+        metadata = None
+        for b in f.namelist():
+            if os.path.basename(b) == name:
+                if metadata is not None:
+                    LOG.warning(f"Found two '{name}' if {path}")
+                metadata = b
+
+    if metadata is not None:
+        with zipfile.ZipFile(path, "r") as f:
+            return json.load(f.open(metadata, "r"))
+    else:
+        raise ValueError(f"Could not find {name} in {path}")
+
+
+def save_metadata(path, metadata, name=DEFAULT_NAME):
+    """Save metadata to a checkpoint file
+
+    Parameters
+    ----------
+    path : str
+        The path to the checkpoint file
+    metadata : JSON
+        A JSON serializable object
+    name : str, optional
+        The name of the metadata file in the zip archive
+    """
+    with zipfile.ZipFile(path, "a") as zipf:
+        base, _ = os.path.splitext(os.path.basename(path))
+        zipf.writestr(
+            f"{base}/{name}",
+            json.dumps(metadata),
+        )

anemoi/utils/config.py

@@ -0,0 +1,94 @@
+# (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.
+
+
+import logging
+import os
+
+try:
+    import tomllib  # Only available since 3.11
+except ImportError:
+    import tomli as tomllib
+
+
+LOG = logging.getLogger(__name__)
+
+
+class DotDict(dict):
+    """A dictionary that allows access to its keys as attributes.
+
+    >>> d = DotDict({"a": 1, "b": {"c": 2}})
+    >>> d.a
+    1
+    >>> d.b.c
+    2
+    >>> d.b = 3
+    >>> d.b
+    3
+
+    The class is recursive, so nested dictionaries are also DotDicts.
+
+    The DotDict class has the same constructor as the dict class.
+
+    >>> d = DotDict(a=1, b=2)
+
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        for k, v in self.items():
+            if isinstance(v, dict):
+                self[k] = DotDict(v)
+
+    def __getattr__(self, attr):
+        try:
+            return self[attr]
+        except KeyError:
+            raise AttributeError(attr)
+
+    def __setattr__(self, attr, value):
+        if isinstance(value, dict):
+            value = DotDict(value)
+        self[attr] = value
+
+    def __repr__(self) -> str:
+        return f"DotDict({super().__repr__()})"
+
+
+CONFIG = None
+
+
+def load_config():
+    """Load the configuration from `~/.anemoi.toml`.
+
+    Returns
+    -------
+    DotDict
+        The configuration
+    """
+    global CONFIG
+    if CONFIG is not None:
+        return CONFIG
+
+    conf = os.path.expanduser("~/.anemoi.toml")
+
+    if os.path.exists(conf):
+
+        with open(conf, "rb") as f:
+            CONFIG = tomllib.load(f)
+    else:
+        CONFIG = {}
+
+    return DotDict(CONFIG)
+
+
+def save_config():
+    """Save the configuration to `~/.anemoi.toml`."""
+
+    conf = os.path.expanduser("~/.anemoi.toml")
+    with open(conf, "w") as f:
+        tomllib.dump(CONFIG, f)

anemoi/utils/dates.py

@@ -0,0 +1,248 @@
+# (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.
+
+
+import calendar
+import datetime
+
+
+def no_time_zone(date):
+    """Remove time zone information from a date.
+
+    Parameters
+    ----------
+    date : datetime.datetime
+        A datetime object.
+
+    Returns
+    -------
+    datetime.datetime
+        The datetime object without time zone information.
+    """
+
+    return date.replace(tzinfo=None)
+
+
+def as_datetime(date):
+    """Convert a date to a datetime object, removing any time zone information.
+
+    Parameters
+    ----------
+    date : datetime.date or datetime.datetime or str
+        The date to convert.
+
+    Returns
+    -------
+    datetime.datetime
+        The datetime object.
+    """
+
+    if isinstance(date, datetime.datetime):
+        return no_time_zone(date)
+
+    if isinstance(date, datetime.date):
+        return no_time_zone(datetime.datetime(date.year, date.month, date.day))
+
+    if isinstance(date, str):
+        return no_time_zone(datetime.datetime.fromisoformat(date))
+
+    raise ValueError(f"Invalid date type: {type(date)}")
+
+
+DOW = {
+    "monday": 0,
+    "tuesday": 1,
+    "wednesday": 2,
+    "thursday": 3,
+    "friday": 4,
+    "saturday": 5,
+    "sunday": 6,
+}
+
+
+MONTH = {
+    "january": 1,
+    "february": 2,
+    "march": 3,
+    "april": 4,
+    "may": 5,
+    "june": 6,
+    "july": 7,
+    "august": 8,
+    "september": 9,
+    "october": 10,
+    "november": 11,
+    "december": 12,
+}
+
+
+def _make_day(day):
+    if day is None:
+        return set(range(1, 32))
+    if not isinstance(day, list):
+        day = [day]
+    return {int(d) for d in day}
+
+
+def _make_week(week):
+    if week is None:
+        return set(range(7))
+    if not isinstance(week, list):
+        week = [week]
+    return {DOW[w.lower()] for w in week}
+
+
+def _make_months(months):
+    if months is None:
+        return set(range(1, 13))
+
+    if not isinstance(months, list):
+        months = [months]
+
+    return {int(MONTH.get(m, m)) for m in months}
+
+
+class DateTimes:
+    """The DateTimes class is an iterator that generates datetime objects within a given range."""
+
+    def __init__(self, start, end, increment=24, *, day_of_month=None, day_of_week=None, calendar_months=None):
+        """_summary_
+
+        Parameters
+        ----------
+        start : _type_
+            _description_
+        end : _type_
+            _description_
+        increment : int, optional
+            _description_, by default 24
+        day_of_month : _type_, optional
+            _description_, by default None
+        day_of_week : _type_, optional
+            _description_, by default None
+        calendar_months : _type_, optional
+            _description_, by default None
+        """
+        self.start = as_datetime(start)
+        self.end = as_datetime(end)
+        self.increment = datetime.timedelta(hours=increment)
+        self.day_of_month = _make_day(day_of_month)
+        self.day_of_week = _make_week(day_of_week)
+        self.calendar_months = _make_months(calendar_months)
+
+    def __iter__(self):
+        date = self.start
+        while date <= self.end:
+            if (
+                (date.weekday() in self.day_of_week)
+                and (date.day in self.day_of_month)
+                and (date.month in self.calendar_months)
+            ):
+
+                yield date
+            date += self.increment
+
+
+class HindcastDatesTimes:
+    """The HindcastDatesTimes class is an iterator that generates datetime objects within a given range."""
+
+    def __init__(self, reference_dates, years=20):
+        """_summary_
+
+        Parameters
+        ----------
+        reference_dates : _type_
+            _description_
+        years : int, optional
+            _description_, by default 20
+        """
+
+        self.reference_dates = reference_dates
+        self.years = (1, years + 1)
+
+    def __iter__(self):
+        for reference_date in self.reference_dates:
+            for year in range(*self.years):
+                if reference_date.month == 2 and reference_date.day == 29:
+                    date = datetime.datetime(reference_date.year - year, 2, 28)
+                else:
+                    date = datetime.datetime(reference_date.year - year, reference_date.month, reference_date.day)
+                yield (date, reference_date)
+
+
+class Year(DateTimes):
+    """Year is defined as the months of January to December."""
+
+    def __init__(self, year, **kwargs):
+        """_summary_
+
+        Parameters
+        ----------
+        year : int
+            _description_
+        """
+        super().__init__(datetime.datetime(year, 1, 1), datetime.datetime(year, 12, 31), **kwargs)
+
+
+class Winter(DateTimes):
+    """Winter is defined as the months of December, January and February."""
+
+    def __init__(self, year, **kwargs):
+        """_summary_
+
+        Parameters
+        ----------
+        year : int
+            _description_
+        """
+        super().__init__(
+            datetime.datetime(year, 12, 1),
+            datetime.datetime(year + 1, 2, calendar.monthrange(year + 1, 2)[1]),
+            **kwargs,
+        )
+
+
+class Spring(DateTimes):
+    """Spring is defined as the months of March, April and May."""
+
+    def __init__(self, year, **kwargs):
+        """_summary_
+
+        Parameters
+        ----------
+        year : int
+            _description_
+        """
+        super().__init__(datetime.datetime(year, 3, 1), datetime.datetime(year, 5, 31), **kwargs)
+
+
+class Summer(DateTimes):
+    """Summer is defined as the months of June, July and August."""
+
+    def __init__(self, year, **kwargs):
+        """_summary_
+
+        Parameters
+        ----------
+        year : int
+            _description_
+        """
+        super().__init__(datetime.datetime(year, 6, 1), datetime.datetime(year, 8, 31), **kwargs)
+
+
+class Autumn(DateTimes):
+    """Autumn is defined as the months of September, October and November."""
+
+    def __init__(self, year, **kwargs):
+        """_summary_
+
+        Parameters
+        ----------
+        year : int
+            _description_
+        """
+        super().__init__(datetime.datetime(year, 9, 1), datetime.datetime(year, 11, 30), **kwargs)

anemoi/utils/grib.py

@@ -0,0 +1,73 @@
+# (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 GRIB parameters.
+
+See https://codes.ecmwf.int/grib/param-db/ for more information.
+
+"""
+
+import re
+
+import requests
+
+
+def _search(name):
+    name = re.escape(name)
+    r = requests.get(f"https://codes.ecmwf.int/parameter-database/api/v1/param/?search=^{name}$&regex=true")
+    r.raise_for_status()
+    results = r.json()
+    if len(results) == 0:
+        raise KeyError(name)
+
+    if len(results) > 1:
+        names = [f'{r.get("id")} ({r.get("name")})' for r in results]
+        raise ValueError(f"{name} is ambiguous: {', '.join(names)}")
+
+    return results[0]
+
+
+def shortname_to_paramid(shortname: str) -> int:
+    """Return the GRIB parameter id given its shortname.
+
+    Parameters
+    ----------
+    shortname : str
+        Parameter shortname.
+
+    Returns
+    -------
+    int
+        Parameter id.
+
+
+    >>> shortname_to_paramid("2t")
+    167
+
+    """
+    return _search(shortname)["id"]
+
+
+def paramid_to_shortname(paramid: int) -> str:
+    """Return the shortname of a GRIB parameter given its id.
+
+    Parameters
+    ----------
+    paramid : int
+        Parameter id.
+
+    Returns
+    -------
+    str
+        Parameter shortname.
+
+
+    >>> paramid_to_shortname(167)
+    '2t'
+
+    """
+    return _search(str(paramid))["shortname"]