Fix MVBS metadata propagation for channel/frequency_nominal and swapped dimensions

echopype/commongrid/api.py

@@ -186,14 +186,33 @@ def compute_MVBS(
     prov_dict["processing_function"] = "commongrid.compute_MVBS"
     ds_MVBS = ds_MVBS.assign_attrs(prov_dict)
 
-    # Preserve the channel order returned by compute_raw_MVBS and align
-    # frequency_nominal to that order.
-    freq = ds_Sv["frequency_nominal"]
-
-    if "channel" in ds_MVBS.dims:
-        ds_MVBS["frequency_nominal"] = freq.sel(channel=ds_MVBS["channel"])
-    else:
-        ds_MVBS["frequency_nominal"] = freq
+    # Reattach the variables that label the primary dimension, aligned to the order
+    # returned by compute_raw_MVBS. That label is frequency_nominal normally, or
+    # channel once dimensions have been swapped, and either may be stored as a
+    # coordinate or as a data variable, so keep whichever role it had in ds_Sv.
+    for name in list(ds_Sv.coords) + list(ds_Sv.data_vars):
+        if name in ds_MVBS.variables:
+            continue
+
+        var = ds_Sv[name]
+        if dim_0 not in var.dims:
+            continue
+
+        # skip anything carrying range information, which is binned separately
+        if any(dim not in (dim_0, "ping_time") for dim in var.dims):
+            continue
+
+        var = var.sel({dim_0: ds_MVBS[dim_0]})
+
+        # these labels do not vary across pings, so reduce rather than carry the
+        # unbinned ping_time over, which would not align with the binned one
+        if "ping_time" in var.dims:
+            var = var.isel(ping_time=0, drop=True)
+
+        if name in ds_Sv.coords:
+            ds_MVBS = ds_MVBS.assign_coords({name: var})
+        else:
+            ds_MVBS[name] = var
 
     ds_MVBS = insert_input_processing_level(ds_MVBS, input_ds=ds_Sv)
 

echopype/tests/calibrate/test_range.py

@@ -0,0 +1,145 @@
+import pytest
+
+import numpy as np
+import xarray as xr
+
+from echopype.calibrate.range import compute_range_EK, range_mod_TVG_EK
+
+pytestmark = pytest.mark.unit
+
+SOUND_SPEED = 1500.0
+SAMPLE_INTERVAL = 1e-4
+TRANSMIT_DURATION = 1e-3
+
+
+def _mock_beam(channel, n_ping=2, n_range_sample=4, with_beam_dim=False):
+    """Minimal Beam_group dataset holding the variables range computation needs."""
+    dims = ["channel", "ping_time", "range_sample"]
+    coords = {
+        "channel": channel,
+        "ping_time": np.arange(n_ping),
+        "range_sample": np.arange(n_range_sample),
+    }
+
+    backscatter_r = xr.DataArray(
+        np.ones((len(channel), n_ping, n_range_sample)), dims=dims, coords=coords
+    )
+    if with_beam_dim:
+        # beam is a coordinate in EK80 beam groups, not a bare dimension
+        backscatter_r = backscatter_r.expand_dims({"beam": ["1", "2"]}, axis=-1)
+
+    # both vary by channel and ping in real data
+    sample_interval = np.full((len(channel), n_ping), SAMPLE_INTERVAL)
+    transmit_duration = np.full((len(channel), n_ping), TRANSMIT_DURATION)
+
+    return xr.Dataset(
+        {
+            "backscatter_r": backscatter_r,
+            "sample_interval": (["channel", "ping_time"], sample_interval),
+            "transmit_duration_nominal": (["channel", "ping_time"], transmit_duration),
+        },
+        coords=coords,
+    )
+
+
+def _mock_tvg_inputs(transceiver_type):
+    """Beam and Vendor_specific datasets plus the unmodified range."""
+    channel = [f"ch{idx}" for idx in range(len(transceiver_type))]
+    beam = _mock_beam(channel)
+    vend = xr.Dataset(
+        {"transceiver_type": ("channel", transceiver_type)},
+        coords={"channel": channel},
+    )
+    range_meter = compute_range_EK("EK60", beam, {"sound_speed": SOUND_SPEED})
+    return beam, vend, range_meter
+
+
+def test_compute_range_EK_values():
+    beam = _mock_beam(["ch1", "ch2"])
+    echo_range = compute_range_EK("EK60", beam, {"sound_speed": SOUND_SPEED})
+
+    expected = beam["range_sample"].data * SAMPLE_INTERVAL * SOUND_SPEED / 2
+    for ch in beam["channel"].data:
+        for ping in beam["ping_time"].data:
+            assert np.allclose(echo_range.sel(channel=ch, ping_time=ping).data, expected)
+
+    assert echo_range.name == "echo_range"
+    assert echo_range.dims == ("channel", "ping_time", "range_sample")
+
+
+def test_compute_range_EK_sound_speed_scaling():
+    beam = _mock_beam(["ch1"])
+    fast = compute_range_EK("EK60", beam, {"sound_speed": SOUND_SPEED})
+    slow = compute_range_EK("EK60", beam, {"sound_speed": SOUND_SPEED / 2})
+    assert np.allclose(slow.data * 2, fast.data)
+
+
+def test_compute_range_EK_nan_backscatter():
+    beam = _mock_beam(["ch1"])
+    beam["backscatter_r"][dict(channel=0, ping_time=0, range_sample=slice(2, None))] = np.nan
+    echo_range = compute_range_EK("EK60", beam, {"sound_speed": SOUND_SPEED})
+
+    assert np.all(np.isnan(echo_range.isel(channel=0, ping_time=0, range_sample=slice(2, None))))
+    assert not np.any(np.isnan(echo_range.isel(channel=0, ping_time=0, range_sample=slice(0, 2))))
+
+
+def test_compute_range_EK_drops_beam_dim():
+    beam = _mock_beam(["ch1"], with_beam_dim=True)
+    echo_range = compute_range_EK("EK80", beam, {"sound_speed": SOUND_SPEED})
+    assert "beam" not in echo_range.dims
+
+
+def test_compute_range_EK_unsupported_sonar_model():
+    beam = _mock_beam(["ch1"])
+    with pytest.raises(ValueError, match="is not supported"):
+        compute_range_EK("AZFP", beam, {"sound_speed": SOUND_SPEED})
+
+
+def test_compute_range_EK_missing_sound_speed():
+    beam = _mock_beam(["ch1"])
+    with pytest.raises(RuntimeError, match="sounds_speed not included"):
+        compute_range_EK("EK60", beam, {})
+
+
+def test_range_mod_TVG_EK_ex60():
+    beam, vend, range_meter = _mock_tvg_inputs(["GPT"])
+    modified = range_mod_TVG_EK(
+        "EK60", beam, vend, range_meter.copy(deep=True), xr.DataArray(SOUND_SPEED)
+    )
+
+    # Ex60 hardware: 2-sample shift at the beginning
+    expected_shift = 2 * SAMPLE_INTERVAL * SOUND_SPEED / 2
+    assert np.allclose(modified.data, range_meter.data - expected_shift)
+
+
+def test_range_mod_TVG_EK_ex80():
+    beam, vend, range_meter = _mock_tvg_inputs(["WBT", "WBT"])
+    modified = range_mod_TVG_EK(
+        "EK80", beam, vend, range_meter.copy(deep=True), xr.DataArray(SOUND_SPEED)
+    )
+
+    # Ex80 hardware: shift by sound_speed * transmit_duration_nominal / 4
+    expected_shift = SOUND_SPEED * TRANSMIT_DURATION / 4
+    assert np.allclose(modified.data, range_meter.data - expected_shift)
+
+
+def test_range_mod_TVG_EK_ex80_mixed_wbt_gpt():
+    beam, vend, range_meter = _mock_tvg_inputs(["GPT", "WBT"])
+    modified = range_mod_TVG_EK(
+        "EK80", beam, vend, range_meter.copy(deep=True), xr.DataArray(SOUND_SPEED)
+    )
+
+    ex60_shift = 2 * SAMPLE_INTERVAL * SOUND_SPEED / 2
+    ex80_shift = SOUND_SPEED * TRANSMIT_DURATION / 4
+
+    # the WBT channel gets the Ex80 correction
+    assert np.allclose(
+        modified.sel(channel="ch1").data, range_meter.sel(channel="ch1").data - ex80_shift
+    )
+
+    # the GPT channel gets the Ex60 correction on top of the Ex80 one already
+    # applied to every channel, so both are subtracted
+    assert np.allclose(
+        modified.sel(channel="ch0").data,
+        range_meter.sel(channel="ch0").data - ex80_shift - ex60_shift,
+    )

echopype/tests/commongrid/test_commongrid_api.py

@@ -155,6 +155,62 @@ def test__groupby_x_along_channels(request, range_var, lat_lon):
     # Check that the range_var is in the dimension
     assert f"{range_var}_bins" in sv_mean.dims
 
+@pytest.mark.unit
+@pytest.mark.parametrize("freq_along_ping_time", [False, True])
+def test_compute_MVBS_preserves_frequency_nominal(
+    ds_Sv_echo_range_regular, freq_along_ping_time
+):
+    """frequency_nominal is kept, aligned to the output channel order, per channel."""
+    ds_Sv = ds_Sv_echo_range_regular.assign_coords(
+        channel=["55196-38-1", "55196-120-2", "55196-200-3", "55196-455-4"]
+    )
+    freq = np.array([38000.0, 120000.0, 200000.0, 455000.0])
+
+    if freq_along_ping_time:
+        # some datasets store frequency_nominal as (ping_time, channel)
+        ds_Sv["frequency_nominal"] = xr.DataArray(
+            np.tile(freq, (ds_Sv.sizes["ping_time"], 1)),
+            dims=["ping_time", "channel"],
+            coords={"ping_time": ds_Sv["ping_time"], "channel": ds_Sv["channel"]},
+        )
+    else:
+        ds_Sv["frequency_nominal"] = xr.DataArray(
+            freq, dims=["channel"], coords={"channel": ds_Sv["channel"]}
+        )
+
+    ds_MVBS = ep.commongrid.compute_MVBS(ds_Sv, range_bin="20m", ping_time_bin="20s")
+
+    assert "frequency_nominal" in ds_MVBS.variables
+
+    # one value per channel, regardless of how it was stored on the way in
+    assert ds_MVBS["frequency_nominal"].dims == ("channel",)
+    assert not np.isnan(ds_MVBS["frequency_nominal"].values).any()
+
+    expected = dict(zip(ds_Sv["channel"].values, freq))
+    for ch, actual in zip(ds_MVBS["channel"].values, ds_MVBS["frequency_nominal"].values):
+        assert actual == expected[ch]
+
+
+@pytest.mark.unit
+def test_compute_MVBS_preserves_channel_when_dims_swapped(ds_Sv_echo_range_regular):
+    """After swap_dims_channel_frequency, channel is a variable and must survive too."""
+    ds_Sv = ds_Sv_echo_range_regular
+    ds_Sv["frequency_nominal"] = xr.DataArray(
+        np.array([38000.0, 120000.0, 200000.0, 455000.0]),
+        dims=["channel"],
+        coords={"channel": ds_Sv["channel"]},
+    )
+    ds_Sv_swapped = ep.consolidate.swap_dims_channel_frequency(ds_Sv)
+
+    ds_MVBS = ep.commongrid.compute_MVBS(ds_Sv_swapped, range_bin="20m", ping_time_bin="20s")
+
+    assert "channel" in ds_MVBS.variables
+    assert ds_MVBS["channel"].dims == ("frequency_nominal",)
+
+    expected = ds_Sv_swapped["channel"].sel(frequency_nominal=ds_MVBS["frequency_nominal"])
+    assert np.array_equal(ds_MVBS["channel"].values, expected.values)
+
+
 # NASC Tests
 @pytest.mark.integration
 @pytest.mark.parametrize("compute_mvbs", [True, False])