New version: Fixes to add average filter, noise floor correction, and rename SNR to intensity

highiq/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "1.1"
+__version__ = "1.2"
 from . import calc
 from . import io
 from . import vis

highiq/calc/moments.py

@@ -1,103 +1,103 @@
 import numpy as np
-import tensorflow as tf
+import jax.numpy as jnp
 import xarray as xr
 
 
 def _gpu_calc_power(psd, dV, block_size=200, normed=True):
     shp = psd.shape
     power = np.zeros((shp[0], shp[1]))
     if len(shp) == 3:
-        gpu_array = tf.constant(psd, dtype=tf.float32)
+        gpu_array = jnp.array(psd, dtype=jnp.float32)
         if normed:
             gpu_array = gpu_array * dV
-        gpu_array = tf.math.reduce_sum(gpu_array, axis=2)
-        power = gpu_array.numpy()
+        gpu_array = jnp.sum(gpu_array, axis=2)
+        power = np.array(gpu_array)
     else:
-        gpu_array = tf.constant(psd.values, dtype=tf.float32)
+        gpu_array = jnp.array(psd.values, dtype=jnp.float32)
         if normed:
             gpu_array = gpu_array * dV
-        gpu_array = tf.math.reduce_sum(gpu_array, axis=1)
-        power = gpu_array.numpy()
+        gpu_array = jnp.sum(gpu_array, axis=1)
+        power = np.array(gpu_array)
 
     return power
 
 
 def _gpu_calc_velocity(psd, power, vel_bins, dV):
     shp = psd.shape
-    gpu_array = tf.constant(psd, dtype=tf.float32)
-    power_array = tf.constant(power, dtype=tf.float32)
-    vel_bins_tiled = np.tile(vel_bins, (shp[0], shp[1], 1))
-    gpu_array = 1 / power_array * tf.math.reduce_sum(gpu_array * vel_bins_tiled, axis=2)
-    velocity = gpu_array.numpy()
+    gpu_array = jnp.array(psd, dtype=jnp.float32)
+    power_array = jnp.array(power, dtype=jnp.float32)
+    vel_bins_tiled = jnp.tile(vel_bins, (shp[0], shp[1], 1))
+    gpu_array = 1 / power_array * jnp.sum(gpu_array * vel_bins_tiled, axis=2)
+    velocity = np.array(gpu_array)
     return velocity
 
 
 def _gpu_calc_velocity_dumb(psd, vel_bins):
     dV = np.diff(vel_bins)[0]
     vel_min = vel_bins.min()
-    gpu_array = tf.constant(psd)
-    gpu_array = tf.math.argmax(gpu_array, axis=2)
-    gpu_array = vel_min + tf.cast(gpu_array, tf.float32) * dV
-    velocity = gpu_array.numpy()
+    gpu_array = jnp.array(psd)
+    gpu_array = jnp.argmax(gpu_array, axis=2)
+    gpu_array = vel_min + gpu_array.astype(jnp.float32) * dV
+    velocity = np.array(gpu_array)
     return velocity
 
 
 def _gpu_calc_spectral_width(psd, power, vel_bins, velocity, dV):
     shp = psd.shape
     times = shp[0]
-    specwidth = np.zeros((shp[0], shp[1]))
+    specwidth = jnp.zeros((shp[0], shp[1]))
 
-    gpu_array = tf.constant(psd.values, dtype=tf.float32)
-    power_array = tf.constant(power, dtype=tf.float32)
+    gpu_array = jnp.array(psd.values, dtype=jnp.float32)
+    power_array = jnp.array(power, dtype=jnp.float32)
 
