Ship underwater st test

tests/conftest.py

@@ -4,7 +4,7 @@
 
 
 @pytest.fixture(autouse=True)
-def change_test_dir(request, monkeypatch):
+def test_in_working_dir(request, monkeypatch):
     """
     Set the working directory for each test to the directory of that test.
 

tests/instruments/test_ship_underwater_st.py

@@ -1,29 +1,100 @@
 """Test the simulation of ship salinity temperature measurements."""
 
+import datetime
+
 import numpy as np
+import py
+import xarray as xr
 from parcels import FieldSet
 
 from virtual_ship import Location, Spacetime
 from virtual_ship.instruments.ship_underwater_st import simulate_ship_underwater_st
 
 
-def test_simulate_ship_underwater_st() -> None:
+def test_simulate_ship_underwater_st(tmpdir: py.path.LocalPath) -> None:
+    # depth at which the sampling will be done
     DEPTH = -2
 
+    # arbitrary time offset for the dummy fieldset
+    base_time = datetime.datetime.strptime("1950-01-01", "%Y-%m-%d")
+
+    # variabes we are going to compare between expected and actual observations
+    variables = ["salinity", "temperature", "lat", "lon"]
+
+    # where to sample
+    sample_points = [
+        Spacetime(Location(3, 0), base_time + datetime.timedelta(seconds=0)),
+        Spacetime(Location(7, 0), base_time + datetime.timedelta(seconds=1)),
+    ]
+
+    # expected observations at sample points
+    expected_obs = [
+        {
+            "salinity": 1,
+            "temperature": 2,
+            "lat": 3,
+            "lon": 0,
+            "time": base_time + datetime.timedelta(seconds=0),
+        },
+        {
+            "salinity": 5,
+            "temperature": 6,
+            "lat": 7,
+            "lon": 0,
+            "time": base_time + datetime.timedelta(seconds=1),
+        },
+    ]
+
+    # create fieldset based on the expected observations
     fieldset = FieldSet.from_data(
-        {"U": 0, "V": 0, "S": 0, "T": 0},
+        {
+            "U": np.zeros((2, 2)),
+            "V": np.zeros((2, 2)),
+            "salinity": [
+                [expected_obs[0]["salinity"], 0],
+                [0, expected_obs[1]["salinity"]],
+            ],
+            "temperature": [
+                [expected_obs[0]["temperature"], 0],
+                [0, expected_obs[1]["temperature"]],
+            ],
+        },
         {
             "lon": 0,
-            "lat": 0,
-            "time": [np.datetime64("1950-01-01") + np.timedelta64(632160, "h")],
+            "lat": np.array([expected_obs[0]["lat"], expected_obs[1]["lat"]]),
+            "time": np.array(
+                [
+                    np.datetime64(expected_obs[0]["time"]),
+                    np.datetime64(expected_obs[1]["time"]),
+                ]
+            ),
         },
     )
 
-    sample_points = [Spacetime(Location(0, 0), 0)]
+    out_path = tmpdir.join("out.zarr")
 
     simulate_ship_underwater_st(
         fieldset=fieldset,
-        out_file_name="test",
+        out_path=out_path,
         depth=DEPTH,
         sample_points=sample_points,
     )
+
+    results = xr.open_zarr(out_path)
+
+    # below we assert if output makes sense
+    assert len(results.trajectory) == 1  # expect a singe trajectory
+    traj = results.trajectory.item()
+    assert len(results.sel(trajectory=traj).obs) == len(
+        sample_points
+    )  # expect as many obs as sample points
+
+    # for every obs, check if the variables match the expected observations
+    for obs_i, exp in zip(results.sel(trajectory=traj).obs, expected_obs, strict=True):
+        obs = results.sel(trajectory=traj, obs=obs_i)
+        for var in variables:
+            obs_value = obs[var].values.item()
+            exp_value = exp[var]
+            assert np.isclose(
+                obs_value, exp_value
+            ), f"Observation incorrect {obs_i=} {var=} {obs_value=} {exp_value=}."

virtual_ship/instruments/ship_underwater_st.py

@@ -1,11 +1,14 @@
 """Ship salinity and temperature."""
 
 import numpy as np
-from parcels import FieldSet, JITParticle, ParticleSet, Variable
+import py
+from parcels import FieldSet, ParticleSet, ScipyParticle, Variable
 
 from ..spacetime import Spacetime
 
-_ShipSTParticle = JITParticle.add_variables(
+# we specifically use ScipyParticle because we have many small calls to execute
+# there is some overhead with JITParticle and this ends up being significantly faster
+_ShipSTParticle = ScipyParticle.add_variables(
     [
         Variable("salinity", dtype=np.float32, initial=np.nan),
         Variable("temperature", dtype=np.float32, initial=np.nan),
@@ -15,25 +18,29 @@
 
 # define function sampling Salinity
 def _sample_salinity(particle, fieldset, time):
-    particle.salinity = fieldset.S[time, particle.depth, particle.lat, particle.lon]
+    particle.salinity = fieldset.salinity[
+        time, particle.depth, particle.lat, particle.lon
+    ]
 
 
 # define function sampling Temperature
 def _sample_temperature(particle, fieldset, time):
-    particle.temperature = fieldset.T[time, particle.depth, particle.lat, particle.lon]
+    particle.temperature = fieldset.temperature[
+        time, particle.depth, particle.lat, particle.lon
+    ]
 
 
 def simulate_ship_underwater_st(
     fieldset: FieldSet,
-    out_file_name: str,
+    out_path: str | py.path.LocalPath,
     depth: float,
     sample_points: list[Spacetime],
 ) -> None:
     """
     Use parcels to simulate underway data, measuring salinity and temperature at the given depth along the ship track in a fieldset.
 
     :param fieldset: The fieldset to simulate the sampling in.
