feat: missing features for observations

src/anemoi/datasets/create/sources/observations.py

@@ -0,0 +1,43 @@
+# (C) Copyright 2025 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 pandas as pd
+
+
+def check_dataframe(df):
+    """Check the DataFrame for consistency."""
+    if df.empty:
+        pass
+    if "times" not in df.columns:
+        raise ValueError("The DataFrame must contain a 'times' column.")
+    if not pd.api.types.is_datetime64_any_dtype(df["times"]):
+        raise TypeError("The 'times' column must be of datetime type.")
+    if "latitudes" not in df.columns or "longitudes" not in df.columns:
+        raise ValueError("The DataFrame must contain 'latitudes' and 'longitudes' columns.")
+
+
+class ObservationsSource:
+    def __call__(self, window):
+        raise NotImplementedError("This method should be implemented by subclasses")
+
+    def _check(self, df):
+        check_dataframe(df)
+        return df
+
+
+class ObservationsFilter:
+    def __call__(self, df):
+        """Filter the data based on the given window."""
+        check_dataframe(df)
+        return df
+
+    def _check(self, df):
+        check_dataframe(df)
+        return df

src/anemoi/datasets/data/dataset.py

@@ -136,7 +136,8 @@ def _subset(self, **kwargs: Any) -> "Dataset":
         if not kwargs:
             return self.mutate()
 
-        name = kwargs.pop("name", None)
+        name = kwargs.pop("set_group", None)  # TODO(Florian)
+        name = kwargs.pop("name", name)
         result = self.__subset(**kwargs)
         result._name = name
 
@@ -177,13 +178,18 @@ def __subset(self, **kwargs: Any) -> "Dataset":
             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))
+                    Padded(
+                        self,
+                        start=start,
+                        end=end,
+                        frequency=frequency,
+                        padding=padding,
+                        reason=dict(start=start, end=end, frequency=frequency, padding=padding),
+                    )
                     ._subset(**kwargs)
                     .mutate()
                 )
