Implementing stable 1.0.0 geff release (#167)

* working version with vlen mask saving

* updating inner type computation

Author: JoOkuma

pyproject.toml

@@ -48,7 +48,7 @@ dependencies = [
   "ilpy >= 0.5.1",
   "pyarrow",
   "bidict>=0.23.1",
-  "geff @ git+https://github.com/live-image-tracking-tools/geff.git@b751718f81d107e1fdda2df2afb62253039c137b",
+  "geff>=1.0.0",
 ]
 
 [project.optional-dependencies]

src/tracksdata/graph/_base_graph.py

@@ -5,19 +5,18 @@
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Literal, TypeVar, overload
 
+import geff
 import numpy as np
 import polars as pl
 import rustworkx as rx
-from geff import GeffMetadata
-from geff.metadata_schema import Axis
-from geff.rustworkx.io import read_rx
-from geff.write_arrays import write_arrays
+from geff.core_io import construct_var_len_props, write_arrays
+from geff_spec import Axis, PropMetadata
 from numpy.typing import ArrayLike
 from zarr.storage import StoreLike
 
 from tracksdata.attrs import AttrComparison, NodeAttr
 from tracksdata.constants import DEFAULT_ATTR_KEYS
-from tracksdata.utils._dtypes import column_to_bytes, column_to_numpy
+from tracksdata.utils._dtypes import column_to_numpy, polars_dtype_to_numpy_dtype
 from tracksdata.utils._logging import LOG
 from tracksdata.utils._multiprocessing import multiprocessing_apply
 
@@ -1184,7 +1183,7 @@ def from_geff(
         from tracksdata.graph import IndexedRXGraph
 
         # this performs a roundtrip with the rustworkx graph
-        rx_graph, _ = read_rx(geff_store)
+        rx_graph, _ = geff.read(geff_store, backend="rustworkx")
 
         if not isinstance(rx_graph, rx.PyDiGraph):
             LOG.warning("The graph is not a directed graph, converting to directed graph.")
@@ -1196,6 +1195,16 @@ def from_geff(
             **kwargs,
         )
 
+        if DEFAULT_ATTR_KEYS.MASK in indexed_graph.node_attr_keys:
+            from tracksdata.nodes._mask import Mask
+
+            # unsafe operation, changing graph content inplace
+            for node_attr in indexed_graph.rx_graph.nodes():
+                node_attr[DEFAULT_ATTR_KEYS.MASK] = Mask(
+                    node_attr[DEFAULT_ATTR_KEYS.MASK].astype(bool),
+                    bbox=node_attr[DEFAULT_ATTR_KEYS.BBOX],
+                )
+
         if cls == IndexedRXGraph:
             return indexed_graph
 
@@ -1204,7 +1213,7 @@ def from_geff(
     def to_geff(
         self,
         geff_store: StoreLike,
-        geff_metadata: GeffMetadata | None = None,
+        geff_metadata: geff.GeffMetadata | None = None,
         zarr_format: Literal[2, 3] = 3,
     ) -> None:
         """
@@ -1225,9 +1234,12 @@ def to_geff(
         """
 
         node_attrs = self.node_attrs()
+        node_ids = node_attrs[DEFAULT_ATTR_KEYS.NODE_ID].to_numpy()
+        node_attrs = node_attrs.drop(DEFAULT_ATTR_KEYS.NODE_ID)
 
-        if DEFAULT_ATTR_KEYS.MASK in node_attrs.columns:
-            node_attrs = column_to_bytes(node_attrs, DEFAULT_ATTR_KEYS.MASK)
+        edge_attrs = self.edge_attrs().drop(DEFAULT_ATTR_KEYS.EDGE_ID)
+        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)
 
         if geff_metadata is None:
             axes = [Axis(name=DEFAULT_ATTR_KEYS.T, type="time")]
@@ -1240,37 +1252,47 @@ def to_geff(
             else:
                 track_node_props = None
 
-            geff_metadata = GeffMetadata(
+            node_props_metadata = {
+                k: PropMetadata(
+                    identifier=k,
+                    dtype=polars_dtype_to_numpy_dtype(v.dtype) if k != DEFAULT_ATTR_KEYS.MASK else np.uint64,
+                    varlength=k == DEFAULT_ATTR_KEYS.MASK,
+                )
+                for k, v in node_attrs.to_dict().items()
+            }
+            edge_props_metadata = {
+                k: PropMetadata(identifier=k, dtype=polars_dtype_to_numpy_dtype(v.dtype))
+                for k, v in edge_attrs.to_dict().items()
+            }
+
+            geff_metadata = geff.GeffMetadata(
                 directed=True,
                 axes=axes,
+                node_props_metadata=node_props_metadata,
+                edge_props_metadata=edge_props_metadata,
                 track_node_props=track_node_props,
             )
 
-        edge_attrs = self.edge_attrs().drop(DEFAULT_ATTR_KEYS.EDGE_ID)
-
         node_dict = {
-            k: (column_to_numpy(v), None) for k, v in node_attrs.drop(DEFAULT_ATTR_KEYS.NODE_ID).to_dict().items()
+            k: {"values": column_to_numpy(v), "missing": None}
+            for k, v in node_attrs.to_dict().items()
+            if k != DEFAULT_ATTR_KEYS.MASK
         }
 
-        edge_ids = edge_attrs.select(DEFAULT_ATTR_KEYS.EDGE_SOURCE, DEFAULT_ATTR_KEYS.EDGE_TARGET).to_numpy()
-
-        edge_dict = {
-            k: (column_to_numpy(v), None)
-            for k, v in edge_attrs.drop(
-                DEFAULT_ATTR_KEYS.EDGE_SOURCE,
-                DEFAULT_ATTR_KEYS.EDGE_TARGET,
+        if DEFAULT_ATTR_KEYS.MASK in node_attrs.columns:
+            node_dict[DEFAULT_ATTR_KEYS.MASK] = construct_var_len_props(
+                [mask.mask.astype(np.uint64) for mask in node_attrs[DEFAULT_ATTR_KEYS.MASK]]
             )
-            .to_dict()
-            .items()
-        }
+
+        edge_dict = {k: {"values": column_to_numpy(v), "missing": None} for k, v in edge_attrs.to_dict().items()}
 
         write_arrays(
             geff_store,
-            metadata=geff_metadata,
-            node_ids=node_attrs[DEFAULT_ATTR_KEYS.NODE_ID].to_numpy(),
+            node_ids=node_ids.astype(np.uint64),
             node_props=node_dict,
-            edge_ids=edge_ids,
+            edge_ids=edge_ids.astype(np.uint64),
             edge_props=edge_dict,
+            metadata=geff_metadata,
             zarr_format=zarr_format,
         )
 

src/tracksdata/graph/_test/test_graph_backends.py

@@ -8,7 +8,7 @@
 
 from tracksdata.attrs import EdgeAttr, NodeAttr
 from tracksdata.constants import DEFAULT_ATTR_KEYS
-from tracksdata.graph import BaseGraph, RustWorkXGraph, SQLGraph
+from tracksdata.graph import BaseGraph, IndexedRXGraph, RustWorkXGraph, SQLGraph
 from tracksdata.io._numpy_array import from_array
 from tracksdata.nodes._mask import Mask
 
@@ -1779,7 +1779,7 @@ def test_geff_roundtrip(graph_backend: BaseGraph) -> None:
 
     graph_backend.to_geff(geff_store=output_store)
 
-    geff_graph = RustWorkXGraph.from_geff(output_store)
+    geff_graph = IndexedRXGraph.from_geff(output_store)
 
     assert geff_graph.num_nodes == 3
     assert geff_graph.num_edges == 2
@@ -1789,6 +1789,9 @@ def test_geff_roundtrip(graph_backend: BaseGraph) -> None:
     assert set(graph_backend.node_attr_keys) == set(geff_graph.node_attr_keys)
     assert set(graph_backend.edge_attr_keys) == set(geff_graph.edge_attr_keys)
 
+    for node_id in geff_graph.node_ids():
+        assert geff_graph[node_id].to_dict() == graph_backend[node_id].to_dict()
+
     assert rx.is_isomorphic(
         rx_graph,
         geff_graph.rx_graph,

src/tracksdata/nodes/_mask.py

@@ -1,5 +1,6 @@
 from collections.abc import Sequence
 from functools import cached_property, lru_cache
+from typing import Any
 
 import blosc2
 import numpy as np
@@ -278,6 +279,11 @@ def from_coordinates(
 
         return cls(mask, bbox)
 
+    def __eq__(self, other: Any) -> bool:
+        if not isinstance(other, Mask):
+            return False
+        return np.array_equal(self.bbox, other.bbox) and np.array_equal(self.mask, other.mask)
+
 
 class MaskDiskAttrs(GenericFuncNodeAttrs):
     """

src/tracksdata/utils/_dtypes.py

@@ -48,6 +48,9 @@ def polars_dtype_to_numpy_dtype(polars_dtype: DataType) -> np.dtype:
     np.dtype
         The numpy dtype.
     """
+    if isinstance(polars_dtype, pl.Array | pl.List):
+        polars_dtype = polars_dtype.inner
+
     try:
         return _POLARS_DTYPE_TO_NUMPY_DTYPE[polars_dtype]
     except KeyError as e: