anemoi-datasets 0.5.24__py3-none-any.whl → 0.5.26__py3-none-any.whl

anemoi/datasets/data/complement.py

@@ -7,7 +7,6 @@
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 
-
 import datetime
 import logging
 from abc import abstractmethod
@@ -19,6 +18,7 @@ from typing import Optional
 from typing import Set
 from typing import Tuple
 
+import numpy as np
 from numpy.typing import NDArray
 
 from ..grids import nearest_grid_points
@@ -85,6 +85,7 @@ class Complement(Combined):
         for v in self._source.variables:
             if v not in self._target.variables:
                 self._variables.append(v)
+        LOG.info(f"The following variables will be complemented: {self._variables}")
 
         if not self._variables:
             raise ValueError("Augment: no missing variables")
@@ -96,9 +97,11 @@ class Complement(Combined):
 
     @property
     def statistics(self) -> Dict[str, NDArray[Any]]:
-        """Returns the statistics of the complemented dataset."""
-        index = [self._source.name_to_index[v] for v in self._variables]
-        return {k: v[index] for k, v in self._source.statistics.items()}
+        datasets = [self._source, self._target]
+        return {
+            k: [d.statistics[k][d.name_to_index[i]] for d in datasets for i in d.variables if i in self.variables]
+            for k in datasets[0].statistics
+        }
 
     def statistics_tendencies(self, delta: Optional[datetime.timedelta] = None) -> Dict[str, NDArray[Any]]:
         index = [self._source.name_to_index[v] for v in self._variables]
@@ -120,7 +123,11 @@ class Complement(Combined):
     @property
     def variables_metadata(self) -> Dict[str, Any]:
         """Returns the metadata of the variables to be added to the target dataset."""
-        return {k: v for k, v in self._source.variables_metadata.items() if k in self._variables}
+        # Merge the two dicts first
+        all_meta = {**self._source.variables_metadata, **self._target.variables_metadata}
+
+        # Filter to keep only desired variables
+        return {k: v for k, v in all_meta.items() if k in self._variables}
 
     def check_same_variables(self, d1: Dataset, d2: Dataset) -> None:
         """Checks if the variables in two datasets are the same.
@@ -231,7 +238,7 @@ class ComplementNone(Complement):
 class ComplementNearest(Complement):
     """A class to complement a target dataset with variables from a source dataset using nearest neighbor interpolation."""
 
-    def __init__(self, target: Any, source: Any, max_distance: float = None) -> None:
+    def __init__(self, target: Any, source: Any, max_distance: float = None, k: int = 1) -> None:
         """Initializes the ComplementNearest class.
 
         Parameters
@@ -242,17 +249,25 @@ class ComplementNearest(Complement):
             The source dataset.
         max_distance : float, optional
             The maximum distance for nearest neighbor interpolation, default is None.
+        k : int, optional
+            The number of k closest neighbors to consider for interpolation
         """
         super().__init__(target, source)
 
-        self._nearest_grid_points = nearest_grid_points(
+        self.k = k
+        self._distances, self._nearest_grid_points = nearest_grid_points(
             self._source.latitudes,
             self._source.longitudes,
             self._target.latitudes,
             self._target.longitudes,
             max_distance=max_distance,
+            k=k,
         )
 
+        if k == 1:
+            self._distances = np.expand_dims(self._distances, axis=1)
+            self._nearest_grid_points = np.expand_dims(self._nearest_grid_points, axis=1)
+
     def check_compatibility(self, d1: Dataset, d2: Dataset) -> None:
         """Checks the compatibility of two datasets for nearest neighbor interpolation.
 
@@ -285,7 +300,19 @@ class ComplementNearest(Complement):
         source_data = self._source[index[0], source_index, index[2], ...]
         target_data = source_data[..., self._nearest_grid_points]
 
