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

anemoi/datasets/create/{functions/filters → filters}/rotate_winds.py

@@ -9,44 +9,59 @@
 
 
 from collections import defaultdict
+from typing import Any
+from typing import Dict
+from typing import Optional
 
+import earthkit.data as ekd
 import tqdm
+from anemoi.transform.fields import new_field_from_numpy
+from anemoi.transform.fields import new_fieldlist_from_list
 from anemoi.utils.humanize import plural
-from earthkit.data.indexing.fieldlist import FieldArray
 from earthkit.geo.rotate import rotate_vector
 
-
-class NewDataField:
-    def __init__(self, field, data):
-        self.field = field
-        self.data = data
-
-    def to_numpy(self, *args, **kwargs):
-        return self.data
-
-    def __getattr__(self, name):
-        return getattr(self.field, name)
-
-    def __repr__(self) -> str:
-        return repr(self.field)
+from .legacy import legacy_filter
 
 
+@legacy_filter(__file__)
 def execute(
-    context,
-    input,
-    x_wind,
-    y_wind,
-    source_projection=None,
-    target_projection="+proj=longlat",
-):
+    context: Any,
+    input: ekd.FieldList,
+    x_wind: str,
+    y_wind: str,
+    source_projection: Optional[str] = None,
+    target_projection: str = "+proj=longlat",
+) -> ekd.FieldList:
+    """Rotate wind components from one projection to another.
+
+    Parameters
+    ----------
+    context : Any
+        The context in which the function is executed.
+    input : ekd.FieldList
+        List of input fields.
+    x_wind : str
+        X wind component parameter.
+    y_wind : str
+        Y wind component parameter.
+    source_projection : Optional[str], optional
+        Source projection, by default None.
+    target_projection : str, optional
+        Target projection, by default "+proj=longlat".
+
+    Returns
+    -------
+    ekd.FieldList
+        Array of fields with rotated wind components.
+    """
     from pyproj import CRS
 
     context.trace("🔄", "Rotating winds (extracting winds from ", plural(len(input), "field"))
 
-    result = FieldArray()
+    result = []
 
-    wind_params = (x_wind, y_wind)
-    wind_pairs = defaultdict(dict)
+    wind_params: tuple[str, str] = (x_wind, y_wind)
+    wind_pairs: Dict[tuple, Dict[str, Any]] = defaultdict(dict)
 
     for f in input:
         key = f.metadata(namespace="mars")
@@ -84,7 +99,7 @@ def execute(
             target_projection,
         )
 
-        result.append(NewDataField(x, x_new))
-        result.append(NewDataField(y, y_new))
+        result.append(new_field_from_numpy(x, x_new))
+        result.append(new_field_from_numpy(y, y_new))
 
-    return result
+    return new_fieldlist_from_list(result)

anemoi/datasets/create/{functions/filters → filters}/single_level_dewpoint_to_relative_humidity.py

@@ -9,19 +9,42 @@
 
 
 from collections import defaultdict
+from typing import Any
+from typing import Dict
 
-from earthkit.data.indexing.fieldlist import FieldArray
+import earthkit.data as ekd
+from anemoi.transform.fields import new_field_from_numpy
+from anemoi.transform.fields import new_fieldlist_from_list
 from earthkit.meteo import thermo
 
-from .single_level_specific_humidity_to_relative_humidity import NewDataField
+from .legacy import legacy_filter
 
 
-def execute(context, input, t, td, rh="d"):
-    """Convert relative humidity on single levels to dewpoint"""
-    result = FieldArray()
+@legacy_filter(__file__)
+def execute(context: Any, input: ekd.FieldList, t: str, td: str, rh: str = "d") -> ekd.FieldList:
+    """Convert dewpoint on single levels to relative humidity.
 
-    params = (t, td)
-    pairs = defaultdict(dict)
+    Parameters
+    ----------
+    context : Any
+        The context in which the function is executed.
+    input : List[Any]
+        List of input fields.
+    t : str
+        Temperature parameter.
+    td : str
+        Dewpoint parameter.
+    rh : str, optional
+        Relative humidity parameter. Defaults to "d".
+
+    Returns
+    -------
+    ekd.FieldList
+        Array of fields with relative humidity.
+    """
+    result = []
+    params: tuple[str, str] = (t, td)
+    pairs: Dict[tuple, Dict[str, Any]] = defaultdict(dict)
 
     # Gather all necessary fields
     for f in input:
