Nortek Dual Profile Dataset Rotation (#414)

Discovered this problem working on PacWave data, where the second
dataset, which contains instrument-averaged velocity measurements,
errors out when trying to use dolfyn's 'rotate' functionality.

This PR ensures the second profile from Nortek instrument can be rotated
like that of the first. This is to fix a somewhat hacky fix from several
months back, and doesn't require the user to make additional changes to
the 2nd dataset beyond that what dolfyn already does.

Author: jmcvey3

examples/adcp_example.ipynb

@@ -205,7 +205,7 @@ def water_depth_from_amplitude(ds, thresh=10, nfilt=None) -> None:
     else:
         long_name = "Instrument Depth"
 
-    ds["depth"] = xr.DataArray(
+    ds["depth" + tag[0]] = xr.DataArray(
         d.astype("float32"),
         dims=["time" + tag[0]],
         attrs={"units": "m", "long_name": long_name, "standard_name": "depth"},
@@ -307,7 +307,13 @@ def water_depth_from_pressure(ds, salinity=35) -> None:
     else:
         long_name = "Instrument Depth"
 
-    ds["water_density"] = xr.DataArray(
+    # Use correct coordinate tag
+    if "_" in pressure[0]:
+        tag = "_" + pressure[0].split("_")[-1]
+    else:
+        tag = ""
+
+    ds["water_density" + tag] = xr.DataArray(
         rho.astype("float32"),
         dims=[ds[pressure[0]].dims[0]],
         attrs={
@@ -317,7 +323,7 @@ def water_depth_from_pressure(ds, salinity=35) -> None:
             "description": "Water density from linear approximation of sea water equation of state",
         },
     )
-    ds["depth"] = xr.DataArray(
+    ds["depth" + tag] = xr.DataArray(
         d.astype("float32"),
         dims=[ds[pressure[0]].dims[0]],
         attrs={"units": "m", "long_name": long_name, "standard_name": "depth"},
@@ -368,7 +374,7 @@ def remove_surface_interference(
     `distance > range * cos(beam angle) - cell size`
     """
 
-    if "depth" not in ds.data_vars:
+    if ("depth" not in ds.data_vars) and ("depth_avg" not in ds.data_vars):
         raise KeyError(
             "Depth variable 'depth' does not exist in input dataset."
             "Please calculate 'depth' using the function 'water_depth_from_pressure'"
@@ -388,27 +394,23 @@ def remove_surface_interference(
         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:
+        # Fetch cell size (usually 'cell_size' or 'cell_size_avg')
+        cell_sizes = []
+        if hasattr(ds, "cell_size"):
+            cell_sizes.append("cell_size")
+        if hasattr(ds, "cell_size_avg"):
+            cell_sizes.append("cell_size_avg")
+        if not cell_sizes:
             raise KeyError(
                 "'cell_size` not found in dataset attributes. "
                 "Please supply the ADCP's cell size."
             )
     else:
         cs = [cell_size]
 
+    # Depth variable(s)
+    depths = [cs.replace("cell_size", "depth") for cs in cell_sizes]
+
     if not inplace:
         ds = ds.copy(deep=True)
 
@@ -417,34 +419,50 @@ def remove_surface_interference(
 
     # 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[cs]
-        ) + range_offset
-    else:
-        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 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:])
+    for depth, cs in zip(depths, cell_sizes):
+        if range_offset:
+            range_limit = (
+                (ds[depth] - range_offset) * np.cos(beam_angle) - ds.attrs[cs]
+            ) + range_offset
+        else:
+            range_limit = ds[depth] * np.cos(beam_angle) - ds.attrs[cs]
+
+        # No good way to do this
+        if "_avg" not in depth:
+            # 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 profile measurements for surface interference
+            for var in profile_vars:
+                if "avg" in var:
+                    continue
+                # 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[..., mask] = val
+                except:  # int dtype
+                    a[..., mask] = 0
+                ds[var].values = a
         else:
-            tag = ""
-        mask = ds["range" + tag] > range_limit
-        # Remove values
-        a = ds[var].values
-        try:  # float dtype
-            a[..., mask] = val
-        except:  # int dtype
-            a[..., mask] = 0
-        ds[var].values = a
+            for var in profile_vars:
+                if "avg" in var:
+                    mask = ds["range_avg"] > range_limit
+                    # Remove values
+                    a = ds[var].values
+                    try:  # float dtype
+                        a[..., mask] = val
+                    except:  # int dtype
+                        a[..., mask] = 0
+                    ds[var].values = a
 
     if not inplace:
         return ds

examples/data/dolfyn/test_data/Sig1000_dp_echo2.nc

@@ -113,10 +113,14 @@ def read_signature(
     ds = _create_dataset(out)
     ds = _set_coords(ds, ref_frame=ds.coord_sys)
 
-    if "orientmat" not in ds:
+    if ("orientmat" not in ds) and ("heading" in ds):
         ds["orientmat"] = _euler2orient(
             ds["time"], ds["heading"], ds["pitch"], ds["roll"]
         )
+    elif ("orientmat_avg" not in ds) and ("heading_avg" in ds):
+        ds["orientmat_avg"] = _euler2orient(
+            ds["time_avg"], ds["heading_avg"], ds["pitch_avg"], ds["roll_avg"]
+        )
 
     if declin is not None:
         set_declination(ds, declin, inplace=True)
@@ -729,13 +733,6 @@ def _reduce(data):
         dc["range_avg"] = (np.arange(dv["vel_avg"].shape[1]) + 1) * da[
             "cell_size_avg"
         ] + da["blank_dist_avg"]
-        # Check to see if orientmat_avg was written (only for instruments with an AHRS)
-        if ("orientmat" not in dv) and ("orientmat_avg" in dv):
-            dv["orientmat"] = dv.pop("orientmat_avg")
-        elif "heading" not in dv:  # need these to write orientation matrix
-            dv["heading"] = dv.pop("heading_avg")
-            dv["pitch"] = dv.pop("pitch_avg")
-            dv["roll"] = dv.pop("roll_avg")
         tmat = da["filehead_config"]["XFAVG"]
         da["duty_cycle_interval"] = dci = da["filehead_config"]["PLAN"]["MIAVG"]
         da["duty_cycle_n_burst"] = da["filehead_config"]["AVG"]["NPING"]
@@ -822,10 +819,11 @@ def split_dp_datasets(ds):
         ds2.attrs["rotate_vars"] = rotate_vars2
         # Set orientation matricies
         ds2["beam2inst_orientmat"] = ds["beam2inst_orientmat"]
-        if ds.attrs["has_imu"]:
-            ds2 = ds2.rename({"orientmat_avg": "orientmat"})
-        else:
-            ds2["orientmat"] = ds["orientmat"]
+        # IMU versions have 'orientmat_avg' variable, but non-IMU versions do not
+        if not ds.attrs["has_imu"]:
+            ds2["orientmat_avg"] = _euler2orient(
+                ds2["time_avg"], ds2["heading_avg"], ds2["pitch_avg"], ds2["roll_avg"]
+            )
         # Set original coordinate system
         cy = ds2.attrs["coord_sys_axes_avg"]
         ds2.attrs["coord_sys"] = {"XYZ": "inst", "ENU": "earth", "beam": "beam"}[cy]

examples/data/dolfyn/test_data/Sig100_avg.nc

@@ -246,15 +246,28 @@ def set_declination(ds, declin, inplace=True):
     else:
         rotate2earth = False
 
-    ds["orientmat"].values = np.einsum(
-        "kj...,ij->ki...",
-        ds["orientmat"].values,
-        Rdec,
-    ).astype(np.float32)
+    # Should only be one of these:
+    if "orientmat" in ds:
+        ds["orientmat"].values = np.einsum(
+            "kj...,ij->ki...",
+            ds["orientmat"].values,
+            Rdec,
+        ).astype(np.float32)
+    elif "orientmat_avg" in ds:
+        ds["orientmat_avg"].values = np.einsum(
+            "kj...,ij->ki...",
+            ds["orientmat_avg"].values,
+            Rdec,
+        ).astype(np.float32)
+
     if "heading" in ds:
         heading = ds["heading"] + angle
         heading[heading > 180] -= 360
         ds["heading"].values = heading
+    elif "heading_avg" in ds:
+        heading = ds["heading_avg"] + angle
+        heading[heading > 180] -= 360
+        ds["heading_avg"].values = heading
 
     if rotate2earth:
         rotate2(ds, "earth", inplace=True)

examples/data/dolfyn/test_data/Sig100_raw_avg.nc

@@ -57,13 +57,22 @@ def _inst2earth(adcpo, reverse=False, rotate_vars=None, force=False):
 
     if "orientmat" in adcpo:
         omat = adcpo["orientmat"]
-    else:
+    elif "orientmat_avg" in adcpo:
+        omat = adcpo["orientmat_avg"]
+    elif "time" in adcpo:
         omat = _euler2orient(
             adcpo["time"],
             adcpo["heading"].values,
             adcpo["pitch"].values,
             adcpo["roll"].values,
         )
+    elif "time_avg" in adcpo:
+        omat = _euler2orient(
+            adcpo["time_avg"],
+            adcpo["heading_avg"].values,
+            adcpo["pitch_avg"].values,
+            adcpo["roll_avg"].values,
+        )
 
     # Take the transpose of the orientation to get the inst->earth rotation
     # matrix.

examples/data/dolfyn/test_data/Sig500_dp_ice2.nc

@@ -118,12 +118,15 @@ def test_rotate_earth2inst(self):
         rotate2(td_sig, "inst", inplace=True)
         td_sig_i = load("Sig1000_IMU_rotate_inst2earth.nc")
         rotate2(td_sig_i, "inst", inplace=True)
-        td_sig_avg = load("Sig100_avg")
+        # Just check that these run without error
+        td_sig_avg = load("Sig100_avg.nc")
         rotate2(td_sig_avg, "inst", inplace=True)
+        td_sig_dp1_ice = load("Sig500_dp_ice2.nc")
+        rotate2(td_sig_dp1_ice, "inst", inplace=True)
 
         cd_rdi = load("RDI_test01_rotate_beam2inst.nc")
         cd_wr2 = tr.dat_wr2
-        # ship and inst are considered equivalent in dolfy
+        # ship and inst are considered equivalent in dolfyn
         cd_wr2.attrs["coord_sys"] = "inst"
         cd_awac = load("AWAC_test01_earth2inst.nc")
         cd_sig = load("BenchFile01_rotate_beam2inst.nc")