Other instrument tests

tests/instruments/test_adcp.py

@@ -1,33 +1,119 @@
 """Test the simulation of ADCP instruments."""
 
+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.adcp import simulate_adcp
 
 
-def test_simulate_adcp() -> None:
-    MAX_DEPTH = -1000
-    MIN_DEPTH = -5
-    BIN_SIZE = 24
+def test_simulate_adcp(tmpdir: py.path.LocalPath) -> None:
+    # maximum depth the ADCP can measure
+    MAX_DEPTH = -1000  # -1000
+    # minimum depth the ADCP can measure
+    MIN_DEPTH = -5  # -5
+    # How many samples to take in the complete range between max_depth and min_depth.
+    NUM_BINS = 40
+
+    # arbitrary time offset for the dummy fieldset
+    base_time = datetime.datetime.strptime("1950-01-01", "%Y-%m-%d")
+
+    # where to sample
+    sample_points = [
+        Spacetime(Location(1, 2), base_time + datetime.timedelta(seconds=0)),
+        Spacetime(Location(3, 4), base_time + datetime.timedelta(seconds=1)),
+    ]
+
+    # expected observations at sample points
+    expected_obs = [
+        {
+            "V": {"surface": 5, "max_depth": 6},
+            "U": {"surface": 7, "max_depth": 8},
+            "lat": sample_points[0].location.lat,
+            "lon": sample_points[0].location.lon,
+            "time": base_time + datetime.timedelta(seconds=0),
+        },
+        {
+            "V": {"surface": 9, "max_depth": 10},
+            "U": {"surface": 11, "max_depth": 12},
+            "lat": sample_points[1].location.lat,
+            "lon": sample_points[1].location.lon,
+            "time": base_time + datetime.timedelta(seconds=1),
+        },
+    ]
+
+    # create fieldset based on the expected observations
+    # indices are time, depth, latitude, longitude
+    v = np.zeros((2, 2, 2, 2))
+    v[0, 0, 0, 0] = expected_obs[0]["V"]["max_depth"]
+    v[0, 1, 0, 0] = expected_obs[0]["V"]["surface"]
+    v[1, 0, 1, 1] = expected_obs[1]["V"]["max_depth"]
+    v[1, 1, 1, 1] = expected_obs[1]["V"]["surface"]
+
+    u = np.zeros((2, 2, 2, 2))
+    u[0, 0, 0, 0] = expected_obs[0]["U"]["max_depth"]
+    u[0, 1, 0, 0] = expected_obs[0]["U"]["surface"]
+    u[1, 0, 1, 1] = expected_obs[1]["U"]["max_depth"]
+    u[1, 1, 1, 1] = expected_obs[1]["U"]["surface"]
 
     fieldset = FieldSet.from_data(
-        {"U": 0, "V": 0},
         {
-            "lon": 0,
-            "lat": 0,
-            "time": [np.datetime64("1950-01-01") + np.timedelta64(632160, "h")],
+            "V": v,
+            "U": u,
+        },
+        {
+            "lat": np.array([expected_obs[0]["lat"], expected_obs[1]["lat"]]),
+            "lon": np.array([expected_obs[0]["lon"], expected_obs[1]["lon"]]),
+            "depth": np.array([MAX_DEPTH, MIN_DEPTH]),
+            "time": np.array(
+                [
+                    np.datetime64(expected_obs[0]["time"]),
+                    np.datetime64(expected_obs[1]["time"]),
+                ]
+            ),
         },
     )
 
-    sample_points = [Spacetime(Location(0, 0), 0)]
+    # perform simulation
+    out_path = tmpdir.join("out.zarr")
 
     simulate_adcp(
         fieldset=fieldset,
-        out_file_name="test",
+        out_path=out_path,
         max_depth=MAX_DEPTH,
         min_depth=MIN_DEPTH,
-        bin_size=BIN_SIZE,
+        num_bins=NUM_BINS,
         sample_points=sample_points,
     )
