WIP

Author: VeckoTheGecko

parcels/particlefile.py

@@ -148,7 +148,7 @@ def _write_particle_data(self, *, particle_data, pclass, time_interval, time, in
         #         stacklevel=2,
         #     )
         #     return
-
+        nparticles = len(particle_data["trajectory"])
         vars_to_write = _get_vars_to_write(pclass)
         if indices is None:
             indices_to_write = _to_write_particles(particle_data, time)
@@ -173,7 +173,7 @@ def _write_particle_data(self, *, particle_data, pclass, time_interval, time, in
         store = self.store
         if self.create_new_zarrfile:
             if self.chunks is None:
-                self._chunks = (len(pset), 1)
+                self._chunks = (nparticles, 1)
             if (self._maxids > len(ids)) or (self._maxids > self.chunks[0]):  # type: ignore[index]
                 arrsize = (self._maxids, self.chunks[1])  # type: ignore[index]
             else:
@@ -280,10 +280,10 @@ 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"]), particle_data["time"], where=np.isfinite(particle_data["time"])
             )
             & 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"]), particle_data["time"], where=np.isfinite(particle_data["time"])
             )
             | (
                 (np.isnan(particle_data["dt"]))
@@ -307,7 +307,9 @@ def _convert_particle_data_time_to_float_seconds(particle_data, time_interval):
 def _maybe_convert_time_dtype(dtype: np.dtype | _SAME_AS_FIELDSET_TIME_INTERVAL) -> np.dtype:
     """Convert the dtype of time to float64 if it is not already."""
     if dtype is _SAME_AS_FIELDSET_TIME_INTERVAL.VALUE:
-        return np.dtype(np.float64)
+        return np.dtype(
+            np.uint64
+        )  #! We need to have here some proper mechanism for converting particle data to the data that is to be output to zarr (namely the time needs to be converted to float seconds by subtracting the time_interval.left)
     return dtype
 
 

tests/v4/test_particlefile.py

@@ -9,10 +9,10 @@
 
 import parcels
 from parcels import AdvectionRK4, Field, FieldSet, Particle, ParticleSet, Variable, VectorField
+from parcels._core.utils.time import TimeInterval
 from parcels._datasets.structured.generic import datasets
+from parcels.particle import Particle, create_particle_data
 from parcels.particlefile import ParticleFile
-from parcels.particle import create_particle_data
-from parcels._core.utils.time import TimeInterval
 from parcels.xgrid import XGrid
 from tests.common_kernels import DoNothing
 
@@ -366,4 +366,48 @@ def test_particlefile_init_invalid(store):  # TODO: Add test for read only store
 
 
 @pytest.mark.new
-def test_particlefile_writing(store): ...
+def test_particlefile_write_particle_data(store):
+    nparticles = 100
+
+    pfile = ParticleFile(store, outputdt=np.timedelta64(1, "s"), chunks=(nparticles, 40))
+    pclass = Particle
+
+    left, right = np.datetime64("2020-01-01", "ns"), np.datetime64("2020-01-02", "ns")
+    time_interval = TimeInterval(left=left, right=right)
+
+    initial_lon = np.linspace(0, 1, nparticles)
+    data = create_particle_data(
+        pclass=pclass,
+        nparticles=nparticles,
+        ngrids=4,
+        time_interval=time_interval,
+        initial={
+            "time": np.full(nparticles, fill_value=left),
+            "lon": initial_lon,
+            "dt": np.full(nparticles, fill_value=1.0),
+            "trajectory": np.arange(nparticles),
+        },
+    )
+    np.testing.assert_array_equal(data["time"], left)
+    pfile._write_particle_data(
+        particle_data=data,
+        pclass=pclass,
+        time_interval=time_interval,
+        time=left,
+    )
+    ds = xr.open_zarr(store, decode_cf=False)
+    # ds.time.attrs
+
+    def patched_time_attrs(attrs):
+        units = attrs["units"]
+        units = units[: units.index("T")]
+        attrs["units"] = units
+        return attrs
+
+    ds.time.attrs = patched_time_attrs(ds.time.attrs)
+    breakpoint()
+    ds = xr.decode_cf(ds)
+    assert ds.sizes["trajectory"] == nparticles
+    assert ds.time.dtype == "datetime64[ns]"
+    np.testing.assert_equal(ds["time"].isel(obs=0).values, left)
+    np.testing.assert_equal(ds["lon"].isel(obs=0).values, initial_lon)