Merge pull request #148 from funkelab/empty-graph

Gracefully handle empty graph and graph with no edges

Author: tlambert03

motile/costs/cost.py

@@ -8,7 +8,11 @@
 
 
 class Cost(ABC):
-    """A base class for a cost that can be added to a solver."""
+    """A base class for a cost that can be added to a solver.
+
+    Weights should be initialized in the __init__ and added to a instance
+    variable so that the Solver can discover them.
+    """
 
     @abstractmethod
     def apply(self, solver: Solver) -> None:

motile/costs/features.py

@@ -52,6 +52,14 @@ def _increase_features(self, num_features: int) -> None:
         new_features = np.zeros(shape, dtype=self._values.dtype)
         self._values = np.hstack((self._values, new_features))
 
+    def register_feature(self, feature_index: int) -> None:
+        num_variables, num_features = self._values.shape
+        if feature_index >= num_features:
+            self.resize(
+                num_variables,
+                max(feature_index + 1, num_features),
+            )
+
     def add_feature(
         self, variable_index: int | ilpy.Variable, feature_index: int, value: float
     ) -> None:
@@ -68,10 +76,10 @@ def add_feature(
         num_variables, num_features = self._values.shape
 
         variable_index = int(variable_index)
-        if variable_index >= num_variables or feature_index >= num_features:
+        if variable_index >= num_variables:
             self.resize(
                 max(variable_index + 1, num_variables),
-                max(feature_index + 1, num_features),
+                num_features,
             )
 
         self._values[variable_index, feature_index] += value

motile/costs/weights.py

@@ -24,23 +24,28 @@ def __init__(self) -> None:
         self._weight_indices: dict[Weight, int] = {}
         self._modify_callbacks: list[Callback] = []
 
-    def add_weight(self, weight: Weight, name: Hashable) -> None:
+    def add_weight(self, weight: Weight, name: Hashable) -> int:
         """Add a weight to the container.
 
         Args:
             weight:
                 The :class:`~motile.costs.Weight` to add.
             name:
                 The name of the weight.
+
+        Returns:
+            int: the index of the weight
         """
-        self._weight_indices[weight] = len(self._weights)
+        weight_index = len(self._weights)
+        self._weight_indices[weight] = weight_index
         self._weights.append(weight)
         self._weights_by_name[name] = weight
 
         for callback in self._modify_callbacks:
             weight.register_modify_callback(callback)
 
         self._notify_modified(None, weight.value)
+        return weight_index
 
     def register_modify_callback(self, callback: Callback) -> None:
         """Register ``callback`` to be called when a weight is modified.

motile/solver.py

@@ -93,7 +93,8 @@ def add_cost(self, cost: Cost, name: str | None = None) -> None:
         for var_name, var in cost.__dict__.items():
             if not isinstance(var, Weight):
                 continue
-            self.weights.add_weight(var, (name, var_name))
+            weight_index = self.weights.add_weight(var, (name, var_name))
+            self.features.register_feature(weight_index)
 
         cost.apply(self)
 

tests/test_solver.py

@@ -0,0 +1,49 @@
+import motile
+import networkx as nx
+from motile.costs import EdgeSelection, NodeSelection
+
+
+def test_empty_graph() -> None:
+    """Test that solving an empty graph does not error and returns empty solution."""
+    nx_graph = nx.DiGraph()
+    graph = motile.TrackGraph(nx_graph)
+
+    solver = motile.Solver(graph)
+    solver.add_cost(NodeSelection(constant=-1))
+    solver.add_cost(EdgeSelection(constant=-1))
+
+    # Should not error
+    solver.solve()
+    solution_graph = solver.get_selected_subgraph()
+
+    # Solution should be empty
+    assert len(solution_graph.nodes) == 0
+    assert len(solution_graph.edges) == 0
+
+
+def test_graph_with_no_edges() -> None:
+    """Test that solving a graph with nodes but no edges does not error."""
+    cells = [
+        {"id": 0, "t": 0},
+        {"id": 1, "t": 0},
+        {"id": 2, "t": 1},
+    ]
+
+    nx_graph = nx.DiGraph()
+    nx_graph.add_nodes_from([(cell["id"], cell) for cell in cells])
+    # No edges added
+
+    graph = motile.TrackGraph(nx_graph)
+
+    solver = motile.Solver(graph)
+    solver.add_cost(NodeSelection(constant=-1))
+    solver.add_cost(EdgeSelection(constant=-1))
+
+    # Should not error
+    solver.solve()
+    solution_graph = solver.get_selected_subgraph()
+
+    # All nodes should be selected due to negative cost
+    assert set(solution_graph.nodes.keys()) == {0, 1, 2}
+    # No edges should be selected (none exist)
+    assert len(solution_graph.edges) == 0