From 8a38fb9f9b400db35ae37e184316d52d18a9ef87 Mon Sep 17 00:00:00 2001
From: Geet George <g.george@tudelft.nl>
Date: Tue, 6 Aug 2024 19:41:22 -0100
Subject: [PATCH] add q and T to l3 prep

---
 src/halodrops/helper/__init__.py | 93 ++++++++++++++++++++++++++++++++
 src/halodrops/processor.py       | 25 +++++++++
 2 files changed, 118 insertions(+)

diff --git a/src/halodrops/helper/__init__.py b/src/halodrops/helper/__init__.py
index 33a4891..a35de4b 100644
--- a/src/halodrops/helper/__init__.py
+++ b/src/halodrops/helper/__init__.py
@@ -1,4 +1,6 @@
 import numpy as np
+import metpy.calc as mpcalc
+from metpy.units import units
 
 # Keys in l2_variables should be variable names in aspen_ds attribute of Sonde object
 l2_variables = {
@@ -151,3 +153,94 @@ def get_si_converter_function_based_on_var(var_name):
     if func is None:
         raise ValueError(f"No function named {func_name} found in the module")
     return func
+
+
+def calc_saturation_pressure(temperature_K, method="hardy1998"):
+    """
+    Calculate saturation water vapor pressure
+
+    Input
+    -----
+    temperature_K : array
+        array of temperature in Kevlin or dew point temperature for actual vapor pressure
+    method : str
+        Formula used for calculating the saturation pressure
+            'hardy1998' : ITS-90 Formulations for Vapor Pressure, Frostpoint Temperature,
+                Dewpoint Temperature, and Enhancement Factors in the Range –100 to +100 C,
+                Bob Hardy, Proceedings of the Third International Symposium on Humidity and Moisture,
+                1998 (same as used in Aspen software after May 2018)
+
+    Return
+    ------
+    e_sw : array
+        saturation pressure (Pa)
+
+    Examples
+    --------
+    >>> calc_saturation_pressure([273.15])
+    array([ 611.2129107])
+
+    >>> calc_saturation_pressure([273.15, 293.15, 253.15])
+    array([  611.2129107 ,  2339.26239586,   125.58350529])
+    """
+
+    if method == "hardy1998":
+        g = np.empty(8)
+        g[0] = -2.8365744 * 10**3
+        g[1] = -6.028076559 * 10**3
+        g[2] = 1.954263612 * 10**1
+        g[3] = -2.737830188 * 10 ** (-2)
+        g[4] = 1.6261698 * 10 ** (-5)
+        g[5] = 7.0229056 * 10 ** (-10)
+        g[6] = -1.8680009 * 10 ** (-13)
+        g[7] = 2.7150305
+
+        e_sw = np.zeros_like(temperature_K)
+
+        for t, temp in enumerate(temperature_K):
+            ln_e_sw = np.sum([g[i] * temp ** (i - 2) for i in range(0, 7)]) + g[
+                7
+            ] * np.log(temp)
+            e_sw[t] = np.exp(ln_e_sw)
+        return e_sw
+
+
+def calc_q_from_rh(ds):
+    """
+    Input :
+
+        ds : Dataset
+
+    Output :
+
+        q : Specific humidity values
+
+    Function to estimate specific humidity from the relative humidity, temperature and pressure in the given dataset.
+    """
+    e_s = calc_saturation_pressure(ds.ta.values)
+    w_s = mpcalc.mixing_ratio(e_s * units.Pa, ds.p.values * units.Pa).magnitude
+    w = ds.rh.values * w_s
+    q = w / (1 + w)
+    ds["q"] = (ds.rh.dims, q)
+
+    return ds
+
+
+def calc_theta_from_T(dataset):
+    """
+    Input :
+
+        dataset : Dataset
+
+    Output :
+
+        theta : Potential temperature values
+
+    Function to estimate potential temperature from the temperature and pressure in the given dataset.
+    """
+    theta = mpcalc.potential_temperature(
+        dataset.p.values * units.Pa, dataset.ta.values * units.kelvin
+    ).magnitude
+    ds["theta"] = (ds.ta.dims, theta)
+
+    return ds
diff --git a/src/halodrops/processor.py b/src/halodrops/processor.py
index f055b25..6ab306a 100644
--- a/src/halodrops/processor.py
+++ b/src/halodrops/processor.py
@@ -948,3 +948,28 @@ def add_l2_ds(self, l2_dir: str = None):
         )
 
         return self
+
+    def add_q_and_theta_to_l2_ds(self):
+        """
+        Adds potential temperature and specific humidity to the L2 dataset.
+
+        Parameters
+        ----------
+        None
+
+        Returns
+        -------
+        self : object
+            Returns the sonde object with potential temperature and specific humidity added to the L2 dataset.
+        """
+        if hasattr(self, "_interim_l3_ds"):
+            ds = self._interim_l3_ds
+        else:
+            ds = self.l2_ds
+
+        ds = calc_q_from_rh(ds)
+        ds = calc_theta_from_T(ds)
+
+        object.__setattr__(self, "_interim_l3_ds", ds)
+
+        return self