From be1b30f805144007774548d918d87eb3bd0f30f6 Mon Sep 17 00:00:00 2001
From: hgloeckner <helene.gloeckner@posteo.de>
Date: Mon, 16 Dec 2024 16:02:04 +0100
Subject: [PATCH] add wvel

---
 pydropsonde/circles.py  | 33 ++++++++++++++++++++++++++++++---
 pydropsonde/pipeline.py |  1 +
 2 files changed, 31 insertions(+), 3 deletions(-)

diff --git a/pydropsonde/circles.py b/pydropsonde/circles.py
index 0b53a1b..3271849 100644
--- a/pydropsonde/circles.py
+++ b/pydropsonde/circles.py
@@ -324,7 +324,34 @@ def add_omega(self):
             "long_name": "Area-averaged atmospheric pressure velocity (omega)",
             "units": "hPa hr-1",
         }
-        self.circle_ds = ds.assign(
-            dict(omega_p=(ds.div.dims, omega.values, omega_attrs))
-        )
+        self.circle_ds = ds.assign(dict(omega=(ds.div.dims, omega.values, omega_attrs)))
+        return self
+
+    def add_wvel(self):
+        """
+        Calculate vertical velocity as
+        - int diff dz
+
+        This calculates the vertical velocity from omega
+
+        Returns:
+            self: circle object with updated circle_ds
+        """
+        ds = self.circle_ds
+        alt_dim = self.alt_dim
+        div = ds.div.where(~np.isnan(ds.div), drop=True).sortby(alt_dim)
+        zero_vel = xr.DataArray(data=[0], dims=alt_dim, coords={alt_dim: [0]})
+
+        height = xr.concat([zero_vel, div[alt_dim]], dim=alt_dim)
+        height_diff = height.diff(dim=alt_dim)
+
+        del_w = -div * height_diff.values
+
+        w_vel = del_w.cumsum(dim=alt_dim)
+        wvel_attrs = {
+            "standard_name": "upward_air_velocity",
+            "long_name": "Area-averaged atmospheric vertical velocity",
+            "units": "m s-1",
+        }
+        self.circle_ds = ds.assign(dict(wvel=(ds.omega.dims, w_vel, wvel_attrs)))
         return self
diff --git a/pydropsonde/pipeline.py b/pydropsonde/pipeline.py
index a932477..df01c50 100644
--- a/pydropsonde/pipeline.py
+++ b/pydropsonde/pipeline.py
@@ -610,6 +610,7 @@ def run_pipeline(pipeline: dict, config: configparser.ConfigParser):
             "add_divergence",
             "add_vorticity",
             "add_omega",
+            "add_wvel",
         ],
         "output": "circles",
         "comment": "calculate circle products",