Modern I/O

Zarr v3 for field data, Parquet/Arrow for particle data, and Icechunk for versioned storage. Every backend here is optional — each sits behind its own extra and raises an install hint rather than an ImportError traceback when the dependency is missing.

Entry point	Extra	Purpose
to_zarr / from_zarr	zarr	Single-step field store, Zarr v3
to_zarr_timeseries	zarr	Multi-step store with a leading time dimension
open_virtual	zarr	VirtualiZarr view over existing HDF5, no conversion
to_icechunk_virtual, open_icechunk_repo, icechunk_create_tag, icechunk_ancestry	icechunk	Git-like versioning and ACID commits over Zarr v3
particles_to_arrow / particles_from_arrow	arrow	In-memory ParticleData ↔ Arrow table
particles_to_parquet / particles_from_parquet	arrow	Single-file particle round-trip
particles_to_dataset / particles_from_dataset	arrow	Hive-partitioned dataset, by step and species
query_sql	duckdb	SQL over particle Parquet, with spatial pushdown

The on-disk layouts are specified in Schema § 4: § 4.2 for the Zarr store this module writes, § 4.3 for the Parquet particle layout. Within each Parquet partition, rows are Morton-ordered over (x, y, z) so row-group statistics support spatial predicate pushdown.

io

Zarr v3, Parquet/Arrow, virtual HDF5, and Icechunk I/O for pypic.

Field data: to_zarr / from_zarr / to_zarr_timeseries (Zarr v3), open_virtual (VirtualiZarr), Icechunk versioning.

Particle data: particles_to_arrow / particles_from_arrow (Arrow), particles_to_parquet / particles_from_parquet (single-file Parquet), particles_to_dataset / particles_from_dataset (partitioned Parquet), query_sql (DuckDB over Parquet).

Optional dependencies per feature: pip install "pypic-plasma[zarr]" — Zarr v3, VirtualiZarr pip install "pypic-plasma[icechunk]" — Icechunk versioned storage pip install "pypic-plasma[arrow]" — Arrow/Parquet particle I/O pip install "pypic-plasma[duckdb]" — DuckDB SQL queries

particles_from_arrow(table)

Reconstruct a ParticleData from a PyArrow Table.

Expects the column layout produced by particles_to_arrow. Species metadata is read from table.schema.metadata[b"pypic"].

Parameters:

Name	Type	Description	Default
table	Table		required

Returns:

Type	Description
ParticleData

Examples:

>>> import numpy as np
>>> from pypic.containers import ParticleData
>>> pcl = ParticleData(
...     species_index=0, species_name="electrons",
...     position=np.zeros((5, 3)), velocity=np.ones((5, 3)),
...     n_particles=5, metadata={},
...     weight=np.ones(5), species_charge=-1.0, species_mass=1.0,
... )
>>> back = particles_from_arrow(particles_to_arrow(pcl))
>>> (back.n_particles, back.species_charge, back.species_mass)
(5, -1.0, 1.0)

Source code in src/pypic/io/_arrow.py

def particles_from_arrow(table: pa.Table) -> ParticleData:
    r"""Reconstruct a ``ParticleData`` from a PyArrow Table.

    Expects the column layout produced by ``particles_to_arrow``.
    Species metadata is read from ``table.schema.metadata[b"pypic"]``.

    Parameters
    ----------
    table : pyarrow.Table

    Returns
    -------
    ParticleData

    Examples
    --------
    >>> import numpy as np
    >>> from pypic.containers import ParticleData
    >>> pcl = ParticleData(
    ...     species_index=0, species_name="electrons",
    ...     position=np.zeros((5, 3)), velocity=np.ones((5, 3)),
    ...     n_particles=5, metadata={},
    ...     weight=np.ones(5), species_charge=-1.0, species_mass=1.0,
    ... )
    >>> back = particles_from_arrow(particles_to_arrow(pcl))
    >>> (back.n_particles, back.species_charge, back.species_mass)
    (5, -1.0, 1.0)
    """
    ensure_arrow()

    meta = _decode_species_meta(table.schema.metadata or {})

    position: np.ndarray | None = None
    if all(c in table.column_names for c in _POSITION_COLS):
        pos_arrays = [
            table.column(c).to_numpy(zero_copy_only=False) for c in _POSITION_COLS
        ]
        position = np.column_stack(pos_arrays)

    velocity: np.ndarray | None = None
    if all(c in table.column_names for c in _VELOCITY_COLS):
        vel_arrays = [
            table.column(c).to_numpy(zero_copy_only=False) for c in _VELOCITY_COLS
        ]
        velocity = np.column_stack(vel_arrays)

    particle_id: np.ndarray | None = None
    if "id" in table.column_names:
        particle_id = table.column("id").to_numpy(zero_copy_only=False)

    weight: np.ndarray | None = None
    if "weight" in table.column_names:
        weight_raw = table.column("weight").to_numpy(zero_copy_only=False)
        weight = (
            weight_raw
            if weight_raw.dtype == np.float64
            else weight_raw.astype(np.float64)
        )

    n_particles = len(table)

    species_charge_val = meta.get("species_charge")
    species_mass_val = meta.get("species_mass")

    from pypic.io.metadata import from_json_native

    return ParticleData(
        species_index=meta["species_index"],
        species_name=meta["species_name"],
        position=position,
        velocity=velocity,
        n_particles=n_particles,
        metadata=from_json_native(meta.get("metadata", {})),
        id=particle_id,
        weight=weight,
        species_charge=(
            float(species_charge_val) if species_charge_val is not None else None
        ),
        species_mass=(
            float(species_mass_val) if species_mass_val is not None else None
        ),
    )

particles_to_arrow(data, *, position_dtype=None, velocity_dtype=None)

Convert a ParticleData to a PyArrow Table.

Column layout: x, y, z, vx, vy, vz, weight, id. Columns for unloaded fields (e.g. velocity when data.velocity is None) are omitted. Scalar species_charge/species_mass travel in schema metadata.

Species metadata is stored in table.schema.metadata[b"pypic"] as a JSON dict.

Parameters:

Name	Type	Description	Default
data	ParticleData	Source particle data.	required
position_dtype	str or None	Downcast position columns (e.g. "float32"). Default: preserve source dtype.	None
velocity_dtype	str or None	Downcast velocity columns (e.g. "float32"). Default: preserve source dtype.	None

Returns:

Type	Description
Table

Examples:

>>> import numpy as np
>>> from pypic.containers import ParticleData
>>> pcl = ParticleData(
...     species_index=0, species_name="electrons",
...     position=np.zeros((5, 3)), velocity=np.ones((5, 3)),
...     n_particles=5, metadata={},
...     weight=np.ones(5), species_charge=-1.0, species_mass=1.0,
... )
>>> table = particles_to_arrow(pcl)
>>> table.num_rows
5

Source code in src/pypic/io/_arrow.py

def particles_to_arrow(
    data: ParticleData,
    *,
    position_dtype: str | None = None,
    velocity_dtype: str | None = None,
) -> pa.Table:
    r"""Convert a ``ParticleData`` to a PyArrow Table.

    Column layout: ``x``, ``y``, ``z``, ``vx``, ``vy``, ``vz``,
    ``weight``, ``id``.  Columns for unloaded fields (e.g. velocity
    when ``data.velocity is None``) are omitted.  Scalar
    ``species_charge``/``species_mass`` travel in schema metadata.

    Species metadata is stored in ``table.schema.metadata[b"pypic"]``
    as a JSON dict.

    Parameters
    ----------
    data : ParticleData
        Source particle data.
    position_dtype : str or None
        Downcast position columns (e.g. ``"float32"``).
        Default: preserve source dtype.
    velocity_dtype : str or None
        Downcast velocity columns (e.g. ``"float32"``).
        Default: preserve source dtype.

    Returns
    -------
    pyarrow.Table

    Examples
    --------
    >>> import numpy as np
    >>> from pypic.containers import ParticleData
    >>> pcl = ParticleData(
    ...     species_index=0, species_name="electrons",
    ...     position=np.zeros((5, 3)), velocity=np.ones((5, 3)),
    ...     n_particles=5, metadata={},
    ...     weight=np.ones(5), species_charge=-1.0, species_mass=1.0,
    ... )
    >>> table = particles_to_arrow(pcl)
    >>> table.num_rows
    5
    """
    ensure_arrow()
    import pyarrow as pa

    arrays: dict[str, pa.Array] = {}

    if data.position is not None:
        for i, col_name in enumerate(_POSITION_COLS):
            col = data.position[:, i]
            if position_dtype is not None and col.dtype != np.dtype(position_dtype):
                col = col.astype(position_dtype)
            arrays[col_name] = pa.array(col)

    if data.velocity is not None:
        for i, col_name in enumerate(_VELOCITY_COLS):
            col = data.velocity[:, i]
            if velocity_dtype is not None and col.dtype != np.dtype(velocity_dtype):
                col = col.astype(velocity_dtype)
            arrays[col_name] = pa.array(col)

    if data.id is not None:
        arrays["id"] = pa.array(data.id)

    if data.weight is not None:
        arrays["weight"] = pa.array(data.weight)

    table = pa.table(arrays)
    meta = table.schema.metadata or {}
    meta[b"pypic"] = _encode_species_meta(data)
    return table.replace_schema_metadata(meta)

