anemoi-datasets 0.5.25__py3-none-any.whl → 0.5.27__py3-none-any.whl

anemoi/datasets/create/sources/empty.py

@@ -9,7 +9,6 @@
 
 
 from typing import Any
-from typing import List
 
 import earthkit.data as ekd
 
@@ -17,7 +16,7 @@ from .legacy import legacy_source
 
 
 @legacy_source(__file__)
-def execute(context: Any, dates: List[str], **kwargs: Any) -> ekd.FieldList:
+def execute(context: Any, dates: list[str], **kwargs: Any) -> ekd.FieldList:
     """Executes the loading of an empty data source.
 
     Parameters

anemoi/datasets/create/sources/fdb.py

@@ -0,0 +1,133 @@
+# (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.
+
+from datetime import datetime
+from typing import Any
+
+import earthkit.data as ekd
+from anemoi.transform.fields import new_field_from_grid
+from anemoi.transform.fields import new_fieldlist_from_list
+from anemoi.transform.flavour import RuleBasedFlavour
+from anemoi.transform.grids import grid_registry
+
+from anemoi.datasets.create.typing import DateList
+
+from ..source import Source
+from . import source_registry
+
+
+@source_registry.register("fdb")
+class FdbSource(Source):
+    """FDB data source."""
+
+    emoji = "💽"
+
+    def __init__(
+        self,
+        context,
+        fdb_request: dict[str, Any],
+        fdb_config: dict | None = None,
+        fdb_userconfig: dict | None = None,
+        flavour: dict[str, Any] | None = None,
+        grid_definition: str | None = None,
+        **kwargs: dict[str, Any],
+    ):
+        """Initialise the FDB input.
+
+        Parameters
+        ----------
+        context : dict
+            The context.
+        fdb_request: dict
+            The FDB request parameters.
+        fdb_config : dict, optional
+            The FDB config to use.
+        fdb_userconfig : dict, optional
+            The FDB userconfig to use.
+        flavour: dict, optional
+            The flavour configuration, see `anemoi.transform.flavour.RuleBasedFlavour`.
+        grid_definition : str, optional
+            The grid definition to use, see `anemoi.transform.grids.grid_registry`.
+        kwargs : dict, optional
+            Additional keyword arguments.
+        """
+        super().__init__(context)
+        self.request = fdb_request.copy()
+        self.configs = {"config": fdb_config, "userconfig": fdb_userconfig}
+
+        self.flavour = RuleBasedFlavour(flavour) if flavour else None
+        if grid_definition is not None:
+            self.grid = grid_registry.from_config(grid_definition)
+        else:
+            self.grid = None
+
+        if "step" not in self.request:
+            self.request["step"] = 0
+
+        self.request["param"] = _shortname_to_paramid(fdb_request["param"], kwargs.pop("param_id_map", None))
+
+        # temporary workarounds for FDB use at MeteoSwiss (adoption is ongoing)
+        # thus not documented
+        self.offset_from_date = kwargs.pop("offset_from_date", None)
+
+    def execute(self, dates: DateList) -> ekd.FieldList:
+        """Execute the FDB source.
+
+        Parameters
+        ----------
+        dates : DateList
+            The input dates.
+
+        Returns
+        -------
+        ekd.FieldList
+            The output data.
+        """
+
+        requests = []
+        for date in dates:
+            time_request = _time_request_keys(date, self.offset_from_date)
+            requests.append(self.request | time_request)
+
+        # in some cases (e.g. repeated_dates 'constant' mode), we might have a fully
+        # defined request already and an empty dates list
+        requests = requests or [self.request]
+
+        fl = ekd.from_source("empty")
+        for request in requests:
+            fl += ekd.from_source("fdb", request, **self.configs, read_all=True)
+
+        if self.grid is not None:
+            fl = new_fieldlist_from_list([new_field_from_grid(f, self.grid) for f in fl])
+
+        if self.flavour:
+            fl = self.flavour.map(fl)
+
+        return fl
+
+
+def _time_request_keys(dt: datetime, offset_from_date: bool | None = None) -> str:
+    """Defines the time-related keys for the FDB request."""
+    out = {}
+    out["date"] = dt.strftime("%Y%m%d")
+    if offset_from_date:
+        out["time"] = "0000"
+        out["step"] = int((dt - dt.replace(hour=0, minute=0)).total_seconds() // 3600)
+    else:
+        out["time"] = dt.strftime("%H%M")
+    return out
+
+
+def _shortname_to_paramid(shortname: list[str], param_id_map: dict[str, int] | None = None) -> list[int]:
+    from anemoi.datasets.create.sources.mars import use_grib_paramid
+
+    """Convert a shortname to a parameter ID."""
+    if param_id_map is None:
+        return use_grib_paramid(shortname)
+    return [param_id_map[s] for s in shortname]

anemoi/datasets/create/sources/forcings.py

@@ -8,7 +8,6 @@
 # nor does it submit to any jurisdiction.
 
 from typing import Any
-from typing import List
 
 from earthkit.data import from_source
 
@@ -16,7 +15,7 @@ from .legacy import legacy_source
 
 
 @legacy_source(__file__)
-def forcings(context: Any, dates: List[str], template: str, param: str) -> Any:
+def forcings(context: Any, dates: list[str], template: str, param: str) -> Any:
     """Loads forcing data from a specified source.
 
     Parameters

