Fix Argo floats getting stuck when using local data

docs/user-guide/documentation/example_copernicus_download.ipynb

@@ -113,6 +113,7 @@
     "            minimum_depth=0.5057600140571594,\n",
     "            maximum_depth=5902.0576171875,\n",
     "            output_filename=filename,\n",
+    "            coordinates_selection_method=\"outside\",\n",
     "        )"
    ]
   },
@@ -153,6 +154,7 @@
     "        minimum_depth=0.5057600140571594,\n",
     "        maximum_depth=5902.05810546875,\n",
     "        output_filename=filename,\n",
+    "        coordinates_selection_method=\"outside\",\n",
     "    )"
    ]
   },

docs/user-guide/documentation/pre_download_data.md

@@ -78,6 +78,7 @@ The following assumptions are also made about the data:
    - Or these strings must appear as substrings within the variable names (e.g. `o2_glor` is acceptable for `o2`).
 4. Bathymetry data files must contain a variable named `deptho`.
 5. Pre-downloaded data files must have a `"positive"` attribute for the depth dimension (e.g. `"positive": "down"` or `"positive": "up"`) in order to ensure that the depth dimension is correctly interpreted under-the-hood.
+6. The depth levels in the pre-downloaded data must extend (at least) one level beyond the deepest depth your instruments are configured to (e.g. the Argo float `max_depth_meter`). For example, for data with depth levels at [..., 1684 m, 1942 m, 2225 m, ...] and an Argo float `max_depth_meter` of -2000, data must be downloaded down to (at least) the 2225 m level.
 
 #### Also of note
 

src/virtualship/instruments/argo_float.py

@@ -302,6 +302,16 @@ def simulate(self, measurements, out_path) -> None:
             _grid_edge_margin = 0.0
         grid_shallowest = grid_depths[-1] - _grid_edge_margin
 