+
+    results = xr.open_zarr(out_path)
+
+    # test if output is as expected
+    assert len(results.trajectory) == NUM_BINS
+
+    # for every obs, check if the variables match the expected observations
+    # we only verify at the surface and max depth of the adcp, because in between is tricky
+    for traj, vert_loc in [
+        (results.trajectory[0], "max_depth"),
+        (results.trajectory[-1], "surface"),
+    ]:
+        obs_all = results.sel(trajectory=traj).obs
+        assert len(obs_all) == len(sample_points)
+        for i, (obs_i, exp) in enumerate(zip(obs_all, expected_obs, strict=True)):
+            obs = results.sel(trajectory=traj, obs=obs_i)
+            for var in ["lat", "lon"]:
+                obs_value = obs[var].values.item()
+                exp_value = exp[var]
+                assert np.isclose(
+                    obs_value, exp_value
+                ), f"Observation incorrect {vert_loc=} {obs_i=} {var=} {obs_value=} {exp_value=}."
+            for var in ["V", "U"]:
+                obs_value = obs[var].values.item()
+                exp_value = exp[var][vert_loc]
+                assert np.isclose(
+                    obs_value, exp_value
+                ), f"Observation incorrect {vert_loc=} {i=} {var=} {obs_value=} {exp_value=}."

tests/instruments/test_ship_underwater_st.py

@@ -18,50 +18,50 @@ def test_simulate_ship_underwater_st(tmpdir: py.path.LocalPath) -> None:
     # 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)),
+        Spacetime(Location(1, 2), base_time + datetime.timedelta(seconds=0)),
+        Spacetime(Location(3, 4), base_time + datetime.timedelta(seconds=1)),
     ]
 
     # expected observations at sample points
     expected_obs = [
         {
-            "salinity": 1,
-            "temperature": 2,
-            "lat": 3,
-            "lon": 0,
+            "salinity": 5,
+            "temperature": 6,
+            "lat": sample_points[0].location.lat,
+            "lon": sample_points[0].location.lon,
             "time": base_time + datetime.timedelta(seconds=0),
         },
         {
-            "salinity": 5,
-            "temperature": 6,
-            "lat": 7,
-            "lon": 0,
+            "salinity": 7,
+            "temperature": 8,
+            "lat": sample_points[1].location.lat,
+            "lon": sample_points[1].location.lon,
             "time": base_time + datetime.timedelta(seconds=1),
         },
     ]
 
     # create fieldset based on the expected observations
+    # indices are time, latitude, longitude
+    salinity = np.zeros((2, 2, 2))
+    salinity[0, 0, 0] = expected_obs[0]["salinity"]
+    salinity[1, 1, 1] = expected_obs[1]["salinity"]
+
+    temperature = np.zeros((2, 2, 2))
+    temperature[0, 0, 0] = expected_obs[0]["temperature"]
+    temperature[1, 1, 1] = expected_obs[1]["temperature"]
+
     fieldset = FieldSet.from_data(
         {
-            "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"]],
-            ],
+            "V": np.zeros((2, 2, 2)),
+            "U": np.zeros((2, 2, 2)),
+            "salinity": salinity,
+            "temperature": temperature,
         },
         {
-            "lon": 0,
             "lat": np.array([expected_obs[0]["lat"], expected_obs[1]["lat"]]),
+            "lon": np.array([expected_obs[0]["lon"], expected_obs[1]["lon"]]),
             "time": np.array(
                 [
                     np.datetime64(expected_obs[0]["time"]),
@@ -71,6 +71,7 @@ def test_simulate_ship_underwater_st(tmpdir: py.path.LocalPath) -> None:
         },
     )
 
+    # perform simulation
     out_path = tmpdir.join("out.zarr")
 
     simulate_ship_underwater_st(
@@ -80,21 +81,23 @@ def test_simulate_ship_underwater_st(tmpdir: py.path.LocalPath) -> None:
         sample_points=sample_points,
     )
 
+    # test if output is as expected
     results = xr.open_zarr(out_path)
 
-    # below we assert if output makes sense
-    assert len(results.trajectory) == 1  # expect a singe trajectory
+    assert len(results.trajectory) == 1  # expect a single 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):
+    for i, (obs_i, exp) in enumerate(
+        zip(results.sel(trajectory=traj).obs, expected_obs, strict=True)
+    ):
         obs = results.sel(trajectory=traj, obs=obs_i)
-        for var in variables:
+        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 {obs_i=} {var=} {obs_value=} {exp_value=}."
+            ), f"Observation incorrect {i=} {var=} {obs_value=} {exp_value=}."