@@ -50,6 +73,6 @@ def execute(context, input, t, td, rh="d"):
         td_values = values[td].to_numpy(flatten=True)
         # actual conversion from td --> rh
         rh_values = thermo.relative_humidity_from_dewpoint(t=t_values, td=td_values)
-        result.append(NewDataField(values[td], rh_values, rh))
+        result.append(new_field_from_numpy(values[td], rh_values, param=rh))
 
-    return result
+    return new_fieldlist_from_list(result)

anemoi/datasets/create/{functions/filters → filters}/single_level_relative_humidity_to_dewpoint.py

@@ -7,23 +7,47 @@
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 
-
 from collections import defaultdict
+from typing import Any
+from typing import Dict
+from typing import Tuple
 
+import earthkit.data as ekd
+from anemoi.transform.fields import new_field_from_numpy
+from anemoi.transform.fields import new_fieldlist_from_list
 from earthkit.data.indexing.fieldlist import FieldArray
 from earthkit.meteo import thermo
 
-from .single_level_specific_humidity_to_relative_humidity import NewDataField
+from .legacy import legacy_filter
 
 EPS = 1.0e-4
 
 
-def execute(context, input, t, rh, td="d"):
-    """Convert relative humidity on single levels to dewpoint"""
-    result = FieldArray()
+@legacy_filter(__file__)
+def execute(context: Any, input: ekd.FieldList, t: str, rh: str, td: str = "d") -> FieldArray:
+    """Convert relative humidity on single levels to dewpoint.
+
+    Parameters
+    ----------
+    context : Any
+        The context in which the function is executed.
+    input : List[Any]
+        List of input fields.
+    t : str
+        Temperature parameter.
+    rh : str
+        Relative humidity parameter.
+    td : str, optional
+        Dewpoint parameter. Defaults to "d".
 
-    params = (t, rh)
-    pairs = defaultdict(dict)
+    Returns
+    -------
+    FieldArray
+        Array of fields with dewpoint.
+    """
+    result = []
+    params: Tuple[str, str] = (t, rh)
+    pairs: Dict[Tuple[Any, ...], Dict[str, Any]] = defaultdict(dict)
 
     # Gather all necessary fields
     for f in input:
@@ -55,6 +79,6 @@ def execute(context, input, t, rh, td="d"):
         rh_values[rh_values == 0] = EPS
         # actual conversion from rh --> td
         td_values = thermo.dewpoint_from_relative_humidity(t=t_values, r=rh_values)
-        result.append(NewDataField(values[rh], td_values, td))
+        result.append(new_field_from_numpy(values[rh], td_values, param=td))
 
-    return result
+    return new_fieldlist_from_list(result)

anemoi/datasets/create/{functions/filters → filters}/single_level_relative_humidity_to_specific_humidity.py

@@ -8,18 +8,66 @@
 # nor does it submit to any jurisdiction.
 
 
+from typing import Any
+from typing import Dict
+from typing import List
+
+import earthkit.data as ekd
 import numpy as np
-from earthkit.data.indexing.fieldlist import FieldArray
+from anemoi.transform.fields import new_field_from_numpy
+from anemoi.transform.fields import new_fieldlist_from_list
 from earthkit.meteo import thermo
 
+from .legacy import legacy_filter
 from .single_level_specific_humidity_to_relative_humidity import AutoDict
-from .single_level_specific_humidity_to_relative_humidity import NewDataField
 from .single_level_specific_humidity_to_relative_humidity import pressure_at_height_level
 
 
-def execute(context, input, height, t, rh, sp, new_name="2q", **kwargs):
-    """Convert the single (height) level relative humidity to specific humidity"""
-    result = FieldArray()
+@legacy_filter(__file__)
+def execute(
+    context: Any,
+    input: List[Any],
+    height: float,
+    t: str,
+    rh: str,
+    sp: str,
+    new_name: str = "2q",
+    **kwargs: Dict[str, Any],
+) -> ekd.FieldList:
+    """Convert the single (height) level relative humidity to specific humidity.
+
+    Parameters
+    ----------
+    context : Any
+        The context in which the function is executed.
+    input : List[Any]
+        List of input fields.
+    height : float
+        The height level.
+    t : str
+        Temperature parameter name.
+    rh : str
+        Relative humidity parameter name.
+    sp : str
+        Surface pressure parameter name.
+    new_name : str, optional
+        The new name for the specific humidity field, by default "2q".
+    **kwargs : Dict[str, Any]
+        Additional keyword arguments.
+
+    Returns
+    -------
+    ekd.FieldList
+        The resulting field list with specific humidity fields.
+
+    Raises
+    ------
+    KeyError
+        If mandatory keys are missing.
+    ValueError
+        If there are duplicate fields or missing fields.
+    """
+    result = []
 
     MANDATORY_KEYS = ["A", "B"]
     OPTIONAL_KEYS = ["t_ml", "q_ml"]