anemoi/datasets/create/sources/grib.py

@@ -11,10 +11,6 @@
 import glob
 import logging
 from typing import Any
-from typing import Dict
-from typing import List
-from typing import Optional
-from typing import Union
 
 import earthkit.data as ekd
 from anemoi.transform.fields import new_field_from_grid
@@ -29,7 +25,7 @@ from .legacy import legacy_source
 LOG = logging.getLogger(__name__)
 
 
-def check(ds: Any, paths: List[str], **kwargs: Any) -> None:
+def check(ds: Any, paths: list[str], **kwargs: Any) -> None:
     """Check if the dataset matches the expected number of fields.
 
     Parameters
@@ -55,7 +51,7 @@ def check(ds: Any, paths: List[str], **kwargs: Any) -> None:
         raise ValueError(f"Expected {count} fields, got {len(ds)} (kwargs={kwargs}, paths={paths})")
 
 
-def _expand(paths: List[str]) -> Any:
+def _expand(paths: list[str]) -> Any:
     """Expand the given paths using glob.
 
     Parameters
@@ -80,10 +76,10 @@ def _expand(paths: List[str]) -> Any:
 @legacy_source(__file__)
 def execute(
     context: Any,
-    dates: List[Any],
-    path: Union[str, List[str]],
-    flavour: Optional[Union[str, Dict[str, Any]]] = None,
-    grid_definition: Optional[Dict[str, Any]] = None,
+    dates: list[Any],
+    path: str | list[str],
+    flavour: str | dict[str, Any] | None = None,
+    grid_definition: dict[str, Any] | None = None,
     *args: Any,
     **kwargs: Any,
 ) -> ekd.FieldList:

anemoi/datasets/create/sources/grib_index.py

@@ -10,10 +10,8 @@
 import logging
 import os
 import sqlite3
+from collections.abc import Iterator
 from typing import Any
-from typing import Iterator
-from typing import List
-from typing import Optional
 
 import earthkit.data as ekd
 import tqdm
@@ -36,8 +34,8 @@ class GribIndex:
         self,
         database: str,
         *,
-        keys: Optional[List[str] | str] = None,
-        flavour: Optional[str] = None,
+        keys: list[str] | str | None = None,
+        flavour: str | None = None,
         update: bool = False,
         overwrite: bool = False,
     ) -> None:
@@ -157,7 +155,7 @@ class GribIndex:
         """Commit the current transaction to the database."""
         self.conn.commit()
 
