From 84bd72e88b64f0c93f7b8b570440282407e16068 Mon Sep 17 00:00:00 2001
From: hgloeckner <helene.gloeckner@posteo.de>
Date: Fri, 13 Dec 2024 23:21:23 +0100
Subject: [PATCH] add omega

---
 pydropsonde/circles.py  | 29 +++++++++++++++++++++++++++++
 pydropsonde/pipeline.py |  8 +++++++-
 2 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/pydropsonde/circles.py b/pydropsonde/circles.py
index 591b5dd6..aea17edc 100644
--- a/pydropsonde/circles.py
+++ b/pydropsonde/circles.py
@@ -297,3 +297,32 @@ def add_vorticity(self):
         }
         self.circle_ds = ds.assign(vor=(ds.dudx.dims, vor.values, vor_attrs))
         return self
+
+    def add_omega(self):
+        """
+        Calculate vertical pressure velocity as
+        \int div dp
+
+        This calculates the vertical pressure velocity as described in
+        Bony and Stevens 2019
+
+        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)
+        p = ds.mean_p.where(~np.isnan(ds.div), drop=True).sortby(alt_dim)
+        zero_vel = xr.DataArray(data=[0], dims=alt_dim, coords={alt_dim: [0]})
+        pres_diff = xr.concat([zero_vel, p.diff(dim=alt_dim)], dim=alt_dim)
+        del_omega = -div * pres_diff.values
+        omega = del_omega.cumsum(dim=alt_dim) * 0.01 * 60**2
+        omega_attrs = {
+            "standard_name": "atmosphere_vertical_velocity",
+            "long_name": "Area-averaged atmospheric pressure velocity",
+            "units": "hPa hr-1",
+        }
+        self.circle_ds = ds.assign(
+            dict(omega_p=(ds.div.dims, omega.values, omega_attrs))
+        )
+        return self
diff --git a/pydropsonde/pipeline.py b/pydropsonde/pipeline.py
index 64b75409..250905f2 100644
--- a/pydropsonde/pipeline.py
+++ b/pydropsonde/pipeline.py
@@ -603,7 +603,13 @@ def run_pipeline(pipeline: dict, config: configparser.ConfigParser):
     "calculate_circle_data": {
         "intake": "circles",
         "apply": iterate_Circle_method_over_dict_of_Circle_objects,
-        "functions": ["add_density", "apply_fit2d", "add_divergence", "add_vorticity"],
+        "functions": [
+            "add_density",
+            "apply_fit2d",
+            "add_divergence",
+            "add_vorticity",
+            "add_omega",
+        ],
         "output": "circles",
         "comment": "calculate circle products",
     },