Merge pull request #145 from funkelab/86-merge-cost

Add merge indicator and cost

Author: tlambert03

docs/source/api.rst

@@ -109,6 +109,10 @@ Split
 ^^^^^
   .. autoclass:: Split
 
+Merge
+^^^^^
+  .. autoclass:: Merge
+
 EdgeDistance
 ^^^^^^^^^^^^
   .. autoclass:: EdgeDistance

motile/costs/__init__.py

@@ -4,6 +4,7 @@
 from .edge_distance import EdgeDistance
 from .edge_selection import EdgeSelection
 from .features import Features
+from .merge import Merge
 from .node_selection import NodeSelection
 from .split import Split
 from .weight import Weight
@@ -16,6 +17,7 @@
     "EdgeDistance",
     "EdgeSelection",
     "Features",
+    "Merge",
     "NodeSelection",
     "Split",
     "Weight",

motile/costs/merge.py

@@ -0,0 +1,44 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from ..variables import NodeMerge
+from .cost import Cost
+from .weight import Weight
+
+if TYPE_CHECKING:
+    from motile.solver import Solver
+
+
+class Merge(Cost):
+    """Cost for :class:`~motile.variables.NodeMerge` variables.
+
+    Args:
+        weight:
+            The weight to apply to the cost of each split. Default is ``1``.
+
+        attribute:
+            The name of the attribute to use to look up the cost. Default is
+            ``None``, which means that a constant cost is used.
+
+        constant:
+            A constant cost for each node that has more than one selected
+            parent. Default is ``0``.
+    """
+
+    def __init__(
+        self, weight: float = 1, attribute: str | None = None, constant: float = 0
+    ) -> None:
+        self.weight = Weight(weight)
+        self.constant = Weight(constant)
+        self.attribute = attribute
+
+    def apply(self, solver: Solver) -> None:
+        merge_indicators = solver.get_variables(NodeMerge)
+
+        for node, index in merge_indicators.items():
+            if self.attribute is not None:
+                solver.add_variable_cost(
+                    index, solver.graph.nodes[node][self.attribute], self.weight
+                )
+            solver.add_variable_cost(index, 1.0, self.constant)

motile/variables/__init__.py

@@ -1,6 +1,7 @@
 from .edge_selected import EdgeSelected
 from .node_appear import NodeAppear
 from .node_disappear import NodeDisappear
+from .node_merge import NodeMerge
 from .node_selected import NodeSelected
 from .node_split import NodeSplit
 from .variable import Variable
@@ -9,6 +10,7 @@
     "EdgeSelected",
     "NodeAppear",
     "NodeDisappear",
+    "NodeMerge",
     "NodeSelected",
     "NodeSplit",
     "Variable",

motile/variables/node_merge.py

@@ -0,0 +1,73 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Collection, Iterable
+
+import ilpy
+
+from .edge_selected import EdgeSelected
+from .variable import Variable
+
+if TYPE_CHECKING:
+    from motile._types import Node
+    from motile.solver import Solver
+
+
+class NodeMerge(Variable):
+    r"""Binary variable indicating whether a node has more than one parent.
+
+    (i.e., the node is selected and has more than one selected incoming edge).
+
+    This variable is coupled to the edge selection variables through the
+    following linear constraints:
+
+    .. math::
+
+        2 m_v\; - &\sum_{e\in\text{in_edges}(v)} x_e &\leq&\;\; 0
+
+        (|\text{in_edges}(v)| - 1) m_v\; - &\sum_{e\in\text{in_edges}(v)}
+        x_e &\geq&\;\; -1
+
+    where :math:`x_e` are selection indicators for edge :math:`e`, and
+    :math:`m_v` is the merge indicator for node :math:`v`.
+    """
+
+    @staticmethod
+    def instantiate(solver: Solver) -> Collection[Node]:
+        return solver.graph.nodes
+
+    @staticmethod
+    def instantiate_constraints(solver: Solver) -> Iterable[ilpy.Constraint]:
+        merge_indicators = solver.get_variables(NodeMerge)
+        edge_indicators = solver.get_variables(EdgeSelected)
+
+        for node in solver.graph.nodes:
+            prev_edges = solver.graph.prev_edges[node]
+
+            # Ensure that the following holds:
+            #
+            # merge = 0 <=> sum(prev_selected) <= 1
+            # merge = 1 <=> sum(prev_selected) > 1
+            #
+            # Two linear constraints are needed for that:
+            #
+            # (1) 2 * merge - sum(prev_selected) <= 0
+            # (2) (num_prev - 1) * merge - sum(prev_selected) >= -1
+
+            constraint1 = ilpy.Constraint()
+            constraint2 = ilpy.Constraint()
+
+            constraint1.set_coefficient(merge_indicators[node], 2.0)
+            constraint2.set_coefficient(merge_indicators[node], len(prev_edges) - 1.0)
+
+            for prev_edge in prev_edges:
+                constraint1.set_coefficient(edge_indicators[prev_edge], -1.0)
+                constraint2.set_coefficient(edge_indicators[prev_edge], -1.0)
+
+            constraint1.set_relation(ilpy.Relation.LessEqual)
+            constraint2.set_relation(ilpy.Relation.GreaterEqual)
+
+            constraint1.set_value(0.0)
+            constraint2.set_value(-1.0)
+
+            yield constraint1
+            yield constraint2

tests/test_costs.py

@@ -1,8 +1,10 @@
 import motile
