diff --git a/src/assembly.rs b/src/assembly.rs index 27264259..7b00db9e 100644 --- a/src/assembly.rs +++ b/src/assembly.rs @@ -187,7 +187,7 @@ pub fn recurse_index_search( // Define a closure that handles recursing to a new assembly state based on // the given match. - let recurse_on_match = |i: usize, match_ix: usize| { + let recurse_on_match = |match_ix: usize| { let (h1, h2) = matches.match_fragments(match_ix); if let Some(fragments) = fragments(mol, &intermediate_frags, h1, h2) { @@ -203,7 +203,7 @@ pub fn recurse_index_search( let (child_index, child_states_searched) = recurse_index_search( mol, matches, - &state.update(fragments, i, match_ix, h1.len()), + &state.update(fragments, match_ix, h1.len().max(h2.len())), best_index.clone(), bounds, &mut cache.clone(), @@ -222,14 +222,12 @@ pub fn recurse_index_search( if parallel_mode == ParallelMode::None { matches_to_remove .iter() - .enumerate() - .for_each(|(i, match_ix)| recurse_on_match(i, *match_ix)); + .for_each(|match_ix| recurse_on_match(*match_ix)); } else { matches_to_remove .iter() - .enumerate() .par_bridge() - .for_each(|(i, match_ix)| recurse_on_match(i, *match_ix)); + .for_each(|match_ix| recurse_on_match(*match_ix)); } ( diff --git a/src/matches.rs b/src/matches.rs index a41c74e1..0df976a1 100644 --- a/src/matches.rs +++ b/src/matches.rs @@ -8,11 +8,33 @@ use petgraph::graph::EdgeIndex; use crate::{ bounds::{match_bounds, Bound}, canonize::{canonize, CanonizeMode, Labeling}, - molecule::Molecule, + molecule::{Bond, Element, Molecule}, state::State, utils::{connected_components_under_edges, edge_neighbors}, }; +/// Cached features of a fragment used for heuristic ordering. +#[derive(Debug, Copy, Clone, PartialEq, Eq)] +struct FragFeatures { + len: u16, + bond_single: u16, + bond_double: u16, + bond_triple: u16, + hetero_endpoints: u16, +} + +impl FragFeatures { + fn score(self) -> u32 { + // Deterministic, cheap heuristic: + // - Prefer rarer/higher-order chemistry (triple > double > single) + // - Prefer fragments with more heteroatom participation + (self.bond_triple as u32) * 1000 + + (self.bond_double as u32) * 100 + + (self.bond_single as u32) * 10 + + (self.hetero_endpoints as u32) + } +} + /// A node in the DAG storing fragment information; see [`Matches`]. struct DagNode { /// The fragment (i.e., connected molecular subgraph) this node represents. @@ -20,6 +42,8 @@ struct DagNode { /// The canonical ID of this node's fragment. Two [`DagNode`]s have the /// same canonical ID iff their fragments are isomorphic. canonical_id: usize, + /// Cached features used for heuristic ordering. + features: FragFeatures, /// Indices of this node's children/out-neighbors in the DAG. If u is a /// child of v then u.fragment is v.fragment with an additional edge. children: Vec, @@ -46,16 +70,57 @@ pub struct Matches { impl DagNode { /// Create a new [`DagNode`]. - pub fn new(fragment: BitSet, canonical_id: usize) -> Self { + pub fn new(fragment: BitSet, canonical_id: usize, features: FragFeatures) -> Self { Self { fragment, canonical_id, + features, children: Vec::new(), } } } impl Matches { + fn fragment_features(mol: &Molecule, fragment: &BitSet) -> FragFeatures { + let g = mol.graph(); + + let len: u16 = fragment.len().try_into().unwrap_or(u16::MAX); + let mut bond_single: u16 = 0; + let mut bond_double: u16 = 0; + let mut bond_triple: u16 = 0; + let mut hetero_endpoints: u16 = 0; + + for idx in fragment.iter() { + let e = EdgeIndex::new(idx); + let bond = *g.edge_weight(e).expect("bad edge index"); + match bond { + Bond::Single => bond_single = bond_single.saturating_add(1), + Bond::Double => bond_double = bond_double.saturating_add(1), + Bond::Triple => bond_triple = bond_triple.saturating_add(1), + } + + let (e1, e2) = g.edge_endpoints(e).expect("bad edge endpoints"); + let e1 = g.node_weight(e1).expect("bad node index").element(); + let e2 = g.node_weight(e2).expect("bad node index").element(); + + let is_hetero = |el: Element| el != Element::Carbon && el != Element::Hydrogen; + if is_hetero(e1) { + hetero_endpoints = hetero_endpoints.saturating_add(1); + } + if is_hetero(e2) { + hetero_endpoints = hetero_endpoints.saturating_add(1); + } + } + + FragFeatures { + len, + bond_single, + bond_double, + bond_triple, + hetero_endpoints, + } + } + /// Generate [`Matches`] from the given molecule and canonization mode. pub fn new(mol: &Molecule, canonize_mode: CanonizeMode) -> Self { let num_edges = mol.graph().edge_count(); @@ -71,7 +136,11 @@ impl Matches { // Add the fragment to the DAG. Since all singleton edge fragments // are trivially isomorphic, give them the same canonical ID. - dag.push(DagNode::new(frag, 0)); + dag.push(DagNode::new( + frag.clone(), + 0, + Self::fragment_features(mol, &frag), + )); parent_frag_ixs.push(i); } @@ -129,14 +198,22 @@ impl Matches { // Add matched fragments to the DAG on first match. if frag_has_match.insert(iso_ix1) { let frag1_ix = dag.len(); - dag.push(DagNode::new(frag1.clone(), next_canonical_id)); + dag.push(DagNode::new( + frag1.clone(), + next_canonical_id, + Self::fragment_features(mol, frag1), + )); dag[*frag1_parent_ix].children.push(frag1_ix); iso_to_frag_ix.insert(iso_ix1, frag1_ix); child_frag_ixs.push(frag1_ix); } if frag_has_match.insert(iso_ix2) { let frag2_ix = dag.len(); - dag.push(DagNode::new(frag2.clone(), next_canonical_id)); + dag.push(DagNode::new( + frag2.clone(), + next_canonical_id, + Self::fragment_features(mol, frag2), + )); dag[*frag2_parent_ix].children.push(frag2_ix); iso_to_frag_ix.insert(iso_ix2, frag2_ix); child_frag_ixs.push(frag2_ix); @@ -372,7 +449,38 @@ impl Matches { // Sort removable matches in ascending order of match index (i.e., // those with larger fragments first). - removable_matches.sort(); + // + // We refine this ordering deterministically using cached fragment + // features as a tie-breaker among same-size fragments. + removable_matches.sort_unstable_by(|&a, &b| { + let (a1, a2) = self.matches[a]; + let (b1, b2) = self.matches[b]; + + let a_len = self.dag[a1].features.len.max(self.dag[a2].features.len); + let b_len = self.dag[b1].features.len.max(self.dag[b2].features.len); + + // Descending match length (critical: removal sizes must be non-increasing + // so bounds that rely on `largest_removed` remain sound). + match b_len.cmp(&a_len) { + std::cmp::Ordering::Equal => { + let a_score = self.dag[a1] + .features + .score() + .max(self.dag[a2].features.score()); + let b_score = self.dag[b1] + .features + .score() + .max(self.dag[b2].features.score()); + + // Descending chemical score within same match length. + match b_score.cmp(&a_score) { + std::cmp::Ordering::Equal => a.cmp(&b), // Ascending match_ix. + ord => ord, + } + } + ord => ord, + } + }); (intermediate_frags, removable_matches) } diff --git a/src/state.rs b/src/state.rs index 048bb4e2..21e36156 100644 --- a/src/state.rs +++ b/src/state.rs @@ -11,7 +11,7 @@ pub struct State { /// The current upper bound on the assembly index, i.e., edges(mol) - 1 - /// [edges(subgraphs removed) - #(subgraphs removed)]. index: usize, - /// The indices of previously removed duplicate subgraphs. Used for + /// The indices of previously removed matches (`match_ix`). Used for /// disambiguating the serial order of two states during memoization. removal_order: Vec, /// Size of the largest duplicatable subgraph removed up to this point. @@ -40,19 +40,13 @@ impl State { /// Construct the child [`State`] resulting from removing the specified /// match from this [`State`]. - pub fn update( - &self, - fragments: Vec, - remove_ix: usize, - match_ix: usize, - remove_len: usize, - ) -> Self { + pub fn update(&self, fragments: Vec, match_ix: usize, remove_len: usize) -> Self { Self { fragments, index: self.index - remove_len + 1, removal_order: { let mut clone = self.removal_order.clone(); - clone.push(remove_ix); + clone.push(match_ix); clone }, largest_removed: remove_len,