Refactor remaining functions from Issue #11

eurec4a_environment/plot/profile.py

@@ -1,9 +1,7 @@
-import numpy as np
 import matplotlib.pyplot as plt
 import xarray as xr
 import matplotlib.dates as mdates
 import seaborn as sns
-import typhon
 
 sns.set(
     context="notebook",

eurec4a_environment/variables/boundary_layer/mixed_layer_height.py

@@ -3,8 +3,7 @@
 from scipy.signal import find_peaks
 import statsmodels.api as sm
 
-from ... import get_field
-from ... import nomenclature as nom
+from ... import get_field, nomenclature as nom
 
 
 def calc_peak_RH(
@@ -33,7 +32,7 @@ def calc_peakRH_linearize(
     ds,
     altitude=nom.ALTITUDE,
     rh=nom.RELATIVE_HUMIDITY,
-    time_dim="sounding",
+    time="sounding",
     z_min=200.0,
     z_max=1500.0,
     z_min_lin=50.0,
@@ -46,7 +45,7 @@ def calc_peakRH_linearize(
      assume linear profile as an idealization
      Inputs:
          -- ds: dataset
-         -- altitude, rh, time_dim: variable names
+         -- altitude, rh, time: variable names
          -- z_min and z_max: lower and upper bounds for mixed layer height
          -- z_min_lin and z_max_lin: bounds for linearization of observed RH profile
      Outputs:
@@ -60,12 +59,14 @@ def calc_peakRH_linearize(
     dz = int(da_alt.diff(dim=altitude)[1])  # dz=10m
 
     mixed_layer = np.logical_and(da_alt >= z_min, da_alt <= z_max)
+
     ml_linfit = np.logical_and(
         da_alt >= z_min_lin, da_alt <= z_max_lin
     )  # for linearized RH profile
 
-    for i in range(len(da_rh)):
-        rh_profile = da_rh.isel({time_dim: i}).interpolate_na(dim=altitude)
+    for i in range(len(ds[rh])):
+
+        rh_profile = ds[rh].isel({time: i}).interpolate_na(dim=altitude)
         X_height = rh_profile[ml_linfit][altitude]
         X = sm.add_constant(X_height.values)  # add intercept
         model = sm.OLS(
@@ -115,7 +116,6 @@ def calc_from_gradient(
     Outputs: DataArray containing mixed layer height from gradient method
 
     Note that function is quite slow and should be optimized
-
     """
 
     def calculateHmix_var(

eurec4a_environment/variables/calc_density.py

@@ -6,7 +6,6 @@
 
 from .. import get_field
 from ..constants import Rd, Rv, eps
-from .. import nomenclature as nom
 
 
 def calc_density(

eurec4a_environment/variables/calc_static_stability.py

@@ -0,0 +1,42 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+import xarray as xr
+
+
+def calc_static_stability(ds, pres="p", temp="T", theta="theta", altitude="height"):
+
+    """ Equation from: https://www.ncl.ucar.edu/Document/Functions/Contributed/static_stability.shtml
+        S = -(T/theta)*d(theta)/dp. From Bluestein (1992), pg 197, eqn 4.3.8
+
+        Static stability measures the gravitational resistance of an atmosphere to vertical displacements. 
+        It results from fundamental buoyant adjustments, and so it is determined by the vertical stratification 
+        of density or potential temperature.
+    """
+
+    # !! edit these lines based on units function !!
+    # Celsius to Kelvin
+    if ds[temp].max().values < 100:
+        temp_K = ds[temp] + 273.15
+    else:
+        temp_K = ds[temp]
+
+    # convert Pa to hPa
+    if ds[pres].max().values > 1400:
+        ds[pres] = ds[pres] / 100
+
+    static_stability = -(temp_K / ds[theta]) * (
+        ds[theta].diff(dim=altitude) / ds[pres].diff(dim=altitude)
+    )
+    static_stability = static_stability.transpose(
+        transpose_coords=True
+    )  # same dimension order as original dataset
+
+    dims = list(static_stability.dims)
+    da = xr.DataArray(
+        static_stability, dims=dims, coords={d: static_stability[d] for d in dims}
+    )
+    da.attrs["long_name"] = "static stability"
+    da.attrs["units"] = "K/hPa"
+
+    return da

eurec4a_environment/variables/inversion_layers/inversion_layer_height.py

@@ -0,0 +1,91 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+import numpy as np
+import xarray as xr
+
+
+def calc_inversion_static_stability(
+    ds, time="sounding", altitude="height", z_min=1000, z_max=3000
+):
+
+    """
+    Returns inversion height defined as the height where static stability profiles
+    first exceed 0.1 K/hPa between given heights
+    Adapted from Bony and Stevens 2018, Measuring Area-Averaged Vertical Motions with Dropsondes, p 772
+    To do: combine static stability criterion with one for moisture (i.e. hydrolapse)
+    """
+
+    from eurec4a_environment.variables.calc_static_stability import (
+        calc_static_stability,
+    )
+
+    ds_inversion = ds.sel({altitude: slice(z_min, z_max)})
+    da_static_stability = calc_static_stability(ds_inversion)
+
+    inversion_height_stability = np.zeros(len(da_static_stability[time]))
+    for i in range(len(da_static_stability[time])):
+        stability_sounding = da_static_stability.isel({time: i}).values
+        idx_stability = np.argmax(stability_sounding > 0.1)
+        inversion_height_stability[i] = da_static_stability[altitude][idx_stability]
+
+    dims = list(da_static_stability.dims)
+    # drop the height coord
+    del dims[dims.index(altitude)]
+
+    da = xr.DataArray(
+        inversion_height_stability,
+        dims=dims,
+        coords={d: da_static_stability[d] for d in dims},
+    )
+    da.attrs["long_name"] = "inversion base height, static stability > 0.1 K/hPa"
+    da.attrs["units"] = "m"
+
+    return da
+
+
+def calc_inversion_grad_RH_T(
+    ds,
+    time="sounding",
+    altitude="height",
+    temperature="T",
+    rh="rh",
+    z_min=1000,
+    z_max=2500,
+):
+
+    """
+    Inversion base following Grindinger et al, 1992; Cao et al 2007
+    dT/dz > 0 and dRH/dz < 0
+    To do: find all inversion bases and heights, not just first inversion base
+    """
+
+    ds_inversion = ds.sel({altitude: slice(z_min, z_max)})
+    gradient_RH = ds_inversion[rh].differentiate(coord=altitude)
+    gradient_T = ds_inversion[temperature].differentiate(coord=altitude)
+    inversion_levels = gradient_T[altitude].where(
+        (gradient_RH <= 0) & (gradient_T >= 0)
+    )
+
+    num_iters = len(inversion_levels[time])
+    inv_base_height_gradients = np.zeros(num_iters)
+    for i in range(num_iters):
+        alt_inv_da = inversion_levels.isel({time: i})
+        inversion_alts = alt_inv_da[~np.isnan(alt_inv_da)]
+        if inversion_alts[altitude].size == 0:
+            print(f"sounding {i} has no inversion below {z_max} m")
+            continue
+        inv_base_height_gradients[i] = inversion_alts[0]  # find first inversion base
+    inv_base_height_gradients[inv_base_height_gradients == 0] = np.nan
+
+    dims = list(ds.dims.keys())
+    # drop the height coord
+    del dims[dims.index(altitude)]
+    da = xr.DataArray(
+        inv_base_height_gradients, dims=dims, coords={d: ds[d] for d in dims}
+    )
+    da.attrs["long_name"] = "inversion base height from RH and T gradient"
+    da.attrs["units"] = "m"
+    return da
+
+

tests/test_boundary_layer.py

@@ -25,6 +25,26 @@ def test_mixed_layer_height_RHmax(ds_isentropic_test_profiles):
     assert np.allclose(da_rh_peak, z0)
 
 
+def test_mixed_layer_height_RH_lin():
+    ds = eurec4a_environment.source_data.open_joanne_dataset()
+    ds = ds.isel(sounding=slice(0, 10))
+    da_lin = boundary_layer.mixed_layer_height.calc_peakRH_linearize(ds)
+    assert len(da_lin) == 10
+    assert np.all(da_lin < 1500.0)
+
+
+def test_mixed_layer_height_gradient():
+    ds = eurec4a_environment.source_data.open_joanne_dataset()
+    ds = ds.isel(sounding=slice(0, 10))
+    # adjust threshold based on units 0.4 g/kg/10m or 0.0004 kg/kg/10m
+    da_gradient = boundary_layer.mixed_layer_height.calc_from_gradient(
+        ds, var=nom.SPECIFIC_HUMIDITY, threshold=0.4 / 1000
+    )
+    assert len(da_gradient) == 10
+    assert np.all(da_gradient < 1500.0)
+    assert da_gradient.units == "m"
+
+
 def test_inversion_height_gradient_RH(ds_isentropic_test_profiles):
     ds = ds_isentropic_test_profiles.copy()
     z_INV = 2000.0
@@ -34,4 +54,4 @@ def test_inversion_height_gradient_RH(ds_isentropic_test_profiles):
     da_inv = boundary_layer.inversion_height.find_inversion_height_grad_RH(
         ds=ds, rh=nom.RELATIVE_HUMIDITY
     )
-    assert np.allclose(da_inv, z_INV, atol=20)  ## within 20m
+    assert np.allclose(da_inv, z_INV, atol=20)  ## within 20m

tests/test_inversion_layers.py

@@ -0,0 +1,25 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+import numpy as np
+from eurec4a_environment.variables.inversion_layers import inversion_layer_height
+import eurec4a_environment.source_data
+
+
+def test_inversion_static_stability():
+    ds = eurec4a_environment.source_data.open_joanne_dataset()
+    ds = ds.isel(sounding=slice(0, 10))
+    da_ss = inversion_layer_height.calc_inversion_static_stability(ds)
+    assert len(da_ss) == 10
+    assert np.all(da_ss < 5000.0)
+    assert da_ss.units == "m"
+
+
+def test_inversion_gradient_RH_T():
+    ds = eurec4a_environment.source_data.open_joanne_dataset()
+    ds = ds.isel(sounding=slice(0, 10))
+    # adjust threshold based on units 0.4 g/kg/10m or 0.0004 kg/kg/10m
+    da_gradRH_T = inversion_layer_height.calc_inversion_grad_RH_T(ds)
+    assert len(da_gradRH_T) == 10
+    assert np.all(da_gradRH_T < 5000.0)
+    assert da_gradRH_T.units == "m"

tests/test_static_stability.py

@@ -0,0 +1,15 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+
+from eurec4a_environment.variables.calc_static_stability import calc_static_stability
+import eurec4a_environment.source_data
+
+
+def test_static_stability_calculation():
+    ds = eurec4a_environment.source_data.open_joanne_dataset()
+    da_static_stability = calc_static_stability(ds)
+    # check that we get some sensible numbers
+    assert da_static_stability.units == "K/hPa"
+    assert da_static_stability.max() < 3
+    assert da_static_stability.min() > -3