Add update_platform method to EchoData

echopype/echodata/echodata.py

@@ -312,6 +312,145 @@ def squeeze_non_scalar(n):
                 "this method only supports sonar_model values of AZFP, EK60, and EK80"
             )
 
+    def update_platform(self, extra_platform_data: xr.Dataset, time_dim="time"):
+        """
+        Updates the `EchoData.platform` group with additional external platform data.
+
+        `extra_platform_data` must be an xarray Dataset.
+        The name of the time dimension in `extra_platform_data` is specified by the
+        `time_dim` parameter.
+        Data is extracted from `extra_platform_data` by variable name; only the data
+        in `extra_platform_data` with the following variable names will be used:
+            - `"pitch"`
+            - `"roll"`
+            - `"heave"`
+            - `"latitude"`
+            - `"longitude"`
+            - `"water_level"`
+        The data inserted into the Platform group will be indexed by a dimension named `"time2"`.
+
+        Parameters
+        ----------
+        extra_platform_data : xr.Dataset
+            An `xr.Dataset` containing the additional platform data to be added
+            to the `EchoData.platform` group.
+        time_dim: str, default="time"
+            The name of the time dimension in `extra_platform_data`; used for extracting
+            data from `extra_platform_data`.
+
+        Examples
+        --------
+        >>> ed = echopype.open_raw(raw_file, "EK60")
+        >>> extra_platform_data = xr.open_dataset(extra_platform_data_file)
+        >>> ed.update_platform(extra_platform_data)
+        """
+
+        # # only take data during ping times
+        # start_time, end_time = min(self.beam["ping_time"]), max(self.beam["ping_time"])
+
+        # Handle data stored as a CF Trajectory Discrete Sampling Geometry
+        # http://cfconventions.org/Data/cf-conventions/cf-conventions-1.8/cf-conventions.html#trajectory-data
+        # The Saildrone sample data file follows this convention
+        if (
+            "featureType" in extra_platform_data.attrs
+            and extra_platform_data.attrs["featureType"].lower() == "trajectory"
+        ):
+            for coordvar in extra_platform_data.coords:
+                if (
+                    "cf_role" in extra_platform_data[coordvar].attrs
+                    and extra_platform_data[coordvar].attrs["cf_role"]
+                    == "trajectory_id"
+                ):
+                    trajectory_var = coordvar
+
+            # assumes there's only one trajectory in the dataset (index 0)
+            extra_platform_data = extra_platform_data.sel(
+                {trajectory_var: extra_platform_data[trajectory_var][0]}
+            )
+            extra_platform_data = extra_platform_data.drop_vars(trajectory_var)
+            extra_platform_data = extra_platform_data.swap_dims({"obs": time_dim})
+
+        platform = self.platform.assign_coords(
+            time2=extra_platform_data[time_dim].values
+        )
+        platform["time2"].attrs[
+            "long_name"
+        ] = "time dimension for external platform data"
+
+        dropped_vars = []
+        for var in ["pitch", "roll", "heave", "latitude", "longitude", "water_level"]:
+            if var in platform and platform[var].isnull().all():
+                dropped_vars.append(var)
+        if len(dropped_vars) > 0:
+            warnings.warn(
+                f"some variables will be overwritten by platform data: {', '.join(dropped_vars)}"
+            )
+        platform = platform.drop_vars(
+            ["pitch", "roll", "heave", "latitude", "longitude", "water_level"],
+            errors="ignore",
+        )
+
+        num_obs = len(extra_platform_data[time_dim])
+
+        def mapping_get_multiple(mapping, keys, default=None):
+            for key in keys:
+                if key in mapping:
+                    return mapping[key].data
+            return default
+
+        self.platform = platform.update(
+            {
+                "pitch": (
+                    "time2",
+                    mapping_get_multiple(
+                        extra_platform_data,
+                        ["pitch", "PITCH"],
+                        platform.get("pitch", np.full(num_obs, np.nan)),
+                    ),
+                ),
+                "roll": (
+                    "time2",
+                    mapping_get_multiple(
+                        extra_platform_data,
+                        ["roll", "ROLL"],
+                        platform.get("roll", np.full(num_obs, np.nan)),
+                    ),
+                ),
+                "heave": (
+                    "time2",
+                    mapping_get_multiple(
+                        extra_platform_data,
+                        ["heave", "HEAVE"],
+                        platform.get("heave", np.full(num_obs, np.nan)),
+                    ),
+                ),
+                "latitude": (
+                    "time2",
+                    mapping_get_multiple(
+                        extra_platform_data,
+                        ["lat", "latitude", "LATITUDE"],
+                        default=platform.get("latitude", np.full(num_obs, np.nan)),
+                    ),
+                ),
+                "longitude": (
+                    "time2",
+                    mapping_get_multiple(
+                        extra_platform_data,
+                        ["lon", "longitude", "LONGITUDE"],
+                        default=platform.get("longitude", np.full(num_obs, np.nan)),
+                    ),
+                ),
+                "water_level": (
+                    "time2",
+                    mapping_get_multiple(
+                        extra_platform_data,
+                        ["water_level", "WATER_LEVEL"],
+                        default=platform.get("water_level", np.zeros(num_obs)),
+                    ),
+                ),
+            }
+        )
+
     @classmethod
     def _load_convert(cls, convert_obj):
         new_cls = cls()

echopype/tests/echodata/test_echodata.py

@@ -9,6 +9,7 @@
 
 import pytest
 import xarray as xr
+import numpy as np
 
 ek60_path = TEST_DATA_FOLDER / "ek60"
 ek80_path = TEST_DATA_FOLDER / "ek80"
@@ -146,3 +147,20 @@ def test_compute_range(filepath, sonar_model, azfp_xml_path, azfp_cal_type, ek_w
     else:
         range = ed.compute_range(env_params, azfp_cal_type, ek_waveform_mode, )
         assert isinstance(range, xr.DataArray)
+
+def test_update_platform():
+    saildrone_path = ek80_path / "saildrone"
+    raw_file = saildrone_path / "SD2019_WCS_v05-Phase0-D20190617-T125959-0.raw"
+    extra_platform_data_file = saildrone_path / "saildrone-gen_5-fisheries-acoustics-code-sprint-sd1039-20190617T130000-20190618T125959-1_hz-v1.1595357449818.nc"
+
+    ed = echopype.open_raw(raw_file, "EK80")
+
+    updated = ["pitch", "roll", "latitude", "longitude", "water_level"]
+    for variable in updated:
+        assert np.isnan(ed.platform[variable].values).all()
+
+    extra_platform_data = xr.open_dataset(extra_platform_data_file)
+    ed.update_platform(extra_platform_data)
+
+    for variable in updated:
+        assert not np.isnan(ed.platform[variable].values).all()