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

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

@@ -11,20 +11,41 @@ import datetime
 import logging
 import warnings
 from copy import deepcopy
+from typing import Any
+from typing import Dict
+from typing import Generator
+from typing import List
+from typing import Optional
+from typing import Tuple
+from typing import Union
 
 import earthkit.data as ekd
 import numpy as np
 from earthkit.data.core.temporary import temp_file
 from earthkit.data.readers.grib.output import new_grib_output
+from numpy.typing import NDArray
 
 from anemoi.datasets.create.utils import to_datetime_list
 
+from .legacy import legacy_source
 from .mars import mars
 
 LOG = logging.getLogger(__name__)
 
 
-def _member(field):
+def _member(field: Any) -> int:
+    """Retrieves the member number from the field metadata.
+
+    Parameters
+    ----------
+    field : Any
+        The field from which to retrieve the member number.
+
+    Returns
+    -------
+    int
+        The member number.
+    """
     # Bug in eccodes has number=0 randomly
     number = field.metadata("number", default=0)
     if number is None:
@@ -33,26 +54,70 @@ def _member(field):
 
 
 class Accumulation:
-    def __init__(self, out, /, param, date, time, number, step, frequency, **kwargs):
+    """Class to handle data accumulation for a specific parameter, date, time, and member."""
+
+    buggy_steps: bool = False
+
+    def __init__(
+        self,
+        out: Any,
+        /,
+        param: str,
+        date: int,
+        time: int,
+        number: int,
+        step: List[int],
+        frequency: int,
+        **kwargs: Any,
+    ) -> None:
+        """Initializes an Accumulation instance.
+
+        Parameters
+        ----------
+        out : Any
+            Output object for writing data.
+        param : str
+            Parameter name.
+        date : int
+            Date of the accumulation.
+        time : int
+            Time of the accumulation.
+        number : int
+            Member number.
+        step : List[int]
+            List of steps.
+        frequency : int
+            Frequency of accumulation.
+        **kwargs : Any
+            Additional keyword arguments.
+        """
         self.out = out
         self.param = param
         self.date = date
         self.time = time
         self.steps = step
         self.number = number
-        self.values = None
+        self.values: Optional[NDArray[None]] = None
         self.seen = set()
-        self.startStep = None
-        self.endStep = None
+        self.startStep: Optional[int] = None
+        self.endStep: Optional[int] = None
         self.done = False
         self.frequency = frequency
         self._check = None
 
     @property
-    def key(self):
+    def key(self) -> Tuple[str, int, int, List[int], int]:
+        """Returns the key for the accumulation."""
         return (self.param, self.date, self.time, self.steps, self.number)
 
-    def check(self, field):
+    def check(self, field: Any) -> None:
+        """Checks the field metadata against the accumulation parameters.
+
+        Parameters
+        ----------
+        field : Any
+            The field to check.
+        """
         if self._check is None:
             self._check = field.metadata(namespace="mars")
 
@@ -82,8 +147,14 @@ class Accumulation:
             if k not in ("step",):
                 assert self._check[k] == mars[k], (k, self._check[k], mars[k])
 
