Feature/filters for carra (ecmwf#149)

* Add wz_to_w filter

* Add orog_to_z filter

* Add sum filter

* Update CHANGELOG.md

* Add documentation of new filters

* Fix line lengths in documentation

CHANGELOG.md

@@ -20,6 +20,7 @@ Keep it human-readable, your future self will thank you!
 
 - Call filters from anemoi-transform
 - make test optional when adls is not installed Pull request #110
+- Add wz_to_w, orog_to_z, and sum filters (#149)
 
 ## [0.5.8](https://github.com/ecmwf/anemoi-datasets/compare/0.5.7...0.5.8) - 2024-10-26
 

docs/building/filters.rst

@@ -15,8 +15,11 @@ Filters are used to modify the data or metadata in a dataset.
    :maxdepth: 1
 
    filters/select
+   filters/orog_to_z
    filters/rename
    filters/rotate_winds
+   filters/sum
    filters/unrotate_winds
+   filters/wz_to_w
    filters/noop
    filters/empty

docs/building/filters/orog_to_z.rst

@@ -0,0 +1,17 @@
+###########
+ orog_to_z
+###########
+
+The ``orog_to_z`` filter converts orography (in meters) to surface
+geopotential height (m^2/s^2) using the equation:
+
+.. math::
+
+   z &= g \cdot \textrm{orog}\\
+   g &= 9.80665\ m \cdot s^{-1}
+
+This filter needs to follow a source that provides orography, which is
+replaced by surface geopotential height.
+
+.. literalinclude:: yaml/orog_to_z.yaml
+   :language: yaml

docs/building/filters/sum.rst

@@ -0,0 +1,13 @@
+#####
+ sum
+#####
+
+The ``sum`` filter computes the sum over multiple variables. This can be
+useful for computing total precipitation from its components (snow,
+rain) or summing the components of total column integrated water. This
+filter needs to follow a source that provides the list of variables to
+be summed up. These variables are removed by the filter and replaced by
+a single summed variable.
+
+.. literalinclude:: yaml/sum.yaml
+   :language: yaml

docs/building/filters/wz_to_w.rst

@@ -0,0 +1,12 @@
+#########
+ wz_to_w
+#########
+
+The ``wz_to_w`` filter converts geometric vertical velocity (provided in
+m/s) to vertical velocity in pressure coordinates (Pa/s). This filter
+needs to follow a source that provides geometric vertical velocity.
+Geometric vertical velocity is removed by the filter and pressure
+vertical velocity is added.
+
+.. literalinclude:: yaml/wz_to_w.yaml
+   :language: yaml

docs/building/filters/yaml/orog_to_z.yaml

@@ -0,0 +1,10 @@
+input:
+  pipe:
+    - source: # mars, grib, netcdf, etc.
+      # source attributes here
+      # ...
+      # Must load an orography variable
+
+    - orog_to_z:
+        orog: orog  # Name of orography (input) variable
+        z: z        # Name of z (output) variable

docs/building/filters/yaml/sum.yaml

@@ -0,0 +1,13 @@
+input:
+  pipe:
+    - source: # mars, grib, netcdf, etc.
+      # source attributes here
+      # ...
+      # Must load the variables to be summed
+
+    - sum:
+        params:  # List of input variables
+          variable1
+          variable2
+          variable3
+        output: variable_total    # Name of output variable

docs/building/filters/yaml/wz_to_w.yaml

@@ -0,0 +1,10 @@
+input:
+  pipe:
+    - source: # mars, grib, netcdf, etc.
+      # source attributes here
+      # ...
+      # Must load geometric vertical velocity
+
+    - wz_to_w:
+        wz: wz  # Name of geometric vertical velocity (input) variable
+        x: z    # Name of pressure vertical velocity (output) variable

src/anemoi/datasets/create/functions/filters/orog_to_z.py

@@ -0,0 +1,58 @@
+# (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 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
+
+    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
+        return value
+
+    def __getattr__(self, name):
+        return getattr(self.field, name)
+
+
+def execute(context, input, orog, z="z"):
+    """Convert orography [m] to z (geopotential height)"""
+    result = FieldArray()
+
+    processed_fields = defaultdict(dict)
+
+    for f in input:
+        key = f.metadata(namespace="mars")
+        param = key.pop("param")
+        if param == orog:
+            key = tuple(key.items())
+
+            if param in processed_fields[key]:
+                raise ValueError(f"Duplicate field {param} for {key}")
+
+            output = f.to_numpy(flatten=True) * 9.80665
+            result.append(NewDataField(f, output, z))
+        else:
+            result.append(f)
+
+    return result

src/anemoi/datasets/create/functions/filters/sum.py

@@ -0,0 +1,71 @@
+# (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 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
+
+    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
+        return value
+
+    def __getattr__(self, name):
+        return getattr(self.field, name)
+
+
+def execute(context, input, params, output):
+    """Computes the sum over a set of variables"""
+    result = FieldArray()
+
+    needed_fields = defaultdict(dict)
+
+    for f in input:
+        key = f.metadata(namespace="mars")
+        param = key.pop("param")
+        if param in params:
+            key = tuple(key.items())
+
+            if param in needed_fields[key]:
+                raise ValueError(f"Duplicate field {param} for {key}")
+
+            needed_fields[key][param] = f
+        else:
+            result.append(f)
+
+    for keys, values in needed_fields.items():
+
+        if len(values) != len(params):
+            raise ValueError("Missing fields")
+
+        s = None
+        for k, v in values.items():
+            c = v.to_numpy(flatten=True)
+            if s is None:
+                s = c
+            else:
+                s += c
+        result.append(NewDataField(values[list(values.keys())[0]], s, output))
+
+    return result

src/anemoi/datasets/create/functions/filters/wz_to_w.py

@@ -0,0 +1,79 @@
+# (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 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
+
+    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
+        return value
+
+    def __getattr__(self, name):
+        return getattr(self.field, name)
+
+
+def execute(context, input, wz, t, w="w"):
+    """Convert geometric vertical velocity (m/s) to vertical velocity (Pa / s)"""
+    result = FieldArray()
+
+    params = (wz, t)
+    pairs = defaultdict(dict)
+
+    for f in input:
+        key = f.metadata(namespace="mars")
+        param = key.pop("param")
+        if param in params:
+            key = tuple(key.items())
+
+            if param in pairs[key]:
+                raise ValueError(f"Duplicate field {param} for {key}")
+
+            pairs[key][param] = f
+            if param == t:
+                result.append(f)
+        else:
+            result.append(f)
+
+    for keys, values in pairs.items():
+
+        if len(values) != 2:
+            raise ValueError("Missing fields")
+
+        wz_pl = values[wz].to_numpy(flatten=True)
+        t_pl = values[t].to_numpy(flatten=True)
+        pressure = keys[4][1] * 100  # TODO: REMOVE HARDCODED INDICES
+
+        w_pl = wz_to_w(wz_pl, t_pl, pressure)
+        result.append(NewDataField(values[wz], w_pl, w))
+
+    return result
+
+
+def wz_to_w(wz, t, pressure):
+    g = 9.81
+    Rd = 287.058
+
+    return -wz * g * pressure / (t * Rd)