anemoi-datasets 0.3.10__py3-none-any.whl → 0.4.2__py3-none-any.whl

anemoi/datasets/create/functions/sources/xarray/flavour.py

@@ -0,0 +1,330 @@
+# (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.
+#
+
+
+from .coordinates import DateCoordinate
+from .coordinates import LatitudeCoordinate
+from .coordinates import LevelCoordinate
+from .coordinates import LongitudeCoordinate
+from .coordinates import ScalarCoordinate
+from .coordinates import StepCoordinate
+from .coordinates import TimeCoordinate
+from .coordinates import XCoordinate
+from .coordinates import YCoordinate
+from .grid import MeshedGrid
+from .grid import UnstructuredGrid
+
+
+class CoordinateGuesser:
+
+    def __init__(self, ds):
+        self.ds = ds
+        self._cache = {}
+
+    @classmethod
+    def from_flavour(cls, ds, flavour):
+        if flavour is None:
+            return DefaultCoordinateGuesser(ds)
+        else:
+            return FlavourCoordinateGuesser(ds, flavour)
+
+    def guess(self, c, coord):
+        if coord not in self._cache:
+            self._cache[coord] = self._guess(c, coord)
+        return self._cache[coord]
+
+    def _guess(self, c, coord):
+
+        name = c.name
+        standard_name = getattr(c, "standard_name", "").lower()
+        axis = getattr(c, "axis", "")
+        long_name = getattr(c, "long_name", "").lower()
+        units = getattr(c, "units", "")
+
+        d = self._is_longitude(
+            c,
+            axis=axis,
+            name=name,
+            long_name=long_name,
+            standard_name=standard_name,
+            units=units,
+        )
+        if d is not None:
+            return d
+
+        d = self._is_latitude(
+            c,
+            axis=axis,
+            name=name,
+            long_name=long_name,
+            standard_name=standard_name,
+            units=units,
+        )
+        if d is not None:
+            return d
+
+        d = self._is_x(
+            c,
+            axis=axis,
+            name=name,
+            long_name=long_name,
+            standard_name=standard_name,
+            units=units,
+        )
+        if d is not None:
+            return d
+
+        d = self._is_y(
+            c,
+            axis=axis,
+            name=name,
+            long_name=long_name,
+            standard_name=standard_name,
+            units=units,
+        )
+        if d is not None:
+            return d
+
+        d = self._is_time(
+            c,
+            axis=axis,
+            name=name,
+            long_name=long_name,
+            standard_name=standard_name,
+            units=units,
+        )
+        if d is not None:
+            return d
+
+        d = self._is_step(
+            c,
+            axis=axis,
+            name=name,
+            long_name=long_name,
+            standard_name=standard_name,
+            units=units,
+        )
+        if d is not None:
+            return d
+
+        d = self._is_date(
+            c,
+            axis=axis,
+            name=name,
+            long_name=long_name,
+            standard_name=standard_name,
+            units=units,
+        )
+        if d is not None:
+            return d
+
+        d = self._is_level(
+            c,
+            axis=axis,
+            name=name,
+            long_name=long_name,
+            standard_name=standard_name,
+            units=units,
+        )
+        if d is not None:
+            return d
+
+        if c.shape in ((1,), tuple()):
+            return ScalarCoordinate(c)
+
+        raise NotImplementedError(
+            f"Coordinate {coord} not supported\n{axis=}, {name=},"
+            f" {long_name=}, {standard_name=}, units\n\n{c}\n\n{type(c.values)} {c.shape}"
+        )
+
+    def grid(self, coordinates):
+        lat = [c for c in coordinates if c.is_lat]
+        lon = [c for c in coordinates if c.is_lon]
+
+        if len(lat) != 1:
+            raise NotImplementedError(f"Expected 1 latitude coordinate, got {len(lat)}")
+
+        if len(lon) != 1:
+            raise NotImplementedError(f"Expected 1 longitude coordinate, got {len(lon)}")
+
+        lat = lat[0]
+        lon = lon[0]
+
+        if (lat.name, lon.name) in self._cache:
+            return self._cache[(lat.name, lon.name)]
+
+        assert len(lat.variable.shape) == len(lon.variable.shape), (lat.variable.shape, lon.variable.shape)
+        if len(lat.variable.shape) == 1:
+            grid = MeshedGrid(lat, lon)
+        else:
+            grid = UnstructuredGrid(lat, lon)
+
+        self._cache[(lat.name, lon.name)] = grid
+        return grid
+
+
+class DefaultCoordinateGuesser(CoordinateGuesser):
+    def __init__(self, ds):
+        super().__init__(ds)
+
+    def _is_longitude(self, c, *, axis, name, long_name, standard_name, units):
+        if standard_name == "longitude":
+            return LongitudeCoordinate(c)
+
+        if long_name == "longitude" and units == "degrees_east":
+            return LongitudeCoordinate(c)
+
+        if name == "longitude":  # WeatherBench
+            return LongitudeCoordinate(c)
+
+    def _is_latitude(self, c, *, axis, name, long_name, standard_name, units):
+        if standard_name == "latitude":
+            return LatitudeCoordinate(c)
+
+        if long_name == "latitude" and units == "degrees_north":
+            return LatitudeCoordinate(c)
+
+        if name == "latitude":  # WeatherBench
+            return LatitudeCoordinate(c)
+
+    def _is_x(self, c, *, axis, name, long_name, standard_name, units):
+        if standard_name == "projection_x_coordinate":
+            return XCoordinate(c)
+
+        if name == "x":
+            return XCoordinate(c)
+
+    def _is_y(self, c, *, axis, name, long_name, standard_name, units):
+        if standard_name == "projection_y_coordinate":
+            return YCoordinate(c)
+
+        if name == "y":
+            return YCoordinate(c)
+
+    def _is_time(self, c, *, axis, name, long_name, standard_name, units):
+        if standard_name == "time":
+            return TimeCoordinate(c)
+
+        if name == "time":
+            return TimeCoordinate(c)
+
+    def _is_date(self, c, *, axis, name, long_name, standard_name, units):
+        if standard_name == "forecast_reference_time":
+            return DateCoordinate(c)
+        if name == "forecast_reference_time":
+            return DateCoordinate(c)
+
+    def _is_step(self, c, *, axis, name, long_name, standard_name, units):
+        if standard_name == "forecast_period":
+            return StepCoordinate(c)
+
+        if long_name == "time elapsed since the start of the forecast":
+            return StepCoordinate(c)
+
+        if name == "prediction_timedelta":  # WeatherBench
+            return StepCoordinate(c)
+
+    def _is_level(self, c, *, axis, name, long_name, standard_name, units):
+        if standard_name == "atmosphere_hybrid_sigma_pressure_coordinate":
+            return LevelCoordinate(c, "ml")
+
+        if long_name == "height" and units == "m":
+            return LevelCoordinate(c, "height")
+
+        if standard_name == "air_pressure" and units == "hPa":
+            return LevelCoordinate(c, "pl")
+
+        if name == "level":
+            return LevelCoordinate(c, "pl")
+
+        if name == "vertical" and units == "hPa":
+            return LevelCoordinate(c, "pl")
+
+        if standard_name == "depth":
+            return LevelCoordinate(c, "depth")
+
+        if name == "pressure":
+            return LevelCoordinate(c, "pl")
+
+
+class FlavourCoordinateGuesser(CoordinateGuesser):
+    def __init__(self, ds, flavour):
+        super().__init__(ds)
+        self.flavour = flavour
+
+    def _match(self, c, key, values):
+
+        if key not in self.flavour["rules"]:
+            return None
+
+        rules = self.flavour["rules"][key]
+
+        if not isinstance(rules, list):
+            rules = [rules]
+
+        for rule in rules:
+            ok = True
+            for k, v in rule.items():
+                if isinstance(v, str) and values.get(k) != v:
+                    ok = False
+            if ok:
+                return rule
+
+        return None
+
+    def _is_longitude(self, c, *, axis, name, long_name, standard_name, units):
+        if self._match(c, "longitude", locals()):
+            return LongitudeCoordinate(c)
+
+    def _is_latitude(self, c, *, axis, name, long_name, standard_name, units):
+        if self._match(c, "latitude", locals()):
+            return LatitudeCoordinate(c)
+
+    def _is_x(self, c, *, axis, name, long_name, standard_name, units):
+        if self._match(c, "x", locals()):
+            return XCoordinate(c)
+
+    def _is_y(self, c, *, axis, name, long_name, standard_name, units):
+        if self._match(c, "y", locals()):
+            return YCoordinate(c)
+
+    def _is_time(self, c, *, axis, name, long_name, standard_name, units):
+        if self._match(c, "time", locals()):
+            return TimeCoordinate(c)
+
+    def _is_step(self, c, *, axis, name, long_name, standard_name, units):
+        if self._match(c, "step", locals()):
+            return StepCoordinate(c)
+
+    def _is_date(self, c, *, axis, name, long_name, standard_name, units):
+        if self._match(c, "date", locals()):
+            return DateCoordinate(c)
+
+    def _is_level(self, c, *, axis, name, long_name, standard_name, units):
+
+        rule = self._match(c, "level", locals())
+        if rule:
+            # assert False, rule
+            return LevelCoordinate(
+                c,
+                self._levtype(
+                    c,
+                    axis=axis,
+                    name=name,
+                    long_name=long_name,
+                    standard_name=standard_name,
+                    units=units,
+                ),
+            )
+
+    def _levtype(self, c, *, axis, name, long_name, standard_name, units):
+        if "levtype" in self.flavour:
+            return self.flavour["levtype"]
+
+        raise NotImplementedError(f"levtype for {c=}")

