anemoi-datasets 0.5.16__py3-none-any.whl → 0.5.18__py3-none-any.whl

anemoi/datasets/create/sources/xarray_support/variable.py

@@ -0,0 +1,331 @@
+# (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
+import math
+from functools import cached_property
+from typing import Any
+from typing import Dict
+from typing import List
+from typing import Optional
+from typing import Tuple
+
+import numpy as np
+import xarray as xr
+
+from .field import XArrayField
+
+LOG = logging.getLogger(__name__)
+
+
+class Variable:
+    """Represents a variable in an xarray dataset.
+
+    Attributes
+    ----------
+    ds : xr.Dataset
+        The xarray dataset.
+    variable : xr.DataArray
+        The data array representing the variable.
+    coordinates : List[Any]
+        List of coordinates associated with the variable.
+    grid : Any
+        The grid associated with the variable.
+    time : Any
+        The time dimension associated with the variable.
+    metadata : Dict[str, Any]
+        Metadata associated with the variable.
+    """
+
+    def __init__(
+        self,
+        *,
+        ds: xr.Dataset,
+        variable: xr.DataArray,
+        coordinates: List[Any],
+        grid: Any,
+        time: Any,
+        metadata: Dict[str, Any],
+    ):
+        """Initialize the Variable object.
+
+        Parameters
+        ----------
+        ds : xr.Dataset
+            The xarray dataset.
+        variable : xr.DataArray
+            The data array representing the variable.
+        coordinates : List[Any]
+            List of coordinates associated with the variable.
+        grid : Any
+            The grid associated with the variable.
+        time : Any
+            The time dimension associated with the variable.
+        metadata : Dict[str, Any]
+            Metadata associated with the variable.
+        """
+        self.ds = ds
+        self.variable = variable
+
+        self.grid = grid
+        self.coordinates = coordinates
+
+        self._metadata = metadata.copy()
+        self._metadata.update({"variable": variable.name, "param": variable.name})
+
+        self.time = time
+
+        self.shape = tuple(len(c.variable) for c in coordinates if c.is_dim and not c.scalar and not c.is_grid)
+        self.names = {c.variable.name: c for c in coordinates if c.is_dim and not c.scalar and not c.is_grid}
+        self.by_name = {c.variable.name: c for c in coordinates}
+
+        # We need that alias for the time dimension
+        self._aliases = dict(valid_datetime="time")
+
+        self.length = math.prod(self.shape)
+
+    @property
+    def name(self) -> str:
+        """Return the name of the variable."""
+        return str(self.variable.name)
+
+    def __len__(self) -> int:
+        """Return the length of the variable.
+
+        Returns
+        -------
+        int
+            The length of the variable.
+        """
+        return self.length
+
+    @property
+    def grid_mapping(self) -> Optional[Dict[str, Any]]:
+        """Return the grid mapping of the variable."""
+        grid_mapping = self.variable.attrs.get("grid_mapping", None)
+        if grid_mapping is None:
+            return None
+        return self.ds[grid_mapping].attrs
+
+    def grid_points(self) -> Any:
+        """Return the grid points of the variable.
+
+        Returns
+        -------
+        Any
+            The grid points of the variable.
+        """
+        return self.grid.grid_points
+
+    @property
+    def latitudes(self) -> Any:
+        """Return the latitudes of the variable."""
+        return self.grid.latitudes
+
+    @property
+    def longitudes(self) -> Any:
+        """Return the longitudes of the variable."""
+        return self.grid.longitudes
+
+    def __repr__(self) -> str:
+        """Return a string representation of the variable.
+
+        Returns
+        -------
+        str
+            A string representation of the variable.
+        """
+        return "Variable[name=%s,coordinates=%s,metadata=%s]" % (
+            self.variable.name,
+            self.coordinates,
+            self._metadata,
+        )
+
+    def __getitem__(self, i: int) -> "XArrayField":
+        """Get a 2D field from the variable.
+
+        Parameters
+        ----------
+        i : int
+            Index of the field.
+
+        Returns
+        -------
+        XArrayField
+            The 2D field at the specified index.
+
+        Raises
+        ------
+        IndexError
+            If the index is out of range.
+        """
+        if i >= self.length:
+            raise IndexError(i)
+
+        coords = np.unravel_index(i, self.shape)
+        kwargs = {k: v for k, v in zip(self.names, coords)}
+        return XArrayField(self, self.variable.isel(kwargs))
+
+    def sel(self, missing: Dict[str, Any], **kwargs: Any) -> Optional["Variable"]:
+        """Select a subset of the variable based on the given coordinates.
+
+        Parameters
+        ----------
+        missing : Dict[str, Any]
+            Dictionary to store missing coordinates.
+        **kwargs : Any
+            Coordinates to select.
+
+        Returns
+        -------
+        Optional[Variable]
+            The selected subset of the variable.
+        """
+        if not kwargs:
+            return self
+
+        k, v = kwargs.popitem()
+
+        user_provided_k = k
+
+        if k == "valid_datetime":
+            # Ask the Time object to select the valid datetime
+            k = self.time.select_valid_datetime(self)
+            if k is None:
+                return None
+
+        c = self.by_name.get(k)
+
+        # assert c is not None, f"Could not find coordinate {k} in {self.variable.name} {self.coordinates} {list(self.by_name)}"
+
+        if c is None:
+            missing[k] = v
+            return self.sel(missing, **kwargs)
+
+        i = c.index(v)
+        if i is None:
+            if k != user_provided_k:
+                LOG.warning(f"Could not find {user_provided_k}={v} in {c} (alias of {k})")
+            else:
+                LOG.warning(f"Could not find {k}={v} in {c}")
+            return None
+
+        coordinates = [x.reduced(i) if c is x else x for x in self.coordinates]
+
+        metadata = self._metadata.copy()
+        metadata.update({k: v})
+
+        variable = Variable(
+            ds=self.ds,
+            variable=self.variable.isel({k: i}),
+            coordinates=coordinates,
+            grid=self.grid,
+            time=self.time,
+            metadata=metadata,
+        )
+
+        return variable.sel(missing, **kwargs)
+
+    def match(self, **kwargs: Any) -> Tuple[bool, Optional[Dict[str, Any]]]:
+        """Match the variable based on the given metadata.
+
+        Parameters
+        ----------
+        **kwargs : Any
+            Metadata to match.
+
+        Returns
+        -------
+        Tuple[bool, Optional[Dict[str, Any]]]
+            A tuple containing a boolean indicating if the match was successful and the remaining metadata.
+        """
+        if "param" in kwargs:
+            assert "variable" not in kwargs
+            kwargs["variable"] = kwargs.pop("param")
+
+        if "variable" in kwargs:
+            name = kwargs.pop("variable")
+            if not isinstance(name, (list, tuple)):
+                name = [name]
+            if self.variable.name not in name:
+                return False, None
+            return True, kwargs
+        return True, kwargs
+
+
+class FilteredVariable:
+    """Represents a filtered variable based on metadata.
+
+    Attributes
+    ----------
+    variable : Variable
+        The variable to filter.
+    kwargs : Any
+        Metadata to filter the variable.
+    """
+
+    def __init__(self, variable: Variable, **kwargs: Any):
+        """Initialize the FilteredVariable object.
+
+        Parameters
+        ----------
+        variable : Variable
+            The variable to filter.
+        **kwargs : Any
+            Metadata to filter the variable.
+        """
+        self.variable = variable
+        self.kwargs = kwargs
+
+    @cached_property
+    def fields(self) -> List["XArrayField"]:
+        """Filter the fields of a variable based on metadata."""
+        return [
+            field
+            for field in self.variable
+            if all(field.metadata(k, default=None) == v for k, v in self.kwargs.items())
+        ]
+
+    @property
+    def length(self) -> int:
+        """Return the length of the filtered variable."""
+        return len(self.fields)
+
+    def __len__(self) -> int:
+        """Return the length of the filtered variable.
+
+        Returns
+        -------
+        int
+            The length of the filtered variable.
+        """
+        return self.length
+
+    def __getitem__(self, i: int) -> "XArrayField":
+        """Get a field from the filtered variable.
+
+        Parameters
+        ----------
+        i : int
+            Index of the field.
+
+        Returns
+        -------
+        XArrayField
+            The field at the specified index.
+
+        Raises
+        ------
+        IndexError
+            If the index is out of range.
+        """
+        if i >= self.length:
+            raise IndexError(i)
+        return self.fields[i]