-    velocity_array = tf.transpose(np.tile(velocity, (shp[2], 1, 1)), [1, 2, 0])
-    vel_bins_tiled = np.tile(vel_bins, (times, shp[1], 1))
-    gpu_array = tf.math.sqrt(1 / power_array * tf.math.reduce_sum(
+    velocity_array = jnp.transpose(np.tile(velocity, (shp[2], 1, 1)), [1, 2, 0])
+    vel_bins_tiled = jnp.tile(vel_bins, (times, shp[1], 1))
+    gpu_array = jnp.sqrt(1 / power_array * jnp.sum(
                              (vel_bins_tiled - velocity_array)**2 * gpu_array, axis=2))
-    specwidth = gpu_array.numpy()
+    specwidth = np.array(gpu_array)
     return specwidth
 
 
 def _gpu_calc_skewness(psd, power, vel_bins, velocity, spec_width, dV):
     shp = psd.shape
     times = shp[0]
-    gpu_array = tf.constant(psd.values, dtype=tf.float32)
-    power_array = tf.constant(power, dtype=tf.float32)
-    spec_width_array = tf.constant(spec_width, dtype=tf.float32)
+    gpu_array = jnp.array(psd.values, dtype=jnp.float32)
+    power_array = jnp.array(power, dtype=jnp.float32)
+    spec_width_array = jnp.array(spec_width, dtype=jnp.float32)
     power_array *= spec_width_array**3
 
-    velocity_array = tf.transpose(np.tile(velocity, (shp[2], 1, 1)), [1, 2, 0])
-    vel_bins_tiled = np.tile(vel_bins, (times, shp[1], 1))
-    gpu_array = 1 / power_array * tf.math.reduce_sum(
+    velocity_array = jnp.transpose(np.tile(velocity, (shp[2], 1, 1)), [1, 2, 0])
+    vel_bins_tiled = jnp.tile(vel_bins, (times, shp[1], 1))
+    gpu_array = 1 / power_array * jnp.sum(
         (vel_bins_tiled - velocity_array)**3 * gpu_array, axis=2)
-    skewness = gpu_array.numpy()
+    skewness = np.array(gpu_array)
     return skewness
 
 
 def _gpu_calc_kurtosis(psd, power, vel_bins, velocity, spec_width, dV):
     shp = psd.shape
     kurtosis = np.zeros((shp[0], shp[1]))
-    gpu_array = tf.constant(psd.values, dtype=tf.float32)
-    power_array = tf.constant(power, dtype=tf.float32)
-    spec_width_array = tf.constant(spec_width, dtype=tf.float32)
+    gpu_array = jnp.array(psd.values, dtype=jnp.float32)
+    power_array = jnp.array(power, dtype=jnp.float32)
+    spec_width_array = jnp.array(spec_width, dtype=jnp.float32)
     power_array *= spec_width_array**4
-    velocity_array = tf.transpose(np.tile(velocity, (shp[2], 1, 1)), [1, 2, 0])
-    vel_bins_tiled = np.tile(vel_bins, (shp[0], shp[1], 1))
-    gpu_array = 1 / power_array * tf.math.reduce_sum(
+    velocity_array = jnp.transpose(jnp.tile(velocity, (shp[2], 1, 1)), [1, 2, 0])
+    vel_bins_tiled = jnp.tile(vel_bins, (shp[0], shp[1], 1))
+    gpu_array = 1 / power_array * jnp.sum(
         (vel_bins_tiled - velocity_array)**4 * gpu_array, axis=2)
-    kurtosis = gpu_array.numpy()
+    kurtosis = np.array(gpu_array)
     return kurtosis
 
 
-def get_lidar_moments(spectra, snr_thresh=0, block_size_ratio=1.0, which_moments=None):
+def get_lidar_moments(spectra, intensity_thresh=0, block_size_ratio=1.0, which_moments=None):
     """
     This function will retrieve the lidar moments of the Doppler spectra.
 
     Parameters
     ----------
     spectra: ACT Dataset
         The dataset containing the processed Doppler spectral density functions.
-    snr_thresh: float
+    intensity_thresh: float
         The minimum signal to noise ratio to use as an initial mask of noise.
     block_size_ratio: float
         This value is used to determine how much data the GPU will process in one loop. If your
@@ -117,7 +117,7 @@ def get_lidar_moments(spectra, snr_thresh=0, block_size_ratio=1.0, which_moments
     if 'power_spectral_density' not in spectra.variables.keys():
         raise ValueError("You must calculate the power spectra before calculating moments!")
     if which_moments is None:
-        which_moments = ['snr', 'doppler_velocity', 'spectral_width',
+        which_moments = ['intensity', 'doppler_velocity', 'spectral_width',
                          'skewness', 'kurtosis']
     else:
         which_moments = [x.lower() for x in which_moments]
@@ -132,17 +132,12 @@ def get_lidar_moments(spectra, snr_thresh=0, block_size_ratio=1.0, which_moments
         linear_psd_0filled, dV)
     velocity = _gpu_calc_velocity(linear_psd_0filled, power, spectra['vel_bins'].values, dV)
 
-    if 'snr' in which_moments:
+    if 'intensity' in which_moments:
         power_with_noise = dV * spectra['power_spectral_density'].sum(dim='vel_bins')
-        spectra['snr'] = power_with_noise / (dV * len(spectra['vel_bins']))
-        spectra['snr'].attrs['long_name'] = "Signal to Noise Ratio"
-        spectra['snr'].attrs['units'] = ""
-        spectra['intensity'] = spectra['snr'] - 1.
-        spectra['intensity'].attrs['long_name'] = "Intensity (SNR - 1)"
+        spectra['intensity'] = power_with_noise / (dV * len(spectra['vel_bins']))
+        spectra['intensity'].attrs['long_name'] = "Signal to Noise Ratio + 1"
         spectra['intensity'].attrs['units'] = ""
-        spectra['intensity'] = \
-            spectra['intensity'].where(spectra.snr > snr_thresh)
-        spectra.attrs['snr_mask'] = "%f" % snr_thresh
+        spectra.attrs['intensity_mask'] = "%f" % intensity_thresh
 
     if 'doppler_velocity' in which_moments:
         velocity_dumb = _gpu_calc_velocity_dumb(
@@ -159,9 +154,9 @@ def get_lidar_moments(spectra, snr_thresh=0, block_size_ratio=1.0, which_moments
             "Doppler velocity using first moment"
         spectra['doppler_velocity'].attrs['units'] = "m s-1"
         spectra['doppler_velocity_max_peak'] = \
-            spectra['doppler_velocity_max_peak'].where(spectra.snr > snr_thresh)
+            spectra['doppler_velocity_max_peak'].where(spectra.intensity > intensity_thresh)
         spectra['doppler_velocity'] = spectra['doppler_velocity'].where(
-            spectra.snr > snr_thresh)
+            spectra.intensity > intensity_thresh)
 
     if 'spectral_width' in which_moments or 'kurtosis' in which_moments or 'skewness' in which_moments:
         spectral_width = _gpu_calc_spectral_width(
@@ -173,24 +168,24 @@ def get_lidar_moments(spectra, snr_thresh=0, block_size_ratio=1.0, which_moments
             spectral_width, dims=('time', 'range'))
         spectra['spectral_width'].attrs["long_name"] = "Spectral width"
         spectra['spectral_width'].attrs["units"] = "m s-1"
-        if 'snr' in which_moments:
-            spectra['spectral_width'] = spectra['spectral_width'].where(spectra.snr > snr_thresh)
+        if 'intensity' in which_moments:
+            spectra['spectral_width'] = spectra['spectral_width'].where(spectra.intensity > intensity_thresh)
 
     if 'skewness' in which_moments:
         skewness = _gpu_calc_skewness(
             linear_psd, power, spectra['vel_bins'].values, velocity, spectral_width, dV)
         spectra['skewness'] = xr.DataArray(skewness, dims=('time', 'range'))
-        if 'snr' in which_moments:
-            spectra['skewness'] = spectra['skewness'].where(spectra.snr > snr_thresh)
+        if 'intensity' in which_moments:
+            spectra['skewness'] = spectra['skewness'].where(spectra.intensity > intensity_thresh)
         spectra['skewness'].attrs["long_name"] = "Skewness"
         spectra['skewness'].attrs["units"] = "m^3 s^-3"
 
     if 'kurtosis' in which_moments:
         kurtosis = _gpu_calc_kurtosis(
             linear_psd, power, spectra['vel_bins'].values, velocity, spectral_width, dV)
         spectra['kurtosis'] = xr.DataArray(kurtosis, dims=('time', 'range'))
-        if 'snr' in which_moments:
-            spectra['kurtosis'] = spectra['kurtosis'].where(spectra.snr > snr_thresh)
+        if 'intensity' in which_moments:
+            spectra['kurtosis'] = spectra['kurtosis'].where(spectra.intensity > intensity_thresh)
         spectra['kurtosis'].attrs["long_name"] = "Kurtosis"
         spectra['kurtosis'].attrs["units"] = "m^4 s^-4"
 

highiq/calc/spectra.py

@@ -1,8 +1,8 @@
 import numpy as np
-import tensorflow as tf
+import jax.numpy as jnp
 import xarray as xr
 
-from scipy.signal import hann, find_peaks
+from scipy.signal import find_peaks
 from scipy.ndimage import convolve1d
 from pint import UnitRegistry
 
@@ -12,24 +12,25 @@ def _fast_expand(complex_array, factor, num_per_block=200):
     expanded_array = np.zeros((shp[0], shp[1], shp[2] * factor))
     weights = np.tile(np.arange(0, factor) / factor, (shp[0], shp[1], 1))
     for i in range(shp[2]):
-        gpu_array = tf.constant(complex_array[:, :, i], dtype=tf.float32)
+        gpu_array = jnp.zeros(complex_array[:, :, i], dtype=jnp.float32)
         if i < shp[2] - 1:
-            gpu_array2 = tf.constant(complex_array[:, :, i + 1], dtype=tf.float32)
+            gpu_array2 = jnp.zeros(complex_array[:, :, i + 1], dtype=jnp.float32)
             diff_array = gpu_array2 - gpu_array
         else:
-            diff_array = tf.zeros((shp[0], shp[1]))
+            diff_array = jnp.zeros((shp[0], shp[1]))
 
-        rep_array = tf.transpose(
-            np.tile(gpu_array, (factor, 1, 1)), [1, 2, 0])
-        diff_array = tf.transpose(
-            np.tile(diff_array, (factor, 1, 1)), [1, 2, 0])
+        rep_array = jnp.transpose(
+            jnp.tile(gpu_array, (factor, 1, 1)), [1, 2, 0])
+        diff_array = jnp.transpose(
+            jnp.tile(diff_array, (factor, 1, 1)), [1, 2, 0])
         temp_array = diff_array * weights + rep_array
-        expanded_array[:, :, factor * i:factor * (i + 1)] = temp_array.numpy()
+        expanded_array[:, :, factor * i:factor * (i + 1)] = np.array(temp_array)
     return expanded_array
 
 
 def get_psd(spectra, gate_resolution=60., wavelength=None, fs=None, nfft=1024, time_window=None,
-            acf_name='acf', acf_bkg_name='acf_bkg', block_size_ratio=1.0):
+            acf_name='acf', acf_bkg_name='acf_bkg', block_size_ratio=1.0,
+            smooth_window=5):
     """
     This function will get the power spectral density from the autocorrelation function.
 
@@ -59,6 +60,9 @@ def get_psd(spectra, gate_resolution=60., wavelength=None, fs=None, nfft=1024, t
     block_size_ratio: float
         Increase this value to use more GPU memory for processing. Doing this can
         poentially optimize processing.
+    smooth_window: int
+        Apply running average to power spectra to remove small scale noise using
+        this window size.
 
     Returns
     -------
@@ -96,7 +100,7 @@ def get_psd(spectra, gate_resolution=60., wavelength=None, fs=None, nfft=1024, t
         (complex_coeff_in.shape[0], int(complex_coeff_in.shape[1] / num_gates),
             complex_coeff_in.shape[2]), dtype=np.complex128)
     for i in range(complex_coeff.shape[1]):
-        complex_coeff[:, i, :] = np.mean(complex_coeff_in[:, (num_gates * i):(num_gates * i + 1), :], axis=1)
+        complex_coeff[:, i, :] = np.sum(complex_coeff_in[:, (num_gates * i):(num_gates * i + 1), :], axis=1)
     del complex_coeff_in
     freq = np.fft.fftfreq(nfft) * fs
     spectra.attrs['nyquist_velocity'] = "%f m s-1" % (wavelength / (4 * 1 / fs))
@@ -116,30 +120,35 @@ def get_psd(spectra, gate_resolution=60., wavelength=None, fs=None, nfft=1024, t
         (complex_coeff_bkg_in.shape[0], int(complex_coeff_bkg_in.shape[1] / num_gates),
             complex_coeff_bkg_in.shape[2]), dtype=np.complex128)
     for i in range(complex_coeff.shape[1]):
-        complex_coeff_bkg[:, i, :] = np.mean(complex_coeff_bkg_in[:, (num_gates * i):(num_gates * i + 1), :], axis=1)
+        complex_coeff_bkg[:, i, :] = np.sum(complex_coeff_bkg_in[:, (num_gates * i):(num_gates * i + 1), :], axis=1)
     num_lags = complex_coeff_bkg_in.shape[2]
     if nfft < num_lags:
         raise RuntimeError("Number of points in FFT < number of lags in sample!")
     pad_after = int((nfft - num_lags))
     pad_before = 0
-    pad_lengths = tf.constant([[0, 0], [0, 0], [pad_before, pad_after]])
-    frames = tf.pad(complex_coeff, pad_lengths, mode='CONSTANT', constant_values=0)
-
-    power = np.abs(tf.signal.fft(frames).numpy())
+    pad_lengths = [(0, 0), (0, 0), (pad_before, pad_after)]
+    frames = jnp.pad(complex_coeff, pad_lengths, mode='constant', constant_values=0)
+    window = 1 / smooth_window * np.ones(smooth_window) 
+    power = convolve1d(np.abs(jnp.fft.fft(frames)),
+        axis=2, weights=window)
     attrs_dict = {'long_name': 'Range', 'units': 'm'}
     spectra['range'] = xr.DataArray(
         gate_resolution * np.arange(int(frames.shape[1])),
         dims=('range'), attrs=attrs_dict)
     spectra['power'] = xr.DataArray(
         power[:, :, inds_sorted], dims=(('time', 'range', 'vel_bins')))
-    frames = tf.pad(complex_coeff_bkg, pad_lengths, mode='CONSTANT', constant_values=0)
-    power = np.abs(tf.signal.fft(frames).numpy())
+    frames = jnp.pad(complex_coeff_bkg, pad_lengths, mode='constant', constant_values=0)
+    power = convolve1d(np.abs(jnp.fft.fft(frames)),
+        axis=2, weights=window)
 
     spectra['power_bkg'] = xr.DataArray(
         power[:, :, inds_sorted], dims=('time', 'range', 'vel_bins'))
 
     # Subtract background noise
     spectra['power_spectral_density'] = spectra['power'] / spectra['power_bkg']
+
+    # Ground noise floor to 1
+    spectra['power_spectral_density'] = spectra['power_spectral_density'] - spectra['power_spectral_density'].min(axis=2) + 1
     spectra['power_spectral_density'].attrs["long_name"] = "Power spectral density"
     spectra['power_spectral_density'].attrs["units"] = ''
     spectra['range'].attrs['long_name'] = "Range"

highiq/io/arm_data.py

@@ -18,21 +18,21 @@ def load_arm_netcdf(arm_file, **kwargs):
     """
 
     This loads netCDF data that are in the Atmospheric Radiation Measurement standard netCDF format.
-    This is a wrapper around :func:`act.io.armfiles.read_netcdf`
+    This is a wrapper around :func:`act.io.read_netcdf`
 
     Parameters
     ----------
     arm_file: str
         The path to the dataset to load.
 
-    Additional keyword arguments are passed into :func:`act.io.armfiles.read_netcdf`
+    Additional keyword arguments are passed into :func:`act.io.read_netcdf`
 
     Returns
     -------
     ds: ACT Dataset
         Returns the ACT dataset (xarray dataset) that contains the autocorrelation functions.
     """
-    ds = act.io.armfiles.read_netcdf(arm_file, **kwargs)
+    ds = act.io.arm.read_arm_netcdf(arm_file, **kwargs)
     if 'time' not in ds['acf'].dims:
         ds['acf'] = ds['acf'].expand_dims('time')
     if 'time' not in ds['acf_bkg'].dims:

requirements.txt

@@ -2,4 +2,5 @@ numpy
 scipy
 act-atmos
 pint
-tensorflow
+jax
+xarray