-    def write(self, template):
+    def write(self, template: Any) -> None:
+        """Writes the accumulated values to the output.
 
+        Parameters
+        ----------
+        template : Any
+            Template for writing the output.
+        """
         assert self.startStep != self.endStep, (self.startStep, self.endStep)
         if np.all(self.values < 0):
             LOG.warning(
@@ -101,8 +172,16 @@ class Accumulation:
         self.values = None
         self.done = True
 
-    def add(self, field, values):
+    def add(self, field: Any, values: NDArray[Any]) -> None:
+        """Adds values to the accumulation.
 
+        Parameters
+        ----------
+        field : Any
+            The field containing the values.
+        values : np.ndarray
+            The values to add.
+        """
         self.check(field)
 
         step = field.metadata("step")
@@ -131,8 +210,37 @@ class Accumulation:
             self.write(template=field)
 
     @classmethod
-    def mars_date_time_steps(cls, dates, step1, step2, frequency, base_times, adjust_step):
-
+    def mars_date_time_steps(
+        cls,
+        dates: List[datetime.datetime],
+        step1: int,
+        step2: int,
+        frequency: Optional[int],
+        base_times: List[int],
+        adjust_step: bool,
+    ) -> Generator[Tuple[int, int, Tuple[int, ...]], None, None]:
+        """Generates MARS date-time steps.
+
+        Parameters
+        ----------
+        dates : List[datetime.datetime]
+            List of dates.
+        step1 : int
+            First step.
+        step2 : int
+            Second step.
+        frequency : Optional[int]
+            Frequency of accumulation.
+        base_times : List[int]
+            List of base times.
+        adjust_step : bool
+            Whether to adjust the step.
+
+        Returns
+        -------
+        Generator[Tuple[int, int, Tuple[int, ...]], None, None]
+            A generator of MARS date-time steps.
+        """
         # assert step1 > 0, (step1, step2, frequency)
 
         for valid_date in dates:
@@ -152,14 +260,33 @@ class Accumulation:
             yield cls._mars_date_time_step(base_date, step1, step2, add_step, frequency)
 
     def __repr__(self) -> str:
+        """Returns a string representation of the Accumulation instance.
+
+        Returns
+        -------
+        str
+            String representation of the Accumulation instance.
+        """
         return f"{self.__class__.__name__}({self.key})"
 
 
 class AccumulationFromStart(Accumulation):
-    buggy_steps = True
+    """Class to handle data accumulation from the start of the forecast."""
 
-    def compute(self, values, startStep, endStep):
+    buggy_steps = True
 
+    def compute(self, values: NDArray[Any], startStep: int, endStep: int) -> None:
+        """Computes the accumulation from the start.
+
+        Parameters
+        ----------
+        values : np.ndarray
+            The values to accumulate.
+        startStep : int
+            The start step.
+        endStep : int
+            The end step.
+        """
         assert startStep == 0, startStep
 
         if self.values is None:
@@ -186,7 +313,29 @@ class AccumulationFromStart(Accumulation):
                 self.values = np.maximum(self.values, 0)
 
     @classmethod
-    def _mars_date_time_step(cls, base_date, step1, step2, add_step, frequency):
+    def _mars_date_time_step(
+        cls, base_date: datetime.datetime, step1: int, step2: int, add_step: int, frequency: Optional[int]
+    ) -> Tuple[int, int, Tuple[int, ...]]:
+        """Generates a MARS date-time step.
+
+        Parameters
+        ----------
+        base_date : datetime.datetime
+            The base date.
+        step1 : int
+            First step.
+        step2 : int
+            Second step.
+        add_step : int
+            Additional step.
+        frequency : Optional[int]
+            Frequency of accumulation.
+
+        Returns
+        -------
+        Tuple[int, int, Tuple[int, ...]]
+            A tuple representing the MARS date-time step.
+        """
         assert not frequency, frequency
 
         steps = (step1 + add_step, step2 + add_step)
@@ -201,10 +350,22 @@ class AccumulationFromStart(Accumulation):
 
 
 class AccumulationFromLastStep(Accumulation):
-    buggy_steps = False
+    """Class to handle data accumulation from the last step of the forecast."""
 
-    def compute(self, values, startStep, endStep):
+    buggy_steps = False
 
+    def compute(self, values: NDArray[Any], startStep: int, endStep: int) -> None:
+        """Computes the accumulation from the last step.
+
+        Parameters
+        ----------
+        values : np.ndarray
+            The values to accumulate.
+        startStep : int
+            The start step.
+        endStep : int
+            The end step.
+        """
         assert endStep - startStep == self.frequency, (
             startStep,
             endStep,
@@ -227,7 +388,29 @@ class AccumulationFromLastStep(Accumulation):
         self.values += values
 
     @classmethod
-    def _mars_date_time_step(cls, base_date, step1, step2, add_step, frequency):
+    def _mars_date_time_step(
+        cls, base_date: datetime.datetime, step1: int, step2: int, add_step: int, frequency: int
+    ) -> Tuple[int, int, Tuple[int, ...]]:
+        """Generates a MARS date-time step.
+
+        Parameters
+        ----------
+        base_date : datetime.datetime
+            The base date.
+        step1 : int
+            First step.
+        step2 : int
+            Second step.
+        add_step : int
+            Additional step.
+        frequency : int
+            Frequency of accumulation.
+
+        Returns
+        -------
+        Tuple[int, int, Tuple[int, ...]]
+            A tuple representing the MARS date-time step.
+        """
         assert frequency > 0, frequency
         # assert step1 > 0, (step1, step2, frequency, add_step, base_date)
 
@@ -241,20 +424,58 @@ class AccumulationFromLastStep(Accumulation):
         )
 
 
-def _identity(x):
+def _identity(x: Any) -> Any:
+    """Identity function that returns the input as is.
+
+    Parameters
+    ----------
+    x : Any
+        Input value.
+
+    Returns
+    -------
+    Any
+        The input value.
+    """
     return x
 
 
 def _compute_accumulations(
-    context,
-    dates,
-    request,
-    user_accumulation_period=6,
-    data_accumulation_period=None,
-    patch=_identity,
-    base_times=None,
-    use_cdsapi_dataset=None,
-):
+    context: Any,
+    dates: List[datetime.datetime],
+    request: Dict[str, Any],
+    user_accumulation_period: Union[int, Tuple[int, int]] = 6,
+    data_accumulation_period: Optional[int] = None,
+    patch: Any = _identity,
+    base_times: Optional[List[int]] = None,
+    use_cdsapi_dataset: Optional[str] = None,
+) -> Any:
+    """Computes accumulations based on the provided parameters.
+
+    Parameters
+    ----------
+    context : Any
+        Context for the computation.
+    dates : List[datetime.datetime]
+        List of dates.
+    request : Dict[str, Any]
+        Request parameters.
+    user_accumulation_period : Union[int, Tuple[int, int]], optional
+        User-defined accumulation period. Defaults to 6.
+    data_accumulation_period : Optional[int], optional
+        Data accumulation period. Defaults to None.
+    patch : Any, optional
+        Patch function. Defaults to _identity.
+    base_times : Optional[List[int]], optional
+        List of base times. Defaults to None.
+    use_cdsapi_dataset : Optional[str], optional
+        CDSAPI dataset to use. Defaults to None.
+
+    Returns
+    -------
+    Any
+        The computed accumulations.
+    """
     adjust_step = isinstance(user_accumulation_period, int)
 
     if not isinstance(user_accumulation_period, (list, tuple)):
