Merge pull request #118 from funkelab/free_dis_appear

Free appear/disappear in first/last frames

motile/costs/appear.py

@@ -13,6 +13,8 @@
 class Appear(Cost):
     """Cost for :class:`~motile.variables.NodeAppear` variables.
 
+    This is cost is not applied to nodes in the first frame of the graph.
+
     Args:
         weight:
             The weight to apply to the cost of each starting track.
@@ -27,8 +29,6 @@ class Appear(Cost):
         ignore_attribute:
             The name of an optional node attribute that, if it is set and
             evaluates to ``True``, will not set the appear cost for that node.
-            This is useful to allow nodes in the first frame to appear at no
-            cost.
     """
 
     def __init__(
@@ -45,13 +45,16 @@ def __init__(
 
     def apply(self, solver: Solver) -> None:
         appear_indicators = solver.get_variables(NodeAppear)
+        G = solver.graph
 
         for node, index in appear_indicators.items():
             if self.ignore_attribute is not None:
-                if solver.graph.nodes[node].get(self.ignore_attribute, False):
+                if G.nodes[node].get(self.ignore_attribute, False):
                     continue
+            if G.nodes[node][G.frame_attribute] == G.get_frames()[0]:
+                continue
             if self.attribute is not None:
                 solver.add_variable_cost(
-                    index, solver.graph.nodes[node][self.attribute], self.weight
+                    index, G.nodes[node][self.attribute], self.weight
                 )
             solver.add_variable_cost(index, 1.0, self.constant)

motile/costs/disappear.py

@@ -13,15 +13,16 @@
 class Disappear(Cost):
     """Cost for :class:`motile.variables.NodeDisappear` variables.
 
+    This is cost is not applied to nodes in the last frame of the graph.
+
     Args:
         constant (float):
             A constant cost for each node that ends a track.
 
         ignore_attribute:
             The name of an optional node attribute that, if it is set and
             evaluates to ``True``, will not set the disappear cost for that
-            node. This is useful to allow nodes in the last frame to disappear
-            at no cost.
+            node.
     """
 
     def __init__(self, constant: float, ignore_attribute: str | None = None) -> None:
@@ -30,9 +31,12 @@ def __init__(self, constant: float, ignore_attribute: str | None = None) -> None
 
     def apply(self, solver: Solver) -> None:
         disappear_indicators = solver.get_variables(NodeDisappear)
+        G = solver.graph
 
         for node, index in disappear_indicators.items():
             if self.ignore_attribute is not None:
-                if solver.graph.nodes[node].get(self.ignore_attribute, False):
+                if G.nodes[node].get(self.ignore_attribute, False):
                     continue
+            if G.nodes[node][G.frame_attribute] == G.get_frames()[1] - 1:
+                continue
             solver.add_variable_cost(index, 1.0, self.constant)

motile/track_graph.py

@@ -220,10 +220,11 @@ def _convert_nx_hypernode(
 
         return (in_nodes, out_nodes)
 
-    def get_frames(self) -> tuple[int | None, int | None]:
+    def get_frames(self) -> tuple[int, int]:
         """Return tuple with first and last (exclusive) frame this graph has nodes for.
 
         Returns ``(t_begin, t_end)`` where ``t_end`` is exclusive.
+        Returns ``(0, 0)`` for empty graph.
         """
         self._update_metadata()
 
@@ -246,8 +247,8 @@ def _update_metadata(self) -> None:
 
         if not self.nodes:
             self._nodes_by_frame = {}
-            self.t_begin = None
-            self.t_end = None
+            self.t_begin = 0
+            self.t_end = 0
             return
 
         self._nodes_by_frame = {}

tests/test_api.py

@@ -4,6 +4,7 @@
 from motile.constraints import MaxChildren, MaxParents
 from motile.costs import (
     Appear,
+    Disappear,
     EdgeDistance,
     EdgeSelection,
     NodeSelection,
@@ -58,6 +59,7 @@ def test_solver(arlo_graph):
     )
     solver.add_cost(EdgeDistance(position_attribute=("x",), weight=0.5))
     solver.add_cost(Appear(constant=200.0, attribute="score", weight=-1.0))
+    solver.add_cost(Disappear(constant=55.0))
     solver.add_cost(Split(constant=100.0, attribute="score", weight=1.0))
 
     solver.add_constraint(MaxParents(1))
@@ -75,4 +77,4 @@ def test_solver(arlo_graph):
     )
     assert list(subgraph.nodes) == _selected_nodes(solver) == [0, 1, 2, 3, 4, 5]
     cost = solution.get_value()
-    assert cost == -206.0, f"{cost=}"
+    assert cost == -604.0, f"{cost=}"

tests/test_costs.py

@@ -1,59 +1,83 @@
 import motile
-from motile.constraints import MaxChildren, MaxParents
 from motile.costs import (
     Appear,
-    EdgeDistance,
+    Disappear,
     EdgeSelection,
     NodeSelection,
-    Split,
 )
 
 
-def test_ignore_attributes(arlo_graph):
-    graph = arlo_graph
+def test_appear_cost(arlo_graph):
+    solver = motile.Solver(arlo_graph)
 
-    # first solve without ignore attribute:
-
-    solver = motile.Solver(graph)
-    solver.add_cost(NodeSelection(weight=-1.0, attribute="score", constant=-100.0))
+    # Make a slightly negative node cost, and a very positive appear cost and edge
+    # cost. We expect only selecting nodes in the first frame, where by default the
+    # appear cost is ignored, and not selecting any edges
+    solver.add_cost(NodeSelection(weight=0, attribute="score", constant=-1))
+    solver.add_cost(Appear(constant=100))
     solver.add_cost(
-        EdgeSelection(weight=0.5, attribute="prediction_distance", constant=-1.0)
+        EdgeSelection(weight=0, attribute="prediction_distance", constant=100)
     )
-    solver.add_cost(EdgeDistance(position_attribute=("x",), weight=0.5))
-    solver.add_cost(Appear(constant=200.0, attribute="score", weight=-1.0))
-    solver.add_cost(Split(constant=100.0, attribute="score", weight=1.0))
-
-    solver.add_constraint(MaxParents(1))
-    solver.add_constraint(MaxChildren(2))
+    solver.solve()
+    solution_graph = solver.get_selected_subgraph()
+    assert list(solution_graph.nodes.keys()) == [0, 1]
+    assert len(solution_graph.edges) == 0
 
-    solution = solver.solve()
-    no_ignore_value = solution.get_value()
+    ignore_attr = "ignore"
+    # now also ignore the appear cost in the second frame
+    for second_node in arlo_graph.nodes_by_frame(1):
+        arlo_graph.nodes[second_node][ignore_attr] = True
+    # but not the third frame
+    for third_node in arlo_graph.nodes_by_frame(2):
+        arlo_graph.nodes[third_node][ignore_attr] = False
 
-    # solve and ignore appear costs in frame 0
+    solver = motile.Solver(arlo_graph)
 
-    for first_node in graph.nodes_by_frame(0):
-        graph.nodes[first_node]["ignore_appear_cost"] = True
-
-    solver = motile.Solver(graph)
-    solver.add_cost(NodeSelection(weight=-1.0, attribute="score", constant=-100.0))
+    # Resolving should also select nodes in second frame
+    solver.add_cost(NodeSelection(weight=0, attribute="score", constant=-1))
+    solver.add_cost(Appear(constant=100, ignore_attribute=ignore_attr))
     solver.add_cost(
-        EdgeSelection(weight=0.5, attribute="prediction_distance", constant=-1.0)
+        EdgeSelection(weight=0, attribute="prediction_distance", constant=100)
     )
-    solver.add_cost(EdgeDistance(position_attribute="x", weight=0.5))
+    solver.solve()
+    solution_graph = solver.get_selected_subgraph()
+    assert list(solution_graph.nodes.keys()) == [0, 1, 2, 3]
+    assert len(solution_graph.edges) == 0
+
+
+def test_disappear_cost(arlo_graph):
+    solver = motile.Solver(arlo_graph)
+
+    # make a slightly negative node cost, and a positive disappear cost and edge cost
+    # we expect only selecting nodes in the last frame, where by default the disappear
+    # cost is ignored, and not selecting any edges
+    solver.add_cost(NodeSelection(weight=0, attribute="score", constant=-1))
+    solver.add_cost(Disappear(constant=100))
     solver.add_cost(
-        Appear(
-            constant=200.0,
-            attribute="score",
-            weight=-1.0,
-            ignore_attribute="ignore_appear_cost",
-        )
+        EdgeSelection(weight=0, attribute="prediction_distance", constant=100)
     )
-    solver.add_cost(Split(constant=100.0, attribute="score", weight=1.0))
+    solver.solve()
+    solution_graph = solver.get_selected_subgraph()
+    assert list(solution_graph.nodes.keys()) == [4, 5, 6]
+    assert len(solution_graph.edges) == 0
 
-    solver.add_constraint(MaxParents(1))
-    solver.add_constraint(MaxChildren(2))
+    ignore_attr = "ignore"
+    # now also ignore the disappear cost in the second frame
+    for second_node in arlo_graph.nodes_by_frame(1):
+        arlo_graph.nodes[second_node][ignore_attr] = True
+    # but not the first frame
+    for first_node in arlo_graph.nodes_by_frame(0):
+        arlo_graph.nodes[first_node][ignore_attr] = False
 
-    solution = solver.solve()
-    ignore_value = solution.get_value()
+    solver = motile.Solver(arlo_graph)
 
-    assert ignore_value < no_ignore_value
+    # Resolving should also select nodes in second frame
+    solver.add_cost(NodeSelection(weight=0, attribute="score", constant=-1))
+    solver.add_cost(Disappear(constant=100, ignore_attribute=ignore_attr))
+    solver.add_cost(
+        EdgeSelection(weight=0, attribute="prediction_distance", constant=100)
+    )
+    solver.solve()
+    solution_graph = solver.get_selected_subgraph()
+    assert list(solution_graph.nodes.keys()) == [2, 3, 4, 5, 6]
+    assert len(solution_graph.edges) == 0

tests/test_structsvm.py

@@ -75,20 +75,18 @@ def test_structsvm_common_toy_example(toy_graph):
     optimal_weights = solver.weights
 
     np.testing.assert_allclose(
-        optimal_weights[("NodeSelection", "weight")], -4.9771062468440785, rtol=0.01
+        optimal_weights[("NodeSelection", "weight")], -3.27, atol=0.01
     )
     np.testing.assert_allclose(
-        optimal_weights[("NodeSelection", "constant")], -3.60083857250377, rtol=0.01
+        optimal_weights[("NodeSelection", "constant")], 1.78, atol=0.01
     )
     np.testing.assert_allclose(
-        optimal_weights[("EdgeSelection", "weight")], -6.209937259450144, rtol=0.01
+        optimal_weights[("EdgeSelection", "weight")], -3.23, atol=0.01
     )
     np.testing.assert_allclose(
-        optimal_weights[("EdgeSelection", "constant")], -2.4005590483600203, rtol=0.01
-    )
-    np.testing.assert_allclose(
-        optimal_weights[("Appear", "constant")], 32.13305455424766, rtol=0.01
+        optimal_weights[("EdgeSelection", "constant")], 1.06, atol=0.01
     )
+    np.testing.assert_allclose(optimal_weights[("Appear", "constant")], 0.20, atol=0.01)
 
     solver = create_toy_solver(graph)
     solver.weights.from_ndarray(optimal_weights.to_ndarray())
@@ -171,20 +169,18 @@ def test_structsvm_noise():
     logger.debug(solver.features.to_ndarray())
 
     np.testing.assert_allclose(
-        optimal_weights[("NodeSelection", "weight")], -2.7777798708004564, rtol=0.01
-    )
-    np.testing.assert_allclose(
-        optimal_weights[("NodeSelection", "constant")], -1.3883786845544988, rtol=0.01
+        optimal_weights[("NodeSelection", "weight")], -2.77, atol=0.01
     )
     np.testing.assert_allclose(
-        optimal_weights[("EdgeSelection", "weight")], -3.3333338262308043, rtol=0.01
+        optimal_weights[("NodeSelection", "constant")], 0.39, atol=0.01
     )
     np.testing.assert_allclose(
-        optimal_weights[("EdgeSelection", "constant")], -0.9255857897041805, rtol=0.01
+        optimal_weights[("EdgeSelection", "weight")], -3.33, atol=0.01
     )
     np.testing.assert_allclose(
-        optimal_weights[("Appear", "constant")], 19.53720680712646, rtol=0.01
+        optimal_weights[("EdgeSelection", "constant")], 0, atol=0.01
     )
+    np.testing.assert_allclose(optimal_weights[("Appear", "constant")], 0.39, atol=0.01)
 
     def _assert_edges(solver, solution):
         edge_indicators = solver.get_variables(EdgeSelected)
@@ -213,7 +209,3 @@ def _assert_edges(solver, solution):
     logger.debug(solver.get_variables(EdgeSelected))
 
     _assert_edges(solver, solution)
-
-
-if __name__ == "__main__":
-    test_structsvm_noise()