Add save_npz, save_dataframe, and read_metadata provenance writers

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01PEAxNzhDCaTrntX3o1JqMV

Author: wmayner

pyphi/provenance.py

@@ -22,11 +22,15 @@
 from datetime import UTC
 from datetime import datetime
 from pathlib import Path
+from typing import TYPE_CHECKING
 from typing import Any
 
 import numpy as np
 import scipy
 
+if TYPE_CHECKING:
+    import pandas as pd
+
 _PACKAGE_ROOT = Path(__file__).resolve().parent
 
 
@@ -269,6 +273,129 @@ def save_json(
     return path
 
 
+def _metadata_json(
+    prov: Provenance, params: Mapping[str, Any] | None
+) -> tuple[str, str]:
+    """Serialize the provenance record and params as JSON strings."""
+    return (
+        json.dumps(asdict(prov), default=_json_default),
+        json.dumps(dict(params or {}), default=_json_default),
+    )
+
+
+def save_npz(
+    arrays: Mapping[str, np.ndarray],
+    directory: Path | str,
+    name: str,
+    *,
+    params: Mapping[str, Any] | None = None,
+    run_label: str | None = None,
+    seed: int | None = None,
+    note: str | None = None,
+) -> Path:
+    """Write ``arrays`` to a self-describing, non-clobbering ``.npz`` file.
+
+    The arrays are stored with :func:`numpy.savez_compressed`, plus two
+    reserved entries ``_provenance`` and ``_params`` holding JSON strings
+    (read back with :func:`read_metadata`). Array names beginning with
+    ``_`` raise :class:`ValueError`. Filename, versioning, and seed
+    resolution follow :func:`save_json`.
+
+    Returns the written path.
+    """
+    reserved = [key for key in arrays if key.startswith("_")]
+    if reserved:
+        raise ValueError(f"array names beginning with '_' are reserved: {reserved}")
+    prov = _capture_metadata(params, seed, note)
+    prov_json, params_json = _metadata_json(prov, params)
+    path = unique_path(directory, format_stem(name, params, run_label), ".npz")
+    entries = {
+        **arrays,
+        "_provenance": np.array(prov_json),
+        "_params": np.array(params_json),
+    }
+    np.savez_compressed(path, allow_pickle=False, **entries)
+    return path
+
+
+def save_dataframe(
+    df: pd.DataFrame,
+    directory: Path | str,
+    name: str,
+    *,
+    params: Mapping[str, Any] | None = None,
+    run_label: str | None = None,
+    seed: int | None = None,
+    note: str | None = None,
+) -> Path:
+    """Write ``df`` to a self-describing, non-clobbering parquet file.
+
+    The frame is written with its index preserved; ``pyphi_provenance``
+    and ``pyphi_params`` entries (JSON strings) are merged into the
+    parquet schema metadata, so :func:`pandas.read_parquet` reads the
+    data normally and :func:`read_metadata` recovers the metadata.
+    Filename, versioning, and seed resolution follow :func:`save_json`.
+    DataFrame fidelity follows parquet semantics.
+
+    Returns the written path.
+    """
+    import pyarrow as pa
+    import pyarrow.parquet as pq
+
+    prov = _capture_metadata(params, seed, note)
+    prov_json, params_json = _metadata_json(prov, params)
+    path = unique_path(directory, format_stem(name, params, run_label), ".parquet")
+    table = pa.Table.from_pandas(df, preserve_index=True)
+    metadata = dict(table.schema.metadata or {})
+    metadata[b"pyphi_provenance"] = prov_json.encode()
+    metadata[b"pyphi_params"] = params_json.encode()
+    pq.write_table(table.replace_schema_metadata(metadata), path)
+    return path
+
+
+def read_metadata(path: Path | str) -> dict[str, Any]:
+    """Read the provenance and params embedded in a writer's output file.
+
+    Dispatches on the file suffix (``.json``, ``.npz``, or ``.parquet``)
+    and returns ``{"provenance": dict, "params": dict}``. A file without
+    the expected metadata (not produced by :func:`save_json`,
+    :func:`save_npz`, or :func:`save_dataframe`) raises
+    :class:`ValueError`, as does an unrecognized suffix.
+    """
+    path = Path(path)
+    missing = ValueError(f"no pyphi provenance metadata in {path}")
+    if path.suffix == ".json":
+        document = json.loads(path.read_text())
+        try:
+            return {
+                "provenance": document["provenance"],
+                "params": document["params"],
+            }
+        except (KeyError, TypeError):
+            raise missing from None
+    if path.suffix == ".npz":
+        with np.load(path) as npz:
+            try:
+                return {
+                    "provenance": json.loads(str(npz["_provenance"][()])),
+                    "params": json.loads(str(npz["_params"][()])),
+                }
+            except KeyError:
+                raise missing from None
+    if path.suffix == ".parquet":
+        import pyarrow.parquet as pq
+
+        metadata = pq.read_schema(path).metadata or {}
+        try:
+            return {
+                "provenance": json.loads(metadata[b"pyphi_provenance"]),
+                "params": json.loads(metadata[b"pyphi_params"]),
+            }
+        except KeyError:
+            raise missing from None
+    raise ValueError(f"unrecognized suffix {path.suffix!r} for {path}")
+
+
 def _set_provenance(result: Any, prov: Provenance) -> None:
     """Assign ``prov`` to ``result.provenance``, working around frozen results."""
     try:
@@ -319,7 +446,10 @@ def with_provenance(self, **fields: Any) -> HasProvenance:
     "HasProvenance",
     "Provenance",
     "format_stem",
+    "read_metadata",
+    "save_dataframe",
     "save_json",
+    "save_npz",
     "stamp_wall_time",
     "unique_path",
 ]

