Merge pull request #2113 from OceanParcels/adding_unit_tests

Adding unit tests

Author: erikvansebille

parcels/_datasets/structured/generated.py

@@ -0,0 +1,20 @@
+import numpy as np
+import xarray as xr
+
+
+def simple_UV_dataset(dims=(360, 2, 30, 4), maxdepth=1, mesh_type="spherical"):
+    max_lon = 180.0 if mesh_type == "spherical" else 1e6
+
+    return xr.Dataset(
+        {"U": (["time", "depth", "YG", "XG"], np.zeros(dims)), "V": (["time", "depth", "YG", "XG"], np.zeros(dims))},
+        coords={
+            "time": (["time"], xr.date_range("2000", "2001", dims[0]), {"axis": "T"}),
+            "depth": (["depth"], np.linspace(0, maxdepth, dims[1]), {"axis": "Z"}),
+            "YC": (["YC"], np.arange(dims[2]) + 0.5, {"axis": "Y"}),
+            "YG": (["YG"], np.arange(dims[2]), {"axis": "Y", "c_grid_axis_shift": -0.5}),
+            "XC": (["XC"], np.arange(dims[3]) + 0.5, {"axis": "X"}),
+            "XG": (["XG"], np.arange(dims[3]), {"axis": "X", "c_grid_axis_shift": -0.5}),
+            "lat": (["YG"], np.linspace(-90, 90, dims[2]), {"axis": "Y", "c_grid_axis_shift": 0.5}),
+            "lon": (["XG"], np.linspace(-max_lon, max_lon, dims[3]), {"axis": "X", "c_grid_axis_shift": -0.5}),
+        },
+    )

parcels/_index_search.py

@@ -63,8 +63,8 @@ def _search_time_index(field: Field, time: datetime):
         tau = 1
     else:
         tau = (
-            (time - field.data.time[ti]).dt.total_seconds()
-            / (field.data.time[ti + 1] - field.data.time[ti]).dt.total_seconds()
+            (time - field.data.time[ti]).dt.total_seconds().values
+            / (field.data.time[ti + 1] - field.data.time[ti]).dt.total_seconds().values
             if field.data.time[ti] != field.data.time[ti + 1]
             else 0
         )

parcels/_reprs.py

@@ -55,8 +55,7 @@ def particleset_repr(pset: ParticleSet) -> str:
     out = f"""<{type(pset).__name__}>
     fieldset   :
 {textwrap.indent(repr(pset.fieldset), " " * 8)}
-    pclass     : {pset.pclass}
-    repeatdt   : {pset.repeatdt}
+    ptype      : {pset._ptype}
     # particles: {len(pset)}
     particles  : {_format_list_items_multiline(particles, level=2)}
 """

parcels/application_kernels/advection.py

@@ -113,7 +113,7 @@ def AdvectionRK45(particle, fieldset, time):  # pragma: no cover
     Time-step dt is halved if error is larger than fieldset.RK45_tol,
     and doubled if error is smaller than 1/10th of tolerance.
     """
-    dt = min(particle.next_dt, fieldset.RK45_max_dt) / np.timedelta64(1, "s")  # noqa TODO improve API for converting dt to seconds
+    dt = min(particle.next_dt / np.timedelta64(1, "s"), fieldset.RK45_max_dt)  # noqa TODO improve API for converting dt to seconds
     c = [1.0 / 4.0, 3.0 / 8.0, 12.0 / 13.0, 1.0, 1.0 / 2.0]
     A = [
         [1.0 / 4.0, 0.0, 0.0, 0.0, 0.0],

parcels/application_kernels/advectiondiffusion.py

@@ -24,9 +24,10 @@ def AdvectionDiffusionM1(particle, fieldset, time):  # pragma: no cover
     The Wiener increment `dW` is normally distributed with zero
     mean and a standard deviation of sqrt(dt).
     """
+    dt = particle.dt / np.timedelta64(1, "s")  # noqa TODO improve API for converting dt to seconds
     # Wiener increment with zero mean and std of sqrt(dt)
-    dWx = random.normalvariate(0, math.sqrt(math.fabs(particle.dt)))
-    dWy = random.normalvariate(0, math.sqrt(math.fabs(particle.dt)))
+    dWx = random.normalvariate(0, math.sqrt(math.fabs(dt)))
+    dWy = random.normalvariate(0, math.sqrt(math.fabs(dt)))
 
     Kxp1 = fieldset.Kh_zonal[time, particle.depth, particle.lat, particle.lon + fieldset.dres]
     Kxm1 = fieldset.Kh_zonal[time, particle.depth, particle.lat, particle.lon - fieldset.dres]