-    def _get_metadata_keys(self) -> List[str]:
+    def _get_metadata_keys(self) -> list[str]:
         """Retrieve the metadata keys from the database.
 
         Returns
@@ -225,7 +223,7 @@ class GribIndex:
                 LOG.info(f"Path: {self._get_path(existing_record[1])}")
             raise
 
-    def _all_columns(self) -> List[str]:
+    def _all_columns(self) -> list[str]:
         """Retrieve all column names from the grib_index table.
 
         Returns
@@ -241,7 +239,7 @@ class GribIndex:
         self._columns = [col for col in columns if not col.startswith("_")]
         return self._columns
 
-    def _ensure_columns(self, columns: List[str]) -> None:
+    def _ensure_columns(self, columns: list[str]) -> None:
         """Add missing columns to the grib_index table.
 
         Parameters
@@ -324,7 +322,7 @@ class GribIndex:
 
         self._commit()
 
-    def _paramdb(self, category: int, discipline: int) -> Optional[dict]:
+    def _paramdb(self, category: int, discipline: int) -> dict | None:
         """Fetch parameter information from the parameter database.
 
         Parameters
@@ -355,7 +353,7 @@ class GribIndex:
         except Exception as e:
             LOG.warning(f"Failed to fetch information from parameter database: {e}")
 
-    def _param_grib2_info(self, paramId: int) -> List[dict]:
+    def _param_grib2_info(self, paramId: int) -> list[dict]:
         """Fetch GRIB2 parameter information for a given parameter ID.
 
         Parameters
@@ -383,7 +381,7 @@ class GribIndex:
             LOG.warning(f"Failed to fetch information from parameter database: {e}")
         return []
 
-    def _param_id_info(self, paramId: int) -> Optional[dict]:
+    def _param_id_info(self, paramId: int) -> dict | None:
         """Fetch detailed information for a given parameter ID.
 
         Parameters
@@ -412,7 +410,7 @@ class GribIndex:
 
         return None
 
-    def _param_id_unit(self, unitId: int) -> Optional[dict]:
+    def _param_id_unit(self, unitId: int) -> dict | None:
         """Fetch unit information for a given unit ID.
 
         Parameters
@@ -520,7 +518,7 @@ class GribIndex:
             raise ValueError(f"No path found for path_id {path_id}")
         return row[0]
 
-    def retrieve(self, dates: List[Any], **kwargs: Any) -> Iterator[Any]:
+    def retrieve(self, dates: list[Any], **kwargs: Any) -> Iterator[Any]:
         """Retrieve GRIB data from the database.
 
         Parameters
@@ -574,9 +572,9 @@ class GribIndex:
 @legacy_source(__file__)
 def execute(
     context: Any,
-    dates: List[Any],
+    dates: list[Any],
     indexdb: str,
-    flavour: Optional[str] = None,
+    flavour: str | None = None,
     **kwargs: Any,
 ) -> FieldArray:
     """Execute the GRIB data retrieval process.

anemoi/datasets/create/sources/hindcasts.py

@@ -9,9 +9,6 @@
 
 import logging
 from typing import Any
-from typing import Dict
-from typing import List
-from typing import Union
 
 from earthkit.data.core.fieldlist import MultiFieldList
 
@@ -22,7 +19,7 @@ from .legacy import legacy_source
 LOGGER = logging.getLogger(__name__)
 
 
-def _to_list(x: Union[list, tuple, Any]) -> List[Any]:
+def _to_list(x: list | tuple | Any) -> list[Any]:
     """Converts the input to a list if it is not already a list or tuple.
 
     Parameters
@@ -41,7 +38,7 @@ def _to_list(x: Union[list, tuple, Any]) -> List[Any]:
 
 
 @legacy_source(__file__)
