Merge pull request #2348 from Parcels-code/updates_from_virtualship_dev

Fixes/tests from virtualship development at EDITO hackathon

Author: erikvansebille

docs/community/v4-migration.md

@@ -34,3 +34,12 @@ Version 4 of Parcels is unreleased at the moment. The information in this migrat
 
 - Particlefiles should be created by `ParticleFile(...)` instead of `pset.ParticleFile(...)`
 - The `name` argument in `ParticleFile` has been replaced by `store` and can now be a string, a Path or a zarr store.
+
+## Field
+
+- `Field.eval()` returns an array of floats instead of a single float (related to the vectorization)
+- `Field.eval()` does not throw OutOfBounds or other errors
+
+## GridSet
+
+- `GridSet` is now a list, so change `fieldset.gridset.grids[0]` to `fieldset.gridset[0]`.

src/parcels/_core/index_search.py

@@ -44,7 +44,7 @@ def _search_1d_array(
     # TODO v4: We probably rework this to deal with 0D arrays before this point (as we already know field dimensionality)
     if len(arr) < 2:
         return np.zeros(shape=x.shape, dtype=np.int32), np.zeros_like(x)
-    index = np.searchsorted(arr, x, side="right") - 1
+    index = np.clip(np.searchsorted(arr, x, side="right") - 1, 0, len(arr) - 2)
     # Use broadcasting to avoid repeated array access
     arr_index = arr[index]
     arr_next = arr[np.clip(index + 1, 1, len(arr) - 1)]  # Ensure we don't go out of bounds
@@ -57,7 +57,7 @@ def _search_1d_array(
     # bcoord = (x - arr[index]) / dx
 
     index = np.where(x < arr[0], LEFT_OUT_OF_BOUNDS, index)
-    index = np.where(x >= arr[-1], RIGHT_OUT_OF_BOUNDS, index)
+    index = np.where(x > arr[-1], RIGHT_OUT_OF_BOUNDS, index)
 
     return np.atleast_1d(index), np.atleast_1d(bcoord)
 
@@ -85,8 +85,8 @@ def _search_time_index(field: Field, time: datetime):
     if not field.time_interval.is_all_time_in_interval(time):
         _raise_outside_time_interval_error(time, field=None)
 
-    ti = np.searchsorted(field.data.time.data, time, side="right") - 1
-    tau = (time - field.data.time.data[ti]) / (field.data.time.data[ti + 1] - field.data.time.data[ti])
+    ti, tau = _search_1d_array(field.data.time.data, time)
+
     return {"T": {"index": np.atleast_1d(ti), "bcoord": np.atleast_1d(tau)}}
 
 

src/parcels/_core/xgrid.py

@@ -104,9 +104,12 @@ def __init__(self, grid: xgcm.Grid, mesh="flat"):
         if "lat" in ds:
             ds.set_coords("lat")
 
-        if len(set(grid.axes) & {"X", "Y", "Z"}) > 0:  # Only if spatial grid is >0D (see #2054 for further development)
+        if len(set(grid.axes) & {"X", "Y"}) > 0:  # Only if spatial grid is >0D (see #2054 for further development)
             assert_valid_lat_lon(ds["lat"], ds["lon"], grid.axes)
 
+        if "Z" in grid.axes:
+            assert_valid_depth(ds["depth"])
+
         assert_valid_mesh(mesh)
         self._ds = ds
 
@@ -482,6 +485,14 @@ def assert_valid_lat_lon(da_lat, da_lon, axes: _XGCM_AXES):
             )
 
 
+def assert_valid_depth(da_depth):
+    if not np.all(np.diff(da_depth.values) > 0):
+        raise ValueError(
+            f"Depth DataArray {da_depth.name!r} with dims {da_depth.dims} must be strictly increasing. "
+            f'HINT: you may be able to use ds.reindex to flip depth - e.g., ds = ds.reindex({da_depth.name}=ds["{da_depth.name}"][::-1])'
+        )
+
+
 def _convert_center_pos_to_fpoint(
     *, index: int, bcoord: float, xgcm_position: _XGCM_AXIS_POSITION, f_points_xgcm_position: _XGCM_AXIS_POSITION
 ) -> tuple[int, float]:

tests/test_particleset_execute.py

@@ -150,6 +150,45 @@ def test_particleset_endtime_type(fieldset, endtime, expectation):
         pset.execute(endtime=endtime, dt=np.timedelta64(10, "m"), pyfunc=DoNothing)
 
 
+def test_particleset_run_to_endtime(fieldset):
+    starttime = fieldset.time_interval.left
+    endtime = fieldset.time_interval.right
+
+    def SampleU(particles, fieldset):  # pragma: no cover
+        _ = fieldset.U[particles]
+
+    pset = ParticleSet(fieldset, lon=[0.2], lat=[5.0], time=[starttime])
+    pset.execute(SampleU, endtime=endtime, dt=np.timedelta64(1, "D"))
+    assert pset[0].time == endtime
+
+
+def test_particleset_interpolate_on_domainedge(zonal_flow_fieldset):
+    fieldset = zonal_flow_fieldset
+
+    MyParticle = Particle.add_variable(Variable("var"))
+
+    def SampleU(particles, fieldset):  # pragma: no cover
+        particles.var = fieldset.U[particles]
+
+    print(fieldset.U.grid.lon)
+    pset = ParticleSet(fieldset, pclass=MyParticle, lon=fieldset.U.grid.lon[-1], lat=fieldset.U.grid.lat[-1])
+    pset.execute(SampleU, runtime=np.timedelta64(1, "D"), dt=np.timedelta64(1, "D"))
+    np.testing.assert_equal(pset[0].var, 1)
+
+
+def test_particleset_interpolate_outside_domainedge(zonal_flow_fieldset):
+    fieldset = zonal_flow_fieldset
+
+    def SampleU(particles, fieldset):  # pragma: no cover
+        particles.dlon = fieldset.U[particles]
+
+    dlat = 1e-3
+    pset = ParticleSet(fieldset, lon=fieldset.U.grid.lon[-1], lat=fieldset.U.grid.lat[-1] + dlat)
+
+    with pytest.raises(FieldOutOfBoundError):
+        pset.execute(SampleU, runtime=np.timedelta64(1, "D"), dt=np.timedelta64(1, "D"))
+
+
 @pytest.mark.parametrize(
     "dt", [np.timedelta64(1, "s"), np.timedelta64(1, "ms"), np.timedelta64(10, "ms"), np.timedelta64(1, "ns")]
 )
@@ -329,7 +368,6 @@ def MoveRight(particles, fieldset):  # pragma: no cover
 
     def MoveLeft(particles, fieldset):  # pragma: no cover
         inds = np.where(particles.state == StatusCode.ErrorOutOfBounds)
-        print(inds, particles.state)
         particles[inds].dlon -= 1.0
         particles[inds].state = StatusCode.Success
 

tests/test_xgrid.py

@@ -126,6 +126,14 @@ def test_invalid_lon_lat():
         XGrid.from_dataset(ds)
 
 
+def test_invalid_depth():
+    ds = datasets["ds_2d_left"].copy()
+    ds = ds.reindex({"ZG": ds.ZG[::-1]})
+
+    with pytest.raises(ValueError, match="Depth DataArray .* must be strictly increasing*"):
+        XGrid.from_dataset(ds)
+
+
 @pytest.mark.parametrize(
     "ds",
     [