+        # error out when input data depth can't fulfil argo max_depth requiremnet
+        grid_deepest = np.min(grid_depths)
+        deepest_target = min(min(m.max_depth, m.drift_depth) for m in measurements)
+        if len(grid_depths) > 1 and deepest_target < grid_deepest:
+            raise ValueError(
+                f"{self.__class__.__name__} max_depth/drift_depth ({deepest_target}m) is deeper than the deepest level "
+                f"of the input data ({grid_deepest:.2f}m). If running with --from-data, ensure the data was downloaded to "
+                f"(at least) one depth level beyond the Argo float max_depth, or make max_depth shallower."
+            )
+
         # define parcel particles
         argo_float_particleset = ParticleSet(
             fieldset=fieldset,

src/virtualship/instruments/base.py

@@ -346,20 +346,32 @@ def _get_local_ds(self, files: list[Path]) -> xr.Dataset:
         depth_max = self.fetch_spec.depth_max
         both_none = depth_min is None and depth_max is None
 
-        if depth_min == depth_max and not both_none:
-            depth_sel = {
-                "depth": [depth_min],
-                "method": "nearest",
-            }
-        else:
-            depth_sel = {"depth": slice(depth_max, depth_min)}
-
         ds = ds.sel(
             longitude=slice(min_lon, max_lon),
             latitude=slice(min_lat, max_lat),
         )
 
-        ds = ds.sel(**depth_sel)
+        if depth_min == depth_max and not both_none:
+            ds = ds.sel(depth=[depth_min], method="nearest")
+        else:
+            # mirror copernicusmarine's coordinates_selection_method="outside" (as in `_get_copernicus_ds`)
+            # i.e. keep one level beyond each requested bound, so that an instrument's max depth is always inside the fieldset
+            depths = ds["depth"].values
+            deep_idx = (
+                0
+                if depth_max is None
+                else max(int(np.searchsorted(depths, depth_max, side="right")) - 1, 0)
+            )
+            shallow_idx = (
+                len(depths)
+                if depth_min is None
+                else min(
+                    int(np.searchsorted(depths, depth_min, side="left")) + 1,
+                    len(depths),
+                )
+            )
+            ds = ds.isel(depth=slice(deep_idx, shallow_idx))
+
         return ds
 
     def _via_tmp_ds(self, ds: xr.Dataset) -> xr.Dataset:

tests/instruments/test_argo_float.py

@@ -57,25 +57,32 @@ def create_fieldset(
     lat_range=(0.0, 10.0),
     include_salinity=True,
     lifetime_days=0.1,
+    depths=None,
 ):
-    """Create a test fieldset with optional salinity."""
-    v = np.full((2, 2, 2), 1.0)
-    u = np.full((2, 2, 2), 1.0)
-    t = np.full((2, 2, 2), 1.0)
+    """Create a test fieldset with optional salinity and optional depth levels."""
+    dims = ("time", "lat", "lon") if depths is None else ("time", "depth", "lat", "lon")
+    shape = (2, 2, 2) if depths is None else (2, len(depths), 2, 2)
     bathy = np.full((2, 2), -5000.0)
 
     data_vars = {
-        "V": (("time", "lat", "lon"), v),
-        "U": (("time", "lat", "lon"), u),
-        "T": (("time", "lat", "lon"), t),
+        "V": (dims, np.full(shape, 1.0)),
+        "U": (dims, np.full(shape, 1.0)),
+        "T": (dims, np.full(shape, 1.0)),
     }
 
     if include_salinity:
-        data_vars["S"] = (("time", "lat", "lon"), np.full((2, 2, 2), 1.0))
+        data_vars["S"] = (dims, np.full(shape, 1.0))
+
+    depth_coord = (
+        {}
+        if depths is None
+        else {"depth": (("depth"), np.array(depths), {"positive": "up"})}
+    )
 
     ds_fields = xr.Dataset(
         data_vars=data_vars,
         coords={
+            **depth_coord,
             "lon": (("lon"), np.array(lon_range), {"units": "degrees_east"}),
             "lat": (("lat"), np.array(lat_range), {"units": "degrees_north"}),
             "time": (
@@ -343,6 +350,63 @@ def test_argo_fieldoutofbounds_error(tmpdir) -> None:
     )
 
 
+def test_argo_float_reaches_max_depth_and_ascends(tmpdir) -> None:
+    """Argo float should reach its max depth and then ascends."""
+    lifetime_days = 1.0  # time enough for one descent to max depth + ascent
+    fieldset = create_fieldset(
+        lifetime_days=lifetime_days, depths=[-2225.1, -1941.9, -1000.0, -0.5]
+    )
+
+    sensors = [SensorConfig(sensor_type=SensorType.TEMPERATURE)]
+    expedition = create_dummy_expedition(
+        sensors, lifetime=timedelta(days=lifetime_days)
+    )
+    argo_instrument = ArgoFloatInstrument(expedition, None)
+    wp = expedition.schedule.waypoints[0]
+    argo_float = ArgoFloat(
+        spacetime=Spacetime(location=wp.location, time=wp.time),
+        min_depth=0.0,
+        max_depth=MAX_DEPTH,
+        drift_depth=DRIFT_DEPTH,
+        vertical_speed=VERTICAL_SPEED,
+        cycle_days=1,
+        drift_days=0,
+    )
+
+    out_path = tmpdir.join("out.parquet")
+    argo_instrument.load_input_data = lambda: fieldset
+    argo_instrument.simulate([argo_float], out_path)
+
+    results = parcels.read_particlefile(out_path)
+    z = results["z"].to_numpy()
+    phase = results["cycle_phase"].to_numpy()
+
+    np.testing.assert_allclose(z.min(), MAX_DEPTH, atol=1e-3)
+    assert not results["grounded"].to_numpy().any()
+
+    # never sent back up early
+    phase2_dz = np.diff(z)[(phase[:-1] == 2) & (phase[1:] == 2)]
+    assert (phase2_dz <= 0).all()
+
+    # ascent (sampling) starts from max depth
+    assert np.isclose(z[phase == 3].min(), MAX_DEPTH, atol=1e-3)
+    assert np.isfinite(results["temperature"].to_numpy()[phase == 3]).any()
+
+
+def test_argo_max_depth_deeper_than_fieldset_error(tmpdir) -> None:
+    """A max depth deeper than the fieldset depth is rejected at setup."""
+    fieldset = create_fieldset(depths=[-1941.9, -1000.0, -0.5])
+
+    sensors = [SensorConfig(sensor_type=SensorType.TEMPERATURE)]
+    expedition = create_dummy_expedition(sensors)
+    argo_instrument = ArgoFloatInstrument(expedition, None)
+    argo_floats = [create_argo_float(wp) for wp in expedition.schedule.waypoints]
+
+    argo_instrument.load_input_data = lambda: fieldset
+    with pytest.raises(ValueError, match="deeper than the deepest level"):
+        argo_instrument.simulate(argo_floats, tmpdir.join("out.parquet"))
+
+
 def test_argo_float_instrument_type():
     """ArgoFloatInstrument returns the correct InstrumentType and if is underway instrument."""
     sensors = [

tests/instruments/test_base.py

@@ -339,6 +339,52 @@ def test_generate_fieldset_combines_fields(mock_expedition):
     fs_A.__add__.assert_called_once_with(fs_B)
 
 
+@pytest.mark.parametrize(
+    "depth_min, depth_max, expected_depths",
+    [
+        # when max depth is between levels, the next deeper level is kept (i.e. max depth is inside the fieldset)
+        (0.0, -2000.0, [-2225.1, -1941.9, -1684.3, -0.5]),
+        # min depth between levels means the next shallower level is kept
+        (-10.0, -1700.0, [-1941.9, -1684.3, -0.5]),
+        # bounds exactly on levels means no extra levels
+        (-1684.3, -1941.9, [-1941.9, -1684.3]),
+    ],
+)
+def test_get_local_ds_depth_selection_outside(
+    mock_expedition, tmp_path, depth_min, depth_max, expected_depths
+):
+    """Local depth selection keeps one level beyond each bound, like copernicusmarine's coordinates_selection_method='outside'."""
+    depths = np.array([0.5, 1684.3, 1941.9, 2225.1, 2533.3])  # inspired by GLORYS
+    ds = xr.Dataset(
+        data_vars={
+            "thetao": (
+                ["time", "depth", "latitude", "longitude"],
+                np.zeros((1, len(depths), 2, 2)),
+            )
+        },
+        coords={
+            "time": [np.datetime64("2026-01-01")],
+            "depth": ("depth", depths, {"positive": "down"}),
+            "latitude": [0.0, 30.0],
+            "longitude": [-30.0, 0.0],
+        },
+    )
+    file = tmp_path / "phys.nc"
+    ds.to_netcdf(file)
+
+    dummy = DummyInstrument(
+        expedition=mock_expedition,
+        variables={"T": "thetao"},
+        add_bathymetry=False,
+        verbose_progress=False,
+        fetch_spec=FetchSpec(depth_min=depth_min, depth_max=depth_max),
+        from_data=tmp_path,
+    )
+
+    result = dummy._get_local_ds([file])
+    np.testing.assert_allclose(result["depth"].values, expected_depths)
+
+
 def test_load_input_data_error(mock_expedition, monkeypatch):
     dummy = DummyInstrument(
         expedition=mock_expedition,