-    :param out_file_name: The file to write the results to.
+    :param out_path: The path to write the results to.
     :param depth: The depth at which to measure. 0 is water surface, negative is into the water.
     :param sample_points: The places and times to sample at.
     """
@@ -42,26 +49,31 @@ def simulate_ship_underwater_st(
     particleset = ParticleSet.from_list(
         fieldset=fieldset,
         pclass=_ShipSTParticle,
-        lon=0.0,  # initial lat/lon are irrelevant and will be overruled later.
+        lon=0.0,  # initial lat/lon are irrelevant and will be overruled later
         lat=0.0,
         depth=depth,
         time=0,  # same for time
     )
 
     # define output file for the simulation
-    out_file = particleset.ParticleFile(
-        name=out_file_name,
-    )
+    # outputdt set to infinie as we want to just want to write at the end of every call to 'execute'
+    out_file = particleset.ParticleFile(name=out_path, outputdt=np.inf)
 
+    # iterate over each point, manually set lat lon time, then
+    # execute the particle set for one step, performing one set of measurement
     for point in sample_points:
         particleset.lon_nextloop[:] = point.location.lon
         particleset.lat_nextloop[:] = point.location.lat
-        particleset.time_nextloop[:] = point.time
+        particleset.time_nextloop[:] = fieldset.time_origin.reltime(
+            np.datetime64(point.time)
+        )
 
+        # perform one step using the particleset
+        # dt and runtime are set so exactly one step is made.
         particleset.execute(
             [_sample_salinity, _sample_temperature],
             dt=1,
             runtime=1,
             verbose_progress=False,
+            output_file=out_file,
         )
-        out_file.write(particleset, time=particleset[0].time)

virtual_ship/sailship.py

@@ -195,7 +195,7 @@ def sailship(config: VirtualShipConfiguration):
     print("Simulating onboard salinity and temperature measurements.")
     simulate_ship_underwater_st(
         fieldset=config.ship_underwater_st_fieldset,
-        out_file_name=os.path.join("results", "ship_underwater_st.zarr"),
+        out_path=os.path.join("results", "ship_underwater_st.zarr"),
         depth=-2,
         sample_points=ship_underwater_sts,
     )

virtual_ship/spacetime.py

@@ -1,6 +1,7 @@
 """Location class. See class description."""
 
 from dataclasses import dataclass
+from datetime import datetime
 
 from .location import Location
 
@@ -11,4 +12,4 @@ class Spacetime:
     """A location and time."""
 
     location: Location
-    time: float
+    time: datetime