Merge branch 'v4-dev' into deal_with_search_interpolation_errors

Author: erikvansebille

docs/v4/index.md

@@ -7,7 +7,7 @@ The key goals of this update are
 1. to support `Fields` on unstructured grids;
 2. to allow for user-defined interpolation methods (somewhat similar to user-defined kernels);
 3. to make the codebase more modular, easier to extend, and more maintainable;
-4. to align Parcels more with other tools in the [Pangeo ecosystemand](https://www.pangeo.io/#ecosystem), particularly by leveraging `xarray` more; and
+4. to align Parcels more with other tools in the [Pangeo ecosystem](https://www.pangeo.io/#ecosystem), particularly by leveraging `xarray` more; and
 5. to improve the performance of Parcels.
 
 The timeline for the release of Parcels v4 is not yet fixed, but we are aiming for a release of an 'alpha' version in September 2025. This v4-alpha will have support for unstructured grids and user-defined interpolation methods, but is not yet performance-optimised.

parcels/particleset.py

@@ -1,6 +1,7 @@
 import sys
 import warnings
 from collections.abc import Iterable
+from typing import Literal
 
 import numpy as np
 import xarray as xr
@@ -506,59 +507,17 @@ def execute(
         if output_file:
             output_file.metadata["parcels_kernels"] = self._kernel.name
 
-        if (dt is not None) and (not isinstance(dt, np.timedelta64)):
-            raise TypeError("dt must be a np.timedelta64 object")
-        if dt is None or np.isnat(dt):
+        if dt is None:
             dt = np.timedelta64(1, "s")
-        self._data["dt"][:] = dt
-        sign_dt = np.sign(dt).astype(int)
-        if sign_dt not in [-1, 1]:
-            raise ValueError("dt must be a positive or negative np.timedelta64 object")
 
-        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 TypeError("The runtime must be provided when the time_interval is not defined for a fieldset.")
+        if not isinstance(dt, np.timedelta64) or np.isnat(dt) or (sign_dt := np.sign(dt).astype(int)) not in [-1, 1]:
+            raise ValueError(f"dt must be a positive or negative np.timedelta64 object, got {dt=!r}")
 
-            else:
-                if isinstance(runtime, np.timedelta64):
-                    end_time = runtime
-                else:
-                    raise TypeError("The runtime must be a np.timedelta64 object")
+        self._data["dt"][:] = dt
 
-        else:
-            if not np.isnat(self.time_nextloop).any():
-                if sign_dt > 0:
-                    start_time = self.time_nextloop.min()
-                else:
-                    start_time = self.time_nextloop.max()
-            else:
-                if sign_dt > 0:
-                    start_time = self.fieldset.time_interval.left
-                else:
-                    start_time = self.fieldset.time_interval.right
-
-            if runtime is None:
-                if endtime is None:
-                    raise ValueError(
-                        "Must provide either runtime or endtime when time_interval is defined for a fieldset."
-                    )
-                # Ensure that the endtime uses the same type as the start_time
-                if isinstance(endtime, self.fieldset.time_interval.left.__class__):
-                    if sign_dt > 0:
-                        if endtime < self.fieldset.time_interval.left:
-                            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:
-                        if endtime > self.fieldset.time_interval.right:
-                            raise ValueError(
-                                "The endtime must be before the end time of the fieldset.time_interval when dt < 0"
-                            )
-                        end_time = max(endtime, self.fieldset.time_interval.left)
-                else:
-                    raise TypeError("The endtime must be of the same type as the fieldset.time_interval start time.")
-            else:
-                end_time = start_time + runtime * sign_dt
+        start_time, end_time = _get_simulation_start_and_end_times(
+            self.fieldset.time_interval, self._data["time_nextloop"], runtime, endtime, sign_dt
+        )
 
         # Set the time of the particles if it hadn't been set on initialisation
         if np.isnat(self._data["time"]).any():
@@ -619,15 +578,69 @@ def _warn_particle_times_outside_fieldset_time_bounds(release_times: np.ndarray,
 
     if isinstance(time.left, np.datetime64) and isinstance(release_times[0], np.timedelta64):
         release_times = np.array([t + time.left for t in release_times])
-    if np.any(release_times < time.left):
+    if np.any(release_times < time.left) or np.any(release_times > time.right):
         warnings.warn(
             "Some particles are set to be released outside the FieldSet's executable time domain.",
             ParticleSetWarning,
             stacklevel=2,
         )
-    if np.any(release_times > time.right):
-        warnings.warn(
-            "Some particles are set to be released after the fieldset's last time and the fields are not constant in time.",
-            ParticleSetWarning,
-            stacklevel=2,
+
+
+def _get_simulation_start_and_end_times(
+    time_interval: TimeInterval,
+    particle_release_times: np.ndarray,
+    runtime: np.timedelta64 | None,
+    endtime: np.datetime64 | None,
+    sign_dt: Literal[-1, 1],
+) -> tuple[np.datetime64, np.datetime64]:
+    if runtime is not None and endtime is not None:
+        raise ValueError(
+            f"runtime and endtime are mutually exclusive - provide one or the other. Got {runtime=!r}, {endtime=!r}"
         )
+
+    if runtime is None and time_interval is None:
+        raise ValueError("The runtime must be provided when the time_interval is not defined for a fieldset.")
+
+    if sign_dt == 1:
+        first_release_time = particle_release_times.min()
+    else:
+        first_release_time = particle_release_times.max()
+
+    start_time = _get_start_time(first_release_time, time_interval, sign_dt, runtime)
+
+    if endtime is None:
+        if not isinstance(runtime, np.timedelta64):
+            raise ValueError(f"The runtime must be a np.timedelta64 object. Got {type(runtime)}")
+
+        endtime = start_time + sign_dt * runtime
+
+    if time_interval is not None:
+        if type(endtime) != type(time_interval.left):  # noqa: E721
+            raise ValueError(
+                f"The endtime must be of the same type as the fieldset.time_interval start time. Got {endtime=!r} with {time_interval=!r}"
+            )
+        if endtime not in time_interval:
+            msg = (
+                f"Calculated/provided end time of {endtime!r} is not in fieldset time interval {time_interval!r}. Either reduce your runtime, modify your "
+                "provided endtime, or change your release timing."
+                "Important info:\n"
+                f"    First particle release: {first_release_time!r}\n"
+                f"    runtime: {runtime!r}\n"
+                f"    (calculated) endtime: {endtime!r}"
+            )
+            raise ValueError(msg)
+
+    return start_time, endtime
+
+
+def _get_start_time(first_release_time, time_interval, sign_dt, runtime):
+    if time_interval is None:
+        time_interval = TimeInterval(left=np.timedelta64(0, "s"), right=runtime)
+
+    if sign_dt == 1:
+        fieldset_start = time_interval.left
+    else:
+        fieldset_start = time_interval.right
+
+    start_time = first_release_time if not np.isnat(first_release_time) else fieldset_start
+    return start_time

tests/v4/test_index_search.py

@@ -1,12 +1,13 @@
 import numpy as np
 import pytest
+import xarray as xr
 
 from parcels._datasets.structured.generic import datasets
 from parcels._index_search import _search_indices_curvilinear_2d
 from parcels.field import Field
-from parcels.xgrid import (
-    XGrid,
-)
+from parcels.tools.exampledata_utils import download_example_dataset
+from parcels.xgcm import Grid
+from parcels.xgrid import XGrid
 
 
 @pytest.fixture
@@ -51,3 +52,32 @@ def test_grid_indexing_fpoints(field_cone):
             ]
             assert x > np.min(cell_lon) and x < np.max(cell_lon)
             assert y > np.min(cell_lat) and y < np.max(cell_lat)
+
+
+def test_indexing_nemo_curvilinear():
+    data_folder = download_example_dataset("NemoCurvilinear_data")
+    ds = xr.open_mfdataset(
+        data_folder.glob("*.nc4"), combine="nested", data_vars="minimal", coords="minimal", compat="override"
+    )
+    ds = ds.isel({"time_counter": 0, "time": 0, "z_a": 0}, drop=True).rename(
+        {"glamf": "lon", "gphif": "lat", "z": "depth"}
+    )
+    xgcm_grid = Grid(ds, coords={"X": {"left": "x"}, "Y": {"left": "y"}}, periodic=False)
+    grid = XGrid(xgcm_grid)
+
+    # Test points on the NEMO 1/4 degree curvilinear grid
+    lats = np.array([-30, 0, 88])
+    lons = np.array([30, 60, -150])
+
+    yi, eta, xi, xsi = _search_indices_curvilinear_2d(grid, lats, lons)
+
+    # Construct cornerpoints px
+    px = np.array([grid.lon[yi, xi], grid.lon[yi, xi + 1], grid.lon[yi + 1, xi + 1], grid.lon[yi + 1, xi]])
+
+    # Maximum 5 degree difference between px values
+    for i in range(lons.shape[0]):
+        np.testing.assert_allclose(px[1, i], px[:, i], atol=5)
+
+    # Reconstruct lons values from cornerpoints
+    xx = (1 - xsi) * (1 - eta) * px[0] + xsi * (1 - eta) * px[1] + xsi * eta * px[2] + (1 - xsi) * eta * px[3]
+    np.testing.assert_allclose(xx, lons, atol=1e-6)

tests/v4/test_particleset.py

@@ -114,21 +114,6 @@ def test_pset_create_outside_time(fieldset):
         ParticleSet(fieldset, pclass=Particle, lon=[0] * len(time), lat=[0] * len(time), time=time)
 
 
-@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)
-
-
 def test_pset_starttime_not_multiple_dt(fieldset):
     times = [0, 1, 2]
     datetimes = [fieldset.time_interval.left + np.timedelta64(t, "s") for t in times]
@@ -141,38 +126,6 @@ def Addlon(particle, fieldset, time):  # pragma: no cover
     assert np.allclose([p.lon_nextloop for p in pset], [8 - t for t in times])
 
 
-@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)
-
-
 def test_pset_add_explicit(fieldset):
     npart = 11
     lon = np.linspace(0, 1, npart)