@@ -110,6 +158,6 @@ def execute(context, input, height, t, rh, sp, new_name="2q", **kwargs):
         p_sl = pressure_at_height_level(height, q_ml, t_ml, sp_sl, np.array(kwargs["A"]), np.array(kwargs["B"]))
         q_sl = thermo.specific_humidity_from_relative_humidity(t_sl, rh_sl, p_sl)
 
-        result.append(NewDataField(values["sfc"][rh], q_sl, new_name))
+        result.append(new_field_from_numpy(values["sfc"][rh], q_sl, param=new_name))
 
-    return result
+    return new_fieldlist_from_list(result)

anemoi/datasets/create/{functions/filters → filters}/single_level_specific_humidity_to_relative_humidity.py

@@ -8,46 +8,51 @@
 # nor does it submit to any jurisdiction.
 
 
+from typing import Any
+from typing import Dict
+from typing import List
+from typing import Tuple
+from typing import Union
+
+import earthkit.data as ekd
 import numpy as np
-from earthkit.data.indexing.fieldlist import FieldArray
+from anemoi.transform.fields import new_field_from_numpy
+from anemoi.transform.fields import new_fieldlist_from_list
 from earthkit.meteo import constants
 from earthkit.meteo import thermo
+from numpy.typing import NDArray
+
+from .legacy import legacy_filter
 
 
 # Alternative proposed by Baudouin Raoult
 class AutoDict(dict):
-    def __missing__(self, key):
-        value = self[key] = type(self)()
-        return value
-
+    """A dictionary that automatically creates nested dictionaries for missing keys."""
 
-class NewDataField:
-    def __init__(self, field, data, new_name):
-        self.field = field
-        self.data = data
-        self.new_name = new_name
+    def __missing__(self, key: Any) -> Any:
+        """Handle missing keys by creating nested dictionaries.
 
-    def to_numpy(self, *args, **kwargs):
-        return self.data
+        Parameters
+        ----------
+        key : Any
+            The missing key.
 
-    def metadata(self, key=None, **kwargs):
-        if key is None:
-            return self.field.metadata(**kwargs)
-
-        value = self.field.metadata(key, **kwargs)
-        if key == "param":
-            return self.new_name
+        Returns
+        -------
+        Any
+            A new nested dictionary.
+        """
+        value = self[key] = type(self)()
         return value
 
-    def __getattr__(self, name):
-        return getattr(self.field, name)
-
 
-def model_level_pressure(A, B, surface_pressure):
+def model_level_pressure(
+    A: NDArray[Any], B: NDArray[Any], surface_pressure: Union[float, np.ndarray]
+) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
     """Calculates:
      - pressure at the model full- and half-levels
      - delta: depth of log(pressure) at full levels
-     - alpha: alpha term #TODO: more descriptive information
+     - alpha: alpha term #TODO: more descriptive information.
 
     Parameters
     ----------
@@ -55,7 +60,7 @@ def model_level_pressure(A, B, surface_pressure):
         A-coefficients defining the model levels
     B : ndarray
         B-coefficients defining the model levels
-    surface_pressure: number or ndarray
+    surface_pressure : number or ndarray
         surface pressure (Pa)
 
     Returns
@@ -114,9 +119,9 @@ def model_level_pressure(A, B, surface_pressure):
     return p_full_level, p_half_level, delta, alpha
 
 
-def calc_specific_gas_constant(q):
+def calc_specific_gas_constant(q: Union[float, np.ndarray]) -> Union[float, NDArray[Any]]:
     """Calculates the specific gas constant of moist air
-    (specific content of cloud particles and hydrometeors are neglected)
+    (specific content of cloud particles and hydrometeors are neglected).
 
     Parameters
     ----------