-        result = target_data[..., index[3]]
+        epsilon = 1e-8  # prevent division by zero
+        weights = 1.0 / (self._distances + epsilon)
+        weights = weights.astype(target_data.dtype)
+        weights /= weights.sum(axis=1, keepdims=True)  # normalize
+
+        # Reshape weights to broadcast correctly
+        # Add leading singleton dimensions so it matches target_data shape
+        while weights.ndim < target_data.ndim:
+            weights = np.expand_dims(weights, axis=0)
+
+        # Compute weighted average along the last dimension
+        final_point = np.sum(target_data * weights, axis=-1)
+        result = final_point[..., index[3]]
 
         return apply_index_to_slices_changes(result, changes)
 
@@ -330,6 +357,13 @@ def complement_factory(args: Tuple, kwargs: dict) -> Dataset:
         "nearest": ComplementNearest,
     }[interpolation]
 
-    complement = Class(target=target, source=source)._subset(**kwargs)
+    if interpolation == "nearest":
+        k = kwargs.pop("k", "1")
+        complement = Class(target=target, source=source, k=k)._subset(**kwargs)
+
+    else:
+        complement = Class(target=target, source=source)._subset(**kwargs)
+
+    joined = _open_dataset([target, complement])
 
-    return _open_dataset([target, complement], reorder=source.variables)
+    return _open_dataset(joined, reorder=sorted(joined.variables))

anemoi/datasets/data/dataset.py

@@ -179,6 +179,19 @@ class Dataset(ABC, Sized):
         if "start" in kwargs or "end" in kwargs:
             start = kwargs.pop("start", None)
             end = kwargs.pop("end", None)
+            padding = kwargs.pop("padding", None)
+
+            if padding:
+                if padding != "empty":
+                    raise ValueError(f"Only 'empty' padding is supported, got {padding=}")
+                from .padded import Padded
+
+                frequency = kwargs.pop("frequency", self.frequency)
+                return (
+                    Padded(self, start, end, frequency, dict(start=start, end=end, frequency=frequency))
+                    ._subset(**kwargs)
+                    .mutate()
+                )
 
             from .subset import Subset
 
@@ -724,6 +737,9 @@ class Dataset(ABC, Sized):
         """Return the grid shape of the dataset."""
         return (self.shape[-1],)
 
+    def empty_item(self) -> NDArray[Any]:
+        return np.zeros((*self.shape[1:-1], 0), dtype=self.dtype)
+
     def _check(self) -> None:
         """Check for overridden private methods in the dataset."""
         common = Dataset.__dict__.keys() & self.__class__.__dict__.keys()
@@ -1075,3 +1091,16 @@ class Dataset(ABC, Sized):
             The dataset names.
         """
         pass
+
+    def get_latitudes(self, i):
+        return self.get_aux(i)[0]
+
+    def get_longitudes(self, i):
+        return self.get_aux(i)[1]
+
+    def get_timedeltas(self, i):
+        return self.get_aux(i)[2]
+
+    def get_aux(self, i):
+        # need to decide if Fields datasets need to implement this
+        raise NotImplementedError(f"get_aux is not implemented for this dataset, {type(self)}")

anemoi/datasets/data/forwards.py

@@ -330,8 +330,14 @@ class Combined(Forwards):
         ValueError
             If the grids are not the same.
         """
-        if (d1.latitudes != d2.latitudes).any() or (d1.longitudes != d2.longitudes).any():
-            raise ValueError(f"Incompatible grid ({d1} {d2})")
+
+        # note: not a proper implementation, should be handled
+        #       in a more consolidated way ...
+        rtol = 1.0e-7
+        if not np.allclose(d1.latitudes, d2.latitudes, rtol=rtol) or not np.allclose(
+            d1.longitudes, d2.longitudes, rtol=rtol
+        ):
+            raise ValueError(f"Incompatible grid ({d1.longitudes} {d2.longitudes})")
 
     def check_same_shape(self, d1: Dataset, d2: Dataset) -> None:
         """Checks if the shapes of two datasets are the same.

anemoi/datasets/data/misc.py

@@ -11,6 +11,7 @@
 import calendar
 import datetime
 import logging
+import os
 from pathlib import PurePath
 from typing import TYPE_CHECKING
 from typing import Any
@@ -22,6 +23,7 @@ from typing import Union
 
 import numpy as np
 import zarr
+from anemoi.utils.config import load_any_dict_format
 from anemoi.utils.config import load_config as load_settings
 from numpy.typing import NDArray
 
@@ -108,7 +110,10 @@ def round_datetime(d: np.datetime64, dates: NDArray[np.datetime64], up: bool) ->
 
 
 def _as_date(
-    d: Union[int, str, np.datetime64, datetime.date], dates: NDArray[np.datetime64], last: bool
+    d: Union[int, str, np.datetime64, datetime.date],
+    dates: NDArray[np.datetime64],
+    last: bool,
+    frequency: Optional[datetime.timedelta] = None,
 ) -> np.datetime64:
     """Convert a date to a numpy datetime64 object, rounding to the nearest date in a list of dates.
 
