Allow clean functions to handle _avg variables (#377)

Makes the ADCP cleaning functions more robust - updated based on latest
reader updates for dual profiling instruments. Solution for Issue #373

---------

Co-authored-by: akeeste <[email protected]>

Author: jmcvey3

mhkit/dolfyn/adp/clean.py

@@ -157,31 +157,39 @@ def water_depth_from_amplitude(ds, thresh=10, nfilt=None) -> None:
             "Please manually remove 'depth' if it needs to be recalculated."
         )
 
+    # Use "avg" velocty if standard isn't available.
+    # Should not matter which is used.
+    tag = []
+    if hasattr(ds, "vel"):
+        tag += [""]
+    if hasattr(ds, "vel_avg"):
+        tag += ["_avg"]
+
     # This finds the maximum of the echo profile:
-    inds = np.argmax(ds["amp"].values, axis=1)
+    inds = np.argmax(ds["amp" + tag[0]].values, axis=1)
     # This finds the first point that increases (away from the profiler) in
     # the echo profile
-    edf = np.diff(ds["amp"].values.astype(np.int16), axis=1)
+    edf = np.diff(ds["amp" + tag[0]].values.astype(np.int16), axis=1)
     inds2 = (
-        np.max(
+        np.nanmax(
             (edf < 0)
-            * np.arange(ds["vel"].shape[1] - 1, dtype=np.uint8)[None, :, None],
+            * np.arange(ds["vel" + tag[0]].shape[1] - 1, dtype=np.uint8)[None, :, None],
             axis=1,
         )
         + 1
     )
 
     # Calculate the depth of these quantities
-    d1 = ds["range"].values[inds]
-    d2 = ds["range"].values[inds2]
+    d1 = ds["range" + tag[0]].values[inds]
+    d2 = ds["range" + tag[0]].values[inds2]
     # Combine them:
     D = np.vstack((d1, d2))
     # Take the median value as the estimate of the surface:
-    d = np.median(D, axis=0)
+    d = np.nanmedian(D, axis=0)
 
     # Throw out values that do not increase near the surface by *thresh*
-    for ip in range(ds["vel"].shape[1]):
-        itmp = np.min(inds[:, ip])
+    for ip in range(ds["vel" + tag[0]].shape[1]):
+        itmp = np.nanmin(inds[:, ip])
         if (edf[itmp:, :, ip] < thresh).all():
             d[ip] = np.nan
 
@@ -199,7 +207,7 @@ def water_depth_from_amplitude(ds, thresh=10, nfilt=None) -> None:
 
     ds["depth"] = xr.DataArray(
         d.astype("float32"),
-        dims=["time"],
+        dims=["time" + tag[0]],
         attrs={"units": "m", "long_name": long_name, "standard_name": "depth"},
     )
 
@@ -230,7 +238,7 @@ def water_depth_from_pressure(ds, salinity=35) -> None:
     ds : xarray.Dataset
       The full adcp dataset
     salinity: numeric
-      Water salinity in psu. Default = 35
+      Water salinity in PSU. Default = 35
 
     Returns
     -------
@@ -259,16 +267,26 @@ def water_depth_from_pressure(ds, salinity=35) -> None:
             "The variable 'depth' already exists. "
             "Please manually remove 'depth' if it needs to be recalculated."
         )
-    if "pressure" not in ds.data_vars:
+    pressure = [v for v in ds.data_vars if "pressure" in v]
+    if not pressure:
         raise NameError("The variable 'pressure' does not exist.")
-    elif not ds["pressure"].sum():
-        raise ValueError("Pressure data not recorded.")
-    if "temp" not in ds.data_vars:
+    else:
+        for p in pressure:
+            if not ds[p].sum():
+                pressure.remove(p)
+        if not pressure:
+            raise ValueError("Pressure data not recorded.")
+    temp = [
+        v
+        for v in ds.data_vars
+        if (("temp" in v) and ("clock" not in v) and ("press" not in v))
+    ]
+    if not temp:
         raise NameError("The variable 'temp' does not exist.")
 
     # Density calcation
-    P = ds["pressure"].values
-    T = ds["temp"].values  # temperature, degC
+    P = ds[pressure[0]].values  # pressure, dbar
+    T = ds[temp[0]].values  # temperature, degC
     S = salinity  # practical salinity
     rho0 = 1027  # kg/m^3
     T0 = 10  # degC