@@ -355,13 +576,37 @@ def _compute_accumulations(
     return ds
 
 
-def _to_list(x):
+def _to_list(x: Union[List[Any], Tuple[Any], Any]) -> List[Any]:
+    """Converts the input to a list if it is not already a list or tuple.
+
+    Parameters
+    ----------
+    x : Union[List[Any], Tuple[Any], Any]
+        Input value.
+
+    Returns
+    -------
+    List[Any]
+        The input value as a list.
+    """
     if isinstance(x, (list, tuple)):
         return x
     return [x]
 
 
-def _scda(request):
+def _scda(request: Dict[str, Any]) -> Dict[str, Any]:
+    """Modifies the request stream based on the time.
+
+    Parameters
+    ----------
+    request : Dict[str, Any]
+        Request parameters.
+
+    Returns
+    -------
+    Dict[str, Any]
+        The modified request parameters.
+    """
     if request["time"] in (6, 18, 600, 1800):
         request["stream"] = "scda"
     else:
@@ -369,7 +614,28 @@ def _scda(request):
     return request
 
 
-def accumulations(context, dates, use_cdsapi_dataset=None, **request):
+@legacy_source(__file__)
+def accumulations(
+    context: Any, dates: List[datetime.datetime], use_cdsapi_dataset: Optional[str] = None, **request: Any
+) -> Any:
+    """Computes accumulations based on the provided context, dates, and request parameters.
+
+    Parameters
+    ----------
+    context : Any
+        Context for the computation.
+    dates : List[datetime.datetime]
+        List of dates.
+    use_cdsapi_dataset : Optional[str], optional
+        CDSAPI dataset to use. Defaults to None.
+    **request : Any
+        Additional request parameters.
+
+    Returns
+    -------
+    Any
+        The computed accumulations.
+    """
     _to_list(request["param"])
     class_ = request.get("class", "od")
     stream = request.get("stream", "oper")

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

@@ -7,10 +7,35 @@
 # 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 Dict
+from typing import List
+
 from earthkit.data import from_source
 
+from .legacy import legacy_source
+
+
+@legacy_source(__file__)
+def constants(context: Any, dates: List[str], template: Dict[str, Any], param: str) -> Any:
+    """Deprecated function to retrieve constants data.
+
+    Parameters
+    ----------
+    context : Any
+        The context object for tracing.
+    dates : list of str
+        List of dates for which data is required.
+    template : dict of str to Any
+        Template dictionary for the data source.
+    param : str
+        Parameter to retrieve.
 
-def constants(context, dates, template, param):
+    Returns
+    -------
+    Any
+        Data retrieved from the source.
+    """
     from warnings import warn
 
     warn(
@@ -25,4 +50,4 @@ def constants(context, dates, template, param):
     return from_source("forcings", source_or_dataset=template, date=dates, param=param)
 
 
-execute = constants
+execute: Any = constants

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

@@ -8,9 +8,13 @@
 # nor does it submit to any jurisdiction.
 
 
-from .xarray import load_many
+from . import source_registry
+from .xarray import XarraySourceBase
 
 
-def execute(context, dates, path, *args, **kwargs):
+@source_registry.register("eccc_fstd")
+class XarrayECCCSource(XarraySourceBase):
+    """An Xarray data source that uses the `fstd` engine."""
+
+    emoji = "🍁"
     options = {"engine": "fstd"}
-    return load_many("🍁", context, dates, path, *args, options=options, **kwargs)

anemoi/datasets/create/sources/empty.py

@@ -0,0 +1,37 @@
+# (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
+
+
+@legacy_source(__file__)
+def execute(context: Any, dates: List[str], **kwargs: Any) -> ekd.FieldList:
+    """Executes the loading of an empty data source.
+
+    Parameters
+    ----------
+    context : object
+        The context in which the function is executed.
+    dates : list
+        List of dates for which data is to be loaded.
+    **kwargs : dict
+        Additional keyword arguments.
+
+    Returns
+    -------
+    ekd.FieldList
+        Loaded empty data source.
+    """
+    return ekd.from_source("empty")

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

@@ -7,10 +7,34 @@
 # 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
+
 from earthkit.data import from_source
 
+from .legacy import legacy_source
+
+
+@legacy_source(__file__)
+def forcings(context: Any, dates: List[str], template: str, param: str) -> Any:
+    """Loads forcing data from a specified source.
+
+    Parameters
+    ----------
+    context : object
+        The context in which the function is executed.
+    dates : list
+        List of dates for which data is to be loaded.
+    template : str
+        Template for the data source.
+    param : str
+        Parameter for the data source.
 
-def forcings(context, dates, template, param):
+    Returns
+    -------
+    object
+        Loaded forcing data.
+    """
     context.trace("✅", f"from_source(forcings, {template}, {param}")
     return from_source("forcings", source_or_dataset=template, date=dates, param=param)