From 9914db4cbae655f00d6240da1a558bf73c6396db Mon Sep 17 00:00:00 2001
From: Aart Stuurman <aartstuurman@hotmail.com>
Date: Mon, 17 Jun 2024 18:03:11 +0200
Subject: [PATCH] adcp test

---
 tests/instruments/test_adcp.py   | 77 +++++++++++++++++++++++++++++---
 virtual_ship/instruments/adcp.py | 15 +++++--
 2 files changed, 84 insertions(+), 8 deletions(-)

diff --git a/tests/instruments/test_adcp.py b/tests/instruments/test_adcp.py
index 2a3d57f8..e648fd1a 100644
--- a/tests/instruments/test_adcp.py
+++ b/tests/instruments/test_adcp.py
@@ -6,24 +6,72 @@
 
 from virtual_ship import Location, Spacetime
 from virtual_ship.instruments.adcp import simulate_adcp
+import datetime
+import xarray as xr
 
 
 def test_simulate_adcp(tmpdir: py.path.LocalPath) -> None:
+    # maximum depth the ADCP can measure
     MAX_DEPTH = -1000
+    # minimum depth the ADCP can measure
     MIN_DEPTH = -5
+    # how many samples to take in the complete range between max_depth and min_depth
     BIN_SIZE = 24
 
+    # 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 = ["U", "V", "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 = [
+        {
+            "U": 1,
+            "V": 2,
+            "lat": 3,
+            "lon": 0,
+            "time": base_time + datetime.timedelta(seconds=0),
+        },
+        {
+            "U": 5,
+            "V": 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},
+        {
+            "U": [
+                [expected_obs[0]["U"], 0],
+                [0, expected_obs[1]["U"]],
+            ],
+            "V": [
+                [expected_obs[0]["V"], 0],
+                [0, expected_obs[1]["V"]],
+            ],
+        },
         {
             "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_adcp(
@@ -34,3 +82,22 @@ def test_simulate_adcp(tmpdir: py.path.LocalPath) -> None:
         bin_size=BIN_SIZE,
         sample_points=sample_points,
     )
+
+    results = xr.open_zarr(out_path)
+
+    # below we assert if output makes sense
+    EXPECTED_NUM_BINS = len(np.arange(MAX_DEPTH, MIN_DEPTH, BIN_SIZE))
+    assert len(results.trajectory) == EXPECTED_NUM_BINS  # expect a single trajectory
+
+    # TODO test depth sampling
+    for traj in results.trajectory:
+        obs_all = results.sel(trajectory=traj).obs
+        assert len(obs_all) == len(sample_points)
+        for obs_i, exp in zip(obs_all, 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=}."
diff --git a/virtual_ship/instruments/adcp.py b/virtual_ship/instruments/adcp.py
index cfb5bf05..a37e2701 100644
--- a/virtual_ship/instruments/adcp.py
+++ b/virtual_ship/instruments/adcp.py
@@ -60,7 +60,16 @@ def simulate_adcp(
     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,
+        )