@@ -291,7 +309,7 @@ def water_depth_from_pressure(ds, salinity=35) -> None:
 
     ds["water_density"] = xr.DataArray(
         rho.astype("float32"),
-        dims=["time"],
+        dims=[ds[pressure[0]].dims[0]],
         attrs={
             "units": "kg m-3",
             "long_name": "Water Density",
@@ -301,7 +319,7 @@ def water_depth_from_pressure(ds, salinity=35) -> None:
     )
     ds["depth"] = xr.DataArray(
         d.astype("float32"),
-        dims=["time"],
+        dims=[ds[pressure[0]].dims[0]],
         attrs={"units": "m", "long_name": long_name, "standard_name": "depth"},
     )
 
@@ -319,7 +337,7 @@ def nan_beyond_surface(*args, **kwargs):
 
 
 def remove_surface_interference(
-    ds, val=np.nan, beam_angle=None, inplace=False
+    ds, val=np.nan, beam_angle=None, cell_size=None, inplace=False
 ) -> Optional[xr.Dataset]:
     """
     Mask the values of 3D data (vel, amp, corr, echo) that are beyond the surface.
@@ -332,6 +350,8 @@ def remove_surface_interference(
       Specifies the value to set the bad values to. Default is `numpy.nan`
     beam_angle : int
       ADCP beam inclination angle in degrees. Default = dataset.attrs['beam_angle']
+    cell_size : float
+      ADCP beam cellsize in meters. Default = dataset.attrs['cell_size']
     inplace : bool
       When True the existing data object is modified. When False
       a copy is returned. Default = False
@@ -352,51 +372,78 @@ def remove_surface_interference(
         raise KeyError(
             "Depth variable 'depth' does not exist in input dataset."
             "Please calculate 'depth' using the function 'water_depth_from_pressure'"
-            "or 'water_depth_from_amplitude."
+            "or 'water_depth_from_amplitude, or it can be found from the 'dist_bt'"
+            "(bottom track) or 'dist_alt' (altimeter) variables, if available."
         )
 
     if beam_angle is None:
         if hasattr(ds, "beam_angle"):
             beam_angle = np.deg2rad(ds.attrs["beam_angle"])
         else:
-            raise Exception(
+            raise KeyError(
                 "'beam_angle` not found in dataset attributes. "
                 "Please supply the ADCP's beam angle."
             )
     else:
         beam_angle = np.deg2rad(beam_angle)
 
+    if cell_size is None:
+        # Fetch cell size
+        cell_sizes = [
+            a
+            for a in ds.attrs
+            if (
+                ("cell_size" in a)
+                and ("_bt" not in a)
+                and ("_alt" not in a)
+                and ("wave" not in a)
+            )
+        ]
+        if cell_sizes:
+            cs = cell_sizes[0]
+        else:
+            raise KeyError(
+                "'cell_size` not found in dataset attributes. "
+                "Please supply the ADCP's cell size."
+            )
+    else:
+        cs = [cell_size]
+
     if not inplace:
         ds = ds.copy(deep=True)
 
     # Get all variables with 'range' coordinate
     profile_vars = [h for h in ds.keys() if any(s for s in ds[h].dims if "range" in s)]
 
-    # Surface interference distance
     # Apply range_offset if available
     range_offset = __check_for_range_offset(ds)
     if range_offset:
         range_limit = (
-            (ds["depth"] - range_offset) * np.cos(beam_angle) - ds.attrs["cell_size"]
+            (ds["depth"] - range_offset) * np.cos(beam_angle) - ds.attrs[cs]
         ) + range_offset
     else:
-        range_limit = ds["depth"] * np.cos(beam_angle) - ds.attrs["cell_size"]
-
-    bds = ds["range"] > range_limit
+        range_limit = ds["depth"] * np.cos(beam_angle) - ds.attrs[cs]
 
     # Echosounder data needs only be trimmed at water surface
     if "echo" in profile_vars:
         mask_echo = ds["range_echo"] > ds["depth"]
         ds["echo"].values[..., mask_echo] = val
         profile_vars.remove("echo")
 
-    # Correct rest of "range" data for surface interference
+    # Correct profile measurements for surface interference
     for var in profile_vars:
+        # Use correct coordinate tag
+        if "_" in var and ("gd" not in var):
+            tag = "_" + "_".join(var.split("_")[1:])
+        else:
+            tag = ""
+        mask = ds["range" + tag] > range_limit
+        # Remove values
         a = ds[var].values
         try:  # float dtype
-            a[..., bds] = val
+            a[..., mask] = val
         except:  # int dtype
-            a[..., bds] = 0
+            a[..., mask] = 0
         ds[var].values = a
 
     if not inplace:
@@ -434,10 +481,13 @@ def correlation_filter(ds, thresh=50, inplace=False) -> Optional[xr.Dataset]:
         ds = ds.copy(deep=True)
 
     # 4 or 5 beam
+    tag = []
+    if hasattr(ds, "vel"):
+        tag += [""]
     if hasattr(ds, "vel_b5"):
-        tag = ["", "_b5"]
-    else:
-        tag = [""]
+        tag += ["_b5"]
+    if hasattr(ds, "vel_avg"):
+        tag += ["_avg"]
 
     # copy original ref frame
     coord_sys_orig = ds.coord_sys
@@ -456,7 +506,7 @@ def correlation_filter(ds, thresh=50, inplace=False) -> Optional[xr.Dataset]:
             ds[var + tg].attrs["Comments"] = (
                 "Filtered of data with a correlation value below "
                 + str(thresh)
-                + ds.corr.units
+                + ds["corr" + tg].units
             )
 
     rotate2(ds, coord_sys_orig, inplace=True)

mhkit/dolfyn/velocity.py

@@ -177,10 +177,11 @@ def __repr__(
         self,
     ):
         time_string = "{:.2f} {} (started: {})"
-        if "time" not in self or dt642epoch(self["time"][0]) < 1:
+        time = "time" if "time" in self else "time_avg"
+        if time not in self or dt642epoch(self[time][0]) < 1:
             time_string = "-->No Time Information!<--"
         else:
-            tm = self["time"][[0, -1]].values
+            tm = self[time][[0, -1]].values
             dt = dt642date(tm[0])[0]
             delta = (dt642epoch(tm[-1]) - dt642epoch(tm[0])) / (3600 * 24)  # days
             if delta > 1:
@@ -202,7 +203,7 @@ def __repr__(
                 time_string = "-->Error in time info<--"
 
         p = self.ds.attrs
-        t_shape = self["time"].shape
+        t_shape = self[time].shape
         if len(t_shape) > 1:
             shape_string = "({} bins, {} pings @ {}Hz)".format(
                 t_shape[0], t_shape, p.get("fs")

mhkit/tests/dolfyn/test_clean.py

@@ -62,6 +62,7 @@ def test_clean_upADCP(self):
         td_awac = tp.dat_awac.copy(deep=True)
         td_sig = tp.dat_sig_tide.copy(deep=True)
         td_rdi = tp.dat_rdi.copy(deep=True)
+        td_dual = tp.dat_sig_dp2.copy(deep=True)
 
         apm.clean.water_depth_from_pressure(td_awac, salinity=30)
         apm.clean.remove_surface_interference(td_awac, beam_angle=20, inplace=True)
@@ -71,6 +72,11 @@ def test_clean_upADCP(self):
         apm.clean.remove_surface_interference(td_sig, inplace=True)
         td_sig = apm.clean.correlation_filter(td_sig, thresh=50)
 
+        apm.clean.set_range_offset(td_dual, 0.6)
+        apm.clean.water_depth_from_pressure(td_dual, salinity=31)
+        apm.clean.remove_surface_interference(td_dual, inplace=True)
+        td_dual = apm.clean.correlation_filter(td_dual, thresh=50)
+
         # Depth should already be found for this RDI file, but it's bad
         td_rdi["pressure"] /= 10  # set to something reasonable
         td_rdi = td_rdi.drop_vars("depth")
@@ -89,6 +95,7 @@ def test_clean_upADCP(self):
 
     def test_clean_downADCP(self):
         td = tp.dat_sig_ie.copy(deep=True)
+        td_dual = tp.dat_sig_dp2.copy(deep=True)
 
         # First remove bad data
         td["vel"] = apm.clean.val_exceeds_thresh(td.vel, thresh=3)
@@ -100,7 +107,12 @@ def test_clean_downADCP(self):
         # Then clean below seabed
         apm.clean.set_range_offset(td, 0.5)
         apm.clean.water_depth_from_amplitude(td, thresh=10, nfilt=3)
-        td = apm.clean.remove_surface_interference(td)
+        td = apm.clean.remove_surface_interference(td, inplace=False)
+
+        # Technically up-facing but a good check
+        apm.clean.set_range_offset(td_dual, 0.6)
+        apm.clean.water_depth_from_amplitude(td_dual, thresh=10, nfilt=3)
+        td_dual = apm.clean.remove_surface_interference(td_dual, inplace=False)
 
         if make_data:
             save(td, "Sig500_Echo_clean.nc")