From 45171bebacf019b3cb4dae5becc2722edf82547c Mon Sep 17 00:00:00 2001
From: Thomas Nipen <tnipen@gmail.com>
Date: Sat, 30 Nov 2024 21:56:17 +0100
Subject: [PATCH] Add orog_to_z filter

---
 .../create/functions/filters/orog_to_z.py     | 58 +++++++++++++++++++
 1 file changed, 58 insertions(+)
 create mode 100644 src/anemoi/datasets/create/functions/filters/orog_to_z.py

diff --git a/src/anemoi/datasets/create/functions/filters/orog_to_z.py b/src/anemoi/datasets/create/functions/filters/orog_to_z.py
new file mode 100644
index 0000000..ddfc3cc
--- /dev/null
+++ b/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