query_sql(path, sql, *, return_type='particledata')

Execute SQL over a partitioned particle Parquet dataset via DuckDB.

The dataset is available in the query as the particles view. DuckDB automatically handles Hive partition discovery, predicate pushdown, and parallel scanning.

Parameters:

Name	Type	Description	Default
path	str or Path	Root directory of the partitioned Parquet dataset.	required
sql	str	SQL query. The dataset is available as particles. Example: "SELECT * FROM particles WHERE step='000100' AND x > 5"	required
return_type	str	"particledata" (default) converts the result to ParticleData. "arrow" returns a raw pyarrow.Table.	'particledata'

Returns:

Type	Description
ParticleData or Table

See Also

pypic.io.particles_to_dataset : Writes the partitioned dataset this queries.

Source code in src/pypic/io/_duckdb.py

def query_sql(
    path: str | Path,
    sql: str,
    *,
    return_type: str = "particledata",
) -> ParticleData | pa.Table:
    r"""Execute SQL over a partitioned particle Parquet dataset via DuckDB.

    The dataset is available in the query as the ``particles`` view.
    DuckDB automatically handles Hive partition discovery, predicate
    pushdown, and parallel scanning.

    Parameters
    ----------
    path : str or Path
        Root directory of the partitioned Parquet dataset.
    sql : str
        SQL query.  The dataset is available as ``particles``.
        Example: ``"SELECT * FROM particles WHERE step='000100' AND x > 5"``
    return_type : str
        ``"particledata"`` (default) converts the result to
        ``ParticleData``.  ``"arrow"`` returns a raw ``pyarrow.Table``.

    Returns
    -------
    ParticleData or pyarrow.Table

    See Also
    --------
    pypic.io.particles_to_dataset : Writes the partitioned dataset this queries.
    """
    ensure_duckdb()
    ensure_arrow()
    import duckdb

    from pypic.io._arrow import (
        inject_species_meta,
        particles_from_arrow,
    )
    from pypic.io._parquet import _strip_extra_columns

    glob_pattern = str(Path(path) / "**" / "*.parquet")
    escaped = glob_pattern.replace("'", "''")
    con = duckdb.connect()
    try:
        con.execute(
            f"CREATE VIEW particles AS "
            f"SELECT * FROM parquet_scan('{escaped}', "
            f"hive_partitioning=true)"
        )
        result = con.execute(sql)
        arrow_table = result.arrow().read_all()
    finally:
        con.close()

    if return_type == "arrow":
        return arrow_table

    # ParticleData is single-species; DuckDB strips Parquet schema
    # metadata, so recover species_index/charge/mass from one matching
    # Parquet fragment after confirming the SQL result hits exactly
    # one species.  The caller must keep the `species` partition column
    # in the projection (SELECT * does by default); cross-species and
    # aggregate queries should use return_type="arrow".
    if "species" not in arrow_table.column_names:
        msg = (
            "query_sql(return_type='particledata') requires the 'species' "
            "partition column in the result.  Keep `species` in the SELECT "
            "projection, or use return_type='arrow' for scalar / aggregate "
            "queries."
        )
        raise ValueError(msg)

    cols = set(arrow_table.column_names)
    has_position = {"x", "y", "z"}.issubset(cols)
    has_velocity = {"vx", "vy", "vz"}.issubset(cols)
    if not (has_position or has_velocity):
        msg = (
            "query_sql(return_type='particledata') requires a full "
            "position triplet (x, y, z) or velocity triplet (vx, vy, vz) "
            f"in the projection.  Got columns: {sorted(cols)}.  "
            "Use return_type='arrow' for scalar or aggregate queries."
        )
        raise ValueError(msg)

    species_values = {
        v for v in arrow_table.column("species").unique().to_pylist() if v is not None
    }
    if not species_values:
        # Well-formed filter, zero rows.  Recover the intended species when
        # one species is on disk, or when the SQL pins exactly one via a
        # literal ``species='NAME'``.  Anything else is ambiguous and falls
        # through to the ``unknown`` placeholder.
        on_disk = sorted(
            {p.name.split("=", 1)[1] for p in Path(path).glob("step=*/species=*")}
        )
        arrow_table = _strip_extra_columns(arrow_table)
        recovered_species: str | None = None
        if len(on_disk) == 1:
            recovered_species = on_disk[0]
        else:
            pinned = {m.group(1) for m in _SPECIES_LITERAL.finditer(sql)}
            on_disk_pinned = pinned & set(on_disk)
            if len(on_disk_pinned) == 1:
                recovered_species = next(iter(on_disk_pinned))
        if recovered_species is not None:
            payload = _lookup_species_meta(path, recovered_species)
            # Keep species identity (index / charge / mass) but drop
            # per-step ``metadata``: the fragment scanned is *some* step's
            # schema, not the zero-row step asked about.
            arrow_table = inject_species_meta(
                arrow_table,
                int(payload.get("species_index", 0)),
                recovered_species,
                species_charge=payload.get("species_charge"),
                species_mass=payload.get("species_mass"),
            )
        else:
            arrow_table = inject_species_meta(arrow_table, 0, "unknown")
        return particles_from_arrow(arrow_table)
    if len(species_values) > 1:
        names = sorted(species_values)
        msg = (
            f"query_sql matched {len(names)} species ({', '.join(names)}); "
            "ParticleData is a single-species container.  Filter with "
            "WHERE species='...' or use return_type='arrow'."
        )
        raise ValueError(msg)
    species_name = next(iter(species_values))
    payload = _lookup_species_meta(path, species_name)

    arrow_table = _strip_extra_columns(arrow_table)
    arrow_table = inject_species_meta(
        arrow_table,
        int(payload.get("species_index", 0)),
        species_name,
        species_charge=payload.get("species_charge"),
        species_mass=payload.get("species_mass"),
        metadata=payload.get("metadata"),
    )
    return particles_from_arrow(arrow_table)

icechunk_ancestry(path, *, branch='main')

Return the commit history of an Icechunk repository.

Parameters:

Name	Type	Description	Default
path	str or Path	Path to the Icechunk repository.	required
branch	str	Branch whose ancestry to inspect.	'main'

Returns:

Type	Description
list[dict[str, str]]	List of {"id": ..., "message": ...} dicts, most recent first.

Source code in src/pypic/io/_icechunk.py

def icechunk_ancestry(
    path: str | Path,
    *,
    branch: str = "main",
) -> list[dict[str, str]]:
    r"""Return the commit history of an Icechunk repository.

    Parameters
    ----------
    path : str or Path
        Path to the Icechunk repository.
    branch : str
        Branch whose ancestry to inspect.

    Returns
    -------
    list[dict[str, str]]
        List of ``{"id": ..., "message": ...}`` dicts, most recent
        first.
    """
    ensure_icechunk()

    repo = open_icechunk_repo(path)
    return [
        {"id": info.id, "message": info.message}
        for info in repo.ancestry(branch=branch)
    ]

icechunk_create_tag(path, tag, *, snapshot_id=None, branch='main')

Create a named tag in an Icechunk repository.

Parameters:

Name	Type	Description	Default
path	str or Path	Path to the Icechunk repository.	required
tag	str	Tag name (e.g. "v1.0-paper-submission").	required
snapshot_id	str or None	Snapshot to tag. Defaults to the tip of branch.	None
branch	str	Branch whose tip to tag (ignored when snapshot_id is given).	'main'

Source code in src/pypic/io/_icechunk.py

