Merge branch 'v4-dev' into updates_from_virtualship_dev

Author: erikvansebille

.github/workflows/cache-pixi-lock.yml

@@ -2,6 +2,9 @@ name: Generate and cache Pixi lockfile
 
 on:
   workflow_call:
+    inputs:
+      pixi-version:
+        type: string
     outputs:
       cache-id:
         description: "The lock file contents"
@@ -30,7 +33,7 @@ jobs:
       - uses: prefix-dev/setup-pixi@v0.9.0
         if: ${{ !steps.restore.outputs.cache-hit }}
         with:
-          pixi-version: v0.56.0
+          pixi-version: ${{ inputs.pixi-version }}
           run-install: false
       - name: Run pixi lock
         if: ${{ !steps.restore.outputs.cache-hit }}

docs/examples_v3/tutorial_kernelloop.ipynb

@@ -26,7 +26,7 @@
     "\n",
     "When you run a Parcels simulation (i.e. a call to `pset.execute()`), the Kernel loop is the main part of the code that is executed. This part of the code loops through all particles and executes the Kernels that are defined for each particle.\n",
     "\n",
-    "In order to make sure that the displacements of a particle in the different Kernels can be summed, all Kernels add to a _change_ in position (`particles.dlon`, `particles.dlat`, and `particles.ddepth`). This is important, because there are situations where movement kernels would otherwise not commute. Take the example of advecting particles by currents _and_ winds. If the particle would first be moved by the currents and then by the winds, the result could be different from first moving by the winds and then by the currents. Instead, by adding the changes in position, the ordering of the Kernels has no consequence on the particle displacement."
+    "In order to make sure that the displacements of a particle in the different Kernels can be summed, all Kernels add to a _change_ in position (`particles.dlon`, `particles.dlat`, and `particles.dz`). This is important, because there are situations where movement kernels would otherwise not commute. Take the example of advecting particles by currents _and_ winds. If the particle would first be moved by the currents and then by the winds, the result could be different from first moving by the winds and then by the currents. Instead, by adding the changes in position, the ordering of the Kernels has no consequence on the particle displacement."
    ]
   },
   {
@@ -40,29 +40,27 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "Below is a structured overview of the Kernel loop is implemented. Note that this is for longitude only, but the same process is applied for latitude and depth.\n",
+    "Below is a structured overview of the Kernel loop is implemented. Note that this is for `lon` only, but the same process is applied for `lat` and `z`.\n",
     "\n",
-    "1. Initialise an extra Variable `particles.lon=0` and `particles.time_nextloop = particles.time`\n",
+    "1. Initialise an extra Variable `particles.dlon=0`\n",
     "\n",
     "2. Within the Kernel loop, for each particle:<br>\n",
     "\n",
     "    1. Update `particles.lon += particles.dlon`<br>\n",
     "\n",
-    "    2. Set variable `particles.dlon = 0`<br>\n",
+    "    2. Update `particles.time += particles.dt` (except for on the first iteration of the Kernel loop)<br>\n",
     "\n",
-    "    3. Update `particles.time = particles.time_nextloop`\n",
+    "    3. Set variable `particles.dlon = 0`<br>\n",
     "\n",
     "    4. For each Kernel in the list of Kernels:\n",
     "        \n",
     "        1. Execute the Kernel\n",
     "        \n",
     "        2. Update `particles.dlon` by adding the change in longitude, if needed<br>\n",
     "\n",
-    "    5. Update `particles.time_nextloop += particles.dt`<br>\n",
+    "    5. If `outputdt` is a multiple of `particle.time`, write `particle.lon` and `particle.time` to zarr output file<br>\n",
     "\n",
-    "    6. If `outputdt` is a multiple of `particle.time`, write `particle.lon` and `particle.time` to zarr output file<br>\n",
-    "\n",
-    "Besides having commutable Kernels, the main advantage of this implementation is that, when using Field Sampling with e.g. `particle.temp = fieldset.Temp[particle.time, particle.depth, particle.lat, particle.lon]`, the particle location stays the same throughout the entire Kernel loop. Additionally, this implementation ensures that the particle location is the same as the location of the sampled field in the output file."
+    "Besides having commutable Kernels, the main advantage of this implementation is that, when using Field Sampling with e.g. `particle.temp = fieldset.Temp[particle.time, particle.z, particle.lat, particle.lon]`, the particle location stays the same throughout the entire Kernel loop. Additionally, this implementation ensures that the particle location is the same as the location of the sampled field in the output file."
    ]
   },
   {
@@ -155,10 +153,10 @@
    "source": [
     "def wind_kernel(particle, fieldset, time):\n",
     "    particle_dlon += (\n",
-    "        fieldset.UWind[time, particle.depth, particle.lat, particle.lon] * particle.dt\n",
+    "        fieldset.UWind[time, particle.z, particle.lat, particle.lon] * particle.dt\n",
     "    )\n",
     "    particle_dlat += (\n",
-    "        fieldset.VWind[time, particle.depth, particle.lat, particle.lon] * particle.dt\n",
+    "        fieldset.VWind[time, particle.z, particle.lat, particle.lon] * particle.dt\n",
     "    )"
    ]
   },
