Ctd cleanup and test

tests/instruments/test_ctd.py

@@ -1,38 +1,138 @@
-"""Test the simulation of CTD instruments."""
+"""
+Test the simulation of CTD instruments.
 
+Fields are kept static over time and time component of CTD measurements is not tested tested because it's tricky to provide expected measurements.
+"""
+
+import datetime
 from datetime import timedelta
 
 import numpy as np
-from parcels import FieldSet
+import py
+import xarray as xr
+from parcels import Field, FieldSet
 
 from virtual_ship import Location, Spacetime
 from virtual_ship.instruments.ctd import CTD, simulate_ctd
 
 
-def test_simulate_ctds() -> None:
+def test_simulate_ctds(tmpdir: py.path.LocalPath) -> None:
+    # arbitrary time offset for the dummy fieldset
+    base_time = datetime.datetime.strptime("1950-01-01", "%Y-%m-%d")
+
+    # where to cast CTDs
+    ctds = [
+        CTD(
+            spacetime=Spacetime(
+                location=Location(latitude=0, longitude=1),
+                time=base_time + datetime.timedelta(hours=0),
+            ),
+            min_depth=0,
+            max_depth=float("-inf"),
+        ),
+        CTD(
+            spacetime=Spacetime(
+                location=Location(latitude=1, longitude=0),
+                time=base_time + datetime.timedelta(minutes=5),
+            ),
+            min_depth=0,
+            max_depth=float("-inf"),
+        ),
+    ]
+
+    # expected observations for ctds at surface and at maximum depth
+    ctd_exp = [
+        {
+            "surface": {
+                "salinity": 5,
+                "temperature": 6,
+                "lat": ctds[0].spacetime.location.lat,
+                "lon": ctds[0].spacetime.location.lon,
+            },
+            "maxdepth": {
+                "salinity": 7,
+                "temperature": 8,
+                "lat": ctds[0].spacetime.location.lat,
+                "lon": ctds[0].spacetime.location.lon,
+            },
+        },
+        {
+            "surface": {
+                "salinity": 5,
+                "temperature": 6,
+                "lat": ctds[1].spacetime.location.lat,
+                "lon": ctds[1].spacetime.location.lon,
+            },
+            "maxdepth": {
+                "salinity": 7,
+                "temperature": 8,
+                "lat": ctds[1].spacetime.location.lat,
+                "lon": ctds[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))
+    s = np.zeros((2, 2, 2, 2))
+
+    t[:, 1, 0, 1] = ctd_exp[0]["surface"]["temperature"]
+    t[:, 0, 0, 1] = ctd_exp[0]["maxdepth"]["temperature"]
+    t[:, 1, 1, 0] = ctd_exp[1]["surface"]["temperature"]
+    t[:, 0, 1, 0] = ctd_exp[1]["maxdepth"]["temperature"]
+
+    s[:, 1, 0, 1] = ctd_exp[0]["surface"]["salinity"]
+    s[:, 0, 0, 1] = ctd_exp[0]["maxdepth"]["salinity"]
+    s[:, 1, 1, 0] = ctd_exp[1]["surface"]["salinity"]
+    s[:, 0, 1, 0] = ctd_exp[1]["maxdepth"]["salinity"]
+
     fieldset = FieldSet.from_data(
-        {"U": 0, "V": 0, "T": 0, "S": 0, "bathymetry": 100},
+        {"V": v, "U": u, "T": t, "S": s},
         {
-            "lon": 0,
-            "lat": 0,
-            "time": [np.datetime64("1950-01-01") + np.timedelta64(632160, "h")],
+            "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))
 
-    min_depth = -fieldset.U.depth[0]
-    max_depth = -fieldset.U.depth[-1]
-
-    ctds = [
-        CTD(
-            spacetime=Spacetime(location=Location(latitude=0, longitude=0), time=0),
-            min_depth=min_depth,
-            max_depth=max_depth,
-        )
-    ]
+    # perform simulation
+    out_path = tmpdir.join("out.zarr")
 
     simulate_ctd(
         ctds=ctds,
         fieldset=fieldset,
-        out_file_name="test",
+        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(ctds)
+
+    for ctd_i, (traj, exp_bothloc) in enumerate(
+        zip(results.trajectory, ctd_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 ["salinity", "temperature", "lat", "lon"]:
+                obs_value = obs[var].values.item()
+                exp_value = exp[var]
+                assert np.isclose(
+                    obs_value, exp_value
+                ), f"Observation incorrect {ctd_i=} {loc=} {var=} {obs_value=} {exp_value=}."

tests/test_sailship.py

@@ -1,12 +1,39 @@
 """Performs a complete cruise with virtual ship."""
 
+import datetime
+
 import numpy as np
-from parcels import FieldSet
+from parcels import Field, FieldSet
 
 from virtual_ship.sailship import sailship
 from virtual_ship.virtual_ship_configuration import VirtualShipConfiguration
 
 
+def _make_ctd_fieldset() -> FieldSet:
+    # arbitrary time offset for the dummy fieldset
+    base_time = datetime.datetime.strptime("2022-01-01T00:00:00", "%Y-%m-%dT%H:%M:%S")
+
+    u = np.zeros((2, 2, 2, 2))
+    v = np.zeros((2, 2, 2, 2))
+    t = np.zeros((2, 2, 2, 2))
+    s = np.zeros((2, 2, 2, 2))
+
+    fieldset = FieldSet.from_data(
+        {"V": v, "U": u, "T": t, "S": s},
+        {
+            "time": [
+                np.datetime64(base_time + datetime.timedelta(seconds=0)),
+                np.datetime64(base_time + datetime.timedelta(minutes=200)),
+            ],
+            "depth": [0, -1000],
+            "lat": [-40, 90],
+            "lon": [-90, 90],
+        },
+    )
+    fieldset.add_field(Field("bathymetry", [-1000], lon=0, lat=0))
+    return fieldset
+
+
 def test_sailship() -> None:
     adcp_fieldset = FieldSet.from_data(
         {"U": 0, "V": 0},
@@ -18,10 +45,7 @@ def test_sailship() -> None:
         {"lon": 0, "lat": 0},
     )
 
-    ctd_fieldset = FieldSet.from_data(
-        {"U": 0, "V": 0, "S": 0, "T": 0, "bathymetry": 0},
-        {"lon": 0, "lat": 0},
-    )
+    ctd_fieldset = _make_ctd_fieldset()
 
     drifter_fieldset = FieldSet.from_data(
         {"U": 0, "V": 0, "T": 0},

virtual_ship/instruments/ctd.py

@@ -4,6 +4,7 @@
 from datetime import timedelta
 
 import numpy as np
+import py
 from parcels import FieldSet, JITParticle, ParticleSet, Variable
 
 from ..spacetime import Spacetime
@@ -22,8 +23,10 @@ class CTD:
     [
         Variable("salinity", dtype=np.float32, initial=np.nan),
         Variable("temperature", dtype=np.float32, initial=np.nan),
-        Variable("raising", dtype=np.int32, initial=0),
+        Variable("raising", dtype=np.int8, initial=0.0),  # bool. 0 is False, 1 is True.
         Variable("max_depth", dtype=np.float32),
+        Variable("min_depth", dtype=np.float32),
+        Variable("winch_speed", dtype=np.float32),
     ]
 )
 
@@ -37,77 +40,89 @@ def _sample_salinity(particle, fieldset, time):
 
 
 def _ctd_cast(particle, fieldset, time):
-    # Lowering and raising CTD
-    if (
-        -fieldset.bathymetry[time, particle.depth, particle.lat, particle.lon]
-        > particle.max_depth
-    ):
-        maxdepth = (
-            -fieldset.bathymetry[time, particle.depth, particle.lat, particle.lon] + 20
-        )
-    else:
-        maxdepth = particle.max_depth
-    winch_speed = -1.0  # sink and rise speed in m/s
-
+    # lowering
     if particle.raising == 0:
-        # Sinking with winch_speed until near seafloor
-        particle_ddepth = winch_speed * particle.dt
-        if particle.depth <= maxdepth:
+        particle_ddepth = -particle.winch_speed * particle.dt
+        if particle.depth + particle_ddepth < particle.max_depth:
             particle.raising = 1
-
-    if particle.raising == 1:
-        # Rising with winch_speed until depth is -2 m
-        if particle.depth < -2:
-            particle_ddepth = -winch_speed * particle.dt
-            if particle.depth + particle_ddepth >= -2:
-                # to break the loop ...
-                particle.state = 41
-                print("CTD cast finished.")
+            particle_ddepth = -particle_ddepth
+    # raising
+    else:
+        particle_ddepth = particle.winch_speed * particle.dt
+        if particle.depth + particle_ddepth > -particle.min_depth:
+            particle.delete()
 
 
 def simulate_ctd(
-    ctds: list[CTD],
     fieldset: FieldSet,
-    out_file_name: str,
+    out_path: str | py.path.LocalPath,
+    ctds: list[CTD],
     outputdt: timedelta,
 ) -> None:
     """
     Use parcels to simulate a set of CTDs in a fieldset.
 
-    :param ctds: A list of CTDs to simulate.
     :param fieldset: The fieldset to simulate the CTDs in.
-    :param out_file_name: The file to write the results to.
+    :param out_path: The path to write the results to.
+    :param ctds: A list of CTDs to simulate.
     :param outputdt: Interval which dictates the update frequency of file output during simulation
+    :raises ValueError: Whenever provided CTDs, fieldset, are not compatible with this function.
     """
-    lon = [ctd.spacetime.location.lon for ctd in ctds]
-    lat = [ctd.spacetime.location.lat for ctd in ctds]
-    time = [ctd.spacetime.time for ctd in ctds]
+    WINCH_SPEED = 1.0  # sink and rise speed in m/s
+    DT = 10.0  # dt of CTD simulation integrator
+
+    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 ctds should be later than fieldset start time
+    if not all([ctd.spacetime.time >= fieldset_starttime for ctd in ctds]):
+        raise ValueError("CTD deployed before fieldset starts.")
+
+    # depth the ctd will go to. shallowest between ctd max depth and bathymetry.
+    max_depths = [
+        max(
+            ctd.max_depth,
+            fieldset.bathymetry.eval(
+                z=0, y=ctd.spacetime.location.lat, x=ctd.spacetime.location.lon, time=0
+            ),
+        )
+        for ctd in ctds
+    ]
+
+    # CTD depth can not be too shallow, because kernel would break.
+    # This shallow is not useful anyway, no need to support.
+    if not all([max_depth <= -DT * WINCH_SPEED for max_depth in max_depths]):
+        raise ValueError(
+            f"CTD max_depth or bathymetry shallower than maximum {-DT * WINCH_SPEED}"
+        )
 
     # define parcel particles
     ctd_particleset = ParticleSet(
         fieldset=fieldset,
         pclass=_CTDParticle,
-        lon=lon,
-        lat=lat,
+        lon=[ctd.spacetime.location.lon for ctd in ctds],
+        lat=[ctd.spacetime.location.lat for ctd in ctds],
         depth=[ctd.min_depth for ctd in ctds],
-        time=time,
-        max_depth=[ctd.max_depth for ctd in ctds],
+        time=[ctd.spacetime.time for ctd in ctds],
+        max_depth=max_depths,
+        min_depth=[ctd.min_depth for ctd in ctds],
+        winch_speed=[WINCH_SPEED for _ in ctds],
     )
 
     # define output file for the simulation
-    out_file = ctd_particleset.ParticleFile(
-        name=out_file_name,
-        outputdt=outputdt,
-        chunks=(1, 500),
-    )
-
-    # get time when the fieldset ends
-    fieldset_endtime = fieldset.time_origin.fulltime(fieldset.U.grid.time_full[-1])
+    out_file = ctd_particleset.ParticleFile(name=out_path, outputdt=outputdt)
 
     # execute simulation
     ctd_particleset.execute(
         [_sample_salinity, _sample_temperature, _ctd_cast],
         endtime=fieldset_endtime,
-        dt=outputdt,
+        dt=DT,
+        verbose_progress=False,
         output_file=out_file,
     )
+
+    # there should be no particles left, as they delete themselves when they resurface
+    if len(ctd_particleset.particledata) != 0:
+        raise ValueError(
+            "Simulation ended before CTD resurfaced. This most likely means the field time dimension did not match the simulation time span."
+        )