diff --git a/src/tracksdata/graph/_base_graph.py b/src/tracksdata/graph/_base_graph.py index 3a214e35..51f7bf6f 100644 --- a/src/tracksdata/graph/_base_graph.py +++ b/src/tracksdata/graph/_base_graph.py @@ -1868,6 +1868,14 @@ def to_geff( It automatically generates the metadata with: - axes: time (t) and spatial axes ((z), y, x) - tracklet node property: tracklet_id + The graph metadata (`graph.metadata`) is always written to + `geff_metadata.extra["tracksdata"]`, including when the metadata is provided + by the caller. On key collisions the caller's value wins, so an explicit + `extra["tracksdata"]` entry still overrides the graph's metadata. + The caller's object is not modified. + `shape` is the canonical key for the shape of the dense segmentation, it is + read back by `GraphArrayView` and `to_ctc`. Use + `tracksdata.io.read_graph_metadata` to read it back without building a graph. overwrite : bool Whether to overwrite the geff data directory if it exists. zarr_format : Literal[2, 3] @@ -1883,6 +1891,10 @@ def to_geff( edge_ids = edge_attrs.select(DEFAULT_ATTR_KEYS.EDGE_SOURCE, DEFAULT_ATTR_KEYS.EDGE_TARGET).to_numpy() edge_attrs = edge_attrs.drop(DEFAULT_ATTR_KEYS.EDGE_SOURCE, DEFAULT_ATTR_KEYS.EDGE_TARGET) + td_metadata = self.metadata.copy() + td_metadata.update(self._private_metadata_for_copy()) + td_metadata.pop("geff", None) # avoid geff being written multiple times + if geff_metadata is None: axes = [Axis(name=DEFAULT_ATTR_KEYS.T, type="time")] axes.extend( @@ -1917,10 +1929,6 @@ def to_geff( for k, v in edge_attrs.to_dict().items() } - td_metadata = self.metadata.copy() - td_metadata.update(self._private_metadata_for_copy()) - td_metadata.pop("geff", None) # avoid geff being written multiple times - geff_metadata = geff.GeffMetadata( directed=True, axes=axes, @@ -1931,6 +1939,15 @@ def to_geff( "tracksdata": td_metadata, }, ) + else: + # copy so the caller's metadata object is left untouched + geff_metadata = geff_metadata.model_copy(deep=True) + extra = dict(geff_metadata.extra) + # caller-provided entries win, so they can still override the graph's metadata + merged = {**td_metadata, **extra.get("tracksdata", {})} + merged.pop("geff", None) # avoid geff being written multiple times + extra["tracksdata"] = merged + geff_metadata.extra = extra node_dict = { k: {"values": column_to_numpy(v), "missing": None} diff --git a/src/tracksdata/graph/_test/test_graph_backends.py b/src/tracksdata/graph/_test/test_graph_backends.py index 3c8095c6..48b2acea 100644 --- a/src/tracksdata/graph/_test/test_graph_backends.py +++ b/src/tracksdata/graph/_test/test_graph_backends.py @@ -9,6 +9,7 @@ import pytest import rustworkx as rx import sqlalchemy as sa +from geff_spec import GeffMetadata from zarr.storage import MemoryStore from tracksdata.attrs import EdgeAttr, NodeAttr @@ -2861,6 +2862,55 @@ def test_geff_roundtrip(graph_backend: BaseGraph) -> None: ) +def test_geff_roundtrip_custom_metadata(graph_backend: BaseGraph) -> None: + """Graph metadata must survive `to_geff` when the caller supplies its own `GeffMetadata`.""" + + _fill_mock_geff_graph(graph_backend) + graph_backend.metadata["shape"] = (5, 100, 100) + + # a downstream library supplying its own metadata: same props as the auto-generated + # one, but with its own `extra` namespace instead of tracksdata's. + reference_store = MemoryStore() + graph_backend.to_geff(geff_store=reference_store) + custom_metadata = GeffMetadata.read(reference_store) + custom_metadata.extra = {"downstream": {"hello": "world"}} + + output_store = MemoryStore() + graph_backend.to_geff(geff_store=output_store, geff_metadata=custom_metadata) + + written_metadata = GeffMetadata.read(output_store) + # the caller's own namespace is untouched ... + assert written_metadata.extra["downstream"] == {"hello": "world"} + # ... and the graph metadata rode along. `shape` is a list, not a tuple, because + # the extras are serialized as JSON. + assert written_metadata.extra["tracksdata"]["shape"] == [5, 100, 100] + + geff_graph, _ = IndexedRXGraph.from_geff(output_store) + assert geff_graph.metadata["shape"] == [5, 100, 100] + + # the metadata object the caller passed in was not modified + assert custom_metadata.extra == {"downstream": {"hello": "world"}} + + +def test_geff_custom_metadata_overrides_graph_metadata(graph_backend: BaseGraph) -> None: + """On key collisions the caller-supplied `extra["tracksdata"]` wins.""" + + _fill_mock_geff_graph(graph_backend) + graph_backend.metadata["shape"] = (5, 100, 100) + + reference_store = MemoryStore() + graph_backend.to_geff(geff_store=reference_store) + custom_metadata = GeffMetadata.read(reference_store) + custom_metadata.extra = {"tracksdata": {"shape": [1, 2, 3], "extra_key": "value"}} + + output_store = MemoryStore() + graph_backend.to_geff(geff_store=output_store, geff_metadata=custom_metadata) + + geff_graph, _ = IndexedRXGraph.from_geff(output_store) + assert geff_graph.metadata["shape"] == [1, 2, 3] + assert geff_graph.metadata["extra_key"] == "value" + + def test_geff_overwrite(graph_backend: BaseGraph, tmp_path: Path) -> None: """Test that to_geff overwrites existing data in the store.""" _fill_mock_geff_graph(graph_backend) diff --git a/src/tracksdata/io/__init__.py b/src/tracksdata/io/__init__.py index 57313e1a..e502138a 100644 --- a/src/tracksdata/io/__init__.py +++ b/src/tracksdata/io/__init__.py @@ -1,12 +1,17 @@ """Input/output utilities for loading and saving tracking data in various formats.""" from tracksdata.io._ctc import compressed_tracks_table, from_ctc, to_ctc -from tracksdata.io._geff_dtypes import convert_geff_prop_dtype, geff_prop_dtype +from tracksdata.io._geff import ( + convert_geff_prop_dtype, + geff_prop_dtype, + read_graph_metadata, +) __all__ = [ "compressed_tracks_table", "convert_geff_prop_dtype", "from_ctc", "geff_prop_dtype", + "read_graph_metadata", "to_ctc", ] diff --git a/src/tracksdata/io/_geff_dtypes.py b/src/tracksdata/io/_geff.py similarity index 72% rename from src/tracksdata/io/_geff_dtypes.py rename to src/tracksdata/io/_geff.py index b29a4e24..ae954e16 100644 --- a/src/tracksdata/io/_geff_dtypes.py +++ b/src/tracksdata/io/_geff.py @@ -1,16 +1,24 @@ -"""Utilities for inspecting and converting the dtype of properties in geff files. - -The motivating case is segmentation masks: they are binary, but geff files -written by older versions of tracksdata stored the mask ``data`` buffer as -``uint64`` (see https://github.com/royerlab/tracksdata/pull/318). That is 8x -larger than a boolean buffer both on disk and, more importantly, when read into -memory, which can cause out-of-memory errors when loading large datasets. New -files store masks as ``bool`` at write time, so :func:`convert_geff_prop_dtype` -provides a one-time fix for legacy files. - -The helpers are not mask-specific: they read and rewrite the payload dtype of -any geff property (node or edge, fixed- or variable-length). The caller names -the property to act on. +"""Utilities for inspecting and repairing geff datasets without loading the graph. + +:func:`read_graph_metadata` reads the tracksdata graph metadata of a geff dataset. +`BaseGraph.to_geff` writes the graph metadata (`graph.metadata`) into the geff +metadata extras and `BaseGraph.from_geff` hoists it back onto the graph, but callers +that need a value *before* they have a graph object -- for example the ``shape`` of +the dense segmentation, which is required to construct a `GraphArrayView` -- can use +neither. This closes that gap so downstream libraries do not have to know where +tracksdata stores the extras. + +:func:`geff_prop_dtype` and :func:`convert_geff_prop_dtype` inspect and convert the +on-disk dtype of a property. The motivating case is segmentation masks: they are +binary, but geff files written by older versions of tracksdata stored the mask +``data`` buffer as ``uint64`` (see +https://github.com/royerlab/tracksdata/pull/318). That is 8x larger than a boolean +buffer both on disk and, more importantly, when read into memory, which can cause +out-of-memory errors when loading large datasets. New files store masks as ``bool`` +at write time, so :func:`convert_geff_prop_dtype` provides a one-time fix for legacy +files. The helpers are not mask-specific: they read and rewrite the payload dtype of +any geff property (node or edge, fixed- or variable-length). The caller names the +property to act on. """ from __future__ import annotations @@ -18,14 +26,54 @@ import os import shutil from pathlib import Path +from typing import Any import numpy as np import zarr +from geff_spec import GeffMetadata from zarr.storage import StoreLike +from tracksdata.graph._base_graph import BaseGraph from tracksdata.utils._logging import LOG -__all__ = ["convert_geff_prop_dtype", "geff_prop_dtype"] +__all__ = ["convert_geff_prop_dtype", "geff_prop_dtype", "read_graph_metadata"] + +_EXTRA_KEY = "tracksdata" + + +def read_graph_metadata(source: StoreLike | GeffMetadata) -> dict[str, Any]: + """ + Read the tracksdata graph metadata of a geff dataset without loading the graph. + + Returns the same metadata that `graph.metadata` would hold after + `BaseGraph.from_geff`, minus the `geff` key. Note that the values went through a + JSON round-trip, so tuples come back as lists. + + Parameters + ---------- + source : StoreLike | GeffMetadata + The store or path of the geff dataset, or an already parsed `GeffMetadata`. + + Returns + ------- + dict[str, Any] + The graph metadata, empty if the dataset was not written by tracksdata. + + Examples + -------- + ```python + graph.metadata["shape"] = (5, 100, 100) + graph.to_geff("tracks.geff") + + shape = read_graph_metadata("tracks.geff")["shape"] # [5, 100, 100] + ``` + """ + if not isinstance(source, GeffMetadata): + source = GeffMetadata.read(source) + + metadata = source.extra.get(_EXTRA_KEY, {}) + + return {k: v for k, v in metadata.items() if not BaseGraph._is_private_metadata_key(k)} def geff_prop_dtype( diff --git a/src/tracksdata/io/_test/test_geff_dtypes.py b/src/tracksdata/io/_test/test_geff_dtypes.py index 18733028..9f7354f1 100644 --- a/src/tracksdata/io/_test/test_geff_dtypes.py +++ b/src/tracksdata/io/_test/test_geff_dtypes.py @@ -8,7 +8,7 @@ from tracksdata.constants import DEFAULT_ATTR_KEYS from tracksdata.graph import IndexedRXGraph, RustWorkXGraph from tracksdata.io import convert_geff_prop_dtype, geff_prop_dtype -from tracksdata.io._geff_dtypes import _overwrite_array, _set_prop_metadata_dtype +from tracksdata.io._geff import _overwrite_array, _set_prop_metadata_dtype from tracksdata.nodes._mask import Mask MASK_KEY = DEFAULT_ATTR_KEYS.MASK diff --git a/src/tracksdata/io/_test/test_geff_metadata.py b/src/tracksdata/io/_test/test_geff_metadata.py new file mode 100644 index 00000000..2c577fc4 --- /dev/null +++ b/src/tracksdata/io/_test/test_geff_metadata.py @@ -0,0 +1,89 @@ +from pathlib import Path + +import polars as pl +from geff_spec import Axis, GeffMetadata, PropMetadata +from zarr.storage import MemoryStore + +from tracksdata.graph import RustWorkXGraph +from tracksdata.io import read_graph_metadata + +SHAPE = (5, 100, 100) + + +def _make_graph() -> RustWorkXGraph: + graph = RustWorkXGraph() + graph.add_node_attr_key("y", pl.Float64) + graph.add_node_attr_key("x", pl.Float64) + graph.add_node({"t": 0, "y": 1.0, "x": 2.0}) + graph.add_node({"t": 1, "y": 3.0, "x": 4.0}) + graph.metadata["shape"] = SHAPE + return graph + + +def _minimal_geff_metadata() -> GeffMetadata: + """A `GeffMetadata` as a downstream library would build it: no tracksdata extras.""" + return GeffMetadata( + directed=True, + axes=[ + Axis(name="t", type="time"), + Axis(name="y", type="space", scale=0.5), + Axis(name="x", type="space", scale=0.5), + ], + node_props_metadata={ + "t": PropMetadata(identifier="t", dtype="int64"), + "y": PropMetadata(identifier="y", dtype="float64"), + "x": PropMetadata(identifier="x", dtype="float64"), + }, + edge_props_metadata={}, + extra={"downstream": {"hello": "world"}}, + ) + + +def test_read_graph_metadata_from_path(tmp_path: Path) -> None: + """The shape is readable from a store path, without building a graph.""" + graph = _make_graph() + geff_path = tmp_path / "tracks.geff" + graph.to_geff(geff_store=geff_path) + + # tuples become lists through the JSON round-trip + assert read_graph_metadata(geff_path) == {"shape": list(SHAPE)} + + +def test_read_graph_metadata_custom_geff_metadata() -> None: + """The shape survives a write with caller-supplied metadata and is readable back.""" + graph = _make_graph() + store = MemoryStore() + graph.to_geff(geff_store=store, geff_metadata=_minimal_geff_metadata()) + + assert read_graph_metadata(store) == {"shape": list(SHAPE)} + + +def test_read_graph_metadata_from_geff_metadata_instance() -> None: + """An already parsed `GeffMetadata` is accepted, so the store is not reopened.""" + graph = _make_graph() + store = MemoryStore() + graph.to_geff(geff_store=store) + + assert read_graph_metadata(GeffMetadata.read(store)) == {"shape": list(SHAPE)} + + +def test_read_graph_metadata_without_tracksdata_extras() -> None: + """A geff not written by tracksdata yields an empty dict rather than raising.""" + # foreign extras only + assert read_graph_metadata(_minimal_geff_metadata()) == {} + + no_extra = _minimal_geff_metadata() + no_extra.extra = {} + assert read_graph_metadata(no_extra) == {} + + +def test_read_graph_metadata_excludes_private_keys() -> None: + """Private metadata is written to the store but not exposed by the reader.""" + graph = _make_graph() + graph._private_metadata["__private_secret"] = 42 + + store = MemoryStore() + graph.to_geff(geff_store=store) + + assert "__private_secret" in GeffMetadata.read(store).extra["tracksdata"] + assert read_graph_metadata(store) == {"shape": list(SHAPE)}