+import networkx as nx
 from motile.costs import (
     Appear,
     Disappear,
     EdgeSelection,
+    Merge,
     NodeSelection,
 )
 
@@ -81,3 +83,119 @@ def test_disappear_cost(arlo_graph):
     solution_graph = solver.get_selected_subgraph()
     assert list(solution_graph.nodes.keys()) == [2, 3, 4, 5, 6]
     assert len(solution_graph.edges) == 0
+
+
+def test_constant_merge_cost() -> None:
+    """Test that merge cost prevents merges when applied.
+
+    Graph structure:
+        t=0: node 0, node 1
+        t=1: node 2
+        edges: 0->2, 1->2
+
+    With only negative edge selection cost, both edges should be selected
+    (resulting in a merge). Adding a merge cost should prevent the merge,
+    resulting in only one edge being selected.
+    """
+    # Create nodes
+    cells = [
+        {"id": 0, "t": 0},
+        {"id": 1, "t": 0},
+        {"id": 2, "t": 1},
+    ]
+
+    # Create edges (both leading to node 2, creating potential merge)
+    edges = [
+        {"source": 0, "target": 2},
+        {"source": 1, "target": 2},
+    ]
+
+    nx_graph = nx.DiGraph()
+    nx_graph.add_nodes_from([(cell["id"], cell) for cell in cells])
+    nx_graph.add_edges_from([(edge["source"], edge["target"], edge) for edge in edges])
+
+    graph = motile.TrackGraph(nx_graph)
+
+    # First test: without merge cost, both edges should be selected
+    solver = motile.Solver(graph)
+    solver.add_cost(EdgeSelection(constant=-1.0))
+    solver.solve()
+    solution_graph = solver.get_selected_subgraph().to_nx_graph()
+
+    # Should select all nodes and both edges (merge occurs)
+    assert set(solution_graph.nodes.keys()) == {0, 1, 2}
+    assert len(solution_graph.edges) == 2
+    assert solution_graph.has_edge(0, 2)
+    assert solution_graph.has_edge(1, 2)
+
+    # Second test: with merge cost, only one edge should be selected
+    solver = motile.Solver(graph)
+    solver.add_cost(EdgeSelection(constant=-1.0))
+    solver.add_cost(Merge(constant=10.0))  # High cost to prevent merge
+    solver.solve()
+    solution_graph = solver.get_selected_subgraph().to_nx_graph()
+
+    # Should select all nodes but only one edge (no merge)
+    assert set(solution_graph.nodes.keys()) == {0, 1, 2}
+    assert len(solution_graph.edges) == 1
+    # Either edge 0->2 or 1->2 should be selected, but not both
+    assert solution_graph.has_edge(0, 2) or solution_graph.has_edge(1, 2)
+    assert not (solution_graph.has_edge(0, 2) and solution_graph.has_edge(1, 2))
+
+
+def test_variable_merge_cost() -> None:
+    """Test that merge cost can use node attributes to selectively allow merges.
+
+    Graph structure:
+        t=0: node 0, node 1, node 2
+        t=1: node 3 (merge_cost=-1.0), node 4 (merge_cost=5.0)
+        edges: 0->3, 1->3, 1->4, 2->4
+
+    With negative edge selection cost and attribute-based merge cost,
+    only the node with negative merge_cost should have a merge (node 3).
+    Node 4 with positive merge_cost should not have a merge.
+    """
+    # Create nodes - all nodes need merge_cost attribute
+    cells = [
+        {"id": 0, "t": 0, "merge_cost": 0.0},
+        {"id": 1, "t": 0, "merge_cost": 0.0},
+        {"id": 2, "t": 0, "merge_cost": 0.0},
+        {"id": 3, "t": 1, "merge_cost": -1.0},  # negative cost = good merge
+        {"id": 4, "t": 1, "merge_cost": 5.0},  # positive cost = bad merge
+    ]
+
+    # Create edges - two edges to each node in time 1 (potential merges)
+    edges = [
+        {"source": 0, "target": 3},
+        {"source": 1, "target": 3},
+        {"source": 1, "target": 4},
+        {"source": 2, "target": 4},
+    ]
+
+    nx_graph = nx.DiGraph()
+    nx_graph.add_nodes_from([(cell["id"], cell) for cell in cells])
+    nx_graph.add_edges_from([(edge["source"], edge["target"], edge) for edge in edges])
+
+    graph = motile.TrackGraph(nx_graph)
+
+    solver = motile.Solver(graph)
+    solver.add_cost(EdgeSelection(constant=-1.0))
+    solver.add_cost(Merge(attribute="merge_cost", weight=1.0, constant=0.0))
+    solver.solve()
+    solution_graph = solver.get_selected_subgraph().to_nx_graph()
+
+    # Should select all nodes
+    assert set(solution_graph.nodes.keys()) == {0, 1, 2, 3, 4}
+
+    # Node 3 should have a merge (both edges 0->3 and 1->3 selected)
+    # because its merge_cost is negative (-1.0), making the total cost attractive
+    assert solution_graph.has_edge(0, 3)
+    assert solution_graph.has_edge(1, 3)
+
+    # Node 4 should NOT have a merge (only one edge selected)
+    # because its merge_cost is positive (5.0), making the merge too expensive
+    edges_to_4 = [
+        solution_graph.has_edge(1, 4),
+        solution_graph.has_edge(2, 4),
+    ]
+    assert sum(edges_to_4) == 1  # exactly one edge to node 4