@@ -404,6 +410,9 @@ def _select_to_columns(self, vars: str | list[str] | tuple[str] | set) -> list[i
         if not isinstance(vars, (list, tuple)):
             vars = [vars]
 
+        for v in vars:
+            if v not in self.name_to_index:
+                raise ValueError(f"select: unknown variable: {v}, available: {list(self.name_to_index)}")
         return [self.name_to_index[v] for v in vars]
 
     def _drop_to_columns(self, vars: str | Sequence[str]) -> list[int]:

src/anemoi/datasets/data/debug.py

@@ -67,6 +67,7 @@ def __init__(self, dataset: "Dataset", kids: list[Any], **kwargs: Any) -> None:
             Additional keyword arguments.
         """
         self.dataset = dataset
+        assert isinstance(kids, list), "Kids must be a list"
         self.kids = kids
         self.kwargs = kwargs
 

src/anemoi/datasets/data/misc.py

@@ -349,19 +349,7 @@ def _open(a: str | PurePath | dict[str, Any] | list[Any] | tuple[Any, ...]) -> "
     """
     from .dataset import Dataset
     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"])
+    from .stores import dataset_lookup
 
     if isinstance(a, Dataset):
         return a.mutate()
@@ -370,7 +358,22 @@ def _open(a: str | PurePath | dict[str, Any] | list[Any] | tuple[Any, ...]) -> "
         return Zarr(a).mutate()
 
     if isinstance(a, str):
-        return Zarr(zarr_lookup(a)).mutate()
+        path = dataset_lookup(a)
+
+        if path and path.endswith(".zarr") or path.endswith(".zip"):
+            return Zarr(path).mutate()
+
+        if path and path.endswith(".vz"):
+            metadata_path = os.path.join(path, "metadata.json")
+            if os.path.exists(metadata_path):
+                if "backend" not in load_any_dict_format(metadata_path):
+                    raise ValueError(f"Metadata for {path} does not contain 'backend' key")
+
+                from anemoi.datasets.data.records import open_records_dataset
+
+                return open_records_dataset(path)
+
+        raise ValueError(f"Unsupported dataset path: {path}. ")
 
     if isinstance(a, PurePath):
         return _open(str(a)).mutate()
@@ -587,6 +590,18 @@ def _open_dataset(*args: Any, **kwargs: Any) -> "Dataset":
 
     assert len(sets) > 0, (args, kwargs)
 
+    if "set_group" in kwargs:
+        from anemoi.datasets.data.records import FieldsRecords
+
+        set_group = kwargs.pop("set_group")
+        assert len(sets) == 1, "set_group can only be used with a single dataset"
+        dataset = sets[0]
+
+        from anemoi.datasets.data.dataset import Dataset
+
+        if isinstance(dataset, Dataset):  # Fields dataset
+            return FieldsRecords(dataset, **kwargs, name=set_group).mutate()
+
     if len(sets) > 1:
         dataset, kwargs = _concat_or_join(sets, kwargs)
         return dataset._subset(**kwargs)

src/anemoi/datasets/data/observations/__init__.py

@@ -67,7 +67,10 @@ def __len__(self):
         return len(self.dates)
 
     def tree(self):
-        return Node(self)
+        return Node(
+            self,
+            [],
+        )
 
     def __getitem__(self, i):
         if isinstance(i, int):
@@ -232,6 +235,7 @@ def get_aux(self, i):
         assert latitudes.shape == longitudes.shape, f"Expected {latitudes.shape}, got {longitudes.shape}"
         assert timedeltas.shape == latitudes.shape, f"Expected {timedeltas.shape}, got {latitudes.shape}"
 
+        assert timedeltas.dtype == "timedelta64[s]", f"Expected timedelta64[s], got {timedeltas.dtype}"
         return latitudes, longitudes, timedeltas
 
     def getitem(self, i):

src/anemoi/datasets/data/padded.py

@@ -17,6 +17,7 @@
 from anemoi.utils.dates import frequency_to_timedelta
 from numpy.typing import NDArray
 
+from anemoi.datasets.data import MissingDateError
 from anemoi.datasets.data.dataset import Dataset
 from anemoi.datasets.data.dataset import FullIndex
 from anemoi.datasets.data.dataset import Shape
@@ -36,7 +37,15 @@ class Padded(Forwards):
     _after: int = 0
     _inside: int = 0
 
-    def __init__(self, dataset: Dataset, start: str, end: str, frequency: str, reason: dict[str, Any]) -> None:
+    def __init__(
+        self,
+        dataset: Dataset,
+        start: str,
+        end: str,
+        frequency: str,
+        reason: Dict[str, Any],
+        padding: str,
+    ) -> None:
         """Create a padded subset of a dataset.
 
         Attributes:
@@ -46,6 +55,7 @@ def __init__(self, dataset: Dataset, start: str, end: str, frequency: str, reaso
         frequency (str): The frequency of the subset.
         reason (Dict[str, Any]): The reason for the padding.
         """
+        self.padding = padding
 
         self.reason = {k: v for k, v in reason.items() if v is not None}
 
@@ -164,12 +174,20 @@ def _get_tuple(self, n: TupleIndex) -> NDArray[Any]:
         return [self[i] for i in n]
 
     def empty_item(self):
-        return self.dataset.empty_item()
+        if self.padding == "empty":
+            return self.dataset.empty_item()
+        elif self.padding == "raise":
+            raise ValueError("Padding is set to 'raise', cannot return an empty item.")
+        elif self.padding == "missing":
+            raise MissingDateError("Padding is set to 'missing'")
+        assert False, self.padding
 
     def get_aux(self, i: FullIndex) -> NDArray[np.timedelta64]:
         if self._i_out_of_range(i):
-            arr = np.array([], dtype=np.float32)
-            aux = arr, arr, arr
+            lats = np.array([], dtype=np.float32)
+            lons = lats
+            timedeltas = np.ones_like(lons, dtype="timedelta64[s]") * 0
+            aux = lats, lons, timedeltas
         else:
             aux = self.dataset.get_aux(i - self._before)