anemoi/datasets/create/functions/sources/xarray/grid.py

@@ -0,0 +1,46 @@
+# (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.
+#
+
+
+import numpy as np
+
+
+class Grid:
+    def __init__(self, lat, lon):
+        self.lat = lat
+        self.lon = lon
+
+    @property
+    def latitudes(self):
+        return self.grid_points()[0]
+
+    @property
+    def longitudes(self):
+        return self.grid_points()[1]
+
+
+class MeshedGrid(Grid):
+    _cache = None
+
+    def grid_points(self):
+        if self._cache is not None:
+            return self._cache
+        lat = self.lat.variable.values
+        lon = self.lon.variable.values
+
+        lat, lon = np.meshgrid(lat, lon)
+        self._cache = (lat.flatten(), lon.flatten())
+        return self._cache
+
+
+class UnstructuredGrid(Grid):
+    def grid_points(self):
+        lat = self.lat.variable.values.flatten()
+        lon = self.lon.variable.values.flatten()
+        return lat, lon

anemoi/datasets/create/functions/sources/xarray/metadata.py

@@ -0,0 +1,161 @@
+# (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.
+#
+
+import logging
+from functools import cached_property
+
+from earthkit.data.core.geography import Geography
+from earthkit.data.core.metadata import RawMetadata
+from earthkit.data.utils.dates import to_datetime
+from earthkit.data.utils.projections import Projection
+
+LOG = logging.getLogger(__name__)
+
+
+class MDMapping:
+
+    def __init__(self, mapping):
+        self.user_to_internal = mapping
+
+    def from_user(self, kwargs):
+        if isinstance(kwargs, str):
+            return self.user_to_internal.get(kwargs, kwargs)
+        return {self.user_to_internal.get(k, k): v for k, v in kwargs.items()}
+
+    def __len__(self):
+        return len(self.user_to_internal)
+
+    def __repr__(self):
+        return f"MDMapping({self.user_to_internal})"
+
+
+class XArrayMetadata(RawMetadata):
+    LS_KEYS = ["variable", "level", "valid_datetime", "units"]
+    NAMESPACES = ["default", "mars"]
+    MARS_KEYS = ["param", "step", "levelist", "levtype", "number", "date", "time"]
+
+    def __init__(self, field, mapping):
+        self._field = field
+        md = field._md.copy()
+
+        self._mapping = mapping
+        if mapping is None:
+            time_coord = [c for c in field.owner.coordinates if c.is_time]
+            if len(time_coord) == 1:
+                time_key = time_coord[0].name
+            else:
+                time_key = "time"
+        else:
+            time_key = mapping.from_user("valid_datetime")
+        self._time = to_datetime(md.pop(time_key))
+        self._field.owner.time.fill_time_metadata(self._time, md)
+        md["valid_datetime"] = self._time.isoformat()
+
+        super().__init__(md)
+
+    @cached_property
+    def geography(self):
+        return XArrayFieldGeography(self._field, self._field.owner.grid)
+
+    def as_namespace(self, namespace=None):
+        if not isinstance(namespace, str) and namespace is not None:
+            raise TypeError("namespace must be a str or None")
+
+        if namespace == "default" or namespace == "" or namespace is None:
+            return dict(self)
+
+        elif namespace == "mars":
+            return self._as_mars()
+
+    def _as_mars(self):
+        return dict(
+            param=self["variable"],
+            step=self["step"],
+            levelist=self["level"],
+            levtype=self["levtype"],
+            number=self["number"],
+            date=self["date"],
+            time=self["time"],
+        )
+
+    def _base_datetime(self):
+        return self._field.forecast_reference_time
+
+    def _valid_datetime(self):
+        return self._time
+
+    def _get(self, key, **kwargs):
+
+        if key.startswith("mars."):
+            key = key[5:]
+            if key not in self.MARS_KEYS:
+                if kwargs.get("raise_on_missing", False):
+                    raise KeyError(f"Invalid key '{key}' in namespace='mars'")
+                else:
+                    return kwargs.get("default", None)
+
+        if self._mapping is not None:
+            key = self._mapping.from_user(key)
+
+        return super()._get(key, **kwargs)
+
+
+class XArrayFieldGeography(Geography):
+    def __init__(self, field, grid):
+        self._field = field
+        self._grid = grid
+
+    def _unique_grid_id(self):
+        raise NotImplementedError()
+
+    def bounding_box(self):
+        raise NotImplementedError()
+        # return BoundingBox(north=self.north, south=self.south, east=self.east, west=self.west)
+
+    def gridspec(self):
+        raise NotImplementedError()
+
+    def latitudes(self, dtype=None):
+        result = self._grid.latitudes
+        if dtype is not None:
+            return result.astype(dtype)
+        return result
+
+    def longitudes(self, dtype=None):
+        result = self._grid.longitudes
+        if dtype is not None:
+            return result.astype(dtype)
+        return result
+
+    def resolution(self):
+        # TODO: implement resolution
+        return None
+
+    @property
+    def mars_grid(self):
+        # TODO: implement mars_grid
+        return None
+
+    @property
+    def mars_area(self):
+        # TODO: code me
+        # return [self.north, self.west, self.south, self.east]
+        return None
+
+    def x(self, dtype=None):
+        raise NotImplementedError()
+
+    def y(self, dtype=None):
+        raise NotImplementedError()
+
+    def shape(self):
+        return self._field.shape
+
+    def projection(self):
+        return Projection.from_cf_grid_mapping(**self._field.grid_mapping)