@@ -133,8 +138,8 @@ def calc_specific_gas_constant(q):
     return R
 
 
-def relative_geopotential_thickness(alpha, q, T):
-    """Calculates the geopotential thickness w.r.t the surface on model full-levels
+def relative_geopotential_thickness(alpha: NDArray[Any], q: NDArray[Any], T: NDArray[Any]) -> NDArray[Any]:
+    """Calculates the geopotential thickness w.r.t the surface on model full-levels.
 
     Parameters
     ----------
@@ -158,10 +163,12 @@ def relative_geopotential_thickness(alpha, q, T):
     return dphi
 
 
-def pressure_at_height_level(height, q, T, sp, A, B):
+def pressure_at_height_level(
+    height: float, q: NDArray[Any], T: NDArray[Any], sp: NDArray[Any], A: NDArray[Any], B: NDArray[Any]
+) -> Union[float, NDArray[Any]]:
     """Calculates the pressure at a height level given in meters above surface.
     This is done by finding the model level above and below the specified height
-    and interpolating the pressure
+    and interpolating the pressure.
 
     Parameters
     ----------
@@ -225,9 +232,54 @@ def pressure_at_height_level(height, q, T, sp, A, B):
     return p_height
 
 
-def execute(context, input, height, t, q, sp, new_name="2r", **kwargs):
-    """Convert the single (height) level specific humidity to relative humidity"""
-    result = FieldArray()
+@legacy_filter(__file__)
+def execute(
+    context: Any,
+    input: List[Any],
+    height: float,
+    t: str,
+    q: str,
+    sp: str,
+    new_name: str = "2r",
+    **kwargs: Dict[str, Any],
+) -> ekd.FieldList:
+    """Convert the single (height) level specific humidity to relative humidity.
+
+    Parameters
+    ----------
+    context : Any
+        The context for the execution.
+    input : list of Any
+        The input data.
+    height : float
+        The height level in meters.
+    t : str
+        The temperature parameter name.
+    q : str
+        The specific humidity parameter name.
+    sp : str
+        The surface pressure parameter name.
+    new_name : str, optional
+        The new name for the relative humidity parameter, by default "2r".
+    **kwargs : dict
+        Additional keyword arguments.
+        t_ml : str, optional
+            The temperature parameter name for model levels, by default "t".
+        q_ml : str, optional
+            The specific humidity parameter name for model levels, by default "q".
+        A : list of float
+            A-coefficients defining the model levels.
+        B : list of float
+            B-coefficients defining the model levels.
+        keep_q : bool, optional
+            Whether to keep the specific humidity field in the result, by default False.
+
+    Returns
+    -------
+    ekd.FieldList
+        The resulting field array with relative humidity.
+    """
+    result = []
 
     MANDATORY_KEYS = ["A", "B"]
     OPTIONAL_KEYS = ["t_ml", "q_ml"]
@@ -323,12 +375,21 @@ def execute(context, input, height, t, q, sp, new_name="2r", **kwargs):
         td_sl = thermo.dewpoint_from_specific_humidity(q=q_sl, p=p_sl)
         rh_sl = thermo.relative_humidity_from_dewpoint(t=t_sl, td=td_sl)
 
-        result.append(NewDataField(values["sfc"][q], rh_sl, new_name))
+        result.append(new_field_from_numpy(values["sfc"][q], rh_sl, param=new_name))
 
-    return result
+    return new_fieldlist_from_list(result)
 
 
-def test():
+def test() -> None:
+    """Test the conversion from specific humidity to relative humidity.
+
+    This function fetches data from a source, performs the conversion, and prints
+    the mean, median, and maximum differences in dewpoint temperature.
+
+    Returns
+    -------
+    None
+    """
     from earthkit.data import from_source
     from earthkit.data.readers.grib.index import GribFieldList
 

anemoi/datasets/create/filters/speeddir_to_uv.py

@@ -0,0 +1,95 @@
+# (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 collections import defaultdict
+from typing import Any
+from typing import List
+
+import earthkit.data as ekd
+import numpy as np
+from anemoi.transform.fields import new_field_from_numpy
+from anemoi.transform.fields import new_fieldlist_from_list
+from earthkit.meteo.wind.array import polar_to_xy
+
+from .legacy import legacy_filter
+
+
+@legacy_filter(__file__)
+def execute(
+    context: Any,
+    input: List[Any],
+    wind_speed: str,
+    wind_dir: str,
+    u_component: str = "u",
+    v_component: str = "v",
+    in_radians: bool = False,
+) -> ekd.FieldList:
+    """Convert wind speed and direction to u and v components.
+
+    Parameters
+    ----------
+    context : Any
+        The context for the execution.
+    input : List[Any]
+        The input data fields.
+    wind_speed : str
+        The name of the wind speed parameter.
+    wind_dir : str
+        The name of the wind direction parameter.
+    u_component : str, optional
+        The name for the u component. Defaults to "u".
+    v_component : str, optional
+        The name for the v component. Defaults to "v".
+    in_radians : bool, optional
+        Whether the wind direction is in radians. Defaults to False.
+
+    Returns
+    -------
+    ekd.FieldList
+        The resulting field array with u and v components.
+    """
+
+    result = []
+
+    wind_params = (wind_speed, wind_dir)
+    wind_pairs = defaultdict(dict)
+
+    for f in input:
+        key = f.metadata(namespace="mars")
+        param = key.pop("param")
+
+        if param not in wind_params:
+            result.append(f)
+            continue
+
+        key = tuple(key.items())
+
+        if param in wind_pairs[key]:
+            raise ValueError(f"Duplicate wind component {param} for {key}")
+
+        wind_pairs[key][param] = f
+
+    for _, pairs in wind_pairs.items():
+        if len(pairs) != 2:
+            raise ValueError("Missing wind component")
+
+        magnitude = pairs[wind_speed]
+        direction = pairs[wind_dir]
+
+        # assert speed.grid_mapping == dir.grid_mapping
+        if in_radians:
+            direction = np.rad2deg(direction)
+
+        u, v = polar_to_xy(magnitude.to_numpy(flatten=True), direction.to_numpy(flatten=True))
+
+        result.append(new_field_from_numpy(magnitude, u, param=u_component))
+        result.append(new_field_from_numpy(direction, v, param=v_component))
+
+    return new_fieldlist_from_list(result)

anemoi/datasets/create/{functions/filters → filters}/sum.py

@@ -7,39 +7,36 @@
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 
-
 from collections import defaultdict
+from typing import Any
+from typing import Dict
+from typing import Hashable
+from typing import List
+from typing import Tuple
 
-from earthkit.data.indexing.fieldlist import FieldArray
-
-
-class NewDataField:
-    def __init__(self, field, data, new_name):
-        self.field = field
-        self.data = data
-        self.new_name = new_name
-
-    def to_numpy(self, *args, **kwargs):
-        return self.data
+import earthkit.data as ekd
+from anemoi.transform.fields import new_field_from_numpy
+from anemoi.transform.fields import new_fieldlist_from_list
 
-    def metadata(self, key=None, **kwargs):
-        if key is None:
-            return self.field.metadata(**kwargs)
+from .legacy import legacy_filter
 
-        value = self.field.metadata(key, **kwargs)
-        if key == "param":
-            return self.new_name
-        return value
 
-    def __getattr__(self, name):
-        return getattr(self.field, name)
+@legacy_filter(__file__)
+def execute(context: Any, input: ekd.FieldList, params: List[str], output: str) -> ekd.FieldList:
+    """Computes the sum over a set of variables.
 
+    Args:
+        context (Any): The execution context.
+        input (List[Any]): The list of input fields.
+        params (List[str]): The list of parameters to sum over.
+        output (str): The name for the output field.
 
-def execute(context, input, params, output):
-    """Computes the sum over a set of variables"""
-    result = FieldArray()
+    Returns:
+        ekd.FieldList: The resulting FieldArray with summed fields.
+    """
+    result = []
 
-    needed_fields = defaultdict(dict)
+    needed_fields: Dict[Tuple[Hashable, ...], Dict[str, ekd.Field]] = defaultdict(dict)
 
     for f in input:
         key = f.metadata(namespace="mars")
@@ -66,6 +63,6 @@ def execute(context, input, params, output):
                 s = c
             else:
                 s += c
-        result.append(NewDataField(values[list(values.keys())[0]], s, output))
+        result.append(new_field_from_numpy(values[list(values.keys())[0]], s, param=output))
 
-    return result
+    return new_fieldlist_from_list(result)