Add XBT instrument simulation based on CTD

README.md

@@ -33,8 +33,9 @@
 
 VirtualShipParcels is a command line simulator allowing students to plan and conduct a virtual research expedition, receiving measurements as if they were coming from actual oceanographic instruments including:
 
-- ADCP (for currents)
-- CTD (for conductivity, and temperature)
+- ADCP (currents)
+- CTD (conductivity and temperature)
+- XBT (temperature)
 - underwater measurements (salinity and temperature)
 - surface drifters
 - argo float deployments

src/virtualship/instruments/__init__.py

@@ -1,5 +1,5 @@
 """Measurement instrument that can be used with Parcels."""
 
-from . import adcp, argo_float, ctd, drifter, ship_underwater_st
+from . import adcp, argo_float, ctd, drifter, ship_underwater_st, xbt
 
-__all__ = ["adcp", "argo_float", "ctd", "drifter", "ship_underwater_st"]
+__all__ = ["adcp", "argo_float", "ctd", "drifter", "ship_underwater_st", "xbt"]

src/virtualship/instruments/xbt.py

@@ -0,0 +1,137 @@
+"""XBT instrument."""
+
+from dataclasses import dataclass
+from datetime import timedelta
+from pathlib import Path
+
+import numpy as np
+from parcels import FieldSet, JITParticle, ParticleSet, Variable
+
+from ..spacetime import Spacetime
+
+
+@dataclass
+class XBT:
+    """Configuration for a single XBT."""
+
+    spacetime: Spacetime
+    min_depth: float
+    max_depth: float
+    initial_fall_speed: float
+    deceleration_coefficient: float
+
+
+_XBTParticle = JITParticle.add_variables(
+    [
+        Variable("temperature", dtype=np.float32, initial=np.nan),
+        Variable("max_depth", dtype=np.float32),
+        Variable("min_depth", dtype=np.float32),
+        Variable("fall_speed", dtype=np.float32),
+        Variable("initial_fall_speed", dtype=np.float32),
+        Variable("deceleration_coefficient", dtype=np.float32),
+    ]
+)
+
+
+def _sample_temperature(particle, fieldset, time):
+    particle.temperature = fieldset.T[time, particle.depth, particle.lat, particle.lon]
+
+
+def _xbt_cast(particle, fieldset, time):
+    particle_ddepth = -particle.fall_speed * particle.dt
+
+    # update the fall speed from the quadractic fall-rate equation
+    particle.fall_speed = (
+        particle.fall_speed - particle.deceleration_coefficient * particle.dt
+    )
+
+    # delete particle when it reaches the maximum depth
+    if particle.depth + particle_ddepth < particle.max_depth:
+        particle.delete()
+
+
+def simulate_xbt(
+    fieldset: FieldSet,
+    out_path: str | Path,
+    xbts: list[XBT],
+    outputdt: timedelta,
+) -> None:
+    """
+    Use Parcels to simulate a set of XBTs in a fieldset.
+
+    :param fieldset: The fieldset to simulate the XBTs in.
+    :param out_path: The path to write the results to.
+    :param xbts: A list of XBTs to simulate.
+    :param outputdt: Interval which dictates the update frequency of file output during simulation
+    :raises ValueError: Whenever provided XBTs, fieldset, are not compatible with this function.
+    """
+    DT = 10.0  # dt of XBT simulation integrator
+
+    if len(xbts) == 0:
+        print(
+            "No XBTs provided. Parcels currently crashes when providing an empty particle set, so no XBT simulation will be done and no files will be created."
+        )
+        # TODO when Parcels supports it this check can be removed.
+        return
+
+    fieldset_starttime = fieldset.time_origin.fulltime(fieldset.U.grid.time_full[0])
+    fieldset_endtime = fieldset.time_origin.fulltime(fieldset.U.grid.time_full[-1])
+
+    # deploy time for all xbts should be later than fieldset start time
+    if not all(
+        [np.datetime64(xbt.spacetime.time) >= fieldset_starttime for xbt in xbts]
+    ):
+        raise ValueError("XBT deployed before fieldset starts.")
+
+    # depth the xbt will go to. shallowest between xbt max depth and bathymetry.
+    max_depths = [
+        max(
+            xbt.max_depth,
+            fieldset.bathymetry.eval(
+                z=0, y=xbt.spacetime.location.lat, x=xbt.spacetime.location.lon, time=0
+            ),
+        )
+        for xbt in xbts
+    ]
+
+    # initial fall speeds
+    initial_fall_speeds = [xbt.initial_fall_speed for xbt in xbts]
+
+    # XBT depth can not be too shallow, because kernel would break.
+    # This shallow is not useful anyway, no need to support.
+    for max_depth, fall_speed in zip(max_depths, initial_fall_speeds, strict=False):
+        if not max_depth <= -DT * fall_speed:
+            raise ValueError(
+                f"XBT max_depth or bathymetry shallower than maximum {-DT * fall_speed}"
+            )
+
+    # define xbt particles
+    xbt_particleset = ParticleSet(
+        fieldset=fieldset,
+        pclass=_XBTParticle,
+        lon=[xbt.spacetime.location.lon for xbt in xbts],
+        lat=[xbt.spacetime.location.lat for xbt in xbts],
+        depth=[xbt.min_depth for xbt in xbts],
+        time=[xbt.spacetime.time for xbt in xbts],
+        max_depth=max_depths,
+        min_depth=[xbt.min_depth for xbt in xbts],
+        fall_speed=[xbt.initial_fall_speed for xbt in xbts],
+    )
+
+    # define output file for the simulation
+    out_file = xbt_particleset.ParticleFile(name=out_path, outputdt=outputdt)
+
+    # execute simulation
+    xbt_particleset.execute(
+        [_sample_temperature, _xbt_cast],
+        endtime=fieldset_endtime,
+        dt=DT,
+        verbose_progress=False,
+        output_file=out_file,
+    )
+
+    # there should be no particles left, as they delete themselves when they finish profiling
+    if len(xbt_particleset.particledata) != 0:
+        raise ValueError(
+            "Simulation ended before XBT finished profiling. This most likely means the field time dimension did not match the simulation time span."
+        )