@@ -42,8 +43,8 @@ def AdvectionDiffusionM1(particle, fieldset, time):  # pragma: no cover
     by = math.sqrt(2 * fieldset.Kh_meridional[time, particle.depth, particle.lat, particle.lon])
 
     # Particle positions are updated only after evaluating all terms.
-    particle_dlon += u * particle.dt + 0.5 * dKdx * (dWx**2 + particle.dt) + bx * dWx  # noqa
-    particle_dlat += v * particle.dt + 0.5 * dKdy * (dWy**2 + particle.dt) + by * dWy  # noqa
+    particle_dlon += u * dt + 0.5 * dKdx * (dWx**2 + dt) + bx * dWx  # noqa
+    particle_dlat += v * dt + 0.5 * dKdy * (dWy**2 + dt) + by * dWy  # noqa
 
 
 def AdvectionDiffusionEM(particle, fieldset, time):  # pragma: no cover
@@ -59,9 +60,10 @@ def AdvectionDiffusionEM(particle, fieldset, time):  # pragma: no cover
     The Wiener increment `dW` is normally distributed with zero
     mean and a standard deviation of sqrt(dt).
     """
+    dt = particle.dt / np.timedelta64(1, "s")  # noqa TODO improve API for converting dt to seconds
     # Wiener increment with zero mean and std of sqrt(dt)
-    dWx = random.normalvariate(0, math.sqrt(math.fabs(particle.dt)))
-    dWy = random.normalvariate(0, math.sqrt(math.fabs(particle.dt)))
+    dWx = random.normalvariate(0, math.sqrt(math.fabs(dt)))
+    dWy = random.normalvariate(0, math.sqrt(math.fabs(dt)))
 
     u, v = fieldset.UV[time, particle.depth, particle.lat, particle.lon]
 
@@ -78,8 +80,8 @@ def AdvectionDiffusionEM(particle, fieldset, time):  # pragma: no cover
     by = math.sqrt(2 * fieldset.Kh_meridional[time, particle.depth, particle.lat, particle.lon])
 
     # Particle positions are updated only after evaluating all terms.
-    particle_dlon += ax * particle.dt + bx * dWx  # noqa
-    particle_dlat += ay * particle.dt + by * dWy  # noqa
+    particle_dlon += ax * dt + bx * dWx  # noqa
+    particle_dlat += ay * dt + by * dWy  # noqa
 
 
 def DiffusionUniformKh(particle, fieldset, time):  # pragma: no cover
@@ -100,9 +102,10 @@ def DiffusionUniformKh(particle, fieldset, time):  # pragma: no cover
     The Wiener increment `dW` is normally distributed with zero
     mean and a standard deviation of sqrt(dt).
     """
+    dt = particle.dt / np.timedelta64(1, "s")  # noqa TODO improve API for converting dt to seconds
     # Wiener increment with zero mean and std of sqrt(dt)
-    dWx = random.normalvariate(0, math.sqrt(math.fabs(particle.dt)))
-    dWy = random.normalvariate(0, math.sqrt(math.fabs(particle.dt)))
+    dWx = random.normalvariate(0, math.sqrt(math.fabs(dt)))
+    dWy = random.normalvariate(0, math.sqrt(math.fabs(dt)))
 
     bx = math.sqrt(2 * fieldset.Kh_zonal[particle])
     by = math.sqrt(2 * fieldset.Kh_meridional[particle])

parcels/application_kernels/interpolation.py

@@ -7,16 +7,135 @@
 import numpy as np
 
 from parcels.field import Field
+from parcels.tools.statuscodes import (
+    FieldOutOfBoundError,
+)
 
 if TYPE_CHECKING:
     from parcels.uxgrid import _UXGRID_AXES
+    from parcels.xgrid import _XGRID_AXES
 
 __all__ = [
     "UXPiecewiseConstantFace",
     "UXPiecewiseLinearNode",
+    "XBiLinear",
+    "XBiLinearPeriodic",
+    "XTriLinear",
 ]
 
 