anemoi/datasets/create/functions/sources/xarray/time.py

@@ -0,0 +1,98 @@
+# (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.
+#
+
+
+import datetime
+
+
+class Time:
+    @classmethod
+    def from_coordinates(cls, coordinates):
+        time_coordinate = [c for c in coordinates if c.is_time]
+        step_coordinate = [c for c in coordinates if c.is_step]
+        date_coordinate = [c for c in coordinates if c.is_date]
+
+        if len(date_coordinate) == 0 and len(time_coordinate) == 1 and len(step_coordinate) == 1:
+            return ForecasstFromValidTimeAndStep(step_coordinate[0])
+
+        if len(date_coordinate) == 0 and len(time_coordinate) == 1 and len(step_coordinate) == 0:
+            return Analysis()
+
+        if len(date_coordinate) == 0 and len(time_coordinate) == 0 and len(step_coordinate) == 0:
+            return Constant()
+
+        if len(date_coordinate) == 1 and len(time_coordinate) == 1 and len(step_coordinate) == 0:
+            return ForecastFromValidTimeAndBaseTime(date_coordinate[0])
+
+        if len(date_coordinate) == 1 and len(time_coordinate) == 0 and len(step_coordinate) == 1:
+            return ForecastFromBaseTimeAndDate(date_coordinate[0], step_coordinate[0])
+
+        raise NotImplementedError(f"{date_coordinate=} {time_coordinate=} {step_coordinate=}")
+
+
+class Constant(Time):
+
+    def fill_time_metadata(self, time, metadata):
+        metadata["date"] = time.strftime("%Y%m%d")
+        metadata["time"] = time.strftime("%H%M")
+        metadata["step"] = 0
+
+
+class Analysis(Time):
+
+    def fill_time_metadata(self, time, metadata):
+        metadata["date"] = time.strftime("%Y%m%d")
+        metadata["time"] = time.strftime("%H%M")
+        metadata["step"] = 0
+
+
+class ForecasstFromValidTimeAndStep(Time):
+    def __init__(self, step_coordinate):
+        self.step_name = step_coordinate.variable.name
+
+    def fill_time_metadata(self, time, metadata):
+        step = metadata.pop(self.step_name)
+        assert isinstance(step, datetime.timedelta)
+        base = time - step
+
+        hours = step.total_seconds() / 3600
+        assert int(hours) == hours
+
+        metadata["date"] = base.strftime("%Y%m%d")
+        metadata["time"] = base.strftime("%H%M")
+        metadata["step"] = int(hours)
+
+
+class ForecastFromValidTimeAndBaseTime(Time):
+    def __init__(self, date_coordinate):
+        self.date_coordinate = date_coordinate
+
+    def fill_time_metadata(self, time, metadata):
+
+        step = time - self.date_coordinate
+
+        hours = step.total_seconds() / 3600
+        assert int(hours) == hours
+
+        metadata["date"] = self.date_coordinate.single_value.strftime("%Y%m%d")
+        metadata["time"] = self.date_coordinate.single_value.strftime("%H%M")
+        metadata["step"] = int(hours)
+
+
+class ForecastFromBaseTimeAndDate(Time):
+    def __init__(self, date_coordinate, step_coordinate):
+        self.date_coordinate = date_coordinate
+        self.step_coordinate = step_coordinate
+
+    def fill_time_metadata(self, time, metadata):
+        metadata["date"] = time.strftime("%Y%m%d")
+        metadata["time"] = time.strftime("%H%M")
+        hours = metadata[self.step_coordinate.name].total_seconds() / 3600
+        assert int(hours) == hours
+        metadata["step"] = int(hours)