tests/instruments/test_xbt.py

@@ -0,0 +1,131 @@
+"""
+Test the simulation of XBT instruments.
+
+Fields are kept static over time and time component of XBT measurements is not tested tested because it's tricky to provide expected measurements.
+"""
+
+import datetime
+from datetime import timedelta
+
+import numpy as np
+import xarray as xr
+from parcels import Field, FieldSet
+
+from virtualship import Location, Spacetime
+from virtualship.instruments.xbt import XBT, simulate_xbt
+
+
+def test_simulate_xbts(tmpdir) -> None:
+    # arbitrary time offset for the dummy fieldset
+    base_time = datetime.datetime.strptime("1950-01-01", "%Y-%m-%d")
+
+    # where to cast XBTs
+    xbts = [
+        XBT(
+            spacetime=Spacetime(
+                location=Location(latitude=0, longitude=1),
+                time=base_time + datetime.timedelta(hours=0),
+            ),
+            min_depth=0,
+            max_depth=float("-inf"),
+            initial_fall_speed=6.553,
+            deceleration_coefficient=0.00242,
+        ),
+        XBT(
+            spacetime=Spacetime(
+                location=Location(latitude=1, longitude=0),
+                time=base_time,
+            ),
+            min_depth=0,
+            max_depth=float("-inf"),
+            initial_fall_speed=6.553,
+            deceleration_coefficient=0.00242,
+        ),
+    ]
+
+    # expected observations for xbts at surface and at maximum depth
+    xbt_exp = [
+        {
+            "surface": {
+                "temperature": 6,
+                "lat": xbts[0].spacetime.location.lat,
+                "lon": xbts[0].spacetime.location.lon,
+            },
+            "maxdepth": {
+                "temperature": 8,
+                "lat": xbts[0].spacetime.location.lat,
+                "lon": xbts[0].spacetime.location.lon,
+            },
+        },
+        {
+            "surface": {
+                "temperature": 6,
+                "lat": xbts[1].spacetime.location.lat,
+                "lon": xbts[1].spacetime.location.lon,
+            },
+            "maxdepth": {
+                "temperature": 8,
+                "lat": xbts[1].spacetime.location.lat,
+                "lon": xbts[1].spacetime.location.lon,
+            },
+        },
+    ]
+
+    # create fieldset based on the expected observations
+    # indices are time, depth, latitude, longitude
+    u = np.zeros((2, 2, 2, 2))
+    v = np.zeros((2, 2, 2, 2))
+    t = np.zeros((2, 2, 2, 2))
+
+    t[:, 1, 0, 1] = xbt_exp[0]["surface"]["temperature"]
+    t[:, 0, 0, 1] = xbt_exp[0]["maxdepth"]["temperature"]
+    t[:, 1, 1, 0] = xbt_exp[1]["surface"]["temperature"]
+    t[:, 0, 1, 0] = xbt_exp[1]["maxdepth"]["temperature"]
+
+    fieldset = FieldSet.from_data(
+        {"V": v, "U": u, "T": t},
+        {
+            "time": [
+                np.datetime64(base_time + datetime.timedelta(hours=0)),
+                np.datetime64(base_time + datetime.timedelta(hours=1)),
+            ],
+            "depth": [-1000, 0],
+            "lat": [0, 1],
+            "lon": [0, 1],
+        },
+    )
+    fieldset.add_field(Field("bathymetry", [-1000], lon=0, lat=0))
+
+    # perform simulation
+    out_path = tmpdir.join("out.zarr")
+
+    simulate_xbt(
+        xbts=xbts,
+        fieldset=fieldset,
+        out_path=out_path,
+        outputdt=timedelta(seconds=10),
+    )
+
+    # test if output is as expected
+    results = xr.open_zarr(out_path)
+
+    assert len(results.trajectory) == len(xbts)
+
+    for xbt_i, (traj, exp_bothloc) in enumerate(
+        zip(results.trajectory, xbt_exp, strict=True)
+    ):
+        obs_surface = results.sel(trajectory=traj, obs=0)
+        min_index = np.argmin(results.sel(trajectory=traj)["z"].data)
+        obs_maxdepth = results.sel(trajectory=traj, obs=min_index)
+
+        for obs, loc in [
+            (obs_surface, "surface"),
+            (obs_maxdepth, "maxdepth"),
+        ]:
+            exp = exp_bothloc[loc]
+            for var in ["temperature", "lat", "lon"]:
+                obs_value = obs[var].values.item()
+                exp_value = exp[var]
+                assert np.isclose(
+                    obs_value, exp_value
+                ), f"Observation incorrect {xbt_i=} {loc=} {var=} {obs_value=} {exp_value=}."