+def XBiLinear(
+    field: Field,
+    ti: int,
+    position: dict[_XGRID_AXES, tuple[int, float | np.ndarray]],
+    tau: np.float32 | np.float64,
+    t: np.float32 | np.float64,
+    z: np.float32 | np.float64,
+    y: np.float32 | np.float64,
+    x: np.float32 | np.float64,
+):
+    """Bilinear interpolation on a regular grid."""
+    xi, xsi = position["X"]
+    yi, eta = position["Y"]
+    zi, _ = position["Z"]
+
+    data = field.data.data[:, zi, yi : yi + 2, xi : xi + 2]
+    data = (1 - tau) * data[ti, :, :] + tau * data[ti + 1, :, :]
+
+    return (
+        (1 - xsi) * (1 - eta) * data[0, 0]
+        + xsi * (1 - eta) * data[0, 1]
+        + xsi * eta * data[1, 1]
+        + (1 - xsi) * eta * data[1, 0]
+    )
+
+
+def XBiLinearPeriodic(
+    field: Field,
+    ti: int,
+    position: dict[_XGRID_AXES, tuple[int, float | np.ndarray]],
+    tau: np.float32 | np.float64,
+    t: np.float32 | np.float64,
+    z: np.float32 | np.float64,
+    y: np.float32 | np.float64,
+    x: np.float32 | np.float64,
+):
+    """Bilinear interpolation on a regular grid with periodic boundary conditions in horizontal directions."""
+    xi, xsi = position["X"]
+    yi, eta = position["Y"]
+    zi, _ = position["Z"]
+
+    if xi < 0:
+        xi = 0
+        xsi = (x - field.grid.lon[xi]) / (field.grid.lon[xi + 1] - field.grid.lon[xi])
+    if yi < 0:
+        yi = 0
+        eta = (y - field.grid.lat[yi]) / (field.grid.lat[yi + 1] - field.grid.lat[yi])
+
+    data = field.data.data[:, zi, yi : yi + 2, xi : xi + 2]
+    data = (1 - tau) * data[ti, :, :] + tau * data[ti + 1, :, :]
+
+    xsi = 0 if not np.isfinite(xsi) else xsi
+    eta = 0 if not np.isfinite(eta) else eta
+
+    if xsi > 0 and eta > 0:
+        return (
+            (1 - xsi) * (1 - eta) * data[0, 0]
+            + xsi * (1 - eta) * data[0, 1]
+            + xsi * eta * data[1, 1]
+            + (1 - xsi) * eta * data[1, 0]
+        )
+    elif xsi > 0 and eta == 0:
+        return (1 - xsi) * data[0, 0] + xsi * data[0, 1]
+    elif xsi == 0 and eta > 0:
+        return (1 - eta) * data[0, 0] + eta * data[1, 0]
+    else:
+        return data[0, 0]
+
+
+def XTriLinear(
+    field: Field,
+    ti: int,
+    position: dict[_XGRID_AXES, tuple[int, float | np.ndarray]],
+    tau: np.float32 | np.float64,
+    t: np.float32 | np.float64,
+    z: np.float32 | np.float64,
+    y: np.float32 | np.float64,
+    x: np.float32 | np.float64,
+):
+    """Trilinear interpolation on a regular grid."""
+    xi, xsi = position["X"]
+    yi, eta = position["Y"]
+    zi, zeta = position["Z"]
+
+    if zi < 0 or xi < 0 or yi < 0:
+        raise FieldOutOfBoundError
+
+    data = field.data.data[:, zi : zi + 2, yi : yi + 2, xi : xi + 2]
+    data = (1 - tau) * data[ti, :, :, :] + tau * data[ti + 1, :, :, :]
+    if zeta > 0:
+        data = (1 - zeta) * data[0, :, :] + zeta * data[1, :, :]
+    else:
+        data = data[0, :, :]
+
+    xsi = 0 if not np.isfinite(xsi) else xsi
+    eta = 0 if not np.isfinite(eta) else eta
+
+    if xsi > 0 and eta > 0:
+        return (
+            (1 - xsi) * (1 - eta) * data[0, 0]
+            + xsi * (1 - eta) * data[0, 1]
+            + xsi * eta * data[1, 1]
+            + (1 - xsi) * eta * data[1, 0]
+        )
+    elif xsi > 0 and eta == 0:
+        return (1 - xsi) * data[0, 0] + xsi * data[0, 1]
+    elif xsi == 0 and eta > 0:
+        return (1 - eta) * data[0, 0] + eta * data[1, 0]
+    else:
+        return data[0, 0]
+
+
 def UXPiecewiseConstantFace(
     field: Field,
     ti: int,

parcels/kernel.py

@@ -12,8 +12,6 @@
 
 from parcels.application_kernels.advection import (
     AdvectionAnalytical,
-    AdvectionRK4_3D,
-    AdvectionRK4_3D_CROCO,
     AdvectionRK45,
 )
 from parcels.basegrid import GridType
@@ -124,9 +122,10 @@ def __init__(
         # Derive meta information from pyfunc, if not given
         self.check_fieldsets_in_kernels(pyfunc)
 
-        if (pyfunc is AdvectionRK4_3D) and fieldset.U.gridindexingtype == "croco":
-            pyfunc = AdvectionRK4_3D_CROCO
-            self.funcname = "AdvectionRK4_3D_CROCO"
+        # # TODO will be implemented when we support CROCO again
+        # if (pyfunc is AdvectionRK4_3D) and fieldset.U.gridindexingtype == "croco":
+        #     pyfunc = AdvectionRK4_3D_CROCO
+        #     self.funcname = "AdvectionRK4_3D_CROCO"
 
         if funcvars is not None:  # TODO v4: check if needed from here onwards
             self.funcvars = funcvars
@@ -385,13 +384,12 @@ def evaluate_particle(self, p, endtime):
                 return p
 
             pre_dt = p.dt
-            # TODO implement below later again
-            # try:  # Use next_dt from AdvectionRK45 if it is set
-            #     if abs(endtime - p.time_nextloop) < abs(p.next_dt) - 1e-6:
-            #         p.next_dt = abs(endtime - p.time_nextloop) * sign_dt
-            # except AttributeError:
-            if sign_dt * (endtime - p.time_nextloop) <= p.dt:
-                p.dt = sign_dt * (endtime - p.time_nextloop)
+            try:  # Use next_dt from AdvectionRK45 if it is set
+                if abs(endtime - p.time_nextloop) < abs(p.next_dt) - np.timedelta64(1000, "ns"):
+                    p.next_dt = sign_dt * (endtime - p.time_nextloop)
+            except KeyError:
+                if sign_dt * (endtime - p.time_nextloop) <= p.dt:
+                    p.dt = sign_dt * (endtime - p.time_nextloop)
             res = self._pyfunc(p, self._fieldset, p.time_nextloop)
 
             if res is None:

parcels/particleset.py

@@ -109,7 +109,9 @@ def __init__(
         assert lon.size == lat.size and lon.size == depth.size, "lon, lat, depth don't all have the same lenghts"
 
         if time is None or len(time) == 0:
-            time = np.datetime64("NaT", "ns")  # do not set a time yet (because sign_dt not known)
+            time = type(fieldset.U.data.time[0].values)(
+                "NaT", "ns"
+            )  # do not set a time yet (because sign_dt not known)
         elif type(time[0]) in [np.datetime64, np.timedelta64]:
             pass  # already in the right format
         else:
@@ -156,7 +158,7 @@ def __init__(
                 if isinstance(v.initial, attrgetter):
                     initial = v.initial(self)
                 else:
-                    initial = v.initial * np.ones(len(trajectory_ids), dtype=v.dtype)
+                    initial = [np.array(v.initial, dtype=v.dtype)] * len(trajectory_ids)
                 self._data[v.name] = initial
 
         # update initial values provided on ParticleSet creation
@@ -843,16 +845,20 @@ def _warn_outputdt_release_desync(outputdt: float, starttime: float, release_tim
 
 
 def _warn_particle_times_outside_fieldset_time_bounds(release_times: np.ndarray, time: TimeInterval):
-    if np.any(release_times):
-        if np.any(release_times < time.left):
-            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,
-            )
+    if np.isnat(release_times).all():
+        return
+
+    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):
+        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,
+        )

parcels/xgrid.py

@@ -9,6 +9,7 @@
 from parcels import xgcm
 from parcels._index_search import _search_indices_curvilinear_2d
 from parcels.basegrid import BaseGrid
+from parcels.tools.statuscodes import FieldOutOfBoundError, FieldOutOfBoundSurfaceError
 
 _XGRID_AXES = Literal["X", "Y", "Z"]
 _XGRID_AXES_ORDERING: Sequence[_XGRID_AXES] = "ZYX"
@@ -271,6 +272,15 @@ def search(self, z, y, x, ei=None):
         ds = self.xgcm_grid._ds
 
         zi, zeta = _search_1d_array(ds.depth.values, z)
+        if zi == -1:
+            if zeta < 0:
+                raise FieldOutOfBoundError(
+                    f"Depth {z} is out of bounds for the grid with depth values {ds.depth.values}."
+                )
+            elif zeta > 1:
+                raise FieldOutOfBoundSurfaceError(
+                    f"Depth {z} is out of the surface for the grid with depth values {ds.depth.values}."
+                )
 
         if ds.lon.ndim == 1:
             yi, eta = _search_1d_array(ds.lat.values, y)
@@ -453,6 +463,9 @@ def _search_1d_array(
     float
         Barycentric coordinate.
     """
+    # 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 0, 0.0
     i = np.argmin(arr <= x) - 1
     bcoord = (x - arr[i]) / (arr[i + 1] - arr[i])
     return i, bcoord