ADD: Harr wavelet example

act/retrievals/__init__.py

@@ -16,6 +16,7 @@
         'pbl_lidar': [
             'calculate_gradient_pbl',
             'calculate_modified_gradient_pbl',
+            'calculate_wavelet_pbl',
             'calculate_tucker_method_pbl',
         ],
         'radiation': [

act/retrievals/pbl_lidar.py

@@ -9,6 +9,13 @@
 import xarray as xr
 from scipy.signal import find_peaks
 
+try:
+    import pywt
+
+    PYWAVELETS_AVAILABLE = True
+except ImportError:
+    PYWAVELETS_AVAILABLE = False
+
 try:
     from statsmodels.tsa.stattools import acf
 except:
@@ -194,6 +201,135 @@ def calculate_modified_gradient_pbl(
     return ds
 
 
+def calculate_wavelet_pbl(
+    ds,
+    var_name='wind_speed',
+    range_name='height',
+    scale=60.0,
+    continuity_window=2,
+    min_height=100,
+    max_height=None,
+):
+    """
+    Estimation of the Planetary Boundary Layer (PBL) height from a ceilometer
+    or Doppler lidar through a Haar wavelet covariance transform. The dataset
+    is averaged into 5-minute periods, and each vertical profile is decomposed
+    with a Haar wavelet. The PBL height at each time is taken to be the range
+    at which the wavelet approximation coefficients show their sharpest
+    transition. A continuity check then replaces PBL height estimates that
+    jump more than 150 m above their neighbors with the local baseline.
+
+    Note:
+    This retrieval method should be applied under a cloud-free, well-mixed PBL condition.
+
+    It is not expected perform well in cloud capped boundary layers.
+    Additional PRs will be included within the near future to address more PBL
+    environmental conditions.
+
+    Parameters
+    ----------
+    ds : xarray.Dataset
+        Dataset containing the zenith-pointing ceilometer or Doppler lidar data.
+    var_name : str
+        Variable in the dataset to compute the wavelet transform on (e.g.,
+        backscatter intensity or vertical velocity).
+    range_name : str
+        Name of the range/height coordinate in the dataset.
+    scale : float
+        Approximate spatial scale, in the same units as range_name, over which
+        the Haar wavelet decomposition is performed. This sets the decomposition level.
+    continuity_window : int
+        Number of neighboring time steps on each side of a given time to
+        average over when checking for, and smoothing out, discontinuous PBL
+        height estimates.
+    min_height : float
+        Minimum allowed PBL height in the units of range_name. Excludes
+        near-surface noise from the search for the sharpest transition.
+    max_height : float or None
+        Maximum allowed PBL height in the units of range_name. Use this to
+        exclude elevated cloud or aerosol layers above the PBL from the
+        search. If None, no upper bound is applied.
+
+    Returns
+    -------
+    ds : xarray.Dataset
+        Dataset resampled to 5-minute periods with new variables
+        `wavelet_backscatter`, containing the Haar wavelet approximation
+        coefficients, and `pbl_wavelet`, containing the estimated PBL heights.
+
+    References
+    ----------
+    Brooks, I. M. (2003). Finding boundary layer top using
+        wavelet covariance transform. Journal of Atmospheric and Oceanic
+        Technology, 20(8), 1092-1105.
+        https://doi.org/10.1175/1520-0426(2003)20%3C1092:FBLTUB%3E2.0.CO;2
+
+    Cohn, S. A., & Angevine, W. M. (2000). Boundary layer height and
+        entrainment zone thickness measured by lidars and wind-profiling
+        radars. Journal of Applied Meteorology, 39(8), 1233-1247.
+        https://doi.org/10.1175/1520-0450(2000)039%3C1233:BLHAEZ%3E2.0.CO;2
+    """
+    if not PYWAVELETS_AVAILABLE:
+        raise ImportError('PyWavelets needs to be installed to use this feature.')
+
+    ds = ds.resample(time='5min').mean()
+    range_resolution = ds[range_name].values[1] - ds[range_name].values[0]
+    level = int(scale / range_resolution) - 1
+
+    coeffs = pywt.wavedec(ds[var_name].values, 'haar', level=level)
+    cA = coeffs[0]
+
+    resampled_range = ds[range_name].values[:: 2**level]
+    resampled_time = ds.time.values
+
+    ds['resampled_range'] = resampled_range
+    ds['resampled_time'] = resampled_time
+    ds = ds.set_coords(['resampled_range', 'resampled_time'])
+    ds['wavelet_backscatter'] = (('resampled_time', 'resampled_range'), cA)
+
+    range_mask = ds.resampled_range >= min_height
+    if max_height is not None:
+        range_mask = range_mask & (ds.resampled_range <= max_height)
+    wavelet_valid = ds.wavelet_backscatter.where(range_mask, drop=True)
+
+    max_gradient = wavelet_valid.diff('resampled_range').max('resampled_range')
+    pbl_heights = []
+    for t in range(len(ds.resampled_time)):
+        profile = wavelet_valid.isel(resampled_time=t)
+        try:
+            pbl_height = profile.where(
+                profile.diff('resampled_range') == max_gradient.isel(resampled_time=t),
+                drop=True,
+            ).resampled_range.values[0]
+        except IndexError:
+            pbl_height = np.nan
+        pbl_heights.append(pbl_height)
+
+    pbl_heights = np.array(pbl_heights, dtype=float)
+    for i in range(continuity_window, len(pbl_heights) - continuity_window):
+        neighbors = np.concatenate(
+            [
+                pbl_heights[i - continuity_window : i],
+                pbl_heights[i + 1 : i + continuity_window + 1],
+            ]
+        )
+        baseline = np.nanmean(neighbors)
+        if pbl_heights[i] > baseline + 150:
+            pbl_heights[i] = baseline
+
+    ds['pbl_wavelet'] = xr.DataArray(pbl_heights, dims='resampled_time')
+    ds['pbl_wavelet'].attrs[
+        'description'
+    ] = 'Planetary Boundary Layer Estimate via Haar wavelet covariance transform'
+    ds['pbl_wavelet'].attrs['input_parameter'] = var_name
+    if hasattr(ds[range_name], 'units'):
+        ds['pbl_wavelet'].attrs['units'] = ds[range_name].attrs['units']
+    else:
+        ds['pbl_wavelet'].attrs['units'] = 'meters'
+
+    return ds
+
+
 def calculate_tucker_method_pbl(
     ds,
     velocity="radial_velocity",
@@ -264,10 +400,6 @@ def calculate_tucker_method_pbl(
     min_gate_height : float
         Minimum height of the range gate to be considered for PBL height determination.
         This is to avoid surface noise and spurious low-level signals. The default value is 100 meters.
-
-    Returns
-    -------
-    ds : xarray.Dataset
         Original dataset with the following variables added:
         pbl_tucker : PBL height for each averaging interval.
         tucker_atmospheric_variance : Atmospheric variance profile for each
@@ -348,4 +480,5 @@ def calculate_tucker_method_pbl(
         "autocorrelation of radial velocity"
     )
     ds["tucker_atmospheric_variance"].attrs["units"] = 'm^2/s^2'
+
     return ds

continuous_integration/environment_actions.yml

@@ -29,6 +29,7 @@ dependencies:
   - metpy
   - arm_pyart
   - openpyxl
+  - pywavelets
   - pip
   - pip:
     - moviepy

docs/environment_docs.yml

@@ -27,6 +27,7 @@ dependencies:
   - myst-nb
   - nbsphinx
   - sgp4
+  - pywavelets
   - pip
   - pip:
     - mpl2nc

examples/retrievals/plot_wavelet_pbl.py

@@ -0,0 +1,54 @@
+"""
+Planetary Boundary Layer Height Wavelet Method Retrieval
+---------------------------------------------------------
+
+This example shows how to estimate the planetary boundary layer
+height via a Haar wavelet covariance transform retrieval
+
+Author: Robert Jackson
+"""
+
+import matplotlib.pyplot as plt
+from arm_test_data import DATASETS
+
+import act
+
+# Read Ceilometer data for an example
+filename_ceil = DATASETS.fetch('sgpceilC1.b1.20190101.000000.nc')
+ds = act.io.arm.read_arm_netcdf(filename_ceil)
+
+# Apply corrections to the dataset
+ds = act.corrections.correct_ceil(ds, var_name='backscatter')
+
+# Estimate PBL Height via a Haar wavelet covariance transform,
+# limiting the search to below 2000 m to exclude elevated cloud layers
+ds = act.retrievals.pbl_lidar.calculate_wavelet_pbl(
+    ds, var_name='backscatter', range_name='range', scale=60.0, max_height=1500.0
+)
+
+# Plot the pbl height estimates
+display = act.plotting.TimeSeriesDisplay(ds, figsize=(10, 5))
+
+# plot the CL backscatter before overlaying the Wavelet Method PBL Height
+display.plot(
+    'backscatter',
+    cmap='ChaseSpectral',
+    vmin=-6,
+    vmax=6,
+    set_title='SGP Ceilometer PBL Height Estimate via Wavelet Method',
+)
+
+# overlay the PBL Height estimate. We will compute a 10 minute running average to smooth the estimate.
+# The rolling function is used to compute a running average of the PBL height estimate over a 10 minute window,
+# with a minimum of 3 valid data points required for the average to be computed.
+# The center=True argument ensures that the average is centered on the current time point.
+display.axes[0].plot(
+    ds['resampled_time'].values,
+    ds['pbl_wavelet'].rolling(resampled_time=10, center=True, min_periods=3).mean().values,
+    color='w',
+    linewidth=2,
+    label='Wavelet PBL Height Estimate',
+)
+# shorten the range
+display.set_yrng([0, 2000])
+plt.show()

tests/retrievals/test_pbl_lidar.py

@@ -1,8 +1,16 @@
 import numpy as np
+import pytest
 from arm_test_data import DATASETS
 
 import act
 
+try:
+    import pywt  # noqa
+
+    PYWAVELETS_AVAILABLE = True
+except ImportError:
+    PYWAVELETS_AVAILABLE = False
+
 
 def test_calculate_gradient_pbl():
     # Read and apply connections
@@ -48,3 +56,45 @@ def test_calculate_modified_gradient_pbl():
     # test attributes
     assert ds['pbl_mod_gradient'].attrs["input_parameter"] == "backscatter"
     assert ds['pbl_mod_gradient'].attrs["units"] == "m"
+
+
+@pytest.mark.skipif(not PYWAVELETS_AVAILABLE, reason='PyWavelets is not installed.')
+def test_calculate_wavelet_pbl():
+    # Read and apply connections
+    ds = act.io.arm.read_arm_netcdf(DATASETS.fetch('sgpceilC1.b1.20190101.000000.nc'))
+    ds = act.corrections.correct_ceil(ds, var_name='backscatter')
+
+    # Call the Retrieval
+    ds = act.retrievals.pbl_lidar.calculate_wavelet_pbl(
+        ds, var_name='backscatter', range_name='range', scale=60.0
+    )
+
+    # create a subset for testing
+    subset = ds.sel(resampled_time=slice("2019-01-01T11:30:00", "2019-01-01T11:40:00"))
+    # Test the mean of the profile for the subset time
+    np.testing.assert_array_almost_equal(subset.pbl_wavelet.mean(), 4569.961, decimal=3)
+    # Test the minimum PBL Height during the period
+    np.testing.assert_almost_equal(subset.pbl_wavelet.min(), 3435.0, 1)
+
+    # test attributes
+    assert ds['pbl_wavelet'].attrs["input_parameter"] == "backscatter"
+    assert ds['pbl_wavelet'].attrs["units"] == "m"
+
+
+@pytest.mark.skipif(not PYWAVELETS_AVAILABLE, reason='PyWavelets is not installed.')
+def test_calculate_wavelet_pbl_max_height():
+    # Read and apply connections
+    ds = act.io.arm.read_arm_netcdf(DATASETS.fetch('sgpceilC1.b1.20190101.000000.nc'))
+    ds = act.corrections.correct_ceil(ds, var_name='backscatter')
+
+    # Call the Retrieval, constraining the search to below 1000 m
+    ds = act.retrievals.pbl_lidar.calculate_wavelet_pbl(
+        ds, var_name='backscatter', range_name='range', scale=60.0, max_height=1000.0
+    )
+
+    # create a subset for testing
+    subset = ds.sel(resampled_time=slice("2019-01-01T11:30:00", "2019-01-01T11:40:00"))
+    # Test that the upper bound is respected and no longer picks up the
+    # higher-altitude layers found without max_height set
+    assert subset.pbl_wavelet.max() <= 1000.0
+    np.testing.assert_array_almost_equal(subset.pbl_wavelet.mean(), 615.0, decimal=3)