wip

src/anemoi/transform/__init__.py

@@ -4,6 +4,3 @@
 # 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 ._version import __version__

src/anemoi/transform/data/__init__.py

@@ -0,0 +1,10 @@
+# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# 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.
+
+
+class Data:
+    pass

src/anemoi/transform/filters/__init__.py

@@ -0,0 +1,107 @@
+# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# 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 abc import ABC
+from abc import abstractmethod
+from collections import defaultdict
+from typing import Any
+
+import earthkit.data as ekd
+import numpy as np
+from earthkit.data import FieldList
+from earthkit.meteo.wind.array import polar_to_xy
+from earthkit.meteo.wind.array import xy_to_polar
+
+
+class Filter(ABC):
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}()"
+
+    @abstractmethod
+    def forward(self, x: ekd.FieldList) -> ekd.FieldList:
+        pass
+
+    @abstractmethod
+    def backward(self, x: ekd.FieldList) -> ekd.FieldList:
+        pass
+
+    def reverse(self) -> "Filter":
+        return ReversedFilter(self)
+
+
+class ReversedFilter(Filter):
+    """Swap the forward and backward methods of a filter."""
+
+    def __init__(self, filter) -> None:
+        self.filter = filter
+
+    def __repr__(self) -> str:
+        return f"Reversed({self.filter})"
+
+    def forward(self, x: ekd.FieldList) -> ekd.FieldList:
+        return self.filter.backward(x)
+
+    def backward(self, x: ekd.FieldList) -> ekd.FieldList:
+        return self.filter.forward(x)
+
+
+class TransformFilter(Filter):
+    """A filter to convert only some fields.
+    The fields are matched by their metadata.
+    """
+
+    def _transform(self, data, transform, *group_by):
+
+        result = []
+
+        groups = defaultdict(dict)
+
+        for f in data:
+            key = f.metadata(namespace="mars")
+            param = key.pop("param")
+
+            if param not in group_by:
+                result.append(f)
+                continue
+
+            key = tuple(key.items())
+
+            if param in groups[key]:
+                raise ValueError(f"Duplicate component {param} for {key}")
+
+            groups[key][param] = f
+
+        for _, group in groups.items():
+            if len(group) != len(group_by):
+                raise ValueError("Missing component")
+
+            for f in transform(*[group[p] for p in group_by]):
+                result.append(f)
+
+        return self.new_fieldlits_from_list(result)
+
+    def new_field_from_numpy(self, array, *, template, param):
+        """Create a new field from a numpy array."""
+        md = template.metadata().override(shortName=param)
+        return FieldList.from_array(array, md)[0]
+
+    def new_fieldlits_from_list(self, fields):
+        from earthkit.data.indexing.fieldlist import FieldArray
+
+        return FieldArray(fields)
+
+    @abstractmethod
+    def forward_transform(self, *fields):
+        """To be implemented by subclasses."""
+        pass
+
+    @abstractmethod
+    def backward_transform(self, *fields):
+        """To be implemented by subclasses."""
+        pass

src/anemoi/transform/filters/uv_to_ddff.py

@@ -0,0 +1,103 @@
+# (C) Copyright 2024 ECMWF.
+#
+# 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 numpy as np
+from earthkit.meteo.wind.array import polar_to_xy
+from earthkit.meteo.wind.array import xy_to_polar
+
+from anemoi.transform.filters import TransformFilter
+
+
+class WindComponents(TransformFilter):
+    """A filter to convert wind speed and direction to U and V components,
+    and back.
+    """
+
+    def __init__(
+        self,
+        *,
+        u_component="u",
+        v_component="v",
+        wind_speed="ws",
+        wind_direction="wdir",
+        convention="meteo",
+        radians=False,
+    ):
+        self.u_component = u_component
+        self.v_component = v_component
+        self.wind_speed = wind_speed
+        self.wind_direction = wind_direction
+        self.convention = convention
+        self.radians = radians
+
+    def forward(self, data):
+        return self._transform(
+            data,
+            self.forward_transform,
+            self.u_component,
+            self.v_component,
+        )
+
+    def backward(self, data):
+        return self._transform(
+            data,
+            self.backward_transform,
+            self.wind_speed,
+            self.wind_direction,
+        )
+
+    def forward_transform(self, u, v):
+        """U/V to DD/FF"""
+
+        speed, direction = xy_to_polar(
+            u.to_numpy(),
+            v.to_numpy(),
+            convention=self.convention,
+        )
+        if self.radians:
+            direction = np.deg2rad(direction)
+
+        yield self.new_field_from_numpy(speed, template=u, param=self.wind_speed)
+        yield self.new_field_from_numpy(direction, template=v, param=self.wind_direction)
+
+    def backward_transform(self, speed, direction):
+        """DD/FF to U/V"""
+
+        if self.radians:
+            direction = np.rad2deg(direction)
+
+        u, v = polar_to_xy(
+            speed.to_numpy(),
+            direction.to_numpy(),
+            convention=self.convention,
+        )
+
+        yield self.new_field_from_numpy(u, template=speed, param=self.u_component)
+        yield self.new_field_from_numpy(v, template=direction, param=self.v_component)
+
+
+if __name__ == "__main__":
+    from earthkit.data import from_source
+
+    ################
+    data = from_source(
+        "mars",
+        param=["u", "v", "t", "q"],
+        grid=[1, 1],
+        date="20200101/to/20200105",
+        levelist=[1000, 850, 500],
+    )
+    for f in data:
+        print(f)
+
+    ################
+    data = WindComponents().forward(data)
+    for f in data:
+        print(f)

src/anemoi/transform/sources/__init__.py

@@ -0,0 +1,6 @@
+# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# 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.