anemoi-datasets 0.5.15__py3-none-any.whl → 0.5.17__py3-none-any.whl

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

@@ -1,472 +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 logging
-
-from .coordinates import DateCoordinate
-from .coordinates import EnsembleCoordinate
-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 UnsupportedCoordinate
-from .coordinates import XCoordinate
-from .coordinates import YCoordinate
-from .coordinates import is_scalar
-from .grid import MeshedGrid
-from .grid import MeshProjectionGrid
-from .grid import UnstructuredGrid
-from .grid import UnstructuredProjectionGrid
-
-LOG = logging.getLogger(__name__)
-
-
-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
-
-        d = self._is_number(
-            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)
-
-        LOG.warning(
-            f"Coordinate {coord} not supported\n{axis=}, {name=},"
-            f" {long_name=}, {standard_name=}, units\n\n{c}\n\n{type(c.values)} {c.shape}"
-        )
-
-        return UnsupportedCoordinate(c)
-
-    def grid(self, coordinates, variable):
-        lat = [c for c in coordinates if c.is_lat]
-        lon = [c for c in coordinates if c.is_lon]
-
-        if len(lat) == 1 and len(lon) == 1:
-            return self._lat_lon_provided(lat, lon, variable)
-
-        x = [c for c in coordinates if c.is_x]
-        y = [c for c in coordinates if c.is_y]
-
-        if len(x) == 1 and len(y) == 1:
-            return self._x_y_provided(x, y, variable)
-
-        raise NotImplementedError(f"Cannot establish grid {coordinates}")
-
-    def _check_dims(self, variable, x_or_lon, y_or_lat):
-
-        x_dims = set(x_or_lon.variable.dims)
-        y_dims = set(y_or_lat.variable.dims)
-        variable_dims = set(variable.dims)
-
-        if not (x_dims <= variable_dims) or not (y_dims <= variable_dims):
-            raise ValueError(
-                f"Dimensions do not match {variable.name}{variable.dims} !="
-                f" {x_or_lon.name}{x_or_lon.variable.dims} and {y_or_lat.name}{y_or_lat.variable.dims}"
-            )
-
-        variable_dims = tuple(v for v in variable.dims if v in (x_dims | y_dims))
-        if x_dims == y_dims:
-            # It's unstructured
-            return variable_dims, True
-
-        if len(x_dims) == 1 and len(y_dims) == 1:
-            # It's a mesh
-            return variable_dims, False
-
-        raise ValueError(
-            f"Cannot establish grid for {variable.name}{variable.dims},"
-            f" {x_or_lon.name}{x_or_lon.variable.dims},"
-            f" {y_or_lat.name}{y_or_lat.variable.dims}"
-        )
-
-    def _lat_lon_provided(self, lat, lon, variable):
-        lat = lat[0]
-        lon = lon[0]
-
-        dim_vars, unstructured = self._check_dims(variable, lon, lat)
-
-        if (lat.name, lon.name, dim_vars) in self._cache:
-            return self._cache[(lat.name, lon.name, dim_vars)]
-
-        if unstructured:
-            grid = UnstructuredGrid(lat, lon, dim_vars)
-        else:
-            grid = MeshedGrid(lat, lon, dim_vars)
-
-        self._cache[(lat.name, lon.name, dim_vars)] = grid
-        return grid
-
-    def _x_y_provided(self, x, y, variable):
-        x = x[0]
-        y = y[0]
-
-        dim_vars, unstructured = self._check_dims(variable, x, y)
-
-        if (x.name, y.name, dim_vars) in self._cache:
-            return self._cache[(x.name, y.name, dim_vars)]
-
-        grid_mapping = variable.attrs.get("grid_mapping", None)
-        if grid_mapping is not None:
-            print(f"grid_mapping: {grid_mapping}")
-            print(self.ds[grid_mapping])
-
-        if grid_mapping is None:
-            LOG.warning(f"No 'grid_mapping' attribute provided for '{variable.name}'")
-            LOG.warning("Trying to guess...")
-
-            PROBE = {
-                "prime_meridian_name",
-                "reference_ellipsoid_name",
-                "crs_wkt",
-                "horizontal_datum_name",
-                "semi_major_axis",
-                "spatial_ref",
-                "inverse_flattening",
-                "semi_minor_axis",
-                "geographic_crs_name",
-                "GeoTransform",
-                "grid_mapping_name",
-                "longitude_of_prime_meridian",
-            }
-            candidate = None
-            for v in self.ds.variables:
-                var = self.ds[v]
-                if not is_scalar(var):
-                    continue
-
-                if PROBE.intersection(var.attrs.keys()):
-                    if candidate:
-                        raise ValueError(f"Multiple candidates for 'grid_mapping': {candidate} and {v}")
-                    candidate = v
-
-            if candidate:
-                LOG.warning(f"Using '{candidate}' as 'grid_mapping'")
-                grid_mapping = candidate
-            else:
-                LOG.warning("Could not fine a candidate for 'grid_mapping'")
-
-        if grid_mapping is None:
-            if "crs" in self.ds[variable].attrs:
-                grid_mapping = self.ds[variable].attrs["crs"]
-                LOG.warning(f"Using CRS {grid_mapping} from variable '{variable.name}' attributes")
-
-        if grid_mapping is None:
-            if "crs" in self.ds.attrs:
-                grid_mapping = self.ds.attrs["crs"]
-                LOG.warning(f"Using CRS {grid_mapping} from global attributes")
-
-        grid = None
-        if grid_mapping is not None:
-
-            grid_mapping = dict(self.ds[grid_mapping].attrs)
-
-            if unstructured:
-                grid = UnstructuredProjectionGrid(x, y, grid_mapping)
-            else:
-                grid = MeshProjectionGrid(x, y, grid_mapping)
-
-        if grid is not None:
-            self._cache[(x.name, y.name, dim_vars)] = grid
-            return grid
-
-        LOG.error("Could not fine a candidate for 'grid_mapping'")
-        raise NotImplementedError(f"Unstructured grid {x.name} {y.name}")
-
-
-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)
-
-        # That is the output of `cfgrib` for forecasts
-        if name == "time" and standard_name != "forecast_reference_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 == "vertical" and units == "hPa":
-            return LevelCoordinate(c, "pl")
-
-    def _is_number(self, c, *, axis, name, long_name, standard_name, units):
-        if name in ("realization", "number"):
-            return EnsembleCoordinate(c)
-
-
-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=}")
-
-    def _is_number(self, c, *, axis, name, long_name, standard_name, units):
-        if self._match(c, "number", locals()):
-            return DateCoordinate(c)

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