@@ -242,30 +240,16 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## Caveats"
+    "## Caveat: Avoid updating particle locations directly in Kernels"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "There are a few important considerations to take into account when writing Kernels\n",
-    "\n",
-    "### 1. Avoid updating particle locations directly in Kernels\n",
     "It is better not to update `particle.lon` directly in a Kernel, as it can interfere with the loop above. Assigning a value to `particle.lon` in a Kernel will throw a warning. \n",
     "\n",
-    "Instead, update the local variable `particle.dlon`.\n",
-    "\n",
-    "### 2. Be careful with updating particle variables that do not depend on Fields.\n",
-    "While assigning the interpolated value of a `Field` to a Particle goes well in the loop above, this is not necessarily so for assigning other attributes. For example, a line like `particle.age += particle.dt` is executed directly so may result in the age being `dt` at `time = 0` in the output file. \n",
-    "\n",
-    "A workaround is to either initialise the age to `-dt`, or to increase the `age` only when `particle.time > 0` (using an `np.where` statement).\n",
-    "\n",
-    "\n",
-    "### 3. The last time is not written to file\n",
-    "Because the location at the start of the loop is written at the end of the Kernel loop, the last `particle.time` of the particle is not written to file. This is similar behaviour to e.g. `np.arange(start, stop)`, which also doesn't include the `stop` value itself. \n",
-    "\n",
-    "If you do want to write the last time to file, you can increase the `runtime` or `endtime` by `dt` (although this may cause an OutsideTimeInterval if your run was to the end of the available hydrodynamic data), or you can call `pfile.write_latest_locations(pset, time=pset[0].time_nextloop)`. Note that in the latter case, the particle locations (longitude, latitude and depth) will be updated, but other variables will _not_ be updated as the Kernels are not run again."
+    "Instead, update the local variable `particle.dlon`."
    ]
   },
   {
@@ -355,7 +339,7 @@
    "source": [
     "def KeepInOcean(particle, fieldset, time):\n",
     "    if particle.state == StatusCode.ErrorThroughSurface:\n",
-    "        particle_ddepth = 0.0\n",
+    "        particle_dz = 0.0\n",
     "        particle.state = StatusCode.Success"
    ]
   },

src/parcels/_core/field.py

@@ -218,12 +218,9 @@ def eval(self, time: datetime, z, y, x, particles=None, applyConversion=True):
         else:
             _ei = particles.ei[:, self.igrid]
 
-        tau, ti = _search_time_index(self, time)
-        position = self.grid.search(z, y, x, ei=_ei)
-        _update_particles_ei(particles, position, self)
-        _update_particle_states_position(particles, position)
+        particle_positions, grid_positions = _get_positions(self, time, z, y, x, particles, _ei)
 
-        value = self._interp_method(self, ti, position, tau, time, z, y, x)
+        value = self._interp_method(particle_positions, grid_positions, self)
 
         _update_particle_states_interp_value(particles, value)
 
@@ -304,20 +301,17 @@ def eval(self, time: datetime, z, y, x, particles=None, applyConversion=True):
         else:
             _ei = particles.ei[:, self.igrid]
 
-        tau, ti = _search_time_index(self.U, time)
-        position = self.grid.search(z, y, x, ei=_ei)
-        _update_particles_ei(particles, position, self)
-        _update_particle_states_position(particles, position)
+        particle_positions, grid_positions = _get_positions(self.U, time, z, y, x, particles, _ei)
 
         if self._vector_interp_method is None:
-            u = self.U._interp_method(self.U, ti, position, tau, time, z, y, x)
-            v = self.V._interp_method(self.V, ti, position, tau, time, z, y, x)
+            u = self.U._interp_method(particle_positions, grid_positions, self.U)
+            v = self.V._interp_method(particle_positions, grid_positions, self.V)
             if "3D" in self.vector_type:
-                w = self.W._interp_method(self.W, ti, position, tau, time, z, y, x)
+                w = self.W._interp_method(particle_positions, grid_positions, self.W)
             else:
                 w = 0.0
         else:
-            (u, v, w) = self._vector_interp_method(self, ti, position, tau, time, z, y, x)
+            (u, v, w) = self._vector_interp_method(particle_positions, grid_positions, self)
 
         if applyConversion:
             u = self.U.units.to_target(u, z, y, x)
