Add geff tracking format (#221)

* Add geff tracking format export

* Update ultrack/core/export/exporter.py

Co-authored-by: Jordão Bragantini <jordao.bragantini@czbiohub.org>

* Apply suggestions from code review

Co-authored-by: Jordão Bragantini <jordao.bragantini@czbiohub.org>

* Remove redudant geff tests

---------

Co-authored-by: Jordão Bragantini <jordao.bragantini@czbiohub.org>

Author: ilan-theodoro

pyproject.toml

@@ -50,6 +50,7 @@ dependencies = [
     "websocket >=0.2.1",
     "websockets >=12.0",
     "zarr >=2.15.0,<3.0.0",
+    "geff>=0.2.1,<0.3",
 ]
 
 [project.optional-dependencies]

ultrack/core/_test/test_tracker.py

@@ -1,8 +1,11 @@
+import functools
 from pathlib import Path
 from typing import Tuple
+from unittest.mock import Mock, patch
 
 import networkx as nx
 import numpy as np
+import pandas as pd
 import pytest
 import torch as th
 import torch.nn.functional as F
@@ -173,6 +176,15 @@ def test_outputs(
 
     assert nx.utils.graphs_equal(nx_tracker, nx_original)
 
+    # test to_geff
+    with patch("ultrack.core.export.geff.geff.write_nx") as mock_write_nx:
+        tracker.to_geff("test.geff")
+        mock_write_nx.assert_called_once()
+
+        # Test with overwrite=True
+        tracker.to_geff("test.geff", overwrite=True)
+        assert mock_write_nx.call_count == 2
+
 
 @pytest.mark.parametrize(
     "config_content,timelapse_mock_data,mock_flow_field",

ultrack/core/export/__init__.py

@@ -1,4 +1,6 @@
 from ultrack.core.export.ctc import to_ctc
+from ultrack.core.export.exporter import export_tracks_by_extension
+from ultrack.core.export.geff import to_geff
 from ultrack.core.export.networkx import to_networkx, tracks_layer_to_networkx
 from ultrack.core.export.trackmate import to_trackmate, tracks_layer_to_trackmate
 from ultrack.core.export.tracks_layer import to_tracks_layer

ultrack/core/export/_test/test_exporter.py

@@ -1,10 +1,12 @@
 from pathlib import Path
 
+import pytest
+
 from ultrack import MainConfig, export_tracks_by_extension
 
 
 def test_exporter(tracked_database_mock_data: MainConfig, tmp_path: Path) -> None:
-    file_ext_list = [".xml", ".csv", ".zarr", ".dot", ".json"]
+    file_ext_list = [".xml", ".csv", ".zarr", ".dot", ".json", ".geff"]
     last_modified_time = {}
     for file_ext in file_ext_list:
         tmp_file = tmp_path / f"tracks{file_ext}"
@@ -40,3 +42,28 @@ def test_exporter(tracked_database_mock_data: MainConfig, tmp_path: Path) -> Non
         assert (tmp_path / f"tracks{file_ext}").stat().st_size > 0
 
         assert last_modified_time[str(tmp_file)] != tmp_file.stat().st_mtime
+
+
+def test_geff_zarr_extension_specific(
+    tracked_database_mock_data: MainConfig, tmp_path: Path
+) -> None:
+    """Test specific functionality of .geff.zarr extension in exporter."""
+    geff_file = tmp_path / "tracks.geff"
+
+    # Test that .geff.zarr extension calls the to_geff function
+    export_tracks_by_extension(tracked_database_mock_data, geff_file)
+
+    # Check that file exists and has content
+    assert geff_file.exists()
+    assert geff_file.stat().st_size > 0
+
+    # Test overwrite behavior
+    with pytest.raises(FileExistsError, match="already exists"):
+        export_tracks_by_extension(
+            tracked_database_mock_data, geff_file, overwrite=False
+        )
+
+    # Test that overwrite=True works
+    export_tracks_by_extension(tracked_database_mock_data, geff_file, overwrite=True)
+    assert geff_file.exists()
+    assert geff_file.stat().st_size > 0

ultrack/core/export/_test/test_geff.py

@@ -0,0 +1,81 @@
+from pathlib import Path
+from unittest.mock import MagicMock, Mock, patch
+
+import networkx as nx
+import pytest
+
+from ultrack import MainConfig
+from ultrack.core.export.geff import to_geff
+
+
+def test_to_geff_basic_functionality(
+    tracked_database_mock_data: MainConfig, tmp_path: Path
+):
+    """Test basic functionality of to_geff function and call chain."""
+    output_file = tmp_path / "test_tracks.geff.zarr"
+
+    # Mock both functions to verify the call chain
+    with patch("ultrack.core.export.geff.to_networkx") as mock_to_networkx, patch(
+        "ultrack.core.export.geff.geff.write_nx"
+    ) as mock_write_nx:
+
+        # Set up the mock to return a simple graph
+        mock_graph = nx.Graph()
+        mock_to_networkx.return_value = mock_graph
+
+        to_geff(tracked_database_mock_data, output_file)
+
+        # Verify that to_networkx was called with the config
+        mock_to_networkx.assert_called_once_with(tracked_database_mock_data)
+
+        # Verify that geff.write_nx was called with the graph and filename
+        mock_write_nx.assert_called_once_with(mock_graph, output_file)
+
+
+def test_to_geff_file_overwrite_false(
+    tracked_database_mock_data: MainConfig, tmp_path: Path
+):
+    """Test that FileExistsError is raised when file exists and overwrite=False."""
+    output_file = tmp_path / "test_tracks.geff.zarr"
+
+    # Create a file that already exists
+    output_file.touch()
+
+    # Test that FileExistsError is raised when overwrite=False
+    with pytest.raises(FileExistsError, match="already exists"):
+        to_geff(tracked_database_mock_data, output_file, overwrite=False)
+
+
+def test_to_geff_file_overwrite_true(
+    tracked_database_mock_data: MainConfig, tmp_path: Path
+):
+    """Test that function works when file exists and overwrite=True."""
+    output_file = tmp_path / "test_tracks.geff.zarr"
+
+    # Create a file that already exists
+    output_file.touch()
+
+    # Mock the geff.write_nx function
+    with patch("ultrack.core.export.geff.geff.write_nx") as mock_write_nx:
+        # This should not raise an error
+        to_geff(tracked_database_mock_data, output_file, overwrite=True)
+
+        # Verify that geff.write_nx was called
+        mock_write_nx.assert_called_once()
+
+
+def test_geff_correctness(tracked_database_mock_data: MainConfig, tmp_path: Path):
+    """Test that the geff file is correct."""
+    output_file = tmp_path / "test_tracks.geff.zarr"
+    to_geff(tracked_database_mock_data, output_file)
+
+    # Read the geff file
+    import geff
+
+    geff_nx = geff.read_nx(output_file)
+
+    from ultrack.core.export import to_networkx
+
+    ultrack_nx = to_networkx(tracked_database_mock_data)
+
+    assert nx.is_isomorphic(geff_nx, ultrack_nx)

ultrack/core/export/exporter.py

@@ -1,16 +1,18 @@
 import json
+import logging
 from pathlib import Path
 from typing import Union
 
 import networkx as nx
 
 from ultrack.config import MainConfig
-from ultrack.core.export import (
-    to_networkx,
-    to_trackmate,
-    to_tracks_layer,
-    tracks_to_zarr,
-)
+from ultrack.core.export.geff import to_geff
+from ultrack.core.export.networkx import to_networkx
+from ultrack.core.export.trackmate import to_trackmate
+from ultrack.core.export.tracks_layer import to_tracks_layer
+from ultrack.core.export.zarr import tracks_to_zarr
+
+LOG = logging.getLogger(__name__)
 
 
 def export_tracks_by_extension(
@@ -19,11 +21,12 @@ def export_tracks_by_extension(
     """
     Export tracks to a file given the file extension.
 
-    Supported file extensions are .xml, .csv, .zarr, .dot, and .json.
+    Supported file extensions are .xml, .csv, .zarr, .parquet, .dot, .json, and .geff
     - `.xml` exports to a TrackMate compatible XML file.
     - `.csv` exports to a CSV file.
     - `.parquet` exports to a Parquet file.
     - `.zarr` exports the tracks to dense segments in a `zarr` array format.
+    - `.geff` exports the tracks to a `zarr` format using the geff standard.
     - `.dot` exports to a Graphviz DOT file.
     - `.json` exports to a networkx JSON file.
 
@@ -46,13 +49,16 @@ def export_tracks_by_extension(
         Export tracks to a `zarr` array.
     to_networkx :
         Export tracks to a networkx graph.
+    to_geff :
+        Export tracks to a geff file.
     """
-    if Path(filename).exists() and not overwrite:
+    filename = Path(filename)
+    if filename.exists() and not overwrite:
         raise FileExistsError(
             f"File {filename} already exists. Set `overwrite=True` to overwrite the file"
         )
 
-    file_ext = Path(filename).suffix
+    file_ext = filename.suffix
     if file_ext.lower() == ".xml":
         to_trackmate(config, filename, overwrite=True)
     elif file_ext.lower() == ".csv":
@@ -61,6 +67,8 @@ def export_tracks_by_extension(
     elif file_ext.lower() == ".zarr":
         df, _ = to_tracks_layer(config)
         tracks_to_zarr(config, df, filename, overwrite=True)
+    elif file_ext.lower() == ".geff":
+        to_geff(config, filename, overwrite=overwrite)
     elif file_ext.lower() == ".parquet":
         df, _ = to_tracks_layer(config)
         df.to_parquet(filename)
@@ -75,5 +83,5 @@ def export_tracks_by_extension(
     else:
         raise ValueError(
             f"Unknown file extension: {file_ext}. "
-            "Supported extensions are .xml, .csv, .zarr, .parquet, .dot, and .json."
+            "Supported extensions are .xml, .csv, .zarr, .geff, .parquet, .dot, and .json."
         )

ultrack/core/export/geff.py

@@ -0,0 +1,42 @@
+from pathlib import Path
+from typing import Union
+
+import geff
+import networkx as nx
+
+from ultrack.config import MainConfig
+from ultrack.core.export.networkx import to_networkx
+
+
+def to_geff(
+    config: MainConfig,
+    filename: Union[str, Path],
+    overwrite: bool = False,
+) -> None:
+    """
+    Export tracks to a geff (Graph Exchange File Format) file.
+
+    Parametersmnist
+    ----------
+    config : MainConfig
+        The configuration object.
+    filename : str or Path
+        The name of the file to save the tracks to.
+    overwrite : bool, optional
+        Whether to overwrite the file if it already exists, by default False.
+
+    Raises
+    ------
+    FileExistsError
+        If the file already exists and overwrite is False.
+    """
+    if Path(filename).exists() and not overwrite:
+        raise FileExistsError(
+            f"File {filename} already exists. Set `overwrite=True` to overwrite the file"
+        )
+
+    # Get the networkx graph from the configuration
+    graph = to_networkx(config)
+
+    # Write the graph to geff format
+    geff.write_nx(graph, filename)

ultrack/core/tracker.py

@@ -11,6 +11,7 @@
 from ultrack.config import MainConfig
 from ultrack.core.export import (
     to_ctc,
+    to_geff,
     to_tracks_layer,
     tracks_layer_to_networkx,
     tracks_layer_to_trackmate,
@@ -154,6 +155,11 @@ def to_ctc(self, *args, **kwargs) -> None:
         self._assert_solved()
         to_ctc(config=self.config, *args, **kwargs)
 
+    @functools.wraps(to_geff)
+    def to_geff(self, filename: str, overwrite: bool = False) -> None:
+        self._assert_solved()
+        to_geff(self.config, filename, overwrite=overwrite)
+
     @functools.wraps(to_tracks_layer)
     def to_tracks_layer(self, *args, **kwargs) -> Tuple[pd.DataFrame, Dict]:
         self._assert_solved()