@@ -1,148 +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 logging
-from functools import cached_property
-
-from anemoi.utils.dates import as_datetime
-from earthkit.data.core.geography import Geography
-from earthkit.data.core.metadata import RawMetadata
-from earthkit.data.utils.projections import Projection
-
-LOG = logging.getLogger(__name__)
-
-
-class _MDMapping:
-
-    def __init__(self, variable):
-        self.variable = variable
-        self.time = variable.time
-        # Aliases
-        self.mapping = dict(param="variable")
-        for c in variable.coordinates:
-            for v in c.mars_names:
-                assert v not in self.mapping, f"Duplicate key '{v}' in {c}"
-                self.mapping[v] = c.variable.name
-
-    def _from_user(self, key):
-        return self.mapping.get(key, key)
-
-    def from_user(self, kwargs):
-        return {self._from_user(k): v for k, v in kwargs.items()}
-
-    def __repr__(self):
-        return f"MDMapping({self.mapping})"
-
-    def fill_time_metadata(self, field, md):
-        valid_datetime = self.variable.time.fill_time_metadata(field._md, md)
-        if valid_datetime is not None:
-            md["valid_datetime"] = as_datetime(valid_datetime).isoformat()
-
-
-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):
-        self._field = field
-        md = field._md.copy()
-        self._mapping = _MDMapping(field.owner)
-        self._mapping.fill_time_metadata(field, md)
-        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 {}
-
-    def _base_datetime(self):
-        return self._field.forecast_reference_time
-
-    def _valid_datetime(self):
-        return self._get("valid_datetime")
-
-    def get(self, key, astype=None, **kwargs):
-
-        if key in self._d:
-            if astype is not None:
-                return astype(self._d[key])
-            return self._d[key]
-
-        key = self._mapping._from_user(key)
-
-        return super().get(key, astype=astype, **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/patch.py

@@ -1,44 +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 logging
-
-LOG = logging.getLogger(__name__)
-
-
-def patch_attributes(ds, attributes):
-    for name, value in attributes.items():
-        variable = ds[name]
-        variable.attrs.update(value)
-
-    return ds
-
-
-def patch_coordinates(ds, coordinates):
-    for name in coordinates:
-        ds = ds.assign_coords({name: ds[name]})
-
-    return ds
-
-
-PATCHES = {
-    "attributes": patch_attributes,
-    "coordinates": patch_coordinates,
-}
-
-
-def patch_dataset(ds, patch):
-    for what, values in patch.items():
-        if what not in PATCHES:
-            raise ValueError(f"Unknown patch type {what!r}")
-
-        ds = PATCHES[what](ds, values)
-
-    return ds