@@ -343,45 +337,54 @@ def __getitem__(self, key):
             return _deal_with_errors(error, key, vector_type=self.vector_type)
 
 
-def _update_particles_ei(particles, position, field):
+def _update_particles_ei(particles, grid_positions: dict, field: Field):
     """Update the element index (ei) of the particles"""
     if particles is not None:
         if isinstance(field.grid, XGrid):
             particles.ei[:, field.igrid] = field.grid.ravel_index(
                 {
-                    "X": position["X"][0],
-                    "Y": position["Y"][0],
-                    "Z": position["Z"][0],
+                    "X": grid_positions["X"]["index"],
+                    "Y": grid_positions["Y"]["index"],
+                    "Z": grid_positions["Z"]["index"],
                 }
             )
         elif isinstance(field.grid, UxGrid):
             particles.ei[:, field.igrid] = field.grid.ravel_index(
                 {
-                    "Z": position["Z"][0],
-                    "FACE": position["FACE"][0],
+                    "Z": grid_positions["Z"]["index"],
+                    "FACE": grid_positions["FACE"]["index"],
                 }
             )
 
 
-def _update_particle_states_position(particles, position):
+def _update_particle_states_position(particles, grid_positions: dict):
     """Update the particle states based on the position dictionary."""
     if particles:  # TODO also support uxgrid search
         for dim in ["X", "Y"]:
-            if dim in position:
+            if dim in grid_positions:
                 particles.state = np.maximum(
-                    np.where(position[dim][0] == -1, StatusCode.ErrorOutOfBounds, particles.state), particles.state
+                    np.where(grid_positions[dim]["index"] == -1, StatusCode.ErrorOutOfBounds, particles.state),
+                    particles.state,
                 )
                 particles.state = np.maximum(
-                    np.where(position[dim][0] == GRID_SEARCH_ERROR, StatusCode.ErrorGridSearching, particles.state),
+                    np.where(
+                        grid_positions[dim]["index"] == GRID_SEARCH_ERROR,
+                        StatusCode.ErrorGridSearching,
+                        particles.state,
+                    ),
                     particles.state,
                 )
-        if "Z" in position:
+        if "Z" in grid_positions:
             particles.state = np.maximum(
-                np.where(position["Z"][0] == RIGHT_OUT_OF_BOUNDS, StatusCode.ErrorOutOfBounds, particles.state),
+                np.where(
+                    grid_positions["Z"]["index"] == RIGHT_OUT_OF_BOUNDS, StatusCode.ErrorOutOfBounds, particles.state
+                ),
                 particles.state,
             )
             particles.state = np.maximum(
-                np.where(position["Z"][0] == LEFT_OUT_OF_BOUNDS, StatusCode.ErrorThroughSurface, particles.state),
+                np.where(
+                    grid_positions["Z"]["index"] == LEFT_OUT_OF_BOUNDS, StatusCode.ErrorThroughSurface, particles.state
+                ),
                 particles.state,
             )
 
@@ -469,3 +472,14 @@ def _assert_same_time_interval(fields: list[Field]) -> None:
             raise ValueError(
                 f"Fields must have the same time domain. {fields[0].name}: {reference_time_interval}, {field.name}: {field.time_interval}"
             )
+
+
+def _get_positions(field: Field, time, z, y, x, particles, _ei) -> tuple[dict, dict]:
+    """Initialize and populate particle_positions and grid_positions dictionaries"""
+    particle_positions = {"time": time, "z": z, "lat": y, "lon": x}
+    grid_positions = {}
+    grid_positions.update(_search_time_index(field, time))
+    grid_positions.update(field.grid.search(z, y, x, ei=_ei))
+    _update_particles_ei(particles, grid_positions, field)
+    _update_particle_states_position(particles, grid_positions)
+    return particle_positions, grid_positions

src/parcels/_core/index_search.py

