Feat filter transitives (#625)

* feat: removes transitive proposals

* refactor: relaxed gt generation, less smoothing

---------

Co-authored-by: anna-grim <anna.grim@alleninstitute.org>

Author: anna-grim

src/neuron_proofreader/proposal_graph.py

@@ -104,15 +104,15 @@ def __init__(
 
         # Instance attributes - Proposals
         self.gt_accepts = set()
-        self.max_proposals_per_leaf = max_proposals_per_leaf
         self.merged_ids = set()
-        self.min_size_with_proposals = min_size_with_proposals
         self.n_merges_blocked = 0
         self.node_proposals = defaultdict(set)
         self.proposals = set()
 
         self.proposal_generator = ProposalGenerator(
-            self, filter_transitive_proposals=filter_transitive_proposals
+            self,
+            max_proposals_per_leaf=max_proposals_per_leaf,
+            min_size_with_proposals=min_size_with_proposals
         )
 
         # Graph Loader
@@ -196,6 +196,37 @@ def _add_edge(self, edge_id, attrs):
         self.add_edge(i, j, radius=attrs["radius"], xyz=attrs["xyz"])
         self.xyz_to_edge.update({tuple(xyz): edge_id for xyz in attrs["xyz"]})
 
+    def relabel_nodes(self):
+        """
+        Reassigns contiguous node IDs and update all dependent structures.
+        """
+               # Set node ids
+        old_node_ids = np.array(self.nodes, dtype=int)
+        new_node_ids = np.arange(len(old_node_ids))
+
+        # Set edge ids
+        old_to_new = dict(zip(old_node_ids, new_node_ids))
+        old_edge_ids = list(self.edges)
+        old_irr_edge_ids = self.irreducible.edges
+        edge_attrs = {(i, j): data for i, j, data in self.edges(data=True)}
+
+        # Reset graph
+        self.clear()
+        for (i, j) in old_edge_ids:
+            edge_id = (int(old_to_new[i]), int(old_to_new[j]))
+            self._add_edge(edge_id, edge_attrs[(i, j)])
+
+        self.irreducible.clear()
+        for (i, j) in old_irr_edge_ids:
+            self.irreducible.add_edge(old_to_new[i], old_to_new[j])
+
+        # Update attributes
+        self.node_radius = self.node_radius[old_node_ids]
+        self.node_xyz = self.node_xyz[old_node_ids]
+        self.node_component_id = self.node_component_id[old_node_ids]
+
+        self.reassign_component_ids()
+
     def remove_line_fragment(self, i, j):
         """
         Deletes nodes "i" and "j" from "graph", where these nodes form a connected

src/neuron_proofreader/skeleton_graph.py

@@ -236,11 +236,12 @@ def relabel_nodes(self):
         old_to_new = dict(zip(old_node_ids, new_node_ids))
         old_edge_ids = list(self.edges)
         old_irr_edge_ids = self.irreducible.edges
+        edge_attrs = {(i, j): data for i, j, data in self.edges(data=True)}
 
         # Reset graph
         self.clear()
         for (i, j) in old_edge_ids:
-            self.add_edge(old_to_new[i], old_to_new[j])
+            self.add_edge(old_to_new[i], old_to_new[j], **edge_attrs[(i, j)])
 
         self.irreducible.clear()
         for (i, j) in old_irr_edge_ids:

src/neuron_proofreader/split_proofreading/groundtruth_generation.py

@@ -74,10 +74,6 @@ def run(gt_graph, pred_graph):
         if dist > 8:
             continue
 
-        # Check if nodes are connected via a->b->c and a->c
-        if accepts_graph.is_transitive_connection(proposal):
-            continue
-
         # Check if proposal is structurally consistent
         gt_id = pred_to_gt[id1]
         is_consistent = is_structure_consistent(
@@ -120,7 +116,7 @@ def compute_proposal_proj_dist(gt_graph, pred_graph, proposal):
     for pt in geometry_util.make_line(xyz_i, xyz_j, n_pts):
         dist, _ = gt_graph.kdtree.query(pt)
         proj_dists.append(dist)
-    return np.max(proj_dists)
+    return np.percentile(proj_dists, 90)
 
 
 def find_aligned_component(gt_graph, pred_graph, nodes):
@@ -161,10 +157,10 @@ def find_aligned_component(gt_graph, pred_graph, nodes):
     gt_id = util.find_best(dists)
     dists = np.array(dists[gt_id])
     percent_aligned = len(dists) / n_pts
-    aligned_score = np.mean(dists[dists < np.percentile(dists, 80)])
+    aligned_score = np.percentile(dists, 60)
 
     # Deterine whether aligned
-    if (aligned_score < 4 and gt_id) and percent_aligned > 0.6:
+    if (aligned_score < 7 and gt_id) and percent_aligned > 0.6:
         return gt_id
     else:
         return None

src/neuron_proofreader/split_proofreading/proposal_generation.py

@@ -26,8 +26,9 @@ def __init__(
         self,
         graph,
         allow_nonleaf_targets=False,
-        filter_transitive_proposals=False,
         max_attempts=2,
+        max_proposals_per_leaf=3,
+        min_size_with_proposals=0,
         search_scaling_factor=1.5
     ):
         """
@@ -40,24 +41,24 @@ def __init__(
         allow_nonleaf_targets : bool, optional
             Indication of whether to generate proposals between leaf and nodes
             with degree 2. Default is False.
-        filter_transitive_proposals : bool, optional
-            Indication of whether to filter proposals between fragments with a
-            connection via other proposals and fragments. Default is False.
         max_attempts : int, optional
             Number of attempts made to generate proposals from a node with
             increasing search radii. Default is 2.
+        max_proposals_per_leaf : bool, optional
+            Maximum number of proposals generated at each leaf. Default is 3.
+        min_size_with_proposals : float, optional
+            Minimum fragment path length required for proposals. Default is 0.
         search_scaling_factor : 1.5, optional
             Scaling actor used to enlarge search radius for each search.
             Default is 2.
         """
         # Instance attributes
         self.allow_nonleaf_targets = allow_nonleaf_targets
-        self.filter_transitive_proposals = filter_transitive_proposals
         self.graph = graph
         self.kdtree = None
         self.max_attempts = max_attempts
-        self.max_proposals_per_leaf = graph.max_proposals_per_leaf
-        self.min_size_with_proposals = graph.min_size_with_proposals
+        self.max_proposals_per_leaf = max_proposals_per_leaf
+        self.min_size_with_proposals = min_size_with_proposals
         self.n_proposals_blocked = 0
         self.search_scaling_factor = search_scaling_factor
 
@@ -115,10 +116,6 @@ def __call__(self, initial_radius):
                     # Add proposal
                     proposals.add(pair_id)
                     connections[pair_component_id] = pair_id
-
-        # Filter proposals (if applicable)
-        if self.filter_transitive_proposals:
-            proposals = self._filter_transitive_proposals(proposals)
         return proposals
 
     def find_node_candidates(self, leaf, radius):
@@ -165,30 +162,6 @@ def get_nearby_points(self, leaf, radius):
         pts_dict = self.select_closest_components(pts_dict)
         return [val["xyz"] for val in pts_dict.values()]
 
-    def _filter_transitive_proposals(self, proposals):
-        """
-        Removes proposals that directly connect fragments a and c when an
-        indirect path exists via a → b → c.
-
-        Parameters
-        ----------
-        proposals : List[Frozenset[int]]
-            Proposals to be filtered.
-        """
-        # Initializations
-        graph = deepcopy(self.graph)
-        proposals = get_sorted_proposals(graph, proposals)
-
-        # Main
-        filtered_proposals = list()
-        for proposal in proposals:
-            if not graph.is_transitive_connection(proposal):
-                i, j = proposal
-                graph.add_edge(i, j)
-                graph.add_proposal(i, j)
-                filtered_proposals.append(proposal)
-        return filtered_proposals
-
     # --- Helpers ---
     def get_closer_endpoint(self, edge, xyz):
         """

src/neuron_proofreader/split_proofreading/split_datasets.py

@@ -135,14 +135,13 @@ def load_graph(self, swc_pointer, is_gt=False):
         # Build graph
         graph = ProposalGraph(
             anisotropy=self.config.anisotropy,
-            filter_transitive_proposals=True,
             min_size=self.config.min_size
         )
         graph.load(swc_pointer)
 
         # Filter doubles (if applicable)
         if not is_gt:
-            geometry_util.remove_doubles(graph, 160)
+            geometry_util.remove_doubles(graph, 200)
         return graph
 
     # --- Get Data ---

src/neuron_proofreader/utils/geometry_util.py

@@ -174,7 +174,7 @@ def fit_spline_1d(pts, k=3, s=None):
         Spline fit to the given points.
     """
     t = np.linspace(0, 1, len(pts))
-    s = len(pts) / s if s else len(pts) / 10
+    s = len(pts) / s if s else len(pts) / 15
     return UnivariateSpline(t, pts, k=k, s=s)
 
 
@@ -520,11 +520,13 @@ def remove_doubles(graph, max_length):
     """
     # Initializations
     components = [c for c in nx.connected_components(graph) if len(c) == 2]
+    iterator = np.argsort([len(c) for c in components])
     kdtree = graph.get_kdtree()
-
+    if graph.verbose:
+        iterator = tqdm(iterator, desc="Filter Doubled Fragments")
+        
     # Main
-    desc = "Filter Doubled Fragments"
-    for idx in tqdm(np.argsort([len(c) for c in components]), desc=desc):
+    for idx in iterator:
         i, j = tuple(components[idx])
         if graph.node_component_id[i] in graph.component_id_to_swc_id:
             if graph.edge_length((i, j)) < max_length:
@@ -533,6 +535,7 @@ def remove_doubles(graph, max_length):
                 hits = compute_projections(graph, kdtree, (i, j))
                 if is_double(hits, n_pts):
                     graph.remove_line_fragment(i, j)
+    graph.relabel_nodes()
 
 
 def compute_projections(graph, kdtree, edge):