Fix selection of particlefiles for writing

Looking at time_nextloop instead. We might rework the kernel loop in future (and the responsibilities of time, time_nextloop, as well as their effect on particle writing).

Author: VeckoTheGecko

parcels/particlefile.py

@@ -283,18 +283,22 @@ def _to_write_particles(particle_data, time):
     return np.where(
         (
             np.less_equal(
-                time - np.abs(particle_data["dt"] / 2), particle_data["time"], where=np.isfinite(particle_data["time"])
+                time - np.abs(particle_data["dt"] / 2),
+                particle_data["time_nextloop"],
+                where=np.isfinite(particle_data["time_nextloop"]),
             )
             & np.greater_equal(
-                time + np.abs(particle_data["dt"] / 2), particle_data["time"], where=np.isfinite(particle_data["time"])
+                time + np.abs(particle_data["dt"] / 2),
+                particle_data["time_nextloop"],
+                where=np.isfinite(particle_data["time_nextloop"]),
             )  # check time - dt/2 <= particle_data["time"] <= time + dt/2
             | (
                 (np.isnan(particle_data["dt"]))
-                & np.equal(time, particle_data["time"], where=np.isfinite(particle_data["time"]))
+                & np.equal(time, particle_data["time_nextloop"], where=np.isfinite(particle_data["time_nextloop"]))
             )  # or dt is NaN and time matches particle_data["time"]
         )
         & (np.isfinite(particle_data["trajectory"]))
-        & (np.isfinite(particle_data["time"]))
+        & (np.isfinite(particle_data["time_nextloop"]))
     )[0]
 
 
@@ -303,6 +307,9 @@ def _convert_particle_data_time_to_float_seconds(particle_data, time_interval):
     particle_data = particle_data.copy()
 
     particle_data["time"] = ((particle_data["time"] - time_interval.left) / np.timedelta64(1, "s")).astype(np.float64)
+    particle_data["time_nextloop"] = (
+        (particle_data["time_nextloop"] - time_interval.left) / np.timedelta64(1, "s")
+    ).astype(np.float64)
     particle_data["dt"] = (particle_data["dt"] / np.timedelta64(1, "s")).astype(np.float64)
     return particle_data
 

tests/v4/test_advection.py

@@ -52,6 +52,26 @@ def test_advection_zonal(mesh, npart=10):
         assert (np.diff(pset.lon) < 1.0e-4).all()
 
 
+def test_advection_zonal_with_particlefile(tmp_store):
+    """Particles at high latitude move geographically faster due to the pole correction in `GeographicPolar`."""
+    npart = 10
+    ds = simple_UV_dataset(mesh="flat")
+    ds["U"].data[:] = 1.0
+    grid = XGrid.from_dataset(ds, mesh="flat")
+    U = Field("U", ds["U"], grid, interp_method=XLinear)
+    V = Field("V", ds["V"], grid, interp_method=XLinear)
+    UV = VectorField("UV", U, V)
+    fieldset = FieldSet([U, V, UV])
+
+    pset = ParticleSet(fieldset, lon=np.zeros(npart) + 20.0, lat=np.linspace(0, 80, npart))
+    pfile = pset.ParticleFile(tmp_store, outputdt=np.timedelta64(15, "m"))
+    pset.execute(AdvectionRK4, runtime=np.timedelta64(2, "h"), dt=np.timedelta64(15, "m"), output_file=pfile)
+
+    assert (np.diff(pset.lon) < 1.0e-4).all()
+    ds = xr.open_zarr(tmp_store)
+    breakpoint()
+
+
 def periodicBC(particle, fieldset, time):
     particle.total_dlon += particle.dlon
     particle.lon = np.fmod(particle.lon, 2)

tests/v4/test_particlefile.py

@@ -385,6 +385,7 @@ def test_particlefile_write_particle_data(tmp_store):
         time_interval=time_interval,
         initial={
             "time": np.full(nparticles, fill_value=left),
+            "time_nextloop": np.full(nparticles, fill_value=left),
             "lon": initial_lon,
             "dt": np.full(nparticles, fill_value=1.0),
             "trajectory": np.arange(nparticles),