test/test_provenance_writers.py

@@ -3,6 +3,7 @@
 import json
 
 import numpy as np
+import pandas as pd
 import pytest
 
 from pyphi import provenance
@@ -84,3 +85,55 @@ def test_no_clobber(self, tmp_path):
         second = provenance.save_json({"x": 2}, tmp_path, "study", params={"seed": 1})
         assert second == tmp_path / "study_seed1_v2.json"
         assert json.loads(first.read_text())["data"] == {"x": 1}
+
+
+class TestSaveNpz:
+    def test_arrays_round_trip_with_metadata(self, tmp_path):
+        arrays = {"phis": np.linspace(0, 1, 5), "states": np.eye(2)}
+        path = provenance.save_npz(arrays, tmp_path, "study", params={"seed": 3})
+        assert path == tmp_path / "study_seed3.npz"
+        with np.load(path) as npz:
+            np.testing.assert_array_equal(npz["phis"], arrays["phis"])
+            np.testing.assert_array_equal(npz["states"], arrays["states"])
+        metadata = provenance.read_metadata(path)
+        assert metadata["params"] == {"seed": 3}
+        assert metadata["provenance"]["seed"] == 3
+
+    def test_reserved_names_rejected(self, tmp_path):
+        with pytest.raises(ValueError, match="reserved"):
+            provenance.save_npz({"_provenance": np.zeros(1)}, tmp_path, "study")
+
+
+class TestSaveDataframe:
+    def test_round_trip_and_metadata(self, tmp_path):
+        df = pd.DataFrame(
+            {"phi": [0.1, 0.2], "n": [3, 4]},
+            index=pd.Index(["a", "b"], name="system"),
+        )
+        path = provenance.save_dataframe(df, tmp_path, "study", params={"seed": 9})
+        assert path == tmp_path / "study_seed9.parquet"
+        pd.testing.assert_frame_equal(pd.read_parquet(path), df)
+        metadata = provenance.read_metadata(path)
+        assert metadata["provenance"]["seed"] == 9
+        assert metadata["params"] == {"seed": 9}
+
+
+class TestReadMetadata:
+    def test_json_metadata(self, tmp_path):
+        path = provenance.save_json({"x": 1}, tmp_path, "study", params={"seed": 5})
+        metadata = provenance.read_metadata(path)
+        assert metadata["params"] == {"seed": 5}
+        assert metadata["provenance"]["seed"] == 5
+        assert "pyphi_version" in metadata["provenance"]
+
+    def test_plain_json_rejected(self, tmp_path):
+        path = tmp_path / "plain.json"
+        path.write_text('{"x": 1}')
+        with pytest.raises(ValueError, match="no pyphi provenance"):
+            provenance.read_metadata(path)
+
+    def test_unknown_suffix_rejected(self, tmp_path):
+        path = tmp_path / "file.csv"
+        path.touch()
+        with pytest.raises(ValueError, match="unrecognized suffix"):
+            provenance.read_metadata(path)