Add particlefile_to_v3_zarr() helper

docs/user_guide/getting_started/tutorial_output.ipynb

@@ -147,7 +147,7 @@
     "\n",
     "However, the `parquet` format does not support the [CF-convention for trajectories data](http://cfconventions.org/cf-conventions/v1.6.0/cf-conventions.html#_multidimensional_array_representation_of_trajectories) implemented with the [NCEI trajectory template](https://www.ncei.noaa.gov/data/oceans/ncei/formats/netcdf/v2.0/trajectoryIncomplete.cdl). We have implemented a `parcels.read_particlefile()` function to facilitate reading `parquet` output files, see more information below.\n",
     "\n",
-    "TODO: Add information on conversion functions once https://github.com/Parcels-code/Parcels/issues/2600 is resolved.\n",
+    "There are also utilities to convert to v3 ParticleFile output - see later in this tutorial.\n",
     "```"
    ]
   },
@@ -487,6 +487,30 @@
     "plt.close(fig)\n",
     "anim"
    ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Using v3 output\n",
+    "\n",
+    "We also have some tooling for converting v4 trajectory output to the v3 Zarr output. This is meant as a way to reduce friction for those not wanting to update their plotting scripts, however we recommend writing plotting scripts that work with the new v4 output.\n",
+    "\n",
+    "At the moment, this tooling is not lazy (i.e., it eagerly loads all data into memory in order to do the transformation) so your mileage may vary if you try to convert larger datasets. If you would like to contribute to the performance of this utility, please open an issue or PR."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import xarray as xr\n",
+    "\n",
+    "parcels.particlefile_to_v3_zarr(\"output.parquet\", \"output.zarr\")\n",
+    "ds = xr.open_dataset(\"output.zarr\")\n",
+    "ds"
+   ]
   }
  ],
  "metadata": {

pixi.toml

@@ -160,6 +160,7 @@ mypy = "*"
 lxml = "*" # in CI
 types-tqdm = "*"
 pandas-stubs = "*"
+pyarrow-stubs = "*"
 
 [feature.typing.tasks]
 typing = { cmd = "mypy src/parcels --install-types", description = "Run static type checking with mypy." }

pyproject.toml

@@ -68,6 +68,8 @@ markers = [ # can be skipped by doing `pytest -m "not slow"` etc.
 filterwarnings = [
   "error:.*removed in a future release of Parcels.*:DeprecationWarning", # Have Parcels DeprecationWarnings fail CI (prevents deprecated items being used in internal code)
   "error:::parcels.*",
+  "error::UserWarning",
+  "ignore:This is an alpha version of Parcels v4.*:UserWarning", # TODO v4: Remove when warning is removed
 ]
 
 [tool.ruff]

src/parcels/__init__.py

@@ -13,6 +13,7 @@
 from parcels._xarray import open_raw_zarr
 from parcels._core.particleset import ParticleSet
 from parcels._core.particlefile import ParticleFile, read_particlefile
+from parcels._compat_v3 import particlefile_to_v3_zarr
 from parcels._core.particle import (
     Variable,
     Particle,
@@ -76,6 +77,7 @@
     # Utilities
     "logger",
     "read_particlefile",
+    "particlefile_to_v3_zarr",
     "convert",
     # kernels
     "kernels",

src/parcels/_compat_v3.py

@@ -0,0 +1,110 @@
+from __future__ import annotations
+
+import os
+import warnings
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+import polars as pl
+import pyarrow.parquet as pq
+import xarray as xr
+
+if TYPE_CHECKING:
+    import io
+    from pathlib import Path
+
+
+def _decode_dict_to_utf8(d: dict[Any, Any]) -> dict[Any, Any]:
+    ret = {}
+    for key, item in d.items():
+        if isinstance(key, bytes):
+            key = key.decode("utf8")
+        if isinstance(item, dict):
+            item = _decode_dict_to_utf8(item)
+        if isinstance(item, bytes):
+            item = item.decode("utf8")
+        ret[key] = item
+    return ret
+
+
+def particlefile_to_v3_zarr(from_parquet: str | Path | io.BytesIO, to_zarr: str | os.PathLike) -> None:
+    """Convert a v4 particle file (parquet) to v3-style zarr output.
+
+    Reads the parquet file, renames columns to v3 conventions
+    (``particle_id`` -> ``trajectory``, ``t`` -> ``time``, ``x`` -> ``lon``,
+    ``y`` -> ``lat``), and reshapes the data into a 2D
+    ``(trajectory, obs)`` zarr store.
+
+    Parameters
+    ----------
+    from_parquet : Path
+        Path to the input parquet file.
+    to_zarr : Path
+        Path to the output zarr store. Must have a ``.zarr`` suffix.
+
+    Raises
+    ------
+    ValueError
+        If ``to_zarr`` does not have a ``.zarr`` suffix.
+
+    Notes
+    -----
+    This is not a lazy operation — the entire parquet file is read into memory
+    and pivoted before writing to zarr. For large particle files this may
+    require significant memory. Performance improvements are welcome via PRs.
+    """
+    to_zarr = Path(to_zarr)
+    if to_zarr.suffix != ".zarr":
+        raise ValueError(f"Parameter `to_zarr` must have a '.zarr' suffix. Got {to_zarr=}.")
+    df = pl.read_parquet(from_parquet)
+    table = pq.read_table(from_parquet)
+
+    # TODO: Check for available memory here and fail as a safeguard?
+
+    # Rename columns to v3 conventions
+    rename_map = {"particle_id": "trajectory", "t": "time", "x": "lon", "y": "lat"}
+    try:
+        df = df.rename(rename_map)
+    except pl.exceptions.ColumnNotFoundError as e:
+        e.add_note(
+            f"Expected to have all columns {list(rename_map)} in the output parquet. Got columns {list(df.columns)}."
+        )
+        raise e
+
+    metadata_per_field = {}
+    for parquet_var in table.schema.names:
+        zarr_var = rename_map.get(parquet_var, parquet_var)  # default to parquet name
+        metadata_per_field[zarr_var] = table.field(parquet_var).metadata or {}
+
+    # Group by trajectory, sort by time, and assign observation index
+    df = df.sort("trajectory", "time")
+    df = df.with_columns(
+        pl.col("time").cum_count().over("trajectory").alias("obs") - 1,
+    )
+
+    # Pivot to (trajectory, obs) dimensions
+    trajectories = df["trajectory"].unique().sort()
+    data_vars = [c for c in df.columns if c not in ("trajectory", "obs")]
+
+    ds_dict = {}
+    for var in data_vars:
+        pivoted = df.pivot(on="obs", index="trajectory", values=var, sort_columns=True)
+        value_cols = [c for c in pivoted.columns if c != "trajectory"]
+        ds_dict[var] = (
+            ["trajectory", "obs"],
+            pivoted.select(value_cols).to_numpy(),
+            _decode_dict_to_utf8(metadata_per_field[var]),
+        )
+
+    ds = xr.Dataset(
+        ds_dict,
+        coords={"trajectory": trajectories.to_numpy()},
+        attrs=_decode_dict_to_utf8(table.schema.metadata),
+    )
+    ds = ds.assign_coords({"obs": ds["obs"]})
+
+    with warnings.catch_warnings():
+        warnings.filterwarnings(
+            "ignore", message="Consolidated metadata is currently not part in the Zarr format 3 specification."
+        )
+        ds.to_zarr(to_zarr)

src/parcels/_core/particlefile.py

@@ -29,7 +29,7 @@
 __all__ = ["ParticleFile"]
 
 
-def _get_schema(
+def get_schema(
     particle: parcels.ParticleClass, file_metadata: dict[Any, Any], fset_time_interval: TimeInterval | None
 ) -> pa.Schema:
 
@@ -163,7 +163,7 @@ def write(self, pset: ParticleSet | ParticleSetView, t, fieldset=None, indices=N
             assert not self.path.exists(), "If the file exists, the writer should already be set"
             self._writer = pq.ParquetWriter(
                 self.path,
-                _get_schema(pclass, self.metadata, fieldset.time_interval),
+                get_schema(pclass, self.metadata, fieldset.time_interval),
                 compression=self._compression,
             )
 

src/parcels/_strategies/__init__.py

@@ -8,6 +8,7 @@
     )
     raise err
 
-from . import sgrid, time
+from . import sgrid, time, particle
+from ._core import particlefile_output
 
-__all__ = ["sgrid", "time"]
+__all__ = ["particle", "particlefile_output", "sgrid", "time"]

src/parcels/_strategies/_core.py

@@ -0,0 +1,71 @@
+import io
+
+import hypothesis.strategies as st
+import numpy as np
+import pandas as pd
+import pyarrow as pa
+import pyarrow.parquet as pq
+
+from parcels._core.particle import ParticleClass
+from parcels._core.particlefile import get_schema
+
+from .particle import particle_class
+from .time import time_interval as st_time_interval
+
+__all__ = [
+    "particlefile_output",
+]
+
+
+def _generate_dummy_data(particle: ParticleClass, nparticles=10, nobs=10) -> pd.DataFrame:
+    """Build a pandas dataframe from a particleclass.
+
+    Only variables with ``to_write=True`` are included.
+    """
+    columns: dict[str, np.ndarray] = {}
+    variables = {var.name: var for var in particle.variables if var.to_write}
+    try:
+        particle_id = variables["particle_id"]
+        t = variables["t"]
+    except KeyError as e:
+        e.add_note("This function requires 'particle_id' and 't' to be set")
+
+    nobs_total = nparticles * nobs
+    columns = {}
+    columns["particle_id"] = np.repeat(
+        np.arange(0, nparticles, dtype=particle_id.dtype).reshape((-1, 1)),
+        nobs,
+        axis=1,
+    ).flatten()
+    columns["t"] = np.repeat(
+        np.linspace(0, nparticles * 3, num=nparticles, dtype=t.dtype).reshape((-1, 1)),
+        nobs,
+        axis=1,
+    ).flatten()
+
+    data_vars = set(variables.keys()) - {"particle_id", "t"}
+
+    for name in data_vars:
+        var = variables[name]
+        columns[name] = np.linspace(0, 10000, num=nobs_total, dtype=var.dtype)
+
+    return pd.DataFrame(columns)
+
+
+@st.composite
+def particlefile_output(draw, nobs=None, nparticles=None) -> io.BytesIO:
+    particle = draw(particle_class())
+    time_interval = draw(st_time_interval())
+    if nobs is None:
+        nobs = draw(st.integers(min_value=5, max_value=100))
+    if nparticles is None:
+        nparticles = draw(st.integers(min_value=5, max_value=100))
+
+    df = _generate_dummy_data(particle, nparticles, nobs)
+    schema = get_schema(particle, {}, time_interval)
+    buf = io.BytesIO()
+    pq.write_table(
+        pa.table(df, schema=schema),
+        buf,
+    )
+    return buf

src/parcels/_strategies/particle.py

@@ -0,0 +1,91 @@
+"""Provides Hypothesis strategies for generating Variable, ParticleClass, and related particle data."""
+
+from __future__ import annotations
+
+import numpy as np
+from hypothesis import strategies as st
+
+from parcels._core.particle import Particle, Variable, get_default_particle
+
+__all__ = ["particle_class", "variable", "variable_name"]
+
+# Valid numpy dtypes for Variable
+_VARIABLE_DTYPES = [np.float32, np.float64, np.int32, np.int64, np.bool_]
+
+variable_dtype = st.sampled_from(_VARIABLE_DTYPES).map(np.dtype)
+
+# Names used by the default Particle — generated variables must not collide with these
+_DEFAULT_PARTICLE_NAMES = {var.name for var in Particle.variables}
+
+# Python identifiers that are not keywords (required by _assert_str_and_python_varname)
+variable_name = (
+    st.from_regex(r"[a-z][a-z0-9_]{0,15}", fullmatch=True)
+    .filter(lambda s: s.isidentifier())
+    .filter(lambda s: not __import__("keyword").iskeyword(s))
+    .filter(lambda s: s not in _DEFAULT_PARTICLE_NAMES)
+)
+
+
+@st.composite
+def variable(draw, name=None, dtype=None, to_write=None):
+    """Strategy for generating Variable instances.
+
+    Parameters
+    ----------
+    name : str, optional
+        Fixed variable name. If None, generates a valid Python identifier.
+    dtype : numpy.dtype, optional
+        Fixed dtype. If None, draws from common numpy dtypes.
+    to_write : bool, optional
+        Fixed to_write value. If None, draws True or False.
+    """
+    if name is None:
+        name = draw(variable_name)
+    if dtype is None:
+        dtype = draw(variable_dtype)
+    if to_write is None:
+        to_write = draw(st.booleans())
+
+    if to_write:
+        attrs = draw(
+            st.just({})
+            | st.dictionaries(
+                keys=st.text(min_size=1, max_size=10, alphabet="abcdefghijklmnopqrstuvwxyz_"),
+                values=st.text(min_size=1, max_size=20),
+                max_size=3,
+            )
+        )
+    else:
+        attrs = {}
+
+    return Variable(name=name, dtype=dtype, initial=0, to_write=to_write, attrs=attrs)
+
+
+@st.composite
+def particle_class(draw, min_vars=0, max_vars=5, spatial_dtype=None):
+    """Strategy that extends the default Particle with additional variables.
+
+    This mirrors the predominant use case in Parcels: starting from
+    ``get_default_particle`` and adding custom variables via ``add_variable``.
+
+    Parameters
+    ----------
+    min_vars : int
+        Minimum number of extra variables to add.
+    max_vars : int
+        Maximum number of extra variables to add.
+    spatial_dtype : type, optional
+        np.float32 or np.float64 for the base particle. If None, draws one.
+    """
+    if spatial_dtype is None:
+        spatial_dtype = draw(st.sampled_from([np.float32, np.float64]))
+
+    base = get_default_particle(spatial_dtype)
+
+    n = draw(st.integers(min_value=min_vars, max_value=max_vars))
+    if n == 0:
+        return base
+
+    names = draw(st.lists(variable_name, min_size=n, max_size=n, unique=True))
+    extra_vars = [draw(variable(name=name)) for name in names]
+    return base.add_variable(extra_vars)

tests/__init__.py

@@ -0,0 +1,3 @@
+from tests import mark, utils
+
+__all__ = ["mark", "utils"]

tests/mark.py

@@ -0,0 +1,5 @@
+import pytest
+
+zarr_filterwarning_consolidated_metadata = pytest.mark.filterwarnings(
+    "ignore:Consolidated metadata is currently not part in the Zarr format 3 specification"
+)