Skip to content

Commit 8982c98

Browse files
sebschmisebschmi
committed
Step through anchor successors instead of generating all at once.
1 parent 70355a7 commit 8982c98

8 files changed

Lines changed: 632 additions & 198 deletions

File tree

‎lib_ts_chainalign/src/chain_align.rs‎

Lines changed: 131 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ use lib_tsalign::a_star_aligner::{
1212
template_switch_distance::{EqualCostRange, TemplateSwitchDirection},
1313
};
1414
use log::{debug, trace};
15+
use num_traits::Zero;
1516
use std::{
1617
fmt::Write,
1718
iter,
@@ -34,11 +35,19 @@ use crate::{
3435

3536
mod chainer;
3637

38+
pub struct AlignmentParameters {
39+
/// The step width for generating successors during chaining.
40+
///
41+
/// At most `max_successors` will be generated at a time, but at least all with minimum chaining cost.
42+
pub max_successors: usize,
43+
}
44+
3745
#[expect(clippy::too_many_arguments)]
3846
pub fn align<AlphabetType: Alphabet, Cost: AStarCost>(
3947
sequences: &AlignmentSequences,
4048
start: AlignmentCoordinates,
4149
end: AlignmentCoordinates,
50+
parameters: &AlignmentParameters,
4251
alignment_costs: &AlignmentCosts<Cost>,
4352
rc_fn: &dyn Fn(u8) -> u8,
4453
max_match_run: u32,
@@ -58,6 +67,7 @@ pub fn align<AlphabetType: Alphabet, Cost: AStarCost>(
5867
chaining_cost_function,
5968
&alignment_costs.ts_limits,
6069
k,
70+
parameters.max_successors,
6171
);
6272
let mut astar = AStar::new(context);
6373
let mut chaining_execution_count = 0;
@@ -104,21 +114,55 @@ pub fn align<AlphabetType: Alphabet, Cost: AStarCost>(
104114
}
105115
match identifier {
106116
Identifier::Start => write!(s, "start").unwrap(),
107-
Identifier::Primary { index } => {
108-
write!(s, "P{}", anchors.primary(*index)).unwrap()
117+
Identifier::StartToPrimary { offset } => {
118+
write!(s, "start-to-primary-{offset}").unwrap()
119+
}
120+
Identifier::StartToSecondary { ts_kind, offset } => {
121+
write!(s, "start-to-secondary{}-{offset}", ts_kind.digits())
122+
.unwrap()
109123
}
110-
Identifier::Secondary {
124+
Identifier::PrimaryToPrimary { index, offset } => {
125+
write!(s, "P{}-to-primary-{offset}", anchors.primary(*index))
126+
.unwrap()
127+
}
128+
Identifier::PrimaryToSecondary {
129+
index,
130+
ts_kind,
131+
offset,
132+
} => write!(
133+
s,
134+
"P{}-to-secondary{}-{offset}",
135+
anchors.primary(*index),
136+
ts_kind.digits()
137+
)
138+
.unwrap(),
139+
Identifier::SecondaryToSecondary {
111140
index,
112141
ts_kind,
113142
first_secondary_index,
143+
offset,
114144
} => write!(
115145
s,
116-
"S{}{}->{}",
146+
"S{}{}->{}-to-secondary-{offset}",
117147
ts_kind.digits(),
118148
anchors.secondary(*first_secondary_index, *ts_kind),
119149
anchors.secondary(*index, *ts_kind),
120150
)
121151
.unwrap(),
152+
Identifier::SecondaryToPrimary {
153+
index,
154+
ts_kind,
155+
first_secondary_index,
156+
offset,
157+
} => write!(
158+
s,
159+
"S{}{}->{}-to-primary-{offset}",
160+
ts_kind.digits(),
161+
anchors.secondary(*first_secondary_index, *ts_kind),
162+
anchors.secondary(*index, *ts_kind),
163+
)
164+
.unwrap(),
165+
122166
Identifier::End => write!(s, "end").unwrap(),
123167
}
124168
}
@@ -305,9 +349,29 @@ fn evaluate_chain<Cost: AStarCost>(
305349
let k = usize::try_from(max_match_run + 1).unwrap();
306350
let mut current_upper_bound = Cost::zero();
307351
let mut alignments = Vec::new();
308-
for window in chain.windows(2) {
309-
let from_anchor = window[0];
310-
let to_anchor = window[1];
352+
let mut current_from_index = 0;
353+
354+
loop {
355+
let from_anchor = chain[current_from_index];
356+
let Some((to_anchor_index, to_anchor)) = chain
357+
.iter()
358+
.copied()
359+
.enumerate()
360+
.skip(current_from_index + 1)
361+
.find(|(_, identifier)| match identifier {
362+
Identifier::Start => true,
363+
Identifier::StartToPrimary { .. } => false,
364+
Identifier::StartToSecondary { .. } => false,
365+
Identifier::PrimaryToPrimary { offset, .. }
366+
| Identifier::PrimaryToSecondary { offset, .. }
367+
| Identifier::SecondaryToSecondary { offset, .. }
368+
| Identifier::SecondaryToPrimary { offset, .. } => offset.is_zero(),
369+
Identifier::End => true,
370+
})
371+
else {
372+
break;
373+
};
374+
current_from_index = to_anchor_index;
311375

312376
match (from_anchor, to_anchor) {
313377
(Identifier::Start, Identifier::End) => {
@@ -328,7 +392,11 @@ fn evaluate_chain<Cost: AStarCost>(
328392
}
329393
current_upper_bound += chaining_cost_function.start_to_end();
330394
}
331-
(Identifier::Start, Identifier::Primary { index }) => {
395+
(
396+
Identifier::Start,
397+
Identifier::PrimaryToPrimary { index, .. }
398+
| Identifier::PrimaryToSecondary { index, .. },
399+
) => {
332400
let end = anchors.primary(index).start();
333401
if final_evaluation || !chaining_cost_function.is_primary_from_start_exact(index) {
334402
let alignment = GapAffineAlignment::new(
@@ -355,7 +423,11 @@ fn evaluate_chain<Cost: AStarCost>(
355423
}
356424
current_upper_bound += chaining_cost_function.primary_from_start(index);
357425
}
358-
(Identifier::Start, Identifier::Secondary { index, ts_kind, .. }) => {
426+
(
427+
Identifier::Start,
428+
Identifier::SecondaryToPrimary { index, ts_kind, .. }
429+
| Identifier::SecondaryToSecondary { index, ts_kind, .. },
430+
) => {
359431
let end = anchors.secondary(index, ts_kind).start(ts_kind);
360432
if final_evaluation
361433
|| !chaining_cost_function.is_jump_12_from_start_exact(index, ts_kind)
@@ -386,7 +458,7 @@ fn evaluate_chain<Cost: AStarCost>(
386458
}
387459
current_upper_bound += chaining_cost_function.jump_12_from_start(index, ts_kind);
388460
}
389-
(Identifier::Primary { index }, Identifier::End) => {
461+
(Identifier::PrimaryToPrimary { index, .. }, Identifier::End) => {
390462
let start = anchors.primary(index).end(k);
391463
if final_evaluation || !chaining_cost_function.is_primary_to_end_exact(index) {
392464
let alignment = GapAffineAlignment::new(
@@ -410,7 +482,7 @@ fn evaluate_chain<Cost: AStarCost>(
410482
}
411483
current_upper_bound += chaining_cost_function.primary_to_end(index);
412484
}
413-
(Identifier::Secondary { index, ts_kind, .. }, Identifier::End) => {
485+
(Identifier::SecondaryToPrimary { index, ts_kind, .. }, Identifier::End) => {
414486
let start = anchors.secondary(index, ts_kind).end(ts_kind, k);
415487
if final_evaluation
416488
|| !chaining_cost_function.is_jump_34_to_end_exact(index, ts_kind)
@@ -443,8 +515,15 @@ fn evaluate_chain<Cost: AStarCost>(
443515
current_upper_bound += chaining_cost_function.jump_34_to_end(index, ts_kind);
444516
}
445517
(
446-
Identifier::Primary { index: from_index },
447-
Identifier::Primary { index: to_index },
518+
Identifier::PrimaryToPrimary {
519+
index: from_index, ..
520+
},
521+
Identifier::PrimaryToPrimary {
522+
index: to_index, ..
523+
}
524+
| Identifier::PrimaryToSecondary {
525+
index: to_index, ..
526+
},
448527
) => {
449528
if anchors
450529
.primary(from_index)
@@ -487,8 +566,15 @@ fn evaluate_chain<Cost: AStarCost>(
487566
current_upper_bound += chaining_cost_function.primary(from_index, to_index);
488567
}
489568
(
490-
Identifier::Primary { index: from_index },
491-
Identifier::Secondary {
569+
Identifier::PrimaryToSecondary {
570+
index: from_index, ..
571+
},
572+
Identifier::SecondaryToSecondary {
573+
index: to_index,
574+
ts_kind,
575+
..
576+
}
577+
| Identifier::SecondaryToPrimary {
492578
index: to_index,
493579
ts_kind,
494580
..
@@ -530,12 +616,17 @@ fn evaluate_chain<Cost: AStarCost>(
530616
chaining_cost_function.jump_12(from_index, to_index, ts_kind);
531617
}
532618
(
533-
Identifier::Secondary {
619+
Identifier::SecondaryToSecondary {
534620
index: from_index,
535621
ts_kind,
536622
..
537623
},
538-
Identifier::Secondary {
624+
Identifier::SecondaryToSecondary {
625+
index: to_index,
626+
ts_kind: to_ts_kind,
627+
..
628+
}
629+
| Identifier::SecondaryToPrimary {
539630
index: to_index,
540631
ts_kind: to_ts_kind,
541632
..
@@ -586,12 +677,17 @@ fn evaluate_chain<Cost: AStarCost>(
586677
chaining_cost_function.secondary(from_index, to_index, ts_kind);
587678
}
588679
(
589-
Identifier::Secondary {
680+
Identifier::SecondaryToPrimary {
590681
index: from_index,
591682
ts_kind,
592683
..
593684
},
594-
Identifier::Primary { index: to_index },
685+
Identifier::PrimaryToPrimary {
686+
index: to_index, ..
687+
}
688+
| Identifier::PrimaryToSecondary {
689+
index: to_index, ..
690+
},
595691
) => {
596692
let start = anchors.secondary(from_index, ts_kind).end(ts_kind, k);
597693
let end = anchors.primary(to_index).start();
@@ -630,7 +726,22 @@ fn evaluate_chain<Cost: AStarCost>(
630726
current_upper_bound +=
631727
chaining_cost_function.jump_34(from_index, to_index, ts_kind);
632728
}
633-
(Identifier::End, _) | (_, Identifier::Start) => unreachable!(),
729+
(Identifier::End, _)
730+
| (_, Identifier::Start)
731+
| (
732+
Identifier::SecondaryToPrimary { .. } | Identifier::PrimaryToPrimary { .. },
733+
Identifier::SecondaryToPrimary { .. } | Identifier::SecondaryToSecondary { .. },
734+
)
735+
| (
736+
Identifier::PrimaryToSecondary { .. } | Identifier::SecondaryToSecondary { .. },
737+
Identifier::PrimaryToPrimary { .. }
738+
| Identifier::PrimaryToSecondary { .. }
739+
| Identifier::End,
740+
)
741+
| (Identifier::StartToPrimary { .. } | Identifier::StartToSecondary { .. }, _)
742+
| (_, Identifier::StartToPrimary { .. } | Identifier::StartToSecondary { .. }) => {
743+
unreachable!()
744+
}
634745
}
635746
}
636747

0 commit comments

Comments
 (0)