-def hindcasts(context: Any, dates: List[Any], **request: Dict[str, Any]) -> MultiFieldList:
+def hindcasts(context: Any, dates: list[Any], **request: dict[str, Any]) -> MultiFieldList:
     """Generates hindcast requests based on the provided dates and request parameters.
 
     Parameters

anemoi/datasets/create/sources/legacy.py

@@ -11,8 +11,8 @@
 import inspect
 import logging
 import os
+from collections.abc import Callable
 from typing import Any
-from typing import Callable
 
 from anemoi.datasets.create.input.template import resolve
 

anemoi/datasets/create/sources/mars.py

@@ -9,12 +9,8 @@
 
 import datetime
 import re
+from collections.abc import Generator
 from typing import Any
-from typing import Dict
-from typing import Generator
-from typing import List
-from typing import Optional
-from typing import Union
 
 from anemoi.utils.humanize import did_you_mean
 from earthkit.data import from_source
@@ -27,7 +23,7 @@ from .legacy import legacy_source
 DEBUG = False
 
 
-def to_list(x: Union[list, tuple, Any]) -> list:
+def to_list(x: list | tuple | Any) -> list:
     """Converts the input to a list if it is not already a list or tuple.
 
     Parameters
@@ -46,8 +42,8 @@ def to_list(x: Union[list, tuple, Any]) -> list:
 
 
 def _date_to_datetime(
-    d: Union[datetime.datetime, list, tuple, str],
-) -> Union[datetime.datetime, List[datetime.datetime]]:
+    d: datetime.datetime | list | tuple | str,
+) -> datetime.datetime | list[datetime.datetime]:
     """Converts the input date(s) to datetime objects.
 
     Parameters
@@ -67,7 +63,7 @@ def _date_to_datetime(
     return datetime.datetime.fromisoformat(d)
 
 
-def expand_to_by(x: Union[str, int, list]) -> Union[str, int, list]:
+def expand_to_by(x: str | int | list) -> str | int | list:
     """Expands a range expression to a list of values.
 
     Parameters
@@ -97,7 +93,7 @@ def expand_to_by(x: Union[str, int, list]) -> Union[str, int, list]:
     return x
 
 
-def normalise_time_delta(t: Union[datetime.timedelta, str]) -> datetime.timedelta:
+def normalise_time_delta(t: datetime.timedelta | str) -> datetime.timedelta:
     """Normalizes a time delta string to a datetime.timedelta object.
 
     Parameters
@@ -120,7 +116,7 @@ def normalise_time_delta(t: Union[datetime.timedelta, str]) -> datetime.timedelt
     return t
 
 
-def _normalise_time(t: Union[int, str]) -> str:
+def _normalise_time(t: int | str) -> str:
     """Normalizes a time value to a string in HHMM format.
 
     Parameters
@@ -136,15 +132,15 @@ def _normalise_time(t: Union[int, str]) -> str:
     t = int(t)
     if t < 100:
         t * 100
-    return "{:04d}".format(t)
+    return f"{t:04d}"
 
 
 def _expand_mars_request(
-    request: Dict[str, Any],
+    request: dict[str, Any],
     date: datetime.datetime,
     request_already_using_valid_datetime: bool = False,
     date_key: str = "date",
-) -> List[Dict[str, Any]]:
+) -> list[dict[str, Any]]:
     """Expands a MARS request with the given date and other parameters.
 
     Parameters
@@ -222,11 +218,11 @@ def _expand_mars_request(
 
 
 def factorise_requests(
-    dates: List[datetime.datetime],
-    *requests: Dict[str, Any],
+    dates: list[datetime.datetime],
+    *requests: dict[str, Any],
     request_already_using_valid_datetime: bool = False,
     date_key: str = "date",
-) -> Generator[Dict[str, Any], None, None]:
+) -> Generator[dict[str, Any], None, None]:
     """Factorizes the requests based on the given dates.
 
     Parameters