virtual_ship/instruments/adcp.py

@@ -1,11 +1,14 @@
 """ADCP instrument."""
 
 import numpy as np
-from parcels import FieldSet, JITParticle, ParticleSet, Variable
+import py
+from parcels import FieldSet, ParticleSet, ScipyParticle, Variable
 
 from ..spacetime import Spacetime
 
-_ADCPParticle = 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
+_ADCPParticle = ScipyParticle.add_variables(
     [
         Variable("U", dtype=np.float32, initial=np.nan),
         Variable("V", dtype=np.float32, initial=np.nan),
@@ -21,25 +24,25 @@ def _sample_velocity(particle, fieldset, time):
 
 def simulate_adcp(
     fieldset: FieldSet,
-    out_file_name: str,
+    out_path: str | py.path.LocalPath,
     max_depth: float,
     min_depth: float,
-    bin_size: float,
+    num_bins: int,
     sample_points: list[Spacetime],
 ) -> None:
     """
     Use parcels to simulate an ADCP in a fieldset.
 
     :param fieldset: The fieldset to simulate the ADCP in.
-    :param out_file_name: The file to write the results to.
+    :param out_path: The path to write the results to.
     :param max_depth: Maximum depth the ADCP can measure.
     :param min_depth: Minimum depth the ADCP can measure.
-    :param bin_size: How many samples to take in the complete range between max_depth and min_depth.
+    :param num_bins: How many samples to take in the complete range between max_depth and min_depth.
     :param sample_points: The places and times to sample at.
     """
     sample_points.sort(key=lambda p: p.time)
 
-    bins = np.arange(max_depth, min_depth, bin_size)
+    bins = np.linspace(max_depth, min_depth, num_bins)
     num_particles = len(bins)
     particleset = ParticleSet.from_list(
         fieldset=fieldset,
@@ -53,14 +56,22 @@ def simulate_adcp(
     )
 
     # define output file for the simulation
-    out_file = particleset.ParticleFile(
-        name=out_file_name,
-    )
+    # outputdt set to infinite as we just want to write at the end of every call to 'execute'
+    out_file = particleset.ParticleFile(name=out_path, outputdt=np.inf)
 
     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)
+        )
 
-        particleset.execute([_sample_velocity], dt=1, runtime=1, verbose_progress=False)
-        out_file.write(particleset, time=particleset[0].time)
+        # perform one step using the particleset
+        # dt and runtime are set so exactly one step is made.
+        particleset.execute(
+            [_sample_velocity],
+            dt=1,
+            runtime=1,
+            verbose_progress=False,
+            output_file=out_file,
+        )

virtual_ship/sailship.py

@@ -203,10 +203,11 @@ def sailship(config: VirtualShipConfiguration):
     print("Simulating onboard ADCP.")
     simulate_adcp(
         fieldset=config.adcp_fieldset,
-        out_file_name=os.path.join("results", "adcp.zarr"),
+        out_path=os.path.join("results", "adcp.zarr"),
         max_depth=config.ADCP_settings["max_depth"],
         min_depth=-5,
-        bin_size=config.ADCP_settings["bin_size_m"],
+        num_bins=(-5 - config.ADCP_settings["max_depth"])
+        // config.ADCP_settings["bin_size_m"],
         sample_points=adcps,
     )