@@ -75,14 +75,19 @@ def _search_time_index(field: Field, time: datetime):
     if the sampled value is outside the time value range.
     """
     if field.time_interval is None:
-        return np.zeros(shape=time.shape, dtype=np.float32), np.zeros(shape=time.shape, dtype=np.int32)
+        return {
+            "T": {
+                "index": np.zeros(shape=time.shape, dtype=np.int32),
+                "bcoord": np.zeros(shape=time.shape, dtype=np.float32),
+            }
+        }
 
     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])
-    return np.atleast_1d(tau), np.atleast_1d(ti)
+    return {"T": {"index": np.atleast_1d(ti), "bcoord": np.atleast_1d(tau)}}
 
 
 def curvilinear_point_in_cell(grid, y: np.ndarray, x: np.ndarray, yi: np.ndarray, xi: np.ndarray):

src/parcels/_core/kernel.py

@@ -75,7 +75,7 @@ def __init__(
         self._fieldset = fieldset
         self._ptype = ptype
 
-        self._positionupdate_kernels_added = False
+        self._positionupdate_kernel_added = False
 
         for f in pyfuncs:
             self.check_fieldsets_in_kernels(f)
@@ -111,23 +111,19 @@ def remove_deleted(self, pset):
         if len(indices) > 0:
             pset.remove_indices(indices)
 
-    def add_positionupdate_kernels(self):
+    def add_positionupdate_kernel(self):
         # Adding kernels that set and update the coordinate changes
-        def Setcoords(particles, fieldset):  # pragma: no cover
+        def PositionUpdate(particles, fieldset):  # pragma: no cover
             particles.lon += particles.dlon
             particles.lat += particles.dlat
             particles.z += particles.dz
+            particles.time += particles.dt
 
             particles.dlon = 0
             particles.dlat = 0
             particles.dz = 0
 
-            particles.time = particles.time_nextloop
-
-        def UpdateTime(particles, fieldset):  # pragma: no cover
-            particles.time_nextloop = particles.time + particles.dt
-
-        self._pyfuncs = (Setcoords + self + UpdateTime)._pyfuncs
+        self._pyfuncs = (PositionUpdate + self)._pyfuncs
 
     def check_fieldsets_in_kernels(self, pyfunc):  # TODO v4: this can go into another method? assert_is_compatible()?
         """
@@ -234,14 +230,13 @@ def execute(self, pset, endtime, dt):
 
         pset._data["state"][:] = StatusCode.Evaluate
 
-        if not self._positionupdate_kernels_added:
-            self.add_positionupdate_kernels()
-            self._positionupdate_kernels_added = True
-
         while (len(pset) > 0) and np.any(np.isin(pset.state, [StatusCode.Evaluate, StatusCode.Repeat])):
-            time_to_endtime = compute_time_direction * (endtime - pset.time_nextloop)
+            time_to_endtime = compute_time_direction * (endtime - pset.time)
 
-            if all(time_to_endtime <= 0):
+            evaluate_particles = (np.isin(pset.state, [StatusCode.Success, StatusCode.Evaluate])) & (
+                time_to_endtime >= 0
+            )
+            if not np.any(evaluate_particles):
                 return StatusCode.Success
 
             # adapt dt to end exactly on endtime
@@ -251,7 +246,6 @@ def execute(self, pset, endtime, dt):
                 pset.dt = np.minimum(np.maximum(pset.dt, -time_to_endtime), 0)
 
             # run kernels for all particles that need to be evaluated
-            evaluate_particles = (pset.state == StatusCode.Evaluate) & (pset.dt != 0)
             for f in self._pyfuncs:
                 f(pset[evaluate_particles], self._fieldset)
 
@@ -265,9 +259,9 @@ def execute(self, pset, endtime, dt):
             if not hasattr(self.fieldset, "RK45_tol"):
                 pset._data["dt"][:] = dt
 
-            # Reset particle state for particles that signalled success and have not reached endtime yet
-            particles_to_evaluate = (pset.state == StatusCode.Success) & (time_to_endtime > 0)
-            pset[particles_to_evaluate].state = StatusCode.Evaluate
+            # Set particle state for particles that reached endtime
+            particles_endofloop = (pset.state == StatusCode.Evaluate) & (pset.time == endtime)
+            pset[particles_endofloop].state = StatusCode.EndofLoop
 
             # delete particles that signalled deletion
             self.remove_deleted(pset)
@@ -284,4 +278,9 @@ def execute(self, pset, endtime, dt):
                     else:
                         error_func(pset[inds].z, pset[inds].lat, pset[inds].lon)
 
+            # Only add PositionUpdate kernel at the end of the first execute call to avoid adding dt to time too early
+            if not self._positionupdate_kernel_added:
+                self.add_positionupdate_kernel()
+                self._positionupdate_kernel_added = True
+
         return pset

src/parcels/_core/particle.py

@@ -180,7 +180,6 @@ def get_default_particle(spatial_dtype: np.float32 | np.float64) -> ParticleClas
                 dtype=_SAME_AS_FIELDSET_TIME_INTERVAL.VALUE,
                 attrs={"standard_name": "time", "units": "seconds", "axis": "T"},
             ),
-            Variable("time_nextloop", dtype=_SAME_AS_FIELDSET_TIME_INTERVAL.VALUE, to_write=False),
             Variable(
                 "trajectory",
                 dtype=np.int64,