anemoi/datasets/create/sources/xarray_zarr.py

@@ -0,0 +1,41 @@
+# (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.
+
+from typing import Any
+from typing import List
+
+import earthkit.data as ekd
+
+from .legacy import legacy_source
+from .xarray import load_many
+
+
+@legacy_source(__file__)
+def execute(context: Any, dates: List[str], url: str, *args: Any, **kwargs: Any) -> ekd.FieldList:
+    """Execute the data loading process.
+
+    Parameters
+    ----------
+    context : Any
+        The context in which the execution occurs.
+    dates : List[str]
+        List of dates for which data is to be loaded.
+    url : str
+        The URL from which data is to be loaded.
+    *args : tuple
+        Additional positional arguments.
+    **kwargs : dict
+        Additional keyword arguments.
+
+    Returns
+    -------
+    ekd.FieldList
+        The loaded data.
+    """
+    return load_many("🇿", context, url, *args, **kwargs).execute(dates)

anemoi/datasets/create/{functions/sources → sources}/zenodo.py

@@ -8,25 +8,54 @@
 # nor does it submit to any jurisdiction.
 
 
+from typing import Any
+from typing import Dict
+from typing import List
+
+import earthkit.data as ekd
 from earthkit.data.core.fieldlist import MultiFieldList
 from earthkit.data.sources.url import download_and_cache
 
-from . import iterate_patterns
+from .legacy import legacy_source
+from .patterns import iterate_patterns
 from .xarray import load_one
 
 
-def execute(context, dates, record_id, file_key, *args, **kwargs):
+@legacy_source(__file__)
+def execute(context: Any, dates: Any, record_id: str, file_key: str, *args: Any, **kwargs: Any) -> ekd.FieldList:
+    """Executes the download and processing of files from Zenodo.
+
+    Parameters
+    ----------
+    context : Any
+        The context in which the function is executed.
+    dates : Any
+        The dates for which the data is required.
+    record_id : str
+        The Zenodo record ID.
+    file_key : str
+        The key to identify the file.
+    *args : Any
+        Additional arguments.
+    **kwargs : Any
+        Additional keyword arguments.
+
+    Returns
+    -------
+    MultiFieldList
+        A list of fields loaded from the downloaded files.
+    """
     import requests
 
-    result = []
+    result: List[Any] = []
 
     URLPATTERN = "https://zenodo.org/api/records/{record_id}"
     url = URLPATTERN.format(record_id=record_id)
     r = requests.get(url)
     r.raise_for_status()
-    record = r.json()
+    record: Dict[str, Any] = r.json()
 
-    urls = {}
+    urls: Dict[str, str] = {}
     for file in record["files"]:
         urls[file["key"]] = file["links"]["self"]