diff --git a/docs/user_guide/getting_started/tutorial_output.ipynb b/docs/user_guide/getting_started/tutorial_output.ipynb index e9adfffa50..ea311a2877 100644 --- a/docs/user_guide/getting_started/tutorial_output.ipynb +++ b/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": { diff --git a/pixi.toml b/pixi.toml index a3149893d8..b9b2713142 100644 --- a/pixi.toml +++ b/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." } diff --git a/pyproject.toml b/pyproject.toml index 32aa744c70..71afa62dc4 100644 --- a/pyproject.toml +++ b/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] diff --git a/src/parcels/__init__.py b/src/parcels/__init__.py index 796ed68bb1..0736109463 100644 --- a/src/parcels/__init__.py +++ b/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", diff --git a/src/parcels/_compat_v3.py b/src/parcels/_compat_v3.py new file mode 100644 index 0000000000..1e1ca9299e --- /dev/null +++ b/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) diff --git a/src/parcels/_core/particlefile.py b/src/parcels/_core/particlefile.py index d12f4eb073..d2cfd5a778 100644 --- a/src/parcels/_core/particlefile.py +++ b/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, ) diff --git a/src/parcels/_strategies/__init__.py b/src/parcels/_strategies/__init__.py index 15536e2696..85336d5d1b 100644 --- a/src/parcels/_strategies/__init__.py +++ b/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"] diff --git a/src/parcels/_strategies/_core.py b/src/parcels/_strategies/_core.py new file mode 100644 index 0000000000..874074fa65 --- /dev/null +++ b/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 diff --git a/src/parcels/_strategies/particle.py b/src/parcels/_strategies/particle.py new file mode 100644 index 0000000000..e41e978369 --- /dev/null +++ b/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) diff --git a/tests/__init__.py b/tests/__init__.py index e69de29bb2..74dd17ab10 100644 --- a/tests/__init__.py +++ b/tests/__init__.py @@ -0,0 +1,3 @@ +from tests import mark, utils + +__all__ = ["mark", "utils"] diff --git a/tests/mark.py b/tests/mark.py new file mode 100644 index 0000000000..29980f0820 --- /dev/null +++ b/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" +) diff --git a/tests/test_compat_v3.py b/tests/test_compat_v3.py new file mode 100644 index 0000000000..a6c9625b92 --- /dev/null +++ b/tests/test_compat_v3.py @@ -0,0 +1,61 @@ +import io +import tempfile +from datetime import timedelta +from pathlib import Path + +import numpy as np +import xarray as xr +from hypothesis import example, given, settings + +import parcels._strategies as pst +from parcels import FieldSet, ParticleFile, ParticleSet, StatusCode +from parcels._compat_v3 import particlefile_to_v3_zarr +from parcels._core.particle import Particle +from parcels._datasets.structured.generic import datasets as datasets_structured + + +def example_particlefile() -> io.BytesIO: + ds = datasets_structured["ds_2d_left"].copy() + ds = ds[["U_A_grid", "V_A_grid", "grid"]].rename({"U_A_grid": "U", "V_A_grid": "V"}) + fieldset = FieldSet.from_sgrid_conventions(ds, mesh="flat") + + npart = 10 + pset = ParticleSet(fieldset, pclass=Particle, x=np.zeros(npart), y=np.zeros(npart)) + + def RandomDelete(particles, fieldset): # pragma: no cover + particles.state = np.where( + np.random.rand(len(particles)) < 0.3, + StatusCode.Delete, + particles.state, + ) + + with tempfile.TemporaryDirectory() as tmpdir: + parquet_path = Path(tmpdir) / "output.parquet" + ofile = ParticleFile(parquet_path, outputdt=np.timedelta64(1, "s")) + pset.execute(RandomDelete, runtime=np.timedelta64(5, "s"), dt=np.timedelta64(1, "s"), output_file=ofile) + + buf = io.BytesIO(parquet_path.read_bytes()) + + return buf + + +def assert_valid_v3_particlefile_structure(ds: xr.Dataset): + for var in ["lat", "lon", "z", "time"]: + assert var in ds.variables + + assert set(ds.dims) == {"obs", "trajectory"} + assert set(ds.coords) == {"obs", "trajectory"} + + assert ds["lat"].attrs["axis"] == "Y" # attrs are copied accross correctly + + +@settings(deadline=timedelta(seconds=0.3)) +@example(buf=example_particlefile()) +@given(buf=pst.particlefile_output()) +def test_particlefile_to_v3_zarr(buf): + with tempfile.TemporaryDirectory() as tmpdir: + tmp_zarr = Path(tmpdir) / "output.zarr" + + particlefile_to_v3_zarr(from_parquet=buf, to_zarr=tmp_zarr) + ds = xr.open_zarr(tmp_zarr) + assert_valid_v3_particlefile_structure(ds) diff --git a/tests/test_convert.py b/tests/test_convert.py index 2fdc277b2b..45933eba12 100644 --- a/tests/test_convert.py +++ b/tests/test_convert.py @@ -159,6 +159,7 @@ def test_convert_copernicusmarine_no_currents(caplog): assert caplog.text == "" +@pytest.mark.filterwarnings("ignore:The delft3d_to_sgrid function is experimental") def test_convert_structured_delft3d(): ds = open_remote_dataset("Delft3D_data/Rotterdam_tiny") coords = ds[["XZETA", "YZETA", "SIGMA_C"]] diff --git a/tests/test_fieldset.py b/tests/test_fieldset.py index 88ad21f132..a81cf43297 100644 --- a/tests/test_fieldset.py +++ b/tests/test_fieldset.py @@ -9,13 +9,13 @@ import xarray as xr import parcels.tutorial +import tests from parcels import ParticleFile, ParticleSet, convert, open_raw_zarr from parcels._core.fieldset import FieldSet, _datetime_to_msg from parcels._core.model import _default_vector_field_components from parcels._datasets.structured.generic import datasets as datasets_structured from parcels._datasets.structured.generic import datasets_sgrid from parcels._datasets.unstructured.generic import datasets as datasets_unstructured -from tests import utils ds = datasets_structured["ds_2d_left"] @@ -109,7 +109,7 @@ def test_fieldset_from_structured_generic_datasets(ds): assert len(fieldset.fields) == len(ds.data_vars) - 1 # `-1` for the SGRID metadata for field in fieldset.fields.values(): - utils.assert_valid_field_data(field.data, field.grid) + tests.utils.assert_valid_field_data(field.data, field.grid) assert len(fieldset.gridset) == 1 @@ -415,6 +415,7 @@ def test_fieldset_add_error_on_duplicate_context_values(): fset1 + fset2 +@tests.mark.zarr_filterwarning_consolidated_metadata @pytest.mark.parametrize("skip", [True, False]) def test_zarr_warning_on_fieldset_creation(skip, tmp_path): """Test that creating a FieldSet from a Zarr-backed dataset raises a warning about potential backend changes.""" @@ -475,6 +476,7 @@ def test_fieldset_describe(fieldset_two_models: FieldSet): assert actual == expected +@tests.mark.zarr_filterwarning_consolidated_metadata def test_fieldset_describe_backends(tmp_path): ds_u = parcels.tutorial.open_dataset("NemoNorthSeaORCA025-N006_data/U") ds_v = parcels.tutorial.open_dataset("NemoNorthSeaORCA025-N006_data/V") diff --git a/tests/test_particlefile.py b/tests/test_particlefile.py index d4ed34e342..426935d0b4 100755 --- a/tests/test_particlefile.py +++ b/tests/test_particlefile.py @@ -22,7 +22,7 @@ convert, ) from parcels._core.particle import Particle, get_default_particle -from parcels._core.particlefile import _get_schema +from parcels._core.particlefile import get_schema from parcels._core.utils.time import TimeInterval, timedelta_to_float from parcels._datasets.structured.generated import peninsula_dataset from parcels.interpolators import XLinear @@ -560,7 +560,7 @@ def Update_lon(particles, fieldset): # pragma: no cover ], ) def test_particle_schema(particle): - s = _get_schema(particle, {}, TimeInterval(datetime(2023, 1, 1, 12, 0), datetime(2023, 1, 2, 12, 0))) + s = get_schema(particle, {}, TimeInterval(datetime(2023, 1, 1, 12, 0), datetime(2023, 1, 2, 12, 0))) written_variables = [v for v in particle.variables if v.to_write] diff --git a/tests/test_xarray.py b/tests/test_xarray.py index dec10e835d..f16326fc99 100644 --- a/tests/test_xarray.py +++ b/tests/test_xarray.py @@ -2,14 +2,16 @@ import xarray as xr import zarr +import tests from parcels import open_raw_zarr from parcels._datasets.structured.generic import datasets -@pytest.mark.filterwarnings("ignore:Consolidated metadata is currently not part in the Zarr format 3 specification") +@tests.mark.zarr_filterwarning_consolidated_metadata @pytest.mark.parametrize("ds", [pytest.param(v, id=k) for k, v in datasets.items()]) def test_open_raw_zarr(ds: xr.Dataset, tmp_path): path = tmp_path / "ds.zarr" + ds.to_zarr(path) result = open_raw_zarr(path)