anemoi-datasets 0.5.0__py3-none-any.whl → 0.5.6__py3-none-any.whl

anemoi/datasets/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.5.0'
-__version_tuple__ = version_tuple = (0, 5, 0)
+__version__ = version = '0.5.6'
+__version_tuple__ = version_tuple = (0, 5, 6)

anemoi/datasets/commands/inspect.py

@@ -311,7 +311,7 @@ class Version:
                 print(f"🕰️  Dataset initialized {when(start)}.")
                 if built and latest:
                     speed = (latest - start) / built
-                    eta = datetime.datetime.utcnow() + speed * (total - built)
+                    eta = datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None) + speed * (total - built)
                     print(f"🏁 ETA {when(eta)}.")
         else:
             if latest:

anemoi/datasets/commands/publish.py

@@ -0,0 +1,30 @@
+import logging
+
+from . import Command
+
+LOG = logging.getLogger(__name__)
+
+
+class Publish(Command):
+    """Publish a dataset."""
+
+    # This is a command that is used to publish a dataset.
+    # it is a class, inheriting from Command.
+
+    internal = True
+    timestamp = True
+
+    def add_arguments(self, parser):
+        parser.add_argument("path", help="Path of the dataset to publish.")
+
+    def run(self, args):
+        try:
+            from anemoi.registry import publish_dataset
+        except ImportError:
+            LOG.error("anemoi-registry is not installed. Please install it to use this command.")
+            return
+
+        publish_dataset(args.path)
+
+
+command = Publish

anemoi/datasets/create/__init__.py

@@ -14,6 +14,7 @@ import os
 import time
 import uuid
 import warnings
+from copy import deepcopy
 from functools import cached_property
 
 import numpy as np
@@ -24,9 +25,11 @@ from anemoi.utils.dates import frequency_to_string
 from anemoi.utils.dates import frequency_to_timedelta
 from anemoi.utils.humanize import compress_dates
 from anemoi.utils.humanize import seconds_to_human
+from earthkit.data.core.order import build_remapping
 
 from anemoi.datasets import MissingDateError
 from anemoi.datasets import open_dataset
+from anemoi.datasets.create.input.trace import enable_trace
 from anemoi.datasets.create.persistent import build_storage
 from anemoi.datasets.data.misc import as_first_date
 from anemoi.datasets.data.misc import as_last_date
@@ -308,7 +311,6 @@ class HasElementForDataMixin:
 
 
 def build_input_(main_config, output_config):
-    from earthkit.data.core.order import build_remapping
 
     builder = build_input(
         main_config.input,
@@ -323,6 +325,43 @@ def build_input_(main_config, output_config):
     return builder
 
 
+def tidy_recipe(config: object):
+    """Remove potentially private information in the config"""
+    config = deepcopy(config)
+    if isinstance(config, (tuple, list)):
+        return [tidy_recipe(_) for _ in config]
+    if isinstance(config, (dict, DotDict)):
+        for k, v in config.items():
+            if k.startswith("_"):
+                config[k] = "*** REMOVED FOR SECURITY ***"
+            else:
+                config[k] = tidy_recipe(v)
+    if isinstance(config, str):
+        if config.startswith("_"):
+            return "*** REMOVED FOR SECURITY ***"
+        if config.startswith("s3://"):
+            return "*** REMOVED FOR SECURITY ***"
+        if config.startswith("gs://"):
+            return "*** REMOVED FOR SECURITY ***"
+        if config.startswith("http"):
+            return "*** REMOVED FOR SECURITY ***"
+        if config.startswith("ftp"):
+            return "*** REMOVED FOR SECURITY ***"
+        if config.startswith("file"):
+            return "*** REMOVED FOR SECURITY ***"
+        if config.startswith("ssh"):
+            return "*** REMOVED FOR SECURITY ***"
+        if config.startswith("scp"):
+            return "*** REMOVED FOR SECURITY ***"
+        if config.startswith("rsync"):
+            return "*** REMOVED FOR SECURITY ***"
+        if config.startswith("/"):
+            return "*** REMOVED FOR SECURITY ***"
+        if "@" in config:
+            return "*** REMOVED FOR SECURITY ***"
+    return config
+
+
 class Init(Actor, HasRegistryMixin, HasStatisticTempMixin, HasElementForDataMixin):
     dataset_class = NewDataset
     def __init__(self, path, config, check_name=False, overwrite=False, use_threads=False, statistics_temp_dir=None, progress=None, test=False, cache=None, **kwargs):  # fmt: skip
@@ -409,6 +448,7 @@ class Init(Actor, HasRegistryMixin, HasStatisticTempMixin, HasElementForDataMixi
         metadata.update(self.main_config.get("add_metadata", {}))
 
         metadata["_create_yaml_config"] = self.main_config.get_serialisable_dict()
+        metadata["recipe"] = tidy_recipe(self.main_config.get_serialisable_dict())
 
         metadata["description"] = self.main_config.description
         metadata["licence"] = self.main_config["licence"]
@@ -524,7 +564,7 @@ class Load(Actor, HasRegistryMixin, HasStatisticTempMixin, HasElementForDataMixi
             # assert isinstance(group[0], datetime.datetime), type(group[0])
             LOG.debug(f"Building data for group {igroup}/{self.n_groups}")
 
-            result = self.input.select(dates=group)
+            result = self.input.select(group_of_dates=group)
             assert result.group_of_dates == group, (len(result.group_of_dates), len(group), group)
 
             # There are several groups.
@@ -992,7 +1032,6 @@ def chain(tasks):
 
 def creator_factory(name, trace=None, **kwargs):
     if trace:
-        from anemoi.datasets.create.trace import enable_trace
 
         enable_trace(trace)
 

anemoi/datasets/create/check.py

@@ -140,9 +140,15 @@ class StatisticsValueError(ValueError):
 
 def check_data_values(arr, *, name: str, log=[], allow_nans=False):
 
+    shape = arr.shape
+
     if (isinstance(allow_nans, (set, list, tuple, dict)) and name in allow_nans) or allow_nans:
         arr = arr[~np.isnan(arr)]
 
+    if arr.size == 0:
+        warnings.warn(f"Empty array for {name} ({shape})")
+        return
+
     assert arr.size > 0, (name, *log)
 
     min, max = arr.min(), arr.max()

anemoi/datasets/create/functions/filters/rename.py

@@ -32,7 +32,7 @@ class RenamedFieldMapping:
 
         value = self.field.metadata(key, **kwargs)
         if key == self.what:
-            return self.renaming.get(value, value)
+            return self.renaming.get(self.what, {}).get(value, value)
 
         return value
 
@@ -68,8 +68,7 @@ class RenamedFieldFormat:
 
 
 def execute(context, input, what="param", **kwargs):
-    # print('🍍🍍🍍🍍🍍🍍🍍🍍🍍🍍🍍🍍🍍 ==========', kwargs)
-    if what in kwargs:
+    if what in kwargs and isinstance(kwargs[what], str):
         return FieldArray([RenamedFieldFormat(fs, kwargs[what]) for fs in input])
 
     return FieldArray([RenamedFieldMapping(fs, what, kwargs) for fs in input])

anemoi/datasets/create/functions/sources/__init__.py

@@ -16,6 +16,10 @@ LOG = logging.getLogger(__name__)
 
 
 def _expand(paths):
+
+    if not isinstance(paths, list):
+        paths = [paths]
+
     for path in paths:
         if path.startswith("file://"):
             path = path[7:]
@@ -40,8 +44,10 @@ def iterate_patterns(path, dates, **kwargs):
     given_paths = path if isinstance(path, list) else [path]
 
     dates = [d.isoformat() for d in dates]
+    if len(dates) > 0:
+        kwargs["date"] = dates
 
     for path in given_paths:
-        paths = Pattern(path, ignore_missing_keys=True).substitute(date=dates, **kwargs)
+        paths = Pattern(path, ignore_missing_keys=True).substitute(**kwargs)
         for path in _expand(paths):
             yield path, dates

anemoi/datasets/create/functions/sources/accumulations.py

@@ -375,6 +375,8 @@ def accumulations(context, dates, **request):
         ("od", "elda"): dict(base_times=(6, 18)),
         ("ea", "oper"): dict(data_accumulation_period=1, base_times=(6, 18)),
         ("ea", "enda"): dict(data_accumulation_period=3, base_times=(6, 18)),
+        ("rr", "oper"): dict(data_accumulation_period=3, base_times=(0, 3, 6, 9, 12, 15, 18, 21)),
+        ("l5", "oper"): dict(data_accumulation_period=1, base_times=(0,)),
     }
 
     kwargs = KWARGS.get((class_, stream), {})

anemoi/datasets/create/functions/sources/grib.py

@@ -135,7 +135,7 @@ def execute(context, dates, path, latitudes=None, longitudes=None, *args, **kwar
             s = s.sel(valid_datetime=dates, **kwargs)
             ds = ds + s
 
-    if kwargs:
+    if kwargs and not context.partial_ok:
         check(ds, given_paths, valid_datetime=dates, **kwargs)
 
     if geography is not None:

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

@@ -12,6 +12,7 @@ import logging
 from earthkit.data.core.fieldlist import MultiFieldList
 
 from anemoi.datasets.data.stores import name_to_zarr_store
+from anemoi.datasets.utils.fields import NewMetadataField as NewMetadataField
 
 from .. import iterate_patterns
 from .fieldlist import XarrayFieldList
@@ -41,7 +42,7 @@ def load_one(emoji, context, dates, dataset, options={}, flavour=None, **kwargs)
     We have seen this bug triggered when we run many clients in parallel, for example, when we create a new dataset using `xarray-zarr`.
     """
 
-    context.trace(emoji, dataset, options)
+    context.trace(emoji, dataset, options, kwargs)
 
     if isinstance(dataset, str) and ".zarr" in dataset:
         data = xr.open_zarr(name_to_zarr_store(dataset), **options)
@@ -49,7 +50,11 @@ def load_one(emoji, context, dates, dataset, options={}, flavour=None, **kwargs)
         data = xr.open_dataset(dataset, **options)
 
     fs = XarrayFieldList.from_xarray(data, flavour)
-    result = MultiFieldList([fs.sel(valid_datetime=date, **kwargs) for date in dates])
+
+    if len(dates) == 0:
+        return fs.sel(**kwargs)
+    else:
+        result = MultiFieldList([fs.sel(valid_datetime=date, **kwargs) for date in dates])
 
     if len(result) == 0:
         LOG.warning(f"No data found for {dataset} and dates {dates} and {kwargs}")

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

@@ -56,6 +56,8 @@ class Coordinate:
     is_step = False
     is_date = False
     is_member = False
+    is_x = False
+    is_y = False
 
     def __init__(self, variable):
         self.variable = variable
@@ -66,10 +68,11 @@ class Coordinate:
         return 1 if self.scalar else len(self.variable)
 
     def __repr__(self):
-        return "%s[name=%s,values=%s]" % (
+        return "%s[name=%s,values=%s,shape=%s]" % (
             self.__class__.__name__,
             self.variable.name,
             self.variable.values if self.scalar else len(self),
+            self.variable.shape,
         )
 
     def reduced(self, i):
@@ -225,11 +228,13 @@ class LatitudeCoordinate(Coordinate):
 
 class XCoordinate(Coordinate):
     is_grid = True
+    is_x = True
     mars_names = ("x",)
 
 
 class YCoordinate(Coordinate):
     is_grid = True
+    is_y = True
     mars_names = ("y",)
 
 
@@ -239,3 +244,9 @@ class ScalarCoordinate(Coordinate):
     @property
     def mars_names(self):
         return (self.variable.name,)
+
+
+class UnsupportedCoordinate(Coordinate):
+    @property
+    def mars_names(self):
+        return (self.variable.name,)

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

@@ -72,13 +72,18 @@ class XArrayField(Field):
     def shape(self):
         return self._shape
 
-    def to_numpy(self, flatten=False, dtype=None):
-        values = self.selection.values
+    def to_numpy(self, flatten=False, dtype=None, index=None):
+        if index is not None:
+            values = self.selection[index]
+        else:
+            values = self.selection
 
         assert dtype is None
+
         if flatten:
-            return values.flatten()
-        return values.reshape(self.shape)
+            return values.values.flatten()
+
+        return values  # .reshape(self.shape)
 
     def _make_metadata(self):
         return XArrayMetadata(self)
@@ -113,3 +118,7 @@ class XArrayField(Field):
 
     def __repr__(self):
         return repr(self._metadata)
+
+    def _values(self):
+        # we don't use .values as this will download the data
+        return self.selection

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

@@ -70,10 +70,10 @@ class XarrayFieldList(FieldList):
                 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")
+            variable = ds[name]
+            _skip_attr(variable, "coordinates")
+            _skip_attr(variable, "bounds")
+            _skip_attr(variable, "grid_mapping")
 
         # Select only geographical variables
         for name in ds.data_vars:
@@ -81,14 +81,14 @@ class XarrayFieldList(FieldList):
             if name in skip:
                 continue
 
-            v = ds[name]
+            variable = ds[name]
             coordinates = []
 
-            for coord in v.coords:
+            for coord in variable.coords:
 
                 c = guess.guess(ds[coord], coord)
                 assert c, f"Could not guess coordinate for {coord}"
-                if coord not in v.dims:
+                if coord not in variable.dims:
                     c.is_dim = False
                 coordinates.append(c)
 
@@ -98,17 +98,17 @@ class XarrayFieldList(FieldList):
             if grid_coords < 2:
                 continue
 
-            variables.append(
-                Variable(
-                    ds=ds,
-                    var=v,
-                    coordinates=coordinates,
-                    grid=guess.grid(coordinates),
-                    time=Time.from_coordinates(coordinates),
-                    metadata={},
-                )
+            v = Variable(
+                ds=ds,
+                variable=variable,
+                coordinates=coordinates,
+                grid=guess.grid(coordinates, variable),
+                time=Time.from_coordinates(coordinates),
+                metadata={},
             )
 
+            variables.append(v)
+
         return cls(ds, variables)
 
     def sel(self, **kwargs):

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

@@ -8,6 +8,8 @@
 #
 
 
+import logging
+
 from .coordinates import DateCoordinate
 from .coordinates import EnsembleCoordinate
 from .coordinates import LatitudeCoordinate
@@ -16,10 +18,16 @@ 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:
@@ -150,36 +158,145 @@ class CoordinateGuesser:
         if c.shape in ((1,), tuple()):
             return ScalarCoordinate(c)
 
-        raise NotImplementedError(
+        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}"
         )
 
-    def grid(self, coordinates):
+        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:
-            raise NotImplementedError(f"Expected 1 latitude coordinate, got {len(lat)}")
+        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(lon) != 1:
-            raise NotImplementedError(f"Expected 1 longitude coordinate, got {len(lon)}")
+        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]
 
-        if (lat.name, lon.name) in self._cache:
-            return self._cache[(lat.name, lon.name)]
+        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)]
 
-        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)
+        if unstructured:
+            grid = UnstructuredGrid(lat, lon, dim_vars)
         else:
-            grid = UnstructuredGrid(lat, lon)
+            grid = MeshedGrid(lat, lon, dim_vars)
 
-        self._cache[(lat.name, lon.name)] = grid
+        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]
+
+        _, unstructured = self._check_dims(variable, x, y)
+
+        if x.variable.dims != y.variable.dims:
+            raise ValueError(f"Dimensions do not match {x.name}{x.variable.dims} != {y.name}{y.variable.dims}")
+
+        if (x.name, y.name) in self._cache:
+            return self._cache[(x.name, y.name)]
+
+        if (x.name, y.name) in self._cache:
+            return self._cache[(x.name, y.name)]
+
+        assert len(x.variable.shape) == len(y.variable.shape), (x.variable.shape, y.variable.shape)
+
+        grid_mapping = variable.attrs.get("grid_mapping", None)
+
+        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")
+
+        if grid_mapping is not None:
+            if unstructured:
+                return UnstructuredProjectionGrid(x, y, grid_mapping)
+            else:
+                return MeshProjectionGrid(x, y, grid_mapping)
+
+        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):