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

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

@@ -0,0 +1,171 @@
+# (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 json
+import logging
+
+import yaml
+from earthkit.data.core.fieldlist import FieldList
+
+from .coordinates import is_scalar as is_scalar
+from .field import EmptyFieldList
+from .flavour import CoordinateGuesser
+from .metadata import XArrayMetadata as XArrayMetadata
+from .time import Time
+from .variable import FilteredVariable
+from .variable import Variable
+
+LOG = logging.getLogger(__name__)
+
+
+class XarrayFieldList(FieldList):
+    def __init__(self, ds, variables):
+        self.ds = ds
+        self.variables = variables.copy()
+        self.total_length = sum(v.length for v in variables)
+
+    def __repr__(self):
+        return f"XarrayFieldList({self.total_length})"
+
+    def __len__(self):
+        return self.total_length
+
+    def __getitem__(self, i):
+        k = i
+
+        if i < 0:
+            i = self.total_length + i
+
+        for v in self.variables:
+            if i < v.length:
+                return v[i]
+            i -= v.length
+
+        raise IndexError(k)
+
+    @classmethod
+    def from_xarray(cls, ds, flavour=None):
+        variables = []
+
+        if isinstance(flavour, str):
+            with open(flavour) as f:
+                if flavour.endswith(".yaml") or flavour.endswith(".yml"):
+                    flavour = yaml.safe_load(f)
+                else:
+                    flavour = json.load(f)
+
+        guess = CoordinateGuesser.from_flavour(ds, flavour)
+
+        skip = set()
+
+        def _skip_attr(v, attr_name):
+            attr_val = getattr(v, attr_name, "")
+            if isinstance(attr_val, str):
+                skip.update(attr_val.split(" "))
+
+        for name in ds.data_vars:
+            v = ds[name]
+            _skip_attr(v, "coordinates")
+            _skip_attr(v, "bounds")
+            _skip_attr(v, "grid_mapping")
+
+        # Select only geographical variables
+        for name in ds.data_vars:
+
+            if name in skip:
+                continue
+
+            v = ds[name]
+            coordinates = []
+
+            for coord in v.coords:
+
+                c = guess.guess(ds[coord], coord)
+                assert c, f"Could not guess coordinate for {coord}"
+                if coord not in v.dims:
+                    c.is_dim = False
+                coordinates.append(c)
+
+            grid_coords = sum(1 for c in coordinates if c.is_grid and c.is_dim)
+            assert grid_coords <= 2
+
+            if grid_coords < 2:
+                continue
+
+            variables.append(
+                Variable(
+                    ds=ds,
+                    var=v,
+                    coordinates=coordinates,
+                    grid=guess.grid(coordinates),
+                    time=Time.from_coordinates(coordinates),
+                    metadata={},
+                )
+            )
+
+        return cls(ds, variables)
+
+    def sel(self, **kwargs):
+        """Override the FieldList's sel method
+
+        Returns
+        -------
+        FieldList
+            The new FieldList
+
+        The algorithm is as follows:
+        1 - Use the kwargs to select the variables that match the selection (`param` or `variable`)
+        2 - For each variable, use the remaining kwargs to select the coordinates (`level`, `number`, ...)
+        3 - Some mars like keys, like `date`, `time`, `step` are not found in the coordinates,
+            but added to the metadata of the selected fields. A example is `step` that is added to the
+            metadata of the field. Step 2 may return a variable that contain all the fields that
+            verify at the same `valid_datetime`, with different base `date` and `time` and a different `step`.
+            So we get an extra chance to filter the fields by the metadata.
+        """
+
+        variables = []
+        count = 0
+
+        for v in self.variables:
+
+            v.update_metadata_mapping(kwargs)
+
+            # First, select matching variables
+            # This will consume 'param' or 'variable' from kwargs
+            # and return the rest
+            match, rest = v.match(**kwargs)
+
+            if match:
+                count += 1
+                missing = {}
+
+                # Select from the variable's coordinates (time, level, number, ....)
+                # This may return a new variable with a isel() slice of the selection
+                # or None if the selection is not possible. In this case, missing is updated
+                # with the values of kwargs (rest) that are not relevant for this variable
+                v = v.sel(missing, **rest)
+                if missing:
+                    if v is not None:
+                        # The remaining kwargs are passed used to create a FilteredVariable
+                        # that will select 2D slices based on their metadata
+                        v = FilteredVariable(v, **missing)
+                    else:
+                        LOG.warning(f"Variable {v} has missing coordinates: {missing}")
+
+                if v is not None:
+                    variables.append(v)
+
+        if count == 0:
+            LOG.warning("No variable found for %s", kwargs)
+            LOG.warning("Variables: %s", sorted([v.name for v in self.variables]))
+
+        if not variables:
+            return EmptyFieldList()
+
+        return self.__class__(self.ds, variables)

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)