@@ -268,7 +264,7 @@ def factorise_requests(
         yield r
 
 
-def use_grib_paramid(r: Dict[str, Any]) -> Dict[str, Any]:
+def use_grib_paramid(r: dict[str, Any]) -> dict[str, Any]:
     """Converts the parameter short names to GRIB parameter IDs.
 
     Parameters
@@ -365,11 +361,11 @@ MARS_KEYS = [
 @legacy_source(__file__)
 def mars(
     context: Any,
-    dates: List[datetime.datetime],
-    *requests: Dict[str, Any],
+    dates: list[datetime.datetime],
+    *requests: dict[str, Any],
     request_already_using_valid_datetime: bool = False,
     date_key: str = "date",
-    use_cdsapi_dataset: Optional[str] = None,
+    use_cdsapi_dataset: str | None = None,
     **kwargs: Any,
 ) -> Any:
     """Executes MARS requests based on the given context, dates, and other parameters.

anemoi/datasets/create/sources/netcdf.py

@@ -9,7 +9,6 @@
 
 
 from typing import Any
-from typing import List
 
 import earthkit.data as ekd
 
@@ -18,7 +17,7 @@ from .xarray import load_many
 
 
 @legacy_source(__file__)
-def execute(context: Any, dates: List[str], path: str, *args: Any, **kwargs: Any) -> ekd.FieldList:
+def execute(context: Any, dates: list[str], path: str, *args: Any, **kwargs: Any) -> ekd.FieldList:
     """Execute the loading of multiple NetCDF files.
 
     Parameters

anemoi/datasets/create/sources/opendap.py

@@ -9,8 +9,6 @@
 
 
 from typing import Any
-from typing import Dict
-from typing import List
 
 import earthkit.data as ekd
 
@@ -19,7 +17,7 @@ from .xarray import load_many
 
 
 @legacy_source(__file__)
-def execute(context: Dict[str, Any], dates: List[str], url: str, *args: Any, **kwargs: Any) -> ekd.FieldList:
+def execute(context: dict[str, Any], dates: list[str], url: str, *args: Any, **kwargs: Any) -> ekd.FieldList:
     """Execute the data loading process from an OpenDAP source.
 
     Parameters

anemoi/datasets/create/sources/patterns.py

@@ -9,15 +9,13 @@
 
 import datetime
 import glob
+from collections.abc import Generator
 from typing import Any
-from typing import Generator
-from typing import List
-from typing import Tuple
 
 from earthkit.data.utils.patterns import Pattern
 
 
-def _expand(paths: List[str]) -> Generator[str, None, None]:
+def _expand(paths: list[str]) -> Generator[str, None, None]:
     """Expand the given paths to include all matching file paths.
 
     Parameters
@@ -54,8 +52,8 @@ def _expand(paths: List[str]) -> Generator[str, None, None]:
 
 
 def iterate_patterns(
-    path: str, dates: List[datetime.datetime], **kwargs: Any
-) -> Generator[Tuple[str, List[str]], None, None]:
+    path: str, dates: list[datetime.datetime], **kwargs: Any
+) -> Generator[tuple[str, list[str]], None, None]:
     """Iterate over patterns and expand them with given dates and additional keyword arguments.
 
     Parameters

anemoi/datasets/create/sources/planetary_computer.py

@@ -0,0 +1,44 @@
+# (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 . import source_registry
+from .xarray import XarraySourceBase
+
+
+@source_registry.register("planetary_computer")
+class PlanetaryComputerSource(XarraySourceBase):
+    """An Xarray data source for the planetary_computer."""
+
+    emoji = "🪐"
+
+    def __init__(self, context, data_catalog_id, version="v1", *args, **kwargs: dict):
+
+        import planetary_computer
+        import pystac_client
+
+        self.data_catalog_id = data_catalog_id
+        self.flavour = kwargs.pop("flavour", None)
+        self.patch = kwargs.pop("patch", None)
+        self.options = kwargs.pop("options", {})
+
+        catalog = pystac_client.Client.open(
+            f"https://planetarycomputer.microsoft.com/api/stac/{version}/",
+            modifier=planetary_computer.sign_inplace,
+        )
+        collection = catalog.get_collection(self.data_catalog_id)
+
+        asset = collection.assets["zarr-abfs"]
+
+        if "xarray:storage_options" in asset.extra_fields:
+            self.options["storage_options"] = asset.extra_fields["xarray:storage_options"]
+
+        self.options.update(asset.extra_fields["xarray:open_kwargs"])
+
+        super().__init__(context, url=asset.href, *args, **kwargs)

anemoi/datasets/create/sources/recentre.py

@@ -9,9 +9,6 @@
 
 from copy import deepcopy
 from typing import Any
-from typing import Dict
-from typing import List
-from typing import Union
 
 from anemoi.datasets.compute.recentre import recentre as _recentre
 
@@ -19,7 +16,7 @@ from .legacy import legacy_source
 from .mars import mars
 
 
-def to_list(x: Union[list, tuple, str]) -> List:
+def to_list(x: list | tuple | str) -> list:
     """Converts the input to a list. If the input is a string, it splits it by '/'.
 
     Parameters
@@ -39,7 +36,7 @@ def to_list(x: Union[list, tuple, str]) -> List:
     return [x]
 
 
-def normalise_number(number: Union[list, tuple, str]) -> List[int]:
+def normalise_number(number: list | tuple | str) -> list[int]:
     """Normalises the input number to a list of integers.
 
     Parameters
@@ -63,7 +60,7 @@ def normalise_number(number: Union[list, tuple, str]) -> List[int]:
     return number
 
 
-def normalise_request(request: Dict) -> Dict:
+def normalise_request(request: dict) -> dict:
     """Normalises the request dictionary by converting certain fields to lists.
 
     Parameters
@@ -85,7 +82,7 @@ def normalise_request(request: Dict) -> Dict:
     return request
 
 
-def load_if_needed(context: Any, dates: Any, dict_or_dataset: Union[Dict, Any]) -> Any:
+def load_if_needed(context: Any, dates: Any, dict_or_dataset: dict | Any) -> Any:
     """Loads the dataset if the input is a dictionary, otherwise returns the input.
 
     Parameters
@@ -112,11 +109,11 @@ def load_if_needed(context: Any, dates: Any, dict_or_dataset: Union[Dict, Any])
 def recentre(
     context: Any,
     dates: Any,
-    members: Union[Dict, Any],
-    centre: Union[Dict, Any],
+    members: dict | Any,
+    centre: dict | Any,
     alpha: float = 1.0,
-    remapping: Dict = {},
-    patches: Dict = {},
+    remapping: dict = {},
+    patches: dict = {},
 ) -> Any:
     """Recentres the members dataset using the centre dataset.
 

anemoi/datasets/create/sources/source.py

@@ -9,9 +9,6 @@
 
 from datetime import datetime
 from typing import Any
-from typing import Dict
-from typing import List
-from typing import Optional
 
 from earthkit.data import from_source
 
@@ -21,7 +18,7 @@ from .legacy import legacy_source
 
 
 @legacy_source(__file__)
-def source(context: Optional[Any], dates: List[datetime], **kwargs: Any) -> Any:
+def source(context: Any | None, dates: list[datetime], **kwargs: Any) -> Any:
     """Generates a source based on the provided context, dates, and additional keyword arguments.
 
     Parameters
@@ -52,7 +49,7 @@ execute = source
 if __name__ == "__main__":
     import yaml
 
-    config: Dict[str, Any] = yaml.safe_load(
+    config: dict[str, Any] = yaml.safe_load(
         """
       name: mars
       class: ea
@@ -65,7 +62,7 @@ if __name__ == "__main__":
       time: $from_dates
     """
     )
-    dates: List[str] = yaml.safe_load("[2022-12-30 18:00, 2022-12-31 00:00, 2022-12-31 06:00, 2022-12-31 12:00]")
+    dates: list[str] = yaml.safe_load("[2022-12-30 18:00, 2022-12-31 00:00, 2022-12-31 06:00, 2022-12-31 12:00]")
     dates = to_datetime_list(dates)
 
     for f in source(None, dates, **config):