Merge pull request #52 from funkelab/add_disappear_cost

Add disappear cost

motile/costs/__init__.py

@@ -1,5 +1,6 @@
 from .appear import Appear
 from .costs import Costs
+from .disappear import Disappear
 from .edge_distance import EdgeDistance
 from .edge_selection import EdgeSelection
 from .features import Features
@@ -11,6 +12,7 @@
 __all__ = [
     "Appear",
     "Costs",
+    "Disappear",
     "EdgeDistance",
     "EdgeSelection",
     "Features",

motile/costs/disappear.py

@@ -0,0 +1,29 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from ..variables import NodeDisappear
+from .costs import Costs
+from .weight import Weight
+
+if TYPE_CHECKING:
+    from motile.solver import Solver
+
+
+class Disappear(Costs):
+    """Costs for :class:`motile.variables.NodeDisappear` variables.
+
+    Args:
+
+        constant (float):
+            A constant cost for each node that ends a track.
+    """
+
+    def __init__(self, constant: float) -> None:
+        self.constant = Weight(constant)
+
+    def apply(self, solver: Solver) -> None:
+        disappear_indicators = solver.get_variables(NodeDisappear)
+
+        for index in disappear_indicators.values():
+            solver.add_variable_cost(index, 1.0, self.constant)

motile/variables/__init__.py

@@ -1,12 +1,14 @@
 from .edge_selected import EdgeSelected
 from .node_appear import NodeAppear
+from .node_disappear import NodeDisappear
 from .node_selected import NodeSelected
 from .node_split import NodeSplit
 from .variable import Variable
 
 __all__ = [
     "EdgeSelected",
     "NodeAppear",
+    "NodeDisappear",
     "NodeSelected",
     "NodeSplit",
     "Variable",

motile/variables/node_disappear.py

@@ -0,0 +1,72 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Collection, Iterable
+
+from .edge_selected import EdgeSelected
+from .node_selected import NodeSelected
+from .variable import Variable
+
+if TYPE_CHECKING:
+    import ilpy
+
+    from motile._types import NodeId
+    from motile.solver import Solver
+
+
+class NodeDisappear(Variable["NodeId"]):
+    r"""A binary variable for each node that indicates whether the node is the
+    end of a track (i.e., the node is selected and has no selected outgoing
+    edges).
+
+    This variable is coupled to the node and edge selection variables through
+    the following linear constraints:
+
+    .. math::
+        |\\text{out_edges}(v)|\cdot x_v - &\sum_{e \in \\text{out_edges}(v)} x_e
+        - d_v &\leq&\;\; |\\text{out_edges}(v)| - 1
+
+        |\\text{out_edges}(v)|\cdot x_v - &\sum_{e \in \\text{out_edges}(v)} x_e
+        - d_v\cdot |\\text{out_edges}(v)| &\geq&\;\; 0
+
+    where :math:`x_v` and :math:`x_e` are selection indicators for node
+    :math:`v` and edge :math:`e`, and :math:`d_v` is the disappear indicator for
+    node :math:`v`.
+    """
+
+    @staticmethod
+    def instantiate(solver: Solver) -> Collection[NodeId]:
+        return solver.graph.nodes
+
+    @staticmethod
+    def instantiate_constraints(solver: Solver) -> Iterable[ilpy.Expression]:
+        node_indicators = solver.get_variables(NodeSelected)
+        edge_indicators = solver.get_variables(EdgeSelected)
+        disappear_indicators = solver.get_variables(NodeDisappear)
+
+        for node in solver.graph.nodes:
+            next_edges = solver.graph.next_edges[node]
+            selected = node_indicators[node]
+            disappear = disappear_indicators[node]
+
+            if not next_edges:
+                # special case: no outgoing edges, disappear indicator is equal to
+                # selection indicator
+                yield selected == disappear
+                continue
+
+            # Ensure that the following holds:
+            #
+            # disappear = 1 <=> sum(next_selected) = 0 and selected
+            # disappear = 0 <=> sum(next_selected) > 0 or not selected
+            #
+            # Two linear constraints are needed for that:
+            #
+            # let s = num_next * selected - sum(next_selected)
+            # (1) s - disappear <= num_next - 1
+            # (2) s - disappear * num_next >= 0
+
+            num_next = len(next_edges)
+            s = num_next * selected - sum(edge_indicators[e] for e in next_edges)
+
+            yield s - disappear <= num_next - 1
+            yield s - disappear >= 0

tests/test_api.py

@@ -2,7 +2,7 @@
 
 import motile
 from motile.constraints import MaxChildren, MaxParents
-from motile.costs import Appear, EdgeSelection, NodeSelection, Split
+from motile.costs import Appear, Disappear, EdgeSelection, NodeSelection, Split
 from motile.data import (
     arlo_graph,
     arlo_graph_nx,
@@ -40,7 +40,8 @@ def test_solver(self):
         solver = motile.Solver(graph)
         solver.add_costs(NodeSelection(weight=-1.0, attribute="score", constant=-100.0))
         solver.add_costs(EdgeSelection(weight=1.0, attribute="prediction_distance"))
-        solver.add_costs(Appear(constant=200.0))
+        solver.add_costs(Appear(constant=100.0))
+        solver.add_costs(Disappear(constant=100.0))
         solver.add_costs(Split(constant=100.0))
         solver.add_constraints(MaxParents(1))
         solver.add_constraints(MaxChildren(2))