Add more add_location tests and weak add_location (#1000)

* Add more add_location tests. Tweak add_location to fill in first lat & lon values that were being interpolated to nan

* Use fill_value extrapolate rather than hard-wiring the nan replacement scheme

* Add data interpolation and broadcast tests

* Update echopype/tests/consolidate/test_consolidate_integration.py

Co-authored-by: Wu-Jung Lee <leewujung@gmail.com>

* Update echopype/tests/consolidate/test_consolidate_integration.py

* Update echopype/tests/consolidate/test_consolidate_integration.py

Co-authored-by: Wu-Jung Lee <leewujung@gmail.com>

* Rename coord_var to position_var, for clarity

* Remove location extrapolation, update corresponding test. Rename coord variable names in that test, for clarity

---------

Co-authored-by: Wu-Jung Lee <leewujung@gmail.com>

echopype/consolidate/api.py

@@ -147,28 +147,29 @@ def add_location(ds: xr.Dataset, echodata: EchoData = None, nmea_sentence: Optio
 
     Returns
     -------
-    The input dataset with the the location data added
+    The input dataset with the location data added
     """
 
     def sel_interp(var):
         # NMEA sentence selection
         if nmea_sentence:
-            coord_var = echodata["Platform"][var][
+            position_var = echodata["Platform"][var][
                 echodata["Platform"]["sentence_type"] == nmea_sentence
             ]
         else:
-            coord_var = echodata["Platform"][var]
+            position_var = echodata["Platform"][var]
 
-        if len(coord_var) == 1:
+        if len(position_var) == 1:
             # Propagate single, fixed-location coordinate
             return xr.DataArray(
-                data=coord_var.values[0] * np.ones(len(ds["ping_time"]), dtype=np.float64),
+                data=position_var.values[0] * np.ones(len(ds["ping_time"]), dtype=np.float64),
                 dims=["ping_time"],
-                attrs=coord_var.attrs,
+                attrs=position_var.attrs,
             )
         else:
             # Interpolation. time1 is always associated with location data
-            return coord_var.interp(time1=ds["ping_time"])
+            # Values may be nan if there are ping_time values outside the time1 range
+            return position_var.interp(time1=ds["ping_time"])
 
     if "longitude" not in echodata["Platform"] or echodata["Platform"]["longitude"].isnull().all():
         raise ValueError("Coordinate variables not present or all nan")

echopype/tests/consolidate/test_consolidate_integration.py

@@ -1,4 +1,7 @@
+import math
+import os
 import pathlib
+import tempfile
 
 import pytest
 
@@ -8,8 +11,6 @@
 import scipy.io as io
 import echopype as ep
 from typing import List
-import tempfile
-import os
 
 """
 For future reference:
@@ -177,25 +178,6 @@ def test_add_depth():
     # assert ds_Sv_depth["depth"].attrs == {"long_name": "Depth", "standard_name": "depth"}
 
 
-def test_add_location(test_path):
-    ed = ep.open_raw(
-        test_path["EK60"] / "Winter2017-D20170115-T150122.raw",
-        sonar_model="EK60"
-    )
-    ds = ep.calibrate.compute_Sv(ed)
-
-    def _check_var(ds_test):
-        assert "latitude" in ds_test
-        assert "longitude" in ds_test
-        assert "time1" not in ds_test
-
-    ds_all = ep.consolidate.add_location(ds=ds, echodata=ed)
-    _check_var(ds_all)
-
-    ds_sel = ep.consolidate.add_location(ds=ds, echodata=ed, nmea_sentence="GGA")
-    _check_var(ds_sel)
-
-
 def _create_array_list_from_echoview_mats(paths_to_echoview_mat: List[pathlib.Path]) -> List[np.ndarray]:
     """
     Opens each mat file in ``paths_to_echoview_mat``, selects the first ``ping_time``,
@@ -223,6 +205,115 @@ def _create_array_list_from_echoview_mats(paths_to_echoview_mat: List[pathlib.Pa
     return list_of_mat_arrays
 
 
+@pytest.mark.parametrize(
+    ["location_type", "sonar_model", "path_model", "raw_and_xml_paths", "extras"],
+    [
+        (
+            "empty-location",
+            "EK60",
+            "EK60",
+            ("ooi/CE02SHBP-MJ01C-07-ZPLSCB101_OOI-D20191201-T000000.raw", None),
+            None,
+        ),
+        (
+            "with-track-location",
+            "EK60",
+            "EK60",
+            ("Winter2017-D20170115-T150122.raw", None),
+            None,
+        ),
+        (
+            "fixed-location",
+            "AZFP",
+            "AZFP",
+            ("17082117.01A", "17041823.XML"),
+            {'longitude': -60.0, 'latitude': 45.0, 'salinity': 27.9, 'pressure': 59},
+        ),
+    ],
+)
+def test_add_location(
+        location_type,
+        sonar_model,
+        path_model,
+        raw_and_xml_paths,
+        extras,
+        test_path
+):
+    # Prepare the Sv dataset
+    raw_path = test_path[path_model] / raw_and_xml_paths[0]
+    if raw_and_xml_paths[1]:
+        xml_path = test_path[path_model] / raw_and_xml_paths[1]
+    else:
+        xml_path = None
+
+    ed = ep.open_raw(raw_path, xml_path=xml_path, sonar_model=sonar_model)
+    if location_type == "fixed-location":
+        point_ds = xr.Dataset(
+            {
+                "latitude": (["time"], np.array([float(extras['latitude'])])),
+                "longitude": (["time"], np.array([float(extras['longitude'])])),
+            },
+            coords={
+                "time": (["time"], np.array([ed["Sonar/Beam_group1"]["ping_time"].values.min()]))
+            },
+        )
+        ed.update_platform(point_ds)
+
+    env_params = None
+    # AZFP data require external salinity and pressure
+    if sonar_model == "AZFP":
+        env_params = {
+            "temperature": ed["Environment"]["temperature"].values.mean(),
+            "salinity": extras["salinity"],
+            "pressure": extras["pressure"],
+        }
+
+    ds = ep.calibrate.compute_Sv(echodata=ed, env_params=env_params)
+
+    # add_location tests
+    if location_type == "empty-location":
+        with pytest.raises(Exception) as exc:
+            ep.consolidate.add_location(ds=ds, echodata=ed)
+        assert exc.type is ValueError
+        assert "Coordinate variables not present or all nan" in str(exc.value)
+    else:
+        def _tests(ds_test, location_type, nmea_sentence=None):
+            # lat,lon & time1 existence
+            assert "latitude" in ds_test
+            assert "longitude" in ds_test
+            assert "time1" not in ds_test
+
+            # lat & lon have a single dimension: 'ping_time'
+            assert len(ds_test["longitude"].dims) == 1 and ds_test["longitude"].dims[0] == "ping_time" # noqa
+            assert len(ds_test["latitude"].dims) == 1 and ds_test["latitude"].dims[0] == "ping_time" # noqa
+
+            # Check interpolated or broadcast values
+            if location_type == "with-track-location":
+                for position in ["longitude", "latitude"]:
+                    position_var = ed["Platform"][position]
+                    if nmea_sentence:
+                        position_var = position_var[ed["Platform"]["sentence_type"] == nmea_sentence]
+                    position_interp = position_var.interp(time1=ds_test["ping_time"])
+                    # interpolated values are identical
+                    assert np.allclose(ds_test[position].values, position_interp.values, equal_nan=True) # noqa
+            elif location_type == "fixed-location":
+                for position in ["longitude", "latitude"]:
+                    position_uniq = set(ds_test[position].values)
+                    # contains a single repeated value equal to the value passed to update_platform
+                    assert (
+                            len(position_uniq) == 1 and
+                            math.isclose(list(position_uniq)[0], extras[position])
+                    )
+
+        ds_all = ep.consolidate.add_location(ds=ds, echodata=ed)
+        _tests(ds_all, location_type)
+
+        # the test for nmea_sentence="GGA" is limited to the with-track-location case
+        if location_type == "with-track-location":
+            ds_sel = ep.consolidate.add_location(ds=ds, echodata=ed, nmea_sentence="GGA")
+            _tests(ds_sel, location_type, nmea_sentence="GGA")
+
+
 @pytest.mark.parametrize(
     ("sonar_model", "test_path_key", "raw_file_name", "paths_to_echoview_mat",
      "waveform_mode", "encode_mode", "pulse_compression", "write_Sv_to_file"),