Adding unit tests for particle.time and execute dt, runtime and endtime

Author: erikvansebille

parcels/particleset.py

@@ -1,7 +1,6 @@
 import sys
 import warnings
 from collections.abc import Iterable
-from datetime import date, datetime, timedelta
 
 import numpy as np
 import xarray as xr
@@ -109,19 +108,16 @@ def __init__(
             depth = convert_to_flat_array(depth)
         assert lon.size == lat.size and lon.size == depth.size, "lon, lat, depth don't all have the same lenghts"
 
-        if time.size > 0:
-            time = np.repeat(time, lon.size) if time.size == 1 else time
+        if time is None or len(time) == 0:
+            time = fieldset.U.time.values[0]  # TODO set this to NaT if no time is given
+        time = np.repeat(time, lon.size) if time.size == 1 else time
 
-            if type(time[0]) in [np.datetime64, np.timedelta64]:
-                pass  # already in the right format
-            elif type(time[0]) in [datetime, date]:
-                time = np.array([np.datetime64(t) for t in time])
-            elif type(time[0]) in [timedelta]:
-                time = np.array([np.timedelta64(t) for t in time])
-            else:
-                raise NotImplementedError("particle time must be a datetime, timedelta, or date object")
+        if type(time[0]) in [np.datetime64, np.timedelta64]:
+            pass  # already in the right format
+        else:
+            raise TypeError("particle time must be a datetime, timedelta, or date object")
 
-            assert lon.size == time.size, "time and positions (lon, lat, depth) do not have the same lengths."
+        assert lon.size == time.size, "time and positions (lon, lat, depth) do not have the same lengths."
 
         if fieldset.time_interval:
             _warn_particle_times_outside_fieldset_time_bounds(time, fieldset.time_interval)
@@ -800,13 +796,13 @@ def execute(
         if self.fieldset.time_interval is None:
             start_time = np.timedelta64(0, "s")  # For the execution loop, we need a start time as a timedelta object
             if runtime is None:
-                raise ValueError("The runtime must be provided when the time_interval is not defined for a fieldset.")
+                raise TypeError("The runtime must be provided when the time_interval is not defined for a fieldset.")
 
             else:
                 if isinstance(runtime, np.timedelta64):
                     end_time = runtime
                 else:
-                    raise ValueError("The runtime must be a np.timedelta64 object")
+                    raise TypeError("The runtime must be a np.timedelta64 object")
 
         else:
             start_time = self.fieldset.time_interval.left
@@ -822,7 +818,7 @@ def execute(
                         raise ValueError("The endtime must be after the start time of the fieldset.time_interval")
                     end_time = min(endtime, self.fieldset.time_interval.right)
                 else:
-                    raise ValueError("The endtime must be of the same type as the fieldset.time_interval start time.")
+                    raise TypeError("The endtime must be of the same type as the fieldset.time_interval start time.")
             else:
                 end_time = start_time + runtime
 

tests/v4/test_particleset.py

@@ -1,4 +1,5 @@
-from datetime import timedelta
+from contextlib import nullcontext as does_not_raise
+from datetime import datetime, timedelta
 
 import numpy as np
 import pytest
@@ -11,9 +12,87 @@
     ParticleSet,
     UXPiecewiseConstantFace,
     VectorField,
+    xgcm,
 )
+from parcels._datasets.structured.generic import datasets as datasets_structured
 from parcels._datasets.unstructured.generic import datasets as datasets_unstructured
 from parcels.uxgrid import UxGrid
+from parcels.xgrid import XGrid
+
+
+@pytest.fixture
+def fieldset() -> FieldSet:
+    ds = datasets_structured["ds_2d_left"]
+    grid = XGrid(xgcm.Grid(ds))
+    U = Field("U", ds["U (A grid)"], grid, mesh_type="flat")
+    V = Field("V", ds["V (A grid)"], grid, mesh_type="flat")
+    return FieldSet([U, V])
+
+
+def DoNothing(particle, fieldset, time):
+    pass
+
+
+@pytest.mark.parametrize(
+    "time, expectation",
+    [
+        (np.timedelta64(0, "s"), does_not_raise()),
+        (np.datetime64("2000-01-02T00:00:00"), does_not_raise()),
+        (0.0, pytest.raises(TypeError)),
+        (timedelta(seconds=0), pytest.raises(TypeError)),
+        (datetime(2023, 1, 1, 0, 0, 0), pytest.raises(TypeError)),
+    ],
+)
+def test_particleset_init_time_type(fieldset, time, expectation):
+    with expectation:
+        ParticleSet(fieldset, lon=[0.2], lat=[5.0], time=[time], pclass=Particle)
+
+
+@pytest.mark.parametrize(
+    "dt, expectation",
+    [
+        (np.timedelta64(5, "s"), does_not_raise()),
+        (5.0, pytest.raises(TypeError)),
+        (np.datetime64("2000-01-02T00:00:00"), pytest.raises(TypeError)),
+        (timedelta(seconds=2), pytest.raises(TypeError)),
+    ],
+)
+def test_particleset_dt_type(fieldset, dt, expectation):
+    pset = ParticleSet(fieldset, lon=[0.2], lat=[5.0], depth=[50.0], pclass=Particle)
+    with expectation:
+        pset.execute(runtime=np.timedelta64(10, "s"), dt=dt, pyfunc=DoNothing)
+
+
+@pytest.mark.parametrize(
+    "runtime, expectation",
+    [
+        (np.timedelta64(5, "s"), does_not_raise()),
+        (5.0, pytest.raises(TypeError)),
+        (timedelta(seconds=2), pytest.raises(TypeError)),
+        (np.datetime64("2001-01-02T00:00:00"), pytest.raises(TypeError)),
+        (datetime(2000, 1, 2, 0, 0, 0), pytest.raises(TypeError)),
+    ],
+)
+def test_particleset_runtime_type(fieldset, runtime, expectation):
+    pset = ParticleSet(fieldset, lon=[0.2], lat=[5.0], depth=[50.0], pclass=Particle)
+    with expectation:
+        pset.execute(runtime=runtime, dt=np.timedelta64(10, "s"), pyfunc=DoNothing)
+
+
+@pytest.mark.parametrize(
+    "endtime, expectation",
+    [
+        (np.datetime64("2000-01-02T00:00:00"), does_not_raise()),
+        (5.0, pytest.raises(TypeError)),
+        (np.timedelta64(5, "s"), pytest.raises(TypeError)),
+        (timedelta(seconds=2), pytest.raises(TypeError)),
+        (datetime(2000, 1, 2, 0, 0, 0), pytest.raises(TypeError)),
+    ],
+)
+def test_particleset_endtime_type(fieldset, endtime, expectation):
+    pset = ParticleSet(fieldset, lon=[0.2], lat=[5.0], depth=[50.0], pclass=Particle)
+    with expectation:
+        pset.execute(endtime=endtime, dt=np.timedelta64(10, "m"), pyfunc=DoNothing)
 
 
 @pytest.mark.parametrize("verbose_progress", [True, False])