def icechunk_create_tag(
    path: str | Path,
    tag: str,
    *,
    snapshot_id: str | None = None,
    branch: str = "main",
) -> None:
    r"""Create a named tag in an Icechunk repository.

    Parameters
    ----------
    path : str or Path
        Path to the Icechunk repository.
    tag : str
        Tag name (e.g. ``"v1.0-paper-submission"``).
    snapshot_id : str or None
        Snapshot to tag.  Defaults to the tip of *branch*.
    branch : str
        Branch whose tip to tag (ignored when *snapshot_id* is given).
    """
    ensure_icechunk()

    repo = open_icechunk_repo(path)
    if snapshot_id is None:
        snapshot_id = repo.lookup_branch(branch)
    repo.create_tag(tag, snapshot_id)
    _log.info("Tagged snapshot %s as %r in %s", snapshot_id, tag, path)

open_icechunk_repo(path, *, create=False, authorize_virtual_chunk_access=None)

Open (or create) a local Icechunk repository.

Parameters:

Name	Type	Description	Default
path	str or Path	Directory for the repository.	required
create	bool	When True, create the repository if it does not exist.	False
authorize_virtual_chunk_access	dict or None	Mapping of URL prefix → credentials (None for unauthenticated local file:// URLs). When None (the default), the function auto-detects every VirtualChunkContainer registered with the repo and authorizes each prefix with None credentials — local virtual stores written by to_icechunk_virtual round-trip without further wiring. Pass an empty dict to disable virtual chunk reads, or supply explicit credentials for cloud (s3://, gs://) containers.	None

Returns:

Type	Description
Repository	The repository handle.

Source code in src/pypic/io/_icechunk.py

def open_icechunk_repo(
    path: str | Path,
    *,
    create: bool = False,
    authorize_virtual_chunk_access: dict[str, Any] | None = None,
) -> Any:  # noqa: ANN401
    r"""Open (or create) a local Icechunk repository.

    Parameters
    ----------
    path : str or Path
        Directory for the repository.
    create : bool
        When ``True``, create the repository if it does not exist.
    authorize_virtual_chunk_access : dict or None
        Mapping of URL prefix → credentials (``None`` for unauthenticated
        local ``file://`` URLs).  When ``None`` (the default), the
        function auto-detects every ``VirtualChunkContainer`` registered
        with the repo and authorizes each prefix with ``None``
        credentials — local virtual stores written by
        ``to_icechunk_virtual`` round-trip without further wiring.  Pass
        an empty dict to disable virtual chunk reads, or supply explicit
        credentials for cloud (``s3://``, ``gs://``) containers.

    Returns
    -------
    icechunk.Repository
        The repository handle.
    """
    ensure_icechunk()
    import icechunk

    storage = icechunk.local_filesystem_storage(str(path))

    if authorize_virtual_chunk_access is None:
        # Two-pass open: peek at the persisted RepositoryConfig to
        # discover registered containers, then re-open with auth that
        # matches their url_prefixes.  Containers store their full
        # prefix (e.g. file:///path/to/source/), which is what
        # Icechunk requires — a generic "file://" auth does not match.
        if create:
            probe = icechunk.Repository.open_or_create(storage)
        else:
            probe = icechunk.Repository.open(storage)
        containers = probe.config.virtual_chunk_containers or {}
        authorize_virtual_chunk_access = dict.fromkeys(containers)

    if create:
        return icechunk.Repository.open_or_create(
            storage,
            authorize_virtual_chunk_access=authorize_virtual_chunk_access,
        )
    return icechunk.Repository.open(
        storage,
        authorize_virtual_chunk_access=authorize_virtual_chunk_access,
    )

particles_from_dataset(path, *, step=None, species=None, spatial_box=None, ids=None, id_column='id', energy_min=None, columns=None)

Read from a partitioned Parquet dataset with selective loading.

Uses pyarrow.dataset with Hive partitioning for partition pruning (step, species) and Parquet row-group statistics for predicate pushdown (spatial box, energy threshold, IDs).

Parameters:

Name	Type	Description	Default
path	str or Path	Root directory of the partitioned dataset.	required
step	int or None	Timestep to load (partition pruning).	None
species	str or int or None	Species name or index (partition pruning).	None
spatial_box	tuple or None	((x_min, x_max), (y_min, y_max), (z_min, z_max)) for spatial filtering via predicate pushdown.	None
ids	array or Sequence or None	Values to filter on against id_column. Default id_column="id" matches the integer tracking column; pass id_column="weight" with float64 values to track iPIC3D non-uniform-plasma particles by their unique per-particle weight.	None
id_column	str	Column name to filter ids against. Default "id". For best pushdown effectiveness when id_column != "id", write the dataset with matching sort_by (e.g. sort_by="weight" for id_column="weight").	'id'
energy_min	float or None	Minimum speed |v| threshold for energy filtering.	None
columns	Sequence[str] or None	Column names to load (e.g. ["x", "y", "z"]). id_column is force-included when ids filtering is active.	None

Returns:

Type	Description
ParticleData

See Also

particles_to_dataset : The writer this reverses.

Source code in src/pypic/io/_parquet.py

def particles_from_dataset(
    path: str | Path,
    *,
    step: int | None = None,
    species: str | int | None = None,
    spatial_box: (
        tuple[tuple[float, float], tuple[float, float], tuple[float, float]] | None
    ) = None,
    ids: np.ndarray | Sequence[int] | Sequence[float] | None = None,
    id_column: str = "id",
    energy_min: float | None = None,
    columns: Sequence[str] | None = None,
) -> ParticleData:
    r"""Read from a partitioned Parquet dataset with selective loading.

    Uses ``pyarrow.dataset`` with Hive partitioning for partition
    pruning (step, species) and Parquet row-group statistics for
    predicate pushdown (spatial box, energy threshold, IDs).

    Parameters
    ----------
    path : str or Path
        Root directory of the partitioned dataset.
    step : int or None
        Timestep to load (partition pruning).
    species : str or int or None
        Species name or index (partition pruning).
    spatial_box : tuple or None
        ``((x_min, x_max), (y_min, y_max), (z_min, z_max))`` for
        spatial filtering via predicate pushdown.
    ids : array or Sequence or None
        Values to filter on against ``id_column``.  Default
        ``id_column="id"`` matches the integer tracking column; pass
        ``id_column="weight"`` with float64 values to track iPIC3D
        non-uniform-plasma particles by their unique per-particle
        weight.
    id_column : str
        Column name to filter ``ids`` against.  Default ``"id"``.
        For best pushdown effectiveness when ``id_column != "id"``,
        write the dataset with matching ``sort_by`` (e.g.
        ``sort_by="weight"`` for ``id_column="weight"``).
    energy_min : float or None
        Minimum speed ``|v|`` threshold for energy filtering.
    columns : Sequence[str] or None
        Column names to load (e.g. ``["x", "y", "z"]``).  ``id_column``
        is force-included when ``ids`` filtering is active.

    Returns
    -------
    ParticleData

    See Also
    --------
    particles_to_dataset : The writer this reverses.
    """
    ensure_arrow()
    # Declare partition schema explicitly so step is read as string, not int
    import pyarrow as pa
    import pyarrow.dataset as pads

    part_schema = pa.schema(
        [
            pa.field("step", pa.string()),
            pa.field("species", pa.string()),
        ]
    )
    partitioning = pads.HivePartitioning(part_schema)
    dataset = pads.dataset(str(path), format="parquet", partitioning=partitioning)

    # Partition filter
    part_filter: Any = None
    if step is not None:
        step_str = f"{step:06d}"
        part_filter = pads.field("step") == step_str
    pinned_species_str: str | None = None
    if species is not None:
        pinned_species_str = _resolve_species_str(species, Path(path))
        sp_filter = pads.field("species") == pinned_species_str
        part_filter = sp_filter if part_filter is None else part_filter & sp_filter

    # Row-level filter (predicate pushdown on Parquet statistics)
    row_filter: Any = None
    if spatial_box is not None:
        (x_min, x_max), (y_min, y_max), (z_min, z_max) = spatial_box
        row_filter = (
            (pads.field("x") >= x_min)
            & (pads.field("x") <= x_max)
            & (pads.field("y") >= y_min)
            & (pads.field("y") <= y_max)
            & (pads.field("z") >= z_min)
            & (pads.field("z") <= z_max)
        )
    if energy_min is not None:
        speed_filter = pads.field("speed") >= energy_min
        row_filter = speed_filter if row_filter is None else row_filter & speed_filter
    if ids is not None:
        id_list = list(ids) if not isinstance(ids, list) else ids
        id_filter = pads.field(id_column).isin(id_list)
        row_filter = id_filter if row_filter is None else row_filter & id_filter

    combined_filter: Any = None
    if part_filter is not None and row_filter is not None:
        combined_filter = part_filter & row_filter
    elif part_filter is not None:
        combined_filter = part_filter
    elif row_filter is not None:
        combined_filter = row_filter

    # Column pruning — only force-include the active id_column when filtering
    read_columns: list[str] | None = None
    if columns is not None:
        read_columns = list(columns)
        _require_vector_columns(read_columns)
        if ids is not None and id_column not in read_columns:
            read_columns.append(id_column)

    # Resolve species metadata from per-fragment Parquet schemas — the
    # only authoritative source for the original species_index after
    # Hive partitioning has flattened it to a directory name.  Raises
    # if the filter matched multiple species.
    payload = _matched_species_metadata(dataset, combined_filter)
    # Empty-filter recovery: a pinned species keeps its schema metadata on
    # disk, so adopt species identity instead of ``"unknown"``.  Per-step
    # ``metadata`` is dropped — an arbitrary fragment's time/tag stamped onto
    # a zero-row result is corruption.  Mirrors ``_duckdb.query_sql``.
    if payload.get("species_name") == "unknown" and pinned_species_str is not None:
        pinned_payload = _payload_from_species_dir(Path(path), pinned_species_str)
        if pinned_payload is not None:
            payload = {k: v for k, v in pinned_payload.items() if k != "metadata"}

    table = dataset.to_table(filter=combined_filter, columns=read_columns)
    table = _strip_extra_columns(table)
    table = inject_species_meta(
        table,
        int(payload.get("species_index", 0)),
        str(payload.get("species_name", "unknown")),
        species_charge=payload.get("species_charge"),
        species_mass=payload.get("species_mass"),
        metadata=payload.get("metadata"),
    )
    return particles_from_arrow(table)

particles_from_parquet(path)

Read a single Parquet file into a ParticleData.

Parameters:

Name	Type	Description	Default
path	str or Path	Path to the .parquet file.	required

Returns:

Type	Description
ParticleData

See Also

particles_to_parquet : The writer this reverses.

Source code in src/pypic/io/_parquet.py

def particles_from_parquet(path: str | Path) -> ParticleData:
    r"""Read a single Parquet file into a ``ParticleData``.

    Parameters
    ----------
    path : str or Path
        Path to the ``.parquet`` file.

    Returns
    -------
    ParticleData

    See Also
    --------
    particles_to_parquet : The writer this reverses.
    """
    ensure_arrow()
    import pyarrow.parquet as pq

    table = pq.read_table(str(path))
    table = _strip_extra_columns(table)
    return particles_from_arrow(table)

particles_to_dataset(source, path, *, steps=None, species=None, position_dtype=None, velocity_dtype=None, compression_level=1, row_group_size=_DEFAULT_ROW_GROUP_SIZE, sort_by='position')

Write a partitioned Parquet dataset from a Simulation or iterable.

Layout::

{path}/step=000000/species=electrons/part-00000.parquet
{path}/step=000000/species=ions/part-00000.parquet
{path}/step=000100/species=electrons/part-00000.parquet
...

When an iterable source yields multiple ParticleData chunks for the same (step, species) (streaming pipelines), each chunk is written to a fresh part-NNNNN.parquet within the partition directory rather than overwriting the previous one. Existing part-*.parquet files in the destination are preserved and the counter resumes from the next free index.

Parameters:

Name	Type	Description	Default
source	Simulation or Iterable[tuple[int, str, ParticleData]]	Either a Simulation (reads via source.particles) or an iterable of (step, species_name, ParticleData) tuples for custom HDF5→Parquet pipelines. When passing an iterable, steps and species must be None.	required
path	str or Path	Root directory for the partitioned dataset.	required
steps	Sequence[int] or None	Timestep indices to write (Simulation source only). Defaults to all particle steps.	None
species	Sequence[int | str] or None	Species indices or names (Simulation source only). Defaults to all species.	None
position_dtype	str or None	Downcast options.	None
velocity_dtype	str or None	Downcast options.	None
compression_level	int	zstd level (1 for processing, 3 for archival).	1
row_group_size	int	Target rows per row group.	_DEFAULT_ROW_GROUP_SIZE
sort_by	('position', 'weight')	Pre-write sort order; forwarded to particles_to_parquet.	"position"

Notes

Accepts either a Simulation or an iterable of (step, species, ParticleData) triples, so custom pipelines can feed it directly::

pairs = [(0, "electrons", pcl_e), (0, "ions", pcl_i)]
particles_to_dataset(pairs, out_dir)

See Also

particles_from_dataset : Reads back one partition.

Source code in src/pypic/io/_parquet.py

def particles_to_dataset(
    source: Simulation | Iterable[tuple[int, str, ParticleData]],
    path: str | Path,
    *,
    steps: Sequence[int] | None = None,
    species: Sequence[int | str] | None = None,
    position_dtype: str | None = None,
    velocity_dtype: str | None = None,
    compression_level: int = 1,
    row_group_size: int = _DEFAULT_ROW_GROUP_SIZE,
    sort_by: Literal["position", "weight"] = "position",
) -> None:
    r"""Write a partitioned Parquet dataset from a Simulation or iterable.

    Layout::

        {path}/step=000000/species=electrons/part-00000.parquet
        {path}/step=000000/species=ions/part-00000.parquet
        {path}/step=000100/species=electrons/part-00000.parquet
        ...

    When an iterable source yields multiple ``ParticleData`` chunks
    for the same ``(step, species)`` (streaming pipelines), each chunk
    is written to a fresh ``part-NNNNN.parquet`` within the partition
    directory rather than overwriting the previous one.  Existing
    ``part-*.parquet`` files in the destination are preserved and the
    counter resumes from the next free index.

    Parameters
    ----------
    source : Simulation or Iterable[tuple[int, str, ParticleData]]
        Either a ``Simulation`` (reads via ``source.particles``) or an
        iterable of ``(step, species_name, ParticleData)`` tuples for
        custom HDF5→Parquet pipelines.  When passing an iterable,
        ``steps`` and ``species`` must be ``None``.
    path : str or Path
        Root directory for the partitioned dataset.
    steps : Sequence[int] or None
        Timestep indices to write (Simulation source only).  Defaults
        to all particle steps.
    species : Sequence[int | str] or None
        Species indices or names (Simulation source only).  Defaults
        to all species.
    position_dtype, velocity_dtype : str or None
        Downcast options.
    compression_level : int
        zstd level (1 for processing, 3 for archival).
    row_group_size : int
        Target rows per row group.
    sort_by : {"position", "weight"}
        Pre-write sort order; forwarded to ``particles_to_parquet``.

    Notes
    -----
    Accepts either a `Simulation` or an iterable of
    ``(step, species, ParticleData)`` triples, so custom pipelines can
    feed it directly::

        pairs = [(0, "electrons", pcl_e), (0, "ions", pcl_i)]
        particles_to_dataset(pairs, out_dir)

    See Also
    --------
    particles_from_dataset : Reads back one partition.
    """
    ensure_arrow()
    root = Path(path)
    pairs = _resolve_particle_pairs(source, steps, species)

    # Enumerate within a single call so iterable sources may yield
    # multiple chunks per (step, species) without overwriting; the
    # initial counter is seeded from any existing part files in the
    # destination so concurrent re-runs append rather than clobber.
    part_counters: dict[tuple[int, str], int] = {}
    n_written = 0
    for step, sp_name, pcl in pairs:
        if pcl.n_particles == 0:
            _log.debug("Skipping empty step=%d species=%s", step, sp_name)
            continue
        part_dir = root / f"step={step:06d}" / f"species={sp_name}"
        part_dir.mkdir(parents=True, exist_ok=True)
        key = (step, sp_name)
        if key not in part_counters:
            indices = [
                int(m.group(1))
                for p in part_dir.iterdir()
                if (m := _PART_RE.match(p.name))
            ]
            # Counting files would skip into the wrong slot when prior
            # numbering is sparse (e.g. part-00001 without part-00000).
            part_counters[key] = (max(indices) + 1) if indices else 0
        part_idx = part_counters[key]
        part_counters[key] = part_idx + 1
        part_path = part_dir / f"part-{part_idx:05d}.parquet"
        particles_to_parquet(
            pcl,
            part_path,
            position_dtype=position_dtype,
            velocity_dtype=velocity_dtype,
            compression_level=compression_level,
            row_group_size=row_group_size,
            sort_by=sort_by,
        )
        n_written += 1
    _log.info("Wrote partitioned dataset (%d files) to %s", n_written, root)

particles_to_parquet(data, path, *, position_dtype=None, velocity_dtype=None, compression_level=1, row_group_size=_DEFAULT_ROW_GROUP_SIZE, sort_by='position')

Write a single ParticleData to a Parquet file.

Particles are sorted before writing so that Parquet row-group min/max statistics enable predicate pushdown on read. The default sort_by="position" uses a Morton Z-order curve, optimal for spatial box queries. Use sort_by="weight" when weight serves as the per-particle tracking ID (iPIC3D non-uniform plasma, where each macroparticle's float64 weight is unique); spatial queries become slower in return.

Parameters:

Name	Type	Description	Default
data	ParticleData	Source particle data (single species, single timestep).	required
path	str or Path	Destination .parquet file.	required
position_dtype	str or None	Downcast position columns (e.g. "float32").	None
velocity_dtype	str or None	Downcast velocity columns (e.g. "float32").	None
compression_level	int	zstd compression level (1 for processing, 3 for archival).	1
row_group_size	int	Target rows per row group (500K--1M recommended).	_DEFAULT_ROW_GROUP_SIZE
sort_by	('position', 'weight')	Pre-write sort order. "position" (default) Morton-sorts on x/y/z for spatial pushdown. "weight" ascending-sorts on the weight column for particle-tracking pushdown.	"position"

Examples:

>>> import numpy as np
>>> from pypic.containers import ParticleData
>>> pcl = ParticleData(
...     species_index=0, species_name="e",
...     position=np.zeros((5, 3)), velocity=np.ones((5, 3)),
...     n_particles=5, metadata={},
...     weight=np.ones(5), species_charge=-1.0, species_mass=1.0,
... )
>>> import tempfile
>>> from pathlib import Path
>>> with tempfile.TemporaryDirectory() as tmp:
...     out = Path(tmp) / "pcl.parquet"
...     particles_to_parquet(pcl, out)
...     particles_from_parquet(out).n_particles
5

Source code in src/pypic/io/_parquet.py

def particles_to_parquet(
    data: ParticleData,
    path: str | Path,
    *,
    position_dtype: str | None = None,
    velocity_dtype: str | None = None,
    compression_level: int = 1,
    row_group_size: int = _DEFAULT_ROW_GROUP_SIZE,
    sort_by: Literal["position", "weight"] = "position",
) -> None:
    r"""Write a single ``ParticleData`` to a Parquet file.

    Particles are sorted before writing so that Parquet row-group
    min/max statistics enable predicate pushdown on read.  The default
    ``sort_by="position"`` uses a Morton Z-order curve, optimal for
    spatial box queries.  Use ``sort_by="weight"`` when weight serves
    as the per-particle tracking ID (iPIC3D non-uniform plasma, where
    each macroparticle's float64 weight is unique); spatial queries
    become slower in return.

    Parameters
    ----------
    data : ParticleData
        Source particle data (single species, single timestep).
    path : str or Path
        Destination ``.parquet`` file.
    position_dtype : str or None
        Downcast position columns (e.g. ``"float32"``).
    velocity_dtype : str or None
        Downcast velocity columns (e.g. ``"float32"``).
    compression_level : int
        zstd compression level (1 for processing, 3 for archival).
    row_group_size : int
        Target rows per row group (500K--1M recommended).
    sort_by : {"position", "weight"}
        Pre-write sort order.  ``"position"`` (default) Morton-sorts on
        x/y/z for spatial pushdown.  ``"weight"`` ascending-sorts on the
        weight column for particle-tracking pushdown.

    Examples
    --------
    >>> import numpy as np
    >>> from pypic.containers import ParticleData
    >>> pcl = ParticleData(
    ...     species_index=0, species_name="e",
    ...     position=np.zeros((5, 3)), velocity=np.ones((5, 3)),
    ...     n_particles=5, metadata={},
    ...     weight=np.ones(5), species_charge=-1.0, species_mass=1.0,
    ... )
    >>> import tempfile
    >>> from pathlib import Path
    >>> with tempfile.TemporaryDirectory() as tmp:
    ...     out = Path(tmp) / "pcl.parquet"
    ...     particles_to_parquet(pcl, out)
    ...     particles_from_parquet(out).n_particles
    5
    """
    ensure_arrow()
    import pyarrow.parquet as pq

    table = particles_to_arrow(
        data, position_dtype=position_dtype, velocity_dtype=velocity_dtype
    )
    table = _add_speed_column(table)
    if sort_by == "position":
        table = _morton_sort_table(table)
    elif sort_by == "weight":
        table = _column_sort_table(table, "weight")
    else:
        msg = f"sort_by must be 'position' or 'weight', got {sort_by!r}"
        raise ValueError(msg)

    pq.write_table(
        table,
        str(path),
        compression="zstd",
        compression_level=compression_level,
        use_byte_stream_split=True,
        row_group_size=row_group_size,
    )
    _log.info(
        "Wrote %d particles (%s) to %s", data.n_particles, data.species_name, path
    )

open_virtual(path, *, fields_group='fields', drop_variables=None, config=None)

Create a virtual FieldDataset backed by HDF5 byte ranges.

Uses VirtualiZarr to extract byte-range metadata from an HDF5 file, producing a FieldDataset that lazily reads field data from the original file without copying.

For files following the pypic canonical HDF5 layout (schema.md Section 4), grid and normalization metadata are read automatically from the grid/ and normalization/ groups. For other layouts, pass a config with the required metadata.

Parameters:

Name	Type	Description	Default
path	str or Path	Path to the HDF5 file.	required
fields_group	str or None	HDF5 group containing field datasets. "fields" for the canonical layout. None reads from the root group.	'fields'
drop_variables	list[str] or None	HDF5 dataset names to exclude from the virtual view.	None
config	SimulationConfig or None	Explicit metadata. When provided, overrides any metadata found in the HDF5 file.	None

Returns:

Type	Description
FieldDataset	A lazy-loading dataset backed by virtual references to the original HDF5 file.

Raises:

Type	Description
ValueError	If grid metadata cannot be determined from the file or config.
ImportError	If virtualizarr or icechunk is not installed.

Notes

Virtual references are persisted via Icechunk's native Zarr v3 backend (in-memory store). Both virtualizarr and icechunk are installed by the zarr extra (pip install "pypic-plasma[zarr]").

Source code in src/pypic/io/_virtual.py

def open_virtual(
    path: str | Path,
    *,
    fields_group: str | None = "fields",
    drop_variables: list[str] | None = None,
    config: SimulationConfig | None = None,
) -> FieldDataset:
    r"""Create a virtual FieldDataset backed by HDF5 byte ranges.

    Uses VirtualiZarr to extract byte-range metadata from an HDF5
    file, producing a ``FieldDataset`` that lazily reads field data
    from the original file without copying.

    For files following the pypic canonical HDF5 layout (schema.md
    Section 4), grid and normalization metadata are read automatically
    from the ``grid/`` and ``normalization/`` groups.  For other
    layouts, pass a ``config`` with the required metadata.

    Parameters
    ----------
    path : str or Path
        Path to the HDF5 file.
    fields_group : str or None
        HDF5 group containing field datasets.  ``"fields"`` for the
        canonical layout.  ``None`` reads from the root group.
    drop_variables : list[str] or None
        HDF5 dataset names to exclude from the virtual view.
    config : SimulationConfig or None
        Explicit metadata.  When provided, overrides any metadata
        found in the HDF5 file.

    Returns
    -------
    FieldDataset
        A lazy-loading dataset backed by virtual references to the
        original HDF5 file.

    Raises
    ------
    ValueError
        If grid metadata cannot be determined from the file or config.
    ImportError
        If ``virtualizarr`` or ``icechunk`` is not installed.

    Notes
    -----
    Virtual references are persisted via Icechunk's native Zarr v3
    backend (in-memory store).  Both ``virtualizarr`` and ``icechunk``
    are installed by the ``zarr`` extra (``pip install "pypic-plasma[zarr]"``).
    """
    ensure_virtualizarr()
    ensure_icechunk()

    from pathlib import Path as _Path

    path_str = str(_Path(path).resolve())
    vds = _build_vds(path_str, fields_group, drop_variables)
    ds = _virtual_to_readable(vds, source_dir=str(_Path(path_str).parent))

    if config is not None:
        grid = config.grid
        normalization = config.normalization
        species = config.species
        physics = config.physics
        metadata: Mapping[str, Any] = config.metadata
        frame = config.frame
        transforms: Mapping[str, FrameTransform] | None = config.transforms
    else:
        h5_grid, h5_norm, h5_extra = _read_metadata_from_h5(path_str)
        if h5_grid is None:
            msg = (
                f"No 'grid' group found in {path_str} and no config provided. "
                "Pass config= with grid metadata, or use a canonical HDF5 layout."
            )
            raise ValueError(msg)
        grid = h5_grid
        normalization = h5_norm if h5_norm is not None else Normalization.undeclared()
        species = None
        physics = None
        metadata = h5_extra
        frame = "simulation"
        transforms = None

    dim_names = list(grid.surviving_axis_names)
    ds_dims = list(ds.dims)
    if len(ds_dims) != len(dim_names):
        msg = (
            f"Dimension count mismatch: HDF5 has {len(ds_dims)} dimensions "
            f"{ds_dims} but grid expects {len(dim_names)} {dim_names}."
        )
        raise ValueError(msg)
    rename_map = dict(zip(ds_dims, dim_names, strict=True))
    ds = ds.rename(rename_map)

    coord_arrays = grid.coordinate_arrays()
    coords = {dim_names[i]: coord_arrays[i] for i in range(len(dim_names))}
    ds = ds.assign_coords(coords)

    return FieldDataset(
        ds,
        grid,
        normalization,
        species=species,
        physics=physics,
        metadata=metadata,
        frame=frame,
        transforms=transforms,
    )

to_icechunk_virtual(source, output, *, fields_group='fields', drop_variables=None, config=None, message=None, branch='main')

Persist HDF5 byte-range references to an Icechunk repository.

Writes the virtual references for source's field datasets to a persistent on-disk Icechunk repo at output and attaches pypic metadata as group attrs. Subsequent reads via from_zarr(output) resolve chunks by reading byte ranges from source — no field data is copied.

Parameters:

Name	Type	Description	Default
source	str or Path	Path to the source HDF5 file.	required
output	str or Path	Destination directory for the Icechunk repository.	required
fields_group	str or None	HDF5 group containing field datasets ("fields" for the canonical layout; None reads from the root group).	'fields'
drop_variables	list[str] or None	HDF5 dataset names to exclude from the virtual view.	None
config	SimulationConfig or None	Explicit metadata. When provided, overrides any metadata found in the HDF5 file.	None
message	str or None	Icechunk commit message. Defaults to a generated string.	None
branch	str	Branch to commit to. Defaults to "main".	'main'

Returns:

Type	Description
str	The Icechunk snapshot ID of the new commit.

Raises:

Type	Description
ImportError	If virtualizarr or icechunk is not installed.
ValueError	If grid metadata cannot be determined from source or config.

Notes

Both virtualizarr and icechunk are installed by the zarr extra (pip install "pypic-plasma[zarr]"). Moving or deleting source after the write breaks the virtual refs in output — the on-disk repo is metadata only.

Source code in src/pypic/io/_virtual.py

def to_icechunk_virtual(
    source: str | Path,
    output: str | Path,
    *,
    fields_group: str | None = "fields",
    drop_variables: list[str] | None = None,
    config: SimulationConfig | None = None,
    message: str | None = None,
    branch: str = "main",
) -> str:
    r"""Persist HDF5 byte-range references to an Icechunk repository.

    Writes the virtual references for *source*'s field datasets to a
    persistent on-disk Icechunk repo at *output* and attaches pypic
    metadata as group attrs.  Subsequent reads via ``from_zarr(output)``
    resolve chunks by reading byte ranges from *source* — no field
    data is copied.

    Parameters
    ----------
    source : str or Path
        Path to the source HDF5 file.
    output : str or Path
        Destination directory for the Icechunk repository.
    fields_group : str or None
        HDF5 group containing field datasets (``"fields"`` for the
        canonical layout; ``None`` reads from the root group).
    drop_variables : list[str] or None
        HDF5 dataset names to exclude from the virtual view.
    config : SimulationConfig or None
        Explicit metadata.  When provided, overrides any metadata
        found in the HDF5 file.
    message : str or None
        Icechunk commit message.  Defaults to a generated string.
    branch : str
        Branch to commit to.  Defaults to ``"main"``.

    Returns
    -------
    str
        The Icechunk snapshot ID of the new commit.

    Raises
    ------
    ImportError
        If ``virtualizarr`` or ``icechunk`` is not installed.
    ValueError
        If grid metadata cannot be determined from *source* or *config*.

    Notes
    -----
    Both ``virtualizarr`` and ``icechunk`` are installed by the
    ``zarr`` extra (``pip install "pypic-plasma[zarr]"``).  Moving or deleting
    *source* after the write breaks the virtual refs in *output* — the
    on-disk repo is metadata only.
    """
    ensure_virtualizarr()
    ensure_icechunk()

    import shutil
    from pathlib import Path as _Path

    import icechunk
    import zarr

    source_path = _Path(source).resolve()
    source_str = str(source_path)
    source_dir = str(source_path.parent)
    output_path = _Path(output)
    # Decide *before* the mkdir whether the directory is ours to clean up.
    # Same gating as ``to_zarr``; here the realistic trigger is
    # ``open_virtual`` raising because the source HDF5 lacks ``grid/``.
    created_new = not output_path.exists() or (
        output_path.is_dir() and not any(output_path.iterdir())
    )
    output_path.mkdir(parents=True, exist_ok=True)

    try:
        vds = _build_vds(source_str, fields_group, drop_variables)

        storage = icechunk.local_filesystem_storage(str(output_path))
        repo_config, url_prefix = _virtual_repo_config(source_dir)

        # Merge the new virtual chunk container into any config already
        # persisted for this repo.  ``open_or_create`` accepts
        # ``config=`` but does not union it with previously-saved
        # containers — passing only the current source's container
        # would silently displace every prefix from earlier commits, so
        # later reads of those refs fail with "no virtual chunk
        # container can handle the chunk location".
        try:
            persisted = icechunk.Repository.fetch_config(storage)
        except icechunk.IcechunkError:
            # Fresh repo: fetch_config raises before any commit exists.
            persisted = None
        if persisted is not None:
            existing = persisted.virtual_chunk_containers or {}
            if url_prefix not in existing:
                persisted.set_virtual_chunk_container(
                    _make_virtual_chunk_container(url_prefix, source_dir)
                )
            repo_config = persisted

        all_prefixes: dict[str, Any] = dict.fromkeys(
            repo_config.virtual_chunk_containers or {}
        )
        all_prefixes.setdefault(url_prefix, None)

        repo = icechunk.Repository.open_or_create(
            storage,
            config=repo_config,
            authorize_virtual_chunk_access=all_prefixes,
        )
        # Persist the (possibly augmented) config so ``from_zarr`` —
        # which opens the repo without knowing which containers to
        # authorize — can auto-discover every prefix this repo has
        # ever written against.
        repo.save_config()
        _ensure_branch(repo, branch)
        session = repo.writable_session(branch)
        # Clear the session's working-tree root so virtualizarr's
        # to_icechunk can create a fresh root group.  Required for (a)
        # repeat commits to the same branch and (b) new branches
        # forked from a non-empty main — both inherit the prior root
        # group from the branch tip, and virtualizarr's
        # ``Group.from_store`` raises ContainsGroupError on any
        # pre-existing node.  Prior snapshots stay intact in repo
        # history; only this commit's root is replaced.
        session.store.sync_clear()
        # Write virtual refs under ``/fields`` to match the schema-v2.0
        # Zarr layout (schema.md §4.2): field arrays live under the
        # ``/fields`` child group, pypic metadata sits flat on the root
        # group's attrs.  ``from_zarr`` enforces both.
        vds.vz.to_icechunk(session.store, group="fields")

        # Reuse open_virtual to assemble the canonical FieldDataset
        # attrs.  This re-extracts vds against an in-memory store
        # (cheap — only HDF5 metadata is read) and gives
        # encode_pypic_attrs the same FieldDataset shape from_zarr
        # will reconstruct on read.
        fds = open_virtual(
            source,
            fields_group=fields_group,
            drop_variables=drop_variables,
            config=config,
        )
        group = zarr.open_group(session.store, mode="r+")
        for key, value in encode_pypic_attrs(fds).items():
            group.attrs[key] = value

        snapshot: str = session.commit(
            message if message is not None else "pypic: virtual refs"
        )
    except BaseException:
        if created_new:
            shutil.rmtree(output_path, ignore_errors=True)
        raise
    return snapshot

from_zarr(path, *, branch=None, tag=None, snapshot_id=None)

Read a FieldDataset from a Zarr v3 store.

Returns a lazy-loading dataset by default — field arrays are read from disk on first access. Icechunk repositories are auto-detected; pass branch, tag, or snapshot_id to read a specific version.

Parameters:

Name	Type	Description	Default
path	str or Path	Path to the Zarr store directory (or Icechunk repository).	required
branch	str or None	Icechunk branch to read from.	None
tag	str or None	Icechunk tag to read from.	None
snapshot_id	str or None	Icechunk snapshot ID to read from.	None

Returns:

Type	Description
FieldDataset	Reconstructed dataset with full metadata.

See Also

to_zarr : The writer this reverses.

Source code in src/pypic/io/zarr.py

def from_zarr(
    path: str | Path,
    *,
    branch: str | None = None,
    tag: str | None = None,
    snapshot_id: str | None = None,
) -> FieldDataset:
    r"""Read a FieldDataset from a Zarr v3 store.

    Returns a lazy-loading dataset by default — field arrays are read
    from disk on first access.  Icechunk repositories are auto-detected;
    pass *branch*, *tag*, or *snapshot_id* to read a specific version.

    Parameters
    ----------
    path : str or Path
        Path to the Zarr store directory (or Icechunk repository).
    branch : str or None
        Icechunk branch to read from.
    tag : str or None
        Icechunk tag to read from.
    snapshot_id : str or None
        Icechunk snapshot ID to read from.

    Returns
    -------
    FieldDataset
        Reconstructed dataset with full metadata.

    See Also
    --------
    to_zarr : The writer this reverses.
    """
    has_ref = any(x is not None for x in (branch, tag, snapshot_id))

    if has_ref:
        from pypic.io._icechunk import from_zarr_icechunk

        return from_zarr_icechunk(
            path,
            branch=branch,
            tag=tag,
            snapshot_id=snapshot_id,
        )

    from pypic.io._icechunk import is_icechunk_store

    if is_icechunk_store(path):
        from pypic.io._icechunk import from_zarr_icechunk

        return from_zarr_icechunk(path)

    ensure_zarr()
    ds, root_attrs = _open_store(str(path), f"Zarr store at {path}")
    return _ds_to_field_dataset(ds, root_attrs, f"Zarr store at {path}")

to_zarr(fds, path, *, dtype=None, encoding=None, backend=None, message=None, branch='main', simulation_toml=None)

Write a FieldDataset to a Zarr v3 store.

All pypic metadata (grid, normalization, species, physics, frame, transforms) is serialized as flat keys on the root group's attrs alongside a schema.version discriminator (see schema.md §4.2), so that from_zarr — and any non-pypic consumer — can reconstruct the full object straight from the store.

Parameters:

Name	Type	Description	Default
fds	FieldDataset	The dataset to write.	required
path	str or Path	Destination directory (created if it does not exist).	required
dtype	str or None	When set (e.g. "float32"), all field arrays are downcast to this dtype on write. Default: preserve source dtype.	None
encoding	dict or None	Per-variable encoding overrides merged on top of the defaults (Blosc zstd + bitshuffle). Keys are variable names, values are dicts passed to xr.Dataset.to_zarr(encoding=...).	None
backend	str or None	Storage backend. None (default) for plain Zarr v3, "icechunk" for versioned Icechunk storage.	None
message	str or None	Commit message (Icechunk only).	None
branch	str	Branch to commit to (Icechunk only). Default "main".	'main'
simulation_toml	str or Path or None	Source simulation.toml to stamp verbatim into attrs.simulation_toml. Carries the schema sections the typed FieldDataset boundary drops. None (default) stamps whatever the dataset already knows.	None

Returns:

Type	Description
str or None	Snapshot ID when backend="icechunk", None otherwise.

Examples:

>>> import numpy as np
>>> from pypic.dataset import FieldDataset
>>> from pypic.grid import GridInfo
>>> from pypic.units import Normalization
>>> grid = GridInfo(dimensions=(4, 3, 2), spacing=(1.0, 1.0, 1.0))
>>> fds = FieldDataset.from_arrays(
...     {"B_1": np.ones((4, 3, 2))}, grid, Normalization.identity(),
... )
>>> import tempfile
>>> from pathlib import Path
>>> with tempfile.TemporaryDirectory() as tmp:
...     store = Path(tmp) / "test.zarr"
...     to_zarr(fds, store)
...     sorted(from_zarr(store).field_names())
['B_1']

Source code in src/pypic/io/zarr.py

def to_zarr(
    fds: FieldDataset,
    path: str | Path,
    *,
    dtype: str | None = None,
    encoding: dict[str, dict[str, Any]] | None = None,
    backend: str | None = None,
    message: str | None = None,
    branch: str = "main",
    simulation_toml: str | Path | None = None,
) -> str | None:
    r"""Write a FieldDataset to a Zarr v3 store.

    All pypic metadata (grid, normalization, species, physics, frame,
    transforms) is serialized as flat keys on the root group's attrs
    alongside a ``schema.version`` discriminator (see schema.md §4.2),
    so that ``from_zarr`` — and any non-pypic consumer — can
    reconstruct the full object straight from the store.

    Parameters
    ----------
    fds : FieldDataset
        The dataset to write.
    path : str or Path
        Destination directory (created if it does not exist).
    dtype : str or None
        When set (e.g. ``"float32"``), all field arrays are downcast
        to this dtype on write.  Default: preserve source dtype.
    encoding : dict or None
        Per-variable encoding overrides merged on top of the defaults
        (Blosc zstd + bitshuffle).  Keys are variable names, values
        are dicts passed to ``xr.Dataset.to_zarr(encoding=...)``.
    backend : str or None
        Storage backend.  ``None`` (default) for plain Zarr v3,
        ``"icechunk"`` for versioned Icechunk storage.
    message : str or None
        Commit message (Icechunk only).
    branch : str
        Branch to commit to (Icechunk only).  Default ``"main"``.
    simulation_toml : str or Path or None
        Source ``simulation.toml`` to stamp verbatim into
        ``attrs.simulation_toml``. Carries the schema sections
        the typed FieldDataset boundary drops. ``None``
        (default) stamps whatever the dataset already knows.

    Returns
    -------
    str or None
        Snapshot ID when ``backend="icechunk"``, ``None`` otherwise.

    Examples
    --------
    >>> import numpy as np
    >>> from pypic.dataset import FieldDataset
    >>> from pypic.grid import GridInfo
    >>> from pypic.units import Normalization
    >>> grid = GridInfo(dimensions=(4, 3, 2), spacing=(1.0, 1.0, 1.0))
    >>> fds = FieldDataset.from_arrays(
    ...     {"B_1": np.ones((4, 3, 2))}, grid, Normalization.identity(),
    ... )
    >>> import tempfile
    >>> from pathlib import Path
    >>> with tempfile.TemporaryDirectory() as tmp:
    ...     store = Path(tmp) / "test.zarr"
    ...     to_zarr(fds, store)
    ...     sorted(from_zarr(store).field_names())
    ['B_1']
    """
    if backend == "icechunk":
        from pypic.io._icechunk import to_zarr_icechunk

        return to_zarr_icechunk(
            fds,
            path,
            dtype=dtype,
            encoding=encoding,
            message=message,
            branch=branch,
            simulation_toml=simulation_toml,
        )
    if backend is not None:
        msg = f"Unknown backend: {backend!r}. Use None or 'icechunk'."
        raise ValueError(msg)
    if message is not None:
        msg = "message= requires backend='icechunk'."
        raise ValueError(msg)

    ensure_zarr()
    import shutil
    from pathlib import Path as _Path

    # Cleanup gating.  ``DataTree.to_zarr(mode="w")`` materializes a partial
    # store before attribute serialization can reject, say, non-serializable
    # metadata; a later ``from_zarr`` then reports "No pypic metadata found"
    # instead of the real write error.  Only a path that was fresh or empty
    # when we started is ours to remove — a non-empty one is user data.
    path_obj = _Path(path)
    created_new = not path_obj.exists() or (
        path_obj.is_dir() and not any(path_obj.iterdir())
    )

    ds = fds.xr.copy(deep=False)
    tree = xr.DataTree.from_dict({"fields": ds})
    pypic_attrs = encode_pypic_attrs(fds)
    if simulation_toml is not None:
        pypic_attrs["simulation_toml"] = read_simulation_toml(simulation_toml)
    tree.attrs = pypic_attrs

    ds_encoding = _build_encoding(ds, dtype, encoding)
    try:
        tree.to_zarr(
            str(path),
            mode="w",
            consolidated=True,
            encoding=_datatree_encoding(ds_encoding),
            zarr_format=3,
        )
    except BaseException:
        if created_new:
            shutil.rmtree(path, ignore_errors=True)
        raise
    _log.info("Wrote %d fields to %s", len(ds.data_vars), path)
    return None

to_zarr_timeseries(source, path, *, steps=None, fields=None, dtype=None, encoding=None, backend=None, message=None, branch='main', simulation_toml=None)

Write multiple timesteps to a single Zarr v3 store.

Each field becomes (nt, n1, n2, n3) with time as the first dimension, chunked so that reading one timestep is O(1).

Parameters:

Name	Type	Description	Default
source	Simulation or Iterable[tuple[float | int, FieldDataset]]	Either a Simulation object (reads timesteps via source.read(step)) or an iterable of (time, fds) pairs.	required
path	str or Path	Destination Zarr store directory.	required
steps	Sequence[int] or None	Timestep indices to write (only used when source is a Simulation). Defaults to source.steps.	None
fields	Sequence[str] or None	Field names to include (only used when source is a Simulation). Defaults to all available fields.	None
dtype	str or None	Downcast dtype (e.g. "float32").	None
encoding	dict or None	Per-variable encoding overrides.	None
backend	str or None	Storage backend. None for plain Zarr v3, "icechunk" for versioned Icechunk storage.	None
message	str or None	Commit message (Icechunk only).	None
branch	str	Branch to commit to (Icechunk only). Default "main".	'main'
simulation_toml	str or Path or None	Source simulation.toml to stamp verbatim into attrs.simulation_toml. Carries the schema sections the typed FieldDataset boundary drops. None (default) stamps whatever the dataset already knows.	None

Returns:

Type	Description
str or None	Snapshot ID when backend="icechunk", None otherwise.

Examples:

>>> import numpy as np
>>> from pypic.dataset import FieldDataset
>>> from pypic.grid import GridInfo
>>> from pypic.units import Normalization
>>> grid = GridInfo(dimensions=(4, 3), spacing=(1.0, 1.0))
>>> pairs = [
...     (0.0, FieldDataset.from_arrays(
...         {"B_1": np.ones((4, 3))}, grid, Normalization.identity())),
...     (1.0, FieldDataset.from_arrays(
...         {"B_1": np.ones((4, 3)) * 2}, grid, Normalization.identity())),
... ]
>>> import tempfile
>>> from pathlib import Path
>>> with tempfile.TemporaryDirectory() as tmp:
...     store = Path(tmp) / "ts.zarr"
...     to_zarr_timeseries(pairs, store)
...     from_zarr(store).xr.sizes["time"]
2

Source code in src/pypic/io/zarr.py

def to_zarr_timeseries(
    source: Simulation | Iterable[tuple[float | int, FieldDataset]],
    path: str | Path,
    *,
    steps: Sequence[int] | None = None,
    fields: Sequence[str] | None = None,
    dtype: str | None = None,
    encoding: dict[str, dict[str, Any]] | None = None,
    backend: str | None = None,
    message: str | None = None,
    branch: str = "main",
    simulation_toml: str | Path | None = None,
) -> str | None:
    r"""Write multiple timesteps to a single Zarr v3 store.

    Each field becomes ``(nt, n1, n2, n3)`` with ``time`` as the first
    dimension, chunked so that reading one timestep is O(1).

    Parameters
    ----------
    source : Simulation or Iterable[tuple[float | int, FieldDataset]]
        Either a ``Simulation`` object (reads timesteps via
        ``source.read(step)``) or an iterable of ``(time, fds)``
        pairs.
    path : str or Path
        Destination Zarr store directory.
    steps : Sequence[int] or None
        Timestep indices to write (only used when *source* is a
        ``Simulation``).  Defaults to ``source.steps``.
    fields : Sequence[str] or None
        Field names to include (only used when *source* is a
        ``Simulation``).  Defaults to all available fields.
    dtype : str or None
        Downcast dtype (e.g. ``"float32"``).
    encoding : dict or None
        Per-variable encoding overrides.
    backend : str or None
        Storage backend.  ``None`` for plain Zarr v3,
        ``"icechunk"`` for versioned Icechunk storage.
    message : str or None
        Commit message (Icechunk only).
    branch : str
        Branch to commit to (Icechunk only).  Default ``"main"``.
    simulation_toml : str or Path or None
        Source ``simulation.toml`` to stamp verbatim into
        ``attrs.simulation_toml``. Carries the schema sections
        the typed FieldDataset boundary drops. ``None``
        (default) stamps whatever the dataset already knows.

    Returns
    -------
    str or None
        Snapshot ID when ``backend="icechunk"``, ``None`` otherwise.

    Examples
    --------
    >>> import numpy as np
    >>> from pypic.dataset import FieldDataset
    >>> from pypic.grid import GridInfo
    >>> from pypic.units import Normalization
    >>> grid = GridInfo(dimensions=(4, 3), spacing=(1.0, 1.0))
    >>> pairs = [
    ...     (0.0, FieldDataset.from_arrays(
    ...         {"B_1": np.ones((4, 3))}, grid, Normalization.identity())),
    ...     (1.0, FieldDataset.from_arrays(
    ...         {"B_1": np.ones((4, 3)) * 2}, grid, Normalization.identity())),
    ... ]
    >>> import tempfile
    >>> from pathlib import Path
    >>> with tempfile.TemporaryDirectory() as tmp:
    ...     store = Path(tmp) / "ts.zarr"
    ...     to_zarr_timeseries(pairs, store)
    ...     from_zarr(store).xr.sizes["time"]
    2
    """
    if backend == "icechunk":
        from pypic.io._icechunk import to_zarr_timeseries_icechunk

        return to_zarr_timeseries_icechunk(
            source,
            path,
            steps=steps,
            fields=fields,
            dtype=dtype,
            encoding=encoding,
            message=message,
            branch=branch,
            simulation_toml=simulation_toml,
        )
    if backend is not None:
        msg = f"Unknown backend: {backend!r}. Use None or 'icechunk'."
        raise ValueError(msg)
    if message is not None:
        msg = "message= requires backend='icechunk'."
        raise ValueError(msg)

    ensure_zarr()
    import shutil
    from pathlib import Path as _Path

    pairs = _resolve_timeseries_pairs(source, steps, fields)
    path_str = str(path)
    # Cleanup gating as in ``to_zarr``, including the case where xarray's
    # *first* ``ds.to_zarr(mode='w')`` fails partway and leaves a stub store
    # holding only ``zarr.json``.
    path_obj = _Path(path)
    created_new = not path_obj.exists() or (
        path_obj.is_dir() and not any(path_obj.iterdir())
    )
    try:
        pypic_attrs = _write_timeseries_steps(
            pairs, path_str, dtype=dtype, encoding=encoding
        )
    except BaseException:
        # Any failure inside the loop — including a first-step
        # materialization failure that leaves only ``zarr.json`` behind
        # — reclaims the store when we owned the directory.
        # BaseException covers KeyboardInterrupt as well: a Ctrl-C
        # between appends would otherwise leave the same orphan store.
        if created_new:
            shutil.rmtree(path_str, ignore_errors=True)
        raise

    if pypic_attrs is None:
        msg = "No timesteps to write — source yielded zero items."
        raise ValueError(msg)

    if simulation_toml is not None:
        pypic_attrs["simulation_toml"] = read_simulation_toml(simulation_toml)

    # Restamp the cross-step metadata intersection on the root group.
    # ``_write_timeseries_steps`` already wrote step-1's pypic_attrs
    # via the DataTree, but those carried step-1's metadata only.
    # After all appends, the running intersection (computed during the
    # loop) is the only set of values true at every timestep, so reset
    # the metadata key — and re-consolidate so the root attrs change
    # is visible to ``consolidated="auto"`` readers.
    import zarr

    root = zarr.open_group(path_str, mode="r+")
    root.attrs.update(pypic_attrs)
    zarr.consolidate_metadata(root.store)

    _log.info("Wrote timeseries to %s", path)
    return None