@@ -120,6 +125,8 @@ def _as_date(
         The list of dates.
     last : bool
         Whether to round to the last date.
+    frequency : Optional[datetime.timedelta]
+        The frequency of the dataset.
 
     Returns
     -------
@@ -142,30 +149,49 @@ def _as_date(
         pass
 
     if isinstance(d, int):
+        delta = frequency
+        if delta is None:
+            delta = np.timedelta64(1, "s")
+        delta = np.timedelta64(delta, "s")
+
         if len(str(d)) == 4:
             year = d
             if last:
-                return _as_date(np.datetime64(f"{year:04}-12-31T23:59:59"), dates, last)
+                year = year + 1
+                npdate = np.datetime64(f"{year:04}-01-01T00:00:00")
+                return _as_date(npdate - delta, dates, last, frequency)
             else:
-                return _as_date(np.datetime64(f"{year:04}-01-01T00:00:00"), dates, last)
+                return _as_date(np.datetime64(f"{year:04}-01-01T00:00:00"), dates, last, frequency)
 
         if len(str(d)) == 6:
             year = d // 100
             month = d % 100
             if last:
-                _, last_day = calendar.monthrange(year, month)
-                return _as_date(np.datetime64(f"{year:04}-{month:02}-{last_day:02}T23:59:59"), dates, last)
+                month = month + 1
+                if month > 12:
+                    month = 1
+                    year += 1
+                npdate = np.datetime64(f"{year:04}-{month:02}-01T00:00:00")
+                return _as_date(npdate - delta, dates, last, frequency)
             else:
-                return _as_date(np.datetime64(f"{year:04}-{month:02}-01T00:00:00"), dates, last)
+                return _as_date(np.datetime64(f"{year:04}-{month:02}-01T00:00:00"), dates, last, frequency)
 
         if len(str(d)) == 8:
             year = d // 10000
             month = (d % 10000) // 100
             day = d % 100
             if last:
-                return _as_date(np.datetime64(f"{year:04}-{month:02}-{day:02}T23:59:59"), dates, last)
+                day = day + 1
+                if day > calendar.monthrange(year, month)[1]:
+                    day = 1
+                    month += 1
+                    if month > 12:
+                        month = 1
+                        year += 1
+                npdate = np.datetime64(f"{year:04}-{month:02}-{day:02}T00:00:00")
+                return _as_date(npdate - delta, dates, last, frequency)
             else:
-                return _as_date(np.datetime64(f"{year:04}-{month:02}-{day:02}T00:00:00"), dates, last)
+                return _as_date(np.datetime64(f"{year:04}-{month:02}-{day:02}T00:00:00"), dates, last, frequency)
 
     if isinstance(d, str):
 
@@ -201,19 +227,20 @@ def _as_date(
                 np.datetime64(f"{year:04}-{month:02}-{day:02}T{hour:02}:{minute:02}:{second:02}"),
                 dates,
                 last,
+                frequency,
             )
 
         if "-" in d:
             assert ":" not in d
             bits = d.split("-")
             if len(bits) == 1:
-                return _as_date(int(bits[0]), dates, last)
+                return _as_date(int(bits[0]), dates, last, frequency)
 
             if len(bits) == 2:
-                return _as_date(int(bits[0]) * 100 + int(bits[1]), dates, last)
+                return _as_date(int(bits[0]) * 100 + int(bits[1]), dates, last, frequency)
 
             if len(bits) == 3:
-                return _as_date(int(bits[0]) * 10000 + int(bits[1]) * 100 + int(bits[2]), dates, last)
+                return _as_date(int(bits[0]) * 10000 + int(bits[1]) * 100 + int(bits[2]), dates, last, frequency)
 
         if ":" in d:
             assert len(d) == 5
@@ -225,12 +252,16 @@ def _as_date(
             month = first.month
             day = first.day
 
-            return _as_date(np.datetime64(f"{year:04}-{month:02}-{day:02}T{hour}:00:00"), dates, last)
+            return _as_date(np.datetime64(f"{year:04}-{month:02}-{day:02}T{hour}:00:00"), dates, last, frequency)
 
     raise NotImplementedError(f"Unsupported date: {d} ({type(d)})")
 
 
-def as_first_date(d: Union[int, str, np.datetime64, datetime.date], dates: NDArray[np.datetime64]) -> np.datetime64:
+def as_first_date(
+    d: Union[int, str, np.datetime64, datetime.date],
+    dates: NDArray[np.datetime64],
+    frequency: Optional[datetime.timedelta] = None,
+) -> np.datetime64:
     """Convert a date to the first date in a list of dates.
 
     Parameters
@@ -239,16 +270,22 @@ def as_first_date(d: Union[int, str, np.datetime64, datetime.date], dates: NDArr
         The date to convert.
     dates : NDArray[np.datetime64]
         The list of dates.
+    frequency : Optional[datetime.timedelta]
+        The frequency of the dataset.
 
     Returns
     -------
     np.datetime64
         The first date.
     """
-    return _as_date(d, dates, last=False)
+    return _as_date(d, dates, last=False, frequency=frequency)
 
 
-def as_last_date(d: Union[int, str, np.datetime64, datetime.date], dates: NDArray[np.datetime64]) -> np.datetime64:
+def as_last_date(
+    d: Union[int, str, np.datetime64, datetime.date],
+    dates: NDArray[np.datetime64],
+    frequency: Optional[datetime.timedelta] = None,
+) -> np.datetime64:
     """Convert a date to the last date in a list of dates.
 
     Parameters
@@ -257,13 +294,15 @@ def as_last_date(d: Union[int, str, np.datetime64, datetime.date], dates: NDArra
         The date to convert.
     dates : NDArray[np.datetime64]
         The list of dates.
+    frequency : Optional[datetime.timedelta]
+        The frequency of the dataset.
 
     Returns
     -------
     np.datetime64
         The last date.
     """
-    return _as_date(d, dates, last=True)
+    return _as_date(d, dates, last=True, frequency=frequency)
 
 
 def _concat_or_join(datasets: List["Dataset"], kwargs: Dict[str, Any]) -> Tuple["Dataset", Dict[str, Any]]:
@@ -317,6 +356,18 @@ def _open(a: Union[str, PurePath, Dict[str, Any], List[Any], Tuple[Any, ...]]) -
     from .stores import Zarr
     from .stores import zarr_lookup
 
+    if isinstance(a, str) and len(a.split(".")) in [2, 3]:
+
+        metadata_path = os.path.join(a, "metadata.json")
+        if os.path.exists(metadata_path):
+            metadata = load_any_dict_format(metadata_path)
+            if "backend" not in metadata:
+                raise ValueError(f"Metadata for {a} does not contain 'backend' key")
+
+            from anemoi.datasets.data.records import open_records_dataset
+
+            return open_records_dataset(a, backend=metadata["backend"])
+
     if isinstance(a, Dataset):
         return a.mutate()
 
@@ -454,6 +505,13 @@ def _open_dataset(*args: Any, **kwargs: Any) -> "Dataset":
     for a in args:
         sets.append(_open(a))
 
+    if "observations" in kwargs:
+        from .observations import observations_factory
+
+        assert not sets, sets
+
+        return observations_factory(args, kwargs).mutate()
+
     if "xy" in kwargs:
         # Experimental feature, may be removed
         from .xy import xy_factory