Fixing test_drifter

Author: erikvansebille

src/virtualship/instruments/drifter.py

@@ -38,13 +38,11 @@ def _sample_temperature(particles, fieldset):
 
 
 def _check_lifetime(particles, fieldset):
-    dt = particles.dt / np.timedelta64(1, "s")  # convert dt to seconds
-
-    def has_lifetime(p):
-        p.age = np.where(p.has_lifetime == 1, p.age + dt, p.age)
-        p.state = np.where(p.age >= p.lifetime, StatusCode.Delete, p.state)
-
-    has_lifetime(particles[particles.has_lifetime == 1])
+    for i in range(len(particles)):
+        if particles[i].has_lifetime == 1:
+            particles[i].age += particles[i].dt / np.timedelta64(1, "s")
+            if particles[i].age >= particles[i].lifetime:
+                particles[i].state = StatusCode.Delete
 
 
 def simulate_drifters(
@@ -78,22 +76,22 @@ def simulate_drifters(
         pclass=_DrifterParticle,
         lat=[drifter.spacetime.location.lat for drifter in drifters],
         lon=[drifter.spacetime.location.lon for drifter in drifters],
-        depth=[drifter.depth for drifter in drifters],
+        z=[drifter.depth for drifter in drifters],
         time=[drifter.spacetime.time for drifter in drifters],
         has_lifetime=[1 if drifter.lifetime is not None else 0 for drifter in drifters],
         lifetime=[
-            0 if drifter.lifetime is None else drifter.lifetime.total_seconds()
+            0 if drifter.lifetime is None else drifter.lifetime / np.timedelta64(1, "s")
             for drifter in drifters
         ],
     )
 
     # define output file for the simulation
-    out_file = drifter_particleset.ParticleFile(
-        name=out_path, outputdt=outputdt, chunks=[len(drifter_particleset), 100]
+    out_file = ParticleFile(
+        store=out_path, outputdt=outputdt, chunks=(len(drifter_particleset), 100)
     )
 
     # get earliest between fieldset end time and provide end time
-    fieldset_endtime = fieldset.time_origin.fulltime(fieldset.U.grid.time_full[-1])
+    fieldset_endtime = fieldset.time_interval.right - np.timedelta64(1, "s")  # TODO remove hack stopping 1 second too early when v4 is fixed
     if endtime is None:
         actual_endtime = fieldset_endtime
     elif endtime > fieldset_endtime:
@@ -112,7 +110,7 @@ def simulate_drifters(
     )
 
     # if there are more particles left than the number of drifters with an indefinite endtime, warn the user
-    if len(drifter_particleset.particledata) > len(
+    if len(drifter_particleset) > len(
         [d for d in drifters if d.lifetime is None]
     ):
         print(

tests/instruments/test_drifter.py

@@ -4,48 +4,59 @@
 
 import numpy as np
 import xarray as xr
-from parcels import FieldSet
+from parcels import FieldSet, Field, VectorField, XGrid
 
 from virtualship.instruments.drifter import Drifter, simulate_drifters
 from virtualship.models import Location, Spacetime
 
 
 def test_simulate_drifters(tmpdir) -> None:
     # arbitrary time offset for the dummy fieldset
-    base_time = datetime.datetime.strptime("1950-01-01", "%Y-%m-%d")
+    base_time = np.datetime64("1950-01-01")
 
     CONST_TEMPERATURE = 1.0  # constant temperature in fieldset
 
-    v = np.full((2, 2, 2), 1.0)
-    u = np.full((2, 2, 2), 1.0)
-    t = np.full((2, 2, 2), CONST_TEMPERATURE)
-
-    fieldset = FieldSet.from_data(
-        {"V": v, "U": u, "T": t},
-        {
-            "lon": np.array([0.0, 10.0]),
-            "lat": np.array([0.0, 10.0]),
-            "time": [
-                np.datetime64(base_time + datetime.timedelta(seconds=0)),
-                np.datetime64(base_time + datetime.timedelta(days=3)),
-            ],
+    dims = (2, 2, 2)  # time, lat, lon
+    v = np.full(dims, 1.0)
+    u = np.full(dims, 1.0)
+    t = np.full(dims, CONST_TEMPERATURE)
+
+    time = [base_time, base_time + np.timedelta64(3, "D")]
+    ds = xr.Dataset(
+        {"U": (["time", "YG", "XG"], u), "V": (["time", "YG", "XG"], v), "T": (["time", "YG", "XG"], t)},
+        coords={
+            "time": (["time"], time, {"axis": "T"}),
+            "YC": (["YC"], np.arange(dims[1]) + 0.5, {"axis": "Y"}),
+            "YG": (["YG"], np.arange(dims[1]), {"axis": "Y", "c_grid_axis_shift": -0.5}),
+            "XC": (["XC"], np.arange(dims[2]) + 0.5, {"axis": "X"}),
+            "XG": (["XG"], np.arange(dims[2]), {"axis": "X", "c_grid_axis_shift": -0.5}),
+            "lat": (["YG"], np.linspace(-10, 10, dims[1]), {"axis": "Y", "c_grid_axis_shift": 0.5}),
+            "lon": (["XG"], np.linspace(-10, 10, dims[2]), {"axis": "X", "c_grid_axis_shift": -0.5}),
         },
     )
 
+    grid = XGrid.from_dataset(ds, mesh="spherical")
+    U = Field("U", ds["U"], grid)
+    V = Field("V", ds["V"], grid)
+    T = Field("T", ds["T"], grid)
+    UV = VectorField("UV", U, V)
+    fieldset = FieldSet([U, V, T, UV])
+
+
     # drifters to deploy
     drifters = [
         Drifter(
             spacetime=Spacetime(
                 location=Location(latitude=0, longitude=0),
-                time=base_time + datetime.timedelta(days=0),
+                time=base_time + np.timedelta64(0, "D"),
             ),
             depth=0.0,
-            lifetime=datetime.timedelta(hours=2),
+            lifetime=np.timedelta64(2, "h"),
         ),
         Drifter(
             spacetime=Spacetime(
                 location=Location(latitude=1, longitude=1),
-                time=base_time + datetime.timedelta(hours=20),
+                time=base_time + np.timedelta64(20, "h"),
             ),
             depth=0.0,
             lifetime=None,
@@ -65,7 +76,7 @@ def test_simulate_drifters(tmpdir) -> None:
     )
 
     # test if output is as expected
-    results = xr.open_zarr(out_path)
+    results = xr.open_zarr(out_path, decode_cf=False)  # TODO fix decode_cf when parcels v4 is fixed
 
     assert len(results.trajectory) == len(drifters)