use crate::config::{AlignmentMode, ConsensusMode, PoaConfig};
use crate::error::PoaError;
use crate::types::{BubbleSite, Consensus, CoverageGap, GapKind, GraphStats};
use std::collections::HashMap;
#[cfg(test)]
use std::cell::Cell as StdCell;
#[cfg(test)]
thread_local! {
pub(crate) static SKIP_COUNTER: StdCell<u64> = StdCell::new(0);
pub(crate) static NODE_COUNTER: StdCell<u64> = StdCell::new(0);
}
#[cfg(test)]
pub(crate) fn reset_skip_counters() {
SKIP_COUNTER.with(|c| c.set(0));
NODE_COUNTER.with(|c| c.set(0));
}
#[cfg(test)]
pub(crate) fn skip_rate() -> f64 {
let skips = SKIP_COUNTER.with(|c| c.get());
let nodes = NODE_COUNTER.with(|c| c.get());
if nodes == 0 {
0.0
} else {
skips as f64 / nodes as f64
}
}
const UNSET: i32 = i32::MIN / 2;
const LOCK_EPS: usize = 5;
const MINI_K: usize = 15;
const MINI_W: usize = 10;
const MINI_EPS_BASE: usize = 3;
const MINI_MIN_CHAIN: usize = 15;
struct AlignScratch {
lock_node_j: Vec<(usize, u32)>,
lock_exit_j: Vec<(usize, u32)>,
}
impl AlignScratch {
fn new() -> Self {
Self {
lock_node_j: Vec::new(),
lock_exit_j: Vec::new(),
}
}
fn clear(&mut self) {
self.lock_node_j.clear();
self.lock_exit_j.clear();
}
fn get_node_j(&self, node_idx: usize) -> Option<usize> {
self.lock_node_j
.binary_search_by_key(&node_idx, |&(idx, _)| idx)
.ok()
.map(|pos| self.lock_node_j[pos].1 as usize)
}
fn get_exit_j(&self, node_idx: usize) -> Option<usize> {
self.lock_exit_j
.binary_search_by_key(&node_idx, |&(idx, _)| idx)
.ok()
.map(|pos| self.lock_exit_j[pos].1 as usize)
}
}
#[inline]
fn edge_weight(nodes: &[Node], from: usize, to: usize) -> i32 {
nodes[from]
.out_edges
.iter()
.find(|&&(t, _)| t == to)
.map(|&(_, w)| w)
.unwrap_or(0)
}
#[inline]
fn safe_add(a: i32, b: i32) -> i32 {
if a == UNSET {
UNSET
} else {
a.saturating_add(b)
}
}
struct Node {
base: u8,
out_edges: Vec<(usize, i32)>,
in_edges: Vec<usize>,
coverage: u32,
delete_count: u32,
}
pub struct PoaGraph {
nodes: Vec<Node>,
config: PoaConfig,
n_reads: usize,
reads: Vec<Vec<u8>>,
edge_reads: HashMap<(usize, usize), Vec<u32>>,
warnings: usize,
cached_spine: Vec<(usize, u8, i32)>,
spine_updated_at: usize,
spine_interval: usize,
align_scratch: AlignScratch,
spine_mers: HashMap<u64, u32>,
}
const VIRTUAL: u32 = u32::MAX;
#[derive(Clone, Copy, PartialEq, Eq)]
enum State {
M,
I,
D,
}
#[derive(Clone, Copy)]
struct Cell {
score: i32,
pred_t: u32,
}
impl Cell {
fn unset() -> Self {
Cell {
score: UNSET,
pred_t: 0,
}
}
}
#[derive(Clone, Copy, Debug)]
pub enum AlignOp {
Match(usize),
Insert(u8),
Delete(usize),
}
fn push_node(nodes: &mut Vec<Node>, base: u8) -> usize {
let idx = nodes.len();
nodes.push(Node {
base,
out_edges: Vec::new(),
in_edges: Vec::new(),
coverage: 0,
delete_count: 0,
});
idx
}
fn add_edge(nodes: &mut [Node], from: usize, to: usize) {
nodes[from].out_edges.push((to, 1));
nodes[to].in_edges.push(from);
}
fn increment_or_add_edge(nodes: &mut [Node], from: usize, to: usize) {
for (succ, w) in nodes[from].out_edges.iter_mut() {
if *succ == to {
*w += 1;
return;
}
}
nodes[from].out_edges.push((to, 1));
nodes[to].in_edges.push(from);
}
fn topological_order(nodes: &[Node]) -> (Vec<usize>, Vec<usize>) {
let n = nodes.len();
let mut in_deg: Vec<usize> = nodes.iter().map(|nd| nd.in_edges.len()).collect();
let mut queue: std::collections::VecDeque<usize> = (0..n).filter(|&i| in_deg[i] == 0).collect();
let mut topo: Vec<usize> = Vec::with_capacity(n);
while let Some(u) = queue.pop_front() {
topo.push(u);
for &(v, _) in &nodes[u].out_edges {
in_deg[v] -= 1;
if in_deg[v] == 0 {
queue.push_back(v);
}
}
}
let mut rank_of = vec![0usize; n];
for (t, &node_idx) in topo.iter().enumerate() {
rank_of[node_idx] = t;
}
(topo, rank_of)
}
fn materialize_arm_len(nodes: &[Node], start: usize, max_depth: usize) -> usize {
let mut len = 0;
let mut cur = start;
for _ in 0..max_depth {
if len >= 1 && nodes[cur].in_edges.len() > 1 {
break;
}
len += 1;
match nodes[cur].out_edges.as_slice() {
[(next, _)] => cur = *next,
_ => break,
}
}
len
}
fn compute_bubble_ranges(nodes: &[Node], topo: &[usize]) -> Vec<Option<(usize, usize)>> {
let n = topo.len();
let nn = nodes.len();
let mut ranges: Vec<Option<(usize, usize)>> = vec![None; n];
let mut gen_mark = vec![0u32; nn];
let mut current_gen = 0u32;
let mut t = 0;
while t < n {
if nodes[topo[t]].out_edges.len() >= 2 {
current_gen = current_gen.wrapping_add(1);
let entry_node = topo[t];
gen_mark[entry_node] = current_gen;
let mut outstanding = nodes[entry_node].out_edges.len();
let mut exit_t = None;
for (tt, &node_idx) in topo.iter().enumerate().skip(t + 1) {
let bubble_in = nodes[node_idx]
.in_edges
.iter()
.filter(|&&p| gen_mark[p] == current_gen)
.count();
if bubble_in == 0 {
continue;
}
gen_mark[node_idx] = current_gen;
outstanding -= bubble_in;
if outstanding == 0 {
exit_t = Some(tt);
break;
}
outstanding += nodes[node_idx].out_edges.len();
}
if let Some(et) = exit_t {
ranges[t..=et].fill(Some((t, et)));
t = et + 1;
} else {
ranges[t..n].fill(Some((t, n - 1)));
break;
}
} else {
t += 1;
}
}
ranges
}
fn spine_diff(old: &[(usize, u8, i32)], new: &[(usize, u8, i32)]) -> usize {
if old.is_empty() {
return usize::MAX;
}
let len_diff = old.len().abs_diff(new.len());
let base_diffs = old[..old.len().min(new.len())]
.iter()
.zip(new.iter())
.filter(|(o, n)| o.1 != n.1)
.count();
base_diffs + len_diff
}
const SPINE_STABLE_THRESHOLD: usize = 3;
const SPINE_MAX_INTERVAL: usize = 32;
const LOOKAHEAD_K: usize = 5;
const LOOKAHEAD_MARGIN: i32 = 2;
const ARM_MAX_DEPTH: usize = 4096;
fn collect_arm_nodes(
nodes: &[Node],
rank_of: &[usize],
start_idx: usize,
exit_t: usize,
out: &mut Vec<usize>,
) -> bool {
out.clear();
let mut cur = start_idx;
for _ in 0..ARM_MAX_DEPTH {
if rank_of[cur] == exit_t {
return true;
}
out.push(cur);
match nodes[cur].out_edges.as_slice() {
[(next, _)] => cur = *next,
_ => return false,
}
}
false
}
fn score_arm_prefix(
arm: &[usize],
nodes: &[Node],
query: &[u8],
j: usize,
match_score: i32,
mismatch_score: i32,
) -> i32 {
let len = arm
.len()
.min(LOOKAHEAD_K)
.min(query.len().saturating_sub(j));
arm[..len]
.iter()
.enumerate()
.map(|(i, &nidx)| {
if nodes[nidx].base == query[j + i] {
match_score
} else {
mismatch_score
}
})
.sum()
}
const SLIDE_MIN_CONSEC: usize = 2;
fn slide_lock(
all_arms: &[Vec<usize>],
nodes: &[Node],
query: &[u8],
j_entry: usize,
) -> Option<usize> {
let n_arms = all_arms.len();
let remaining = query.len().saturating_sub(j_entry);
if remaining == 0 || n_arms < 2 {
return None;
}
let min_arm_len = all_arms.iter().map(|a| a.len()).min().unwrap_or(0);
if min_arm_len < LOOKAHEAD_K {
return None;
}
let max_steps = all_arms
.iter()
.map(|a| a.len())
.max()
.unwrap_or(0)
.min(remaining);
let mut alive = vec![true; n_arms];
let mut alive_count = n_arms;
let mut consec_unique = 0usize;
let mut consec_winner: Option<usize> = None;
for i in 0..max_steps {
let any_continuing = (0..n_arms).any(|idx| alive[idx] && all_arms[idx].len() > i);
if any_continuing {
for idx in 0..n_arms {
if alive[idx] && all_arms[idx].len() <= i {
alive[idx] = false;
alive_count -= 1;
}
}
}
if alive_count <= 1 {
break;
}
let q = query[j_entry + i];
let matched: Vec<bool> = (0..n_arms)
.map(|idx| alive[idx] && all_arms[idx].len() > i && nodes[all_arms[idx][i]].base == q)
.collect();
let match_count = matched.iter().filter(|&&m| m).count();
let mismatch_count = (0..n_arms)
.filter(|&idx| alive[idx] && !matched[idx])
.count();
if match_count == 1 && mismatch_count > 0 {
let candidate = matched.iter().position(|&m| m).unwrap();
if consec_winner == Some(candidate) {
consec_unique += 1;
} else {
consec_unique = 1;
consec_winner = Some(candidate);
}
if consec_unique >= SLIDE_MIN_CONSEC {
for idx in 0..n_arms {
if alive[idx] && !matched[idx] {
alive[idx] = false;
alive_count -= 1;
}
}
break;
}
} else {
consec_unique = 0;
consec_winner = None;
}
}
if alive_count == 1 {
alive.iter().position(|&a| a)
} else {
None
}
}
#[inline]
fn encode_base(b: u8) -> Option<u64> {
match b {
b'A' | b'a' => Some(0),
b'C' | b'c' => Some(1),
b'G' | b'g' => Some(2),
b'T' | b't' => Some(3),
_ => None,
}
}
fn compute_minimizers(seq: &[u8], k: usize, w: usize) -> Vec<(u64, usize)> {
let n = seq.len();
if n < k {
return vec![];
}
let n_kmers = n - k + 1;
let mask = if k * 2 < 64 {
(1u64 << (k * 2)) - 1
} else {
u64::MAX
};
let mut kmer_hashes: Vec<Option<u64>> = Vec::with_capacity(n_kmers);
let mut hash: u64 = 0;
let mut valid: usize = 0;
for (i, &b) in seq.iter().enumerate() {
match encode_base(b) {
Some(bits) => {
hash = ((hash << 2) | bits) & mask;
valid += 1;
}
None => {
hash = 0;
valid = 0;
}
}
if i + 1 >= k {
kmer_hashes.push(if valid >= k { Some(hash) } else { None });
}
}
let win = w.min(n_kmers);
let mut result: Vec<(u64, usize)> = Vec::new();
let mut last: Option<(u64, usize)> = None;
for start in 0..=(n_kmers - win) {
let mut best_hash = u64::MAX;
let mut best_pos = start;
for (off, entry) in kmer_hashes[start..start + win].iter().enumerate() {
if let Some(h) = entry {
if *h < best_hash {
best_hash = *h;
best_pos = start + off;
}
}
}
if best_hash == u64::MAX {
continue;
}
let entry = (best_hash, best_pos);
if last != Some(entry) {
last = Some(entry);
result.push(entry);
}
}
result
}
fn build_spine_mers(spine: &[(usize, u8, i32)], k: usize, w: usize) -> HashMap<u64, u32> {
let spine_seq: Vec<u8> = spine.iter().map(|&(_, b, _)| b).collect();
let mers = compute_minimizers(&spine_seq, k, w);
let mut counts: HashMap<u64, (u32, u32)> = HashMap::new();
for (hash, pos) in mers {
let e = counts.entry(hash).or_insert((0, pos as u32));
e.0 += 1;
}
counts
.into_iter()
.filter(|(_, (c, _))| *c == 1)
.map(|(h, (_, p))| (h, p))
.collect()
}
fn lis_indices(vals: &[usize]) -> Vec<usize> {
let n = vals.len();
if n == 0 {
return vec![];
}
let mut tails: Vec<usize> = Vec::new();
let mut tail_idx: Vec<usize> = Vec::new();
let mut pred: Vec<Option<usize>> = vec![None; n];
for i in 0..n {
let v = vals[i];
let p = tails.partition_point(|&t| t < v);
if p == tails.len() {
tails.push(v);
tail_idx.push(i);
} else {
tails[p] = v;
tail_idx[p] = i;
}
pred[i] = if p > 0 { Some(tail_idx[p - 1]) } else { None };
}
let lis_len = tails.len();
let mut chain = Vec::with_capacity(lis_len);
let mut cur = Some(tail_idx[lis_len - 1]);
while let Some(i) = cur {
chain.push(i);
cur = pred[i];
}
chain.reverse();
chain
}
fn build_anchor_chain(
read_mers: &[(u64, usize)],
spine_mers: &HashMap<u64, u32>,
spine: &[(usize, u8, i32)],
rank_of: &[usize],
) -> Vec<(usize, usize)> {
let mut candidates: Vec<(usize, usize)> = read_mers
.iter()
.filter_map(|&(hash, read_pos)| {
spine_mers.get(&hash).map(|&sr| {
let node_idx = spine[sr as usize].0;
(read_pos, rank_of[node_idx])
})
})
.collect();
if candidates.is_empty() {
return vec![];
}
candidates.sort_unstable_by_key(|&(_, t)| t);
let read_positions: Vec<usize> = candidates.iter().map(|&(r, _)| r).collect();
let idxs = lis_indices(&read_positions);
idxs.iter().map(|&i| candidates[i]).collect()
}
fn anchor_j_bounds(
t: usize,
anchors: &[(usize, usize)], l: usize,
) -> Option<(usize, usize)> {
if anchors.is_empty() {
return None;
}
let pos = anchors.partition_point(|&(_, at)| at <= t);
if pos > 0 && anchors[pos - 1].1 == t {
let r = anchors[pos - 1].0;
let lo = r.saturating_sub(MINI_EPS_BASE).max(1);
let hi = (r + MINI_EPS_BASE).min(l);
Some((lo, hi))
} else {
None
}
}
#[allow(clippy::too_many_arguments)]
#[inline]
fn anchor_refine_spine(
j_lo: usize,
j_hi: usize,
t: usize,
anchors: &[(usize, usize)],
l: usize,
is_source: bool,
is_spine: bool,
in_bubble: bool,
is_locked: bool,
j_center: usize,
) -> (usize, usize) {
if anchors.is_empty() || is_source || !is_spine || in_bubble || is_locked {
return (j_lo, j_hi);
}
match anchor_j_bounds(t, anchors, l) {
Some((alo, ahi)) => {
if ahi < j_center {
return (j_lo, j_hi);
}
let nlo = j_lo.max(alo);
let nhi = j_hi.min(ahi);
if nlo <= nhi { (nlo, nhi) } else { (j_lo, j_hi) }
}
None => (j_lo, j_hi),
}
}
const SPINE_MARGIN_MIN: usize = 50;
#[inline(always)]
fn gs(mat: &[Cell], t: usize, j: usize, j_lo: usize, j_hi: usize, rw: usize) -> i32 {
if j == 0 || j < j_lo || j > j_hi {
UNSET
} else {
mat[t * rw + (j - j_lo)].score
}
}
#[inline(always)]
fn gsd(
del: &[Cell],
del0: &[Cell],
t: usize,
j: usize,
j_lo: usize,
j_hi: usize,
rw: usize,
) -> i32 {
if j == 0 {
del0[t].score
} else if j < j_lo || j > j_hi {
UNSET
} else {
del[t * rw + (j - j_lo)].score
}
}
#[allow(clippy::too_many_arguments)]
fn align(
nodes: &[Node],
topo: &[usize],
rank_of: &[usize],
spine: &[(usize, u8, i32)],
query: &[u8],
cfg: &PoaConfig,
scratch: &mut AlignScratch,
anchors: &[(usize, usize)], ) -> Result<Vec<AlignOp>, PoaError> {
scratch.clear();
let n = topo.len();
let l = query.len();
let go = cfg.gap_open;
let ge = cfg.gap_extend;
let semi = cfg.alignment_mode == AlignmentMode::SemiGlobal;
let nn = nodes.len();
let mut on_spine = vec![false; nn];
let mut spine_prev: Vec<Option<usize>> = vec![None; nn];
for (sp, &(node_idx, _, _)) in spine.iter().enumerate() {
on_spine[node_idx] = true;
if sp > 0 {
spine_prev[node_idx] = Some(spine[sp - 1].0);
}
}
let bubble_ranges = compute_bubble_ranges(nodes, topo);
let mut bubble_entry_j = vec![0usize; n];
let mut best_j_per_t = vec![0usize; n];
let mut lookahead_skip = vec![false; nn];
let spine_margin: usize = if cfg.adaptive_band {
let w = cfg.adaptive_band_b + (cfg.adaptive_band_f * l as f32).ceil() as usize;
let w = if cfg.band_width > 0 {
w.max(cfg.band_width)
} else {
w
};
w.max(4)
} else if cfg.band_width > 0 {
cfg.band_width
} else {
SPINE_MARGIN_MIN
};
let row_width = (2 * spine_margin + 2).min(l).max(1);
let mut del0 = vec![Cell::unset(); n];
let bsize = n * row_width;
let mut m = vec![Cell::unset(); bsize];
let mut ins = vec![Cell::unset(); bsize];
let mut del = vec![Cell::unset(); bsize];
let mut j_lo_arr = vec![1usize; n];
let mut j_hi_arr = vec![0usize; n];
let mut arm_scratch: Vec<usize> = Vec::new();
let mut all_arms: Vec<Vec<usize>> = Vec::new();
for (t, &node_idx) in topo.iter().enumerate() {
if lookahead_skip[node_idx] {
continue;
}
let node_base = nodes[node_idx].base;
let is_source = nodes[node_idx].in_edges.is_empty();
let j_pred_max = if is_source {
0
} else {
nodes[node_idx]
.in_edges
.iter()
.map(|&p| best_j_per_t[rank_of[p]])
.max()
.unwrap_or(0)
};
let (j_lo, j_hi) = if is_source {
(1usize, row_width.min(l))
} else if let Some(j_center) = scratch.get_node_j(node_idx) {
let lo = j_center.saturating_sub(LOCK_EPS).max(1);
let hi = (j_center + LOCK_EPS).min(l);
(lo, hi)
} else if let Some(j_center) = scratch.get_exit_j(node_idx) {
let lo = j_center.saturating_sub(spine_margin).max(1);
let hi = (j_center + spine_margin).min(l);
(lo, hi)
} else {
match bubble_ranges[t] {
Some((entry_t, exit_t)) => {
if t == entry_t {
bubble_entry_j[entry_t] = j_pred_max;
}
let bej = bubble_entry_j[entry_t];
let bubble_span = exit_t.saturating_sub(entry_t) + 1;
let lo = bej.saturating_sub(spine_margin).max(1);
let hi = (bej + bubble_span.min(spine_margin) + spine_margin).min(l);
(lo, hi)
}
None => {
let j_center = j_pred_max.saturating_add(1);
let lo = j_center.saturating_sub(spine_margin).max(1);
let hi = (j_center + spine_margin).min(l);
(lo, hi)
}
}
};
let j_hi = j_hi.min(j_lo + row_width - 1);
let (j_lo, j_hi) = anchor_refine_spine(
j_lo,
j_hi,
t,
anchors,
l,
is_source,
on_spine[node_idx],
bubble_ranges[t].is_some(),
scratch.get_node_j(node_idx).is_some() || scratch.get_exit_j(node_idx).is_some(),
j_pred_max.saturating_add(1),
);
j_lo_arr[t] = j_lo;
j_hi_arr[t] = j_hi;
#[cfg(test)]
if !is_source {
NODE_COUNTER.with(|c| c.set(c.get() + 1));
}
if !is_source && on_spine[node_idx] {
if let Some(prev_sp) = spine_prev[node_idx] {
let active_pred_ok = match nodes[node_idx].in_edges.as_slice() {
[p] => *p == prev_sp,
_ => {
nodes[node_idx]
.in_edges
.iter()
.all(|&p| p == prev_sp || lookahead_skip[p])
&& nodes[node_idx].in_edges.contains(&prev_sp)
}
};
if active_pred_ok {
let pt = rank_of[prev_sp];
let bj = best_j_per_t[pt];
if bj < l && node_base == query[bj] {
let m_prev = gs(&m, pt, bj, j_lo_arr[pt], j_hi_arr[pt], row_width);
let i_prev = gs(&ins, pt, bj, j_lo_arr[pt], j_hi_arr[pt], row_width);
let d_prev =
gsd(&del, &del0, pt, bj, j_lo_arr[pt], j_hi_arr[pt], row_width);
let pred_is_source = nodes[topo[pt]].in_edges.is_empty();
let i_ok = i_prev == UNSET || (pred_is_source && m_prev > i_prev);
let d_ok = d_prev == UNSET || (pred_is_source && m_prev > d_prev);
if m_prev != UNSET && i_ok && d_ok {
let do_skip = match nodes[node_idx].out_edges.as_slice() {
[_] => true,
_ if bj + 1 < l => {
let next_q = query[bj + 1];
let mut spine_succ = None;
let mut resolved = true;
for &(s, _) in &nodes[node_idx].out_edges {
if on_spine[s] && spine_prev[s] == Some(node_idx) {
if nodes[s].base == next_q {
spine_succ = Some(s);
} else {
resolved = false;
break;
}
} else if !lookahead_skip[s] && nodes[s].base == next_q {
resolved = false;
break;
}
}
if resolved {
if let Some(ss) = spine_succ {
for &(s, _) in &nodes[node_idx].out_edges {
if s != ss && !lookahead_skip[s] {
lookahead_skip[s] = true;
}
}
true
} else {
false
}
} else {
false
}
}
_ => false,
};
if do_skip {
let score = m_prev + cfg.match_score;
m[t * row_width + (bj + 1 - j_lo)] = Cell {
score,
pred_t: pt as u32,
};
best_j_per_t[t] = bj + 1;
#[cfg(test)]
SKIP_COUNTER.with(|c| c.set(c.get() + 1));
continue;
}
}
}
}
}
}
{
let (mut best, mut best_pred) = (UNSET, 0u32);
if is_source {
let val = go + ge;
if val > best {
best = val;
best_pred = VIRTUAL;
}
}
for &p in &nodes[node_idx].in_edges {
let pt = rank_of[p];
let vd = safe_add(del0[pt].score, ge);
if vd > best {
best = vd;
best_pred = pt as u32;
}
}
if best != UNSET {
del0[t] = Cell {
score: best,
pred_t: best_pred,
};
}
}
let mut row_best_j = 0usize;
let mut row_best_score = UNSET;
for j in j_lo..=j_hi {
let q_base = query[j - 1];
let sc = if node_base == q_base {
cfg.match_score
} else {
cfg.mismatch_score
};
let ixcur = t * row_width + (j - j_lo);
{
let (mut best, mut best_pred) = (UNSET, VIRTUAL);
if j == 1 && (is_source || semi) && sc > best {
best = sc;
best_pred = VIRTUAL;
}
if is_source && j > 1 {
let val = safe_add(go + (j as i32 - 1) * ge, sc);
if val != UNSET && val > best {
best = val;
best_pred = VIRTUAL;
}
}
let mut best_edge_w = i32::MAX;
for &p in &nodes[node_idx].in_edges {
let pt = rank_of[p];
let ew = edge_weight(nodes, p, node_idx);
let vm = safe_add(gs(&m, pt, j - 1, j_lo_arr[pt], j_hi_arr[pt], row_width), sc);
if vm != UNSET && (vm > best || (vm == best && ew > best_edge_w)) {
best = vm;
best_pred = pt as u32;
best_edge_w = ew;
}
let vi = safe_add(
gs(&ins, pt, j - 1, j_lo_arr[pt], j_hi_arr[pt], row_width),
sc,
);
if vi != UNSET && (vi > best || (vi == best && ew > best_edge_w)) {
best = vi;
best_pred = pt as u32;
best_edge_w = ew;
}
let vd = safe_add(
gsd(
&del,
&del0,
pt,
j - 1,
j_lo_arr[pt],
j_hi_arr[pt],
row_width,
),
sc,
);
if vd != UNSET && (vd > best || (vd == best && ew > best_edge_w)) {
best = vd;
best_pred = pt as u32;
best_edge_w = ew;
}
}
if best != UNSET {
m[ixcur] = Cell {
score: best,
pred_t: best_pred,
};
if best > row_best_score {
row_best_score = best;
row_best_j = j;
}
}
}
{
let (mut best, mut best_pred) = (UNSET, VIRTUAL);
if is_source && j == 1 {
let val = go + ge;
if val > best {
best = val;
best_pred = VIRTUAL;
}
}
if j > j_lo {
let ixprev = t * row_width + (j - 1 - j_lo);
let vm = safe_add(m[ixprev].score, go + ge);
if vm != UNSET && vm > best {
best = vm;
best_pred = t as u32;
}
let vi = safe_add(ins[ixprev].score, ge);
if vi != UNSET && vi > best {
best = vi;
best_pred = t as u32;
}
}
if best != UNSET {
ins[ixcur] = Cell {
score: best,
pred_t: best_pred,
};
}
}
{
let (mut best, mut best_pred) = (UNSET, 0u32);
let mut best_edge_w = i32::MAX;
for &p in &nodes[node_idx].in_edges {
let pt = rank_of[p];
let ew = edge_weight(nodes, p, node_idx);
let vm = safe_add(
gs(&m, pt, j, j_lo_arr[pt], j_hi_arr[pt], row_width),
go + ge,
);
if vm != UNSET && (vm > best || (vm == best && ew > best_edge_w)) {
best = vm;
best_pred = pt as u32;
best_edge_w = ew;
}
let vi = safe_add(
gs(&ins, pt, j, j_lo_arr[pt], j_hi_arr[pt], row_width),
go + ge,
);
if vi != UNSET && (vi > best || (vi == best && ew > best_edge_w)) {
best = vi;
best_pred = pt as u32;
best_edge_w = ew;
}
let vd = safe_add(
gsd(&del, &del0, pt, j, j_lo_arr[pt], j_hi_arr[pt], row_width),
ge,
);
if vd != UNSET && (vd > best || (vd == best && ew > best_edge_w)) {
best = vd;
best_pred = pt as u32;
best_edge_w = ew;
}
}
if best != UNSET {
del[ixcur] = Cell {
score: best,
pred_t: best_pred,
};
}
}
}
if row_best_score != UNSET {
best_j_per_t[t] = row_best_j;
} else {
best_j_per_t[t] = if is_source {
1
} else {
j_pred_max.saturating_add(1).min(l)
};
}
if let Some((entry_t, exit_t)) = bubble_ranges[t] {
if t == entry_t && nodes[node_idx].out_edges.len() >= 2 {
let j_entry = best_j_per_t[t];
if j_entry < l {
all_arms.clear();
let mut complex = false;
for &(arm_start, _) in &nodes[node_idx].out_edges {
if collect_arm_nodes(nodes, rank_of, arm_start, exit_t, &mut arm_scratch) {
all_arms.push(arm_scratch.clone());
} else {
complex = true;
break;
}
}
if !complex && all_arms.len() >= 2 {
let min_scored = all_arms
.iter()
.map(|arm| {
arm.len()
.min(LOOKAHEAD_K)
.min(query.len().saturating_sub(j_entry))
})
.min()
.unwrap_or(0);
let winner = if min_scored >= LOOKAHEAD_K {
let scores: Vec<i32> = all_arms
.iter()
.map(|arm| {
score_arm_prefix(
arm,
nodes,
query,
j_entry,
cfg.match_score,
cfg.mismatch_score,
)
})
.collect();
let best = *scores.iter().max().unwrap();
let winner_count = scores.iter().filter(|&&s| s == best).count();
if winner_count == 1 {
let second_best = scores
.iter()
.copied()
.filter(|&s| s != best)
.max()
.unwrap_or(i32::MIN);
if best - second_best >= LOOKAHEAD_MARGIN {
scores.iter().position(|&s| s == best)
} else {
None
}
} else {
None
}
} else {
None
};
let winner =
winner.or_else(|| slide_lock(&all_arms, nodes, query, j_entry));
if let Some(w) = winner {
for (arm_idx, arm) in all_arms.iter().enumerate() {
if arm_idx != w {
for &losing_node in arm {
lookahead_skip[losing_node] = true;
}
}
}
let arm = &all_arms[w];
let arm_len = arm.len();
for (d, &win_node) in arm.iter().enumerate() {
let j_c = (j_entry + d + 1).min(l) as u32;
scratch.lock_node_j.push((win_node, j_c));
}
scratch.lock_node_j.sort_unstable_by_key(|&(idx, _)| idx);
let exit_node = topo[exit_t];
let exit_j = (j_entry + arm_len).min(l) as u32;
match scratch
.lock_exit_j
.binary_search_by_key(&exit_node, |&(idx, _)| idx)
{
Ok(pos) => {
scratch.lock_exit_j[pos].1 =
exit_j.max(scratch.lock_exit_j[pos].1)
}
Err(pos) => scratch.lock_exit_j.insert(pos, (exit_node, exit_j)),
}
}
}
}
}
}
}
let terminal_best = (0..n)
.flat_map(|t| {
let jlo = j_lo_arr[t];
let jhi = j_hi_arr[t];
let sm = gs(&m, t, l, jlo, jhi, row_width);
let si = gs(&ins, t, l, jlo, jhi, row_width);
let sd = gsd(&del, &del0, t, l, jlo, jhi, row_width);
let best_sc = [sm, si, sd].into_iter().filter(|&s| s != UNSET).max();
best_sc.map(|sc| {
let st = if sm == sc {
State::M
} else if si == sc {
State::I
} else {
State::D
};
(t, st, sc)
})
})
.max_by_key(|&(_, _, sc)| sc)
.map(|(t, s, _)| (t, s));
let (best_t, best_state) = match terminal_best {
Some(result) => result,
None => (0, State::M),
};
let mut ops: Vec<AlignOp> = Vec::with_capacity(l + n / 4);
let mut t = best_t;
let mut j = l;
let mut cur_state = best_state;
loop {
let cell = {
let jlo = j_lo_arr[t];
let jhi = j_hi_arr[t];
match cur_state {
State::M => {
if j < jlo || j > jhi {
Cell::unset()
} else {
m[t * row_width + (j - jlo)]
}
}
State::I => {
if j < jlo || j > jhi {
Cell::unset()
} else {
ins[t * row_width + (j - jlo)]
}
}
State::D => {
if j == 0 {
del0[t]
} else if j < jlo || j > jhi {
Cell::unset()
} else {
del[t * row_width + (j - jlo)]
}
}
}
};
if cell.score == UNSET {
break;
}
match cur_state {
State::M => {
ops.push(AlignOp::Match(topo[t]));
if cell.pred_t == VIRTUAL {
for k in (1..j).rev() {
ops.push(AlignOp::Insert(query[k - 1]));
}
break;
}
t = cell.pred_t as usize;
j -= 1;
cur_state = best_prev_state_banded(
&m, &ins, &del, &del0, t, j, &j_lo_arr, &j_hi_arr, row_width,
);
}
State::I => {
ops.push(AlignOp::Insert(query[j - 1]));
j -= 1;
if cell.pred_t == VIRTUAL {
for k in (1..j).rev() {
ops.push(AlignOp::Insert(query[k - 1]));
}
break;
}
let sm = gs(&m, t, j, j_lo_arr[t], j_hi_arr[t], row_width);
let si = gs(&ins, t, j, j_lo_arr[t], j_hi_arr[t], row_width);
cur_state = if sm >= si { State::M } else { State::I };
}
State::D => {
ops.push(AlignOp::Delete(topo[t]));
if cell.pred_t == VIRTUAL {
break;
}
t = cell.pred_t as usize;
cur_state = best_prev_state_banded(
&m, &ins, &del, &del0, t, j, &j_lo_arr, &j_hi_arr, row_width,
);
}
}
if j == 0 && cur_state != State::D {
break;
}
}
ops.reverse();
Ok(ops)
}
#[allow(clippy::too_many_arguments)]
#[inline]
fn best_prev_state_banded(
m: &[Cell],
ins: &[Cell],
del: &[Cell],
del0: &[Cell],
t: usize,
j: usize,
j_lo_arr: &[usize],
j_hi_arr: &[usize],
row_width: usize,
) -> State {
let jlo = j_lo_arr[t];
let jhi = j_hi_arr[t];
let sm = gs(m, t, j, jlo, jhi, row_width);
let si = gs(ins, t, j, jlo, jhi, row_width);
let sd = gsd(del, del0, t, j, jlo, jhi, row_width);
if sm != UNSET && sm >= si && sm >= sd {
State::M
} else if si != UNSET && si >= sd {
State::I
} else {
State::D
}
}
fn add_to_graph(
nodes: &mut Vec<Node>,
edge_reads: &mut HashMap<(usize, usize), Vec<u32>>,
query: &[u8],
ops: &[AlignOp],
read_idx: u32,
) {
let mut prev: Option<usize> = None;
let mut q_idx: usize = 0;
for &op in ops {
match op {
AlignOp::Match(node_idx) => {
let q_base = query[q_idx];
q_idx += 1;
let cur = if nodes[node_idx].base == q_base {
nodes[node_idx].coverage += 1;
if let Some(p) = prev {
increment_or_add_edge(nodes, p, node_idx);
edge_reads.entry((p, node_idx)).or_default().push(read_idx);
}
node_idx
} else {
let new_idx = push_node(nodes, q_base);
nodes[new_idx].coverage = 1;
if let Some(p) = prev {
nodes[p].out_edges.push((new_idx, 1));
nodes[new_idx].in_edges.push(p);
edge_reads.entry((p, new_idx)).or_default().push(read_idx);
}
new_idx
};
prev = Some(cur);
}
AlignOp::Insert(q_base) => {
q_idx += 1;
let new_idx = push_node(nodes, q_base);
nodes[new_idx].coverage = 1;
if let Some(p) = prev {
nodes[p].out_edges.push((new_idx, 1));
nodes[new_idx].in_edges.push(p);
edge_reads.entry((p, new_idx)).or_default().push(read_idx);
}
prev = Some(new_idx);
}
AlignOp::Delete(node_idx) => {
nodes[node_idx].delete_count += 1;
if let Some(p) = prev {
increment_or_add_edge(nodes, p, node_idx);
edge_reads.entry((p, node_idx)).or_default().push(read_idx);
}
prev = Some(node_idx);
}
}
}
}
fn heaviest_path(nodes: &[Node], topo: &[usize], rank_of: &[usize]) -> Vec<(usize, u8, i32)> {
let n = topo.len();
let mut cum: Vec<(i64, Option<usize>, i32)> = vec![(0, None, 0); n];
for t in 0..n {
let node_idx = topo[t];
let curr = cum[t].0;
for &(succ_idx, weight) in &nodes[node_idx].out_edges {
let succ_t = rank_of[succ_idx];
let candidate = curr + (weight - 1) as i64;
if candidate > cum[succ_t].0 {
cum[succ_t] = (candidate, Some(t), weight);
}
}
}
let max_cum = (0..n).map(|t| cum[t].0).max().unwrap_or(0);
let best_t = (0..n).find(|&t| cum[t].0 == max_cum).unwrap_or(0);
let mut path: Vec<(usize, u8, i32)> = Vec::new();
let mut t = best_t;
loop {
let node_idx = topo[t];
let w = if cum[t].1.is_none() {
nodes[node_idx].coverage as i32
} else {
cum[t].2
};
path.push((node_idx, nodes[node_idx].base, w));
match cum[t].1 {
None => break,
Some(pred_t) => t = pred_t,
}
}
path.reverse();
path
}
fn majority_frequency(nodes: &[Node], topo: &[usize], min_cov: u32) -> Vec<(usize, u8, i32)> {
topo.iter()
.copied()
.filter(|&idx| {
let cov = nodes[idx].coverage;
let del = nodes[idx].delete_count;
let total = cov + del;
total >= min_cov && cov * 2 >= total
})
.map(|idx| (idx, nodes[idx].base, nodes[idx].coverage as i32))
.collect()
}
fn compute_stats(nodes: &[Node], min_allele_freq: f64, n_reads: usize) -> GraphStats {
let node_count = nodes.len();
let edge_count: usize = nodes.iter().map(|nd| nd.out_edges.len()).sum();
let coverages: Vec<f64> = nodes.iter().map(|nd| nd.coverage as f64).collect();
let coverage_mean = if node_count == 0 {
0.0
} else {
coverages.iter().sum::<f64>() / node_count as f64
};
let coverage_variance = if node_count == 0 {
0.0
} else {
coverages
.iter()
.map(|&c| (c - coverage_mean).powi(2))
.sum::<f64>()
/ node_count as f64
};
let single_support = nodes.iter().filter(|nd| nd.coverage == 1).count();
let single_support_fraction = if node_count == 0 {
0.0
} else {
single_support as f64 / node_count as f64
};
let mut weights: Vec<f64> = nodes
.iter()
.flat_map(|nd| nd.out_edges.iter().map(|&(_, w)| w as f64))
.collect();
weights.sort_by(|a, b| a.partial_cmp(b).unwrap());
let edge_weight_gini = if weights.len() < 2 {
0.0
} else {
let n = weights.len() as f64;
let sum: f64 = weights.iter().sum();
if sum == 0.0 {
0.0
} else {
let numerator: f64 = weights
.iter()
.enumerate()
.map(|(i, &w)| (2.0 * (i as f64 + 1.0) - n - 1.0) * w)
.sum::<f64>();
numerator / (n * sum)
}
};
let threshold = (n_reads as f64 * min_allele_freq).ceil() as i32;
let mut bubble_count = 0usize;
let mut max_bubble_depth = 0usize;
let mut longest_bubble_span = 0usize;
for nd in nodes {
let qualifying: Vec<(usize, i32)> = nd
.out_edges
.iter()
.filter(|&&(_, w)| w >= threshold)
.map(|&(to, w)| (to, w))
.collect();
if qualifying.len() >= 2 {
bubble_count += 1;
let mut weights: Vec<i32> = qualifying.iter().map(|&(_, w)| w).collect();
weights.sort_unstable_by(|a, b| b.cmp(a));
max_bubble_depth = max_bubble_depth.max(weights[1] as usize);
for &(arm_start, _) in &qualifying {
let span = if nodes[arm_start].in_edges.len() > 1 {
0 } else {
materialize_arm_len(nodes, arm_start, ARM_MAX_DEPTH)
};
longest_bubble_span = longest_bubble_span.max(span);
}
}
}
let mean_column_entropy = {
let mut sum = 0.0f64;
let mut count = 0usize;
for nd in nodes {
let cov = nd.coverage as f64;
let del = nd.delete_count as f64;
let total = cov + del;
if total > 0.0 {
let h = binary_entropy(cov / total);
sum += h;
count += 1;
}
}
if count == 0 { 0.0 } else { sum / count as f64 }
};
GraphStats {
node_count,
edge_count,
bubble_count,
max_bubble_depth,
coverage_mean,
coverage_variance,
edge_weight_gini,
single_support_fraction,
mean_column_entropy,
longest_bubble_span,
median_input_read_len: 0, }
}
fn median_read_len(reads: &[Vec<u8>]) -> usize {
if reads.is_empty() {
return 0;
}
let mut lens: Vec<usize> = reads.iter().map(|r| r.len()).collect();
lens.sort_unstable();
lens[lens.len() / 2]
}
#[inline]
fn binary_entropy(p: f64) -> f64 {
if p <= 0.0 || p >= 1.0 {
0.0
} else {
let q = 1.0 - p;
-(p * p.log2() + q * q.log2())
}
}
fn detect_coverage_gaps(coverage: &[u32]) -> Vec<CoverageGap> {
let first = coverage.iter().position(|&c| c >= 2);
let last = coverage.iter().rposition(|&c| c >= 2);
let (first, last) = match (first, last) {
(Some(f), Some(l)) if f < l => (f, l),
_ => return vec![],
};
let mut gaps = Vec::new();
let mut gap_start: Option<usize> = None;
for (offset, &cov) in coverage[(first + 1)..last].iter().enumerate() {
let i = first + 1 + offset;
if cov < 2 {
gap_start.get_or_insert(i);
} else if let Some(s) = gap_start.take() {
gaps.push(CoverageGap {
start: s,
end: i,
kind: GapKind::Spanning,
});
}
}
if let Some(s) = gap_start {
gaps.push(CoverageGap {
start: s,
end: last,
kind: GapKind::Spanning,
});
}
gaps
}
fn find_bubbles(
nodes: &[Node],
topo: &[usize],
n_reads: usize,
min_allele_freq: f64,
) -> Vec<(usize, Vec<usize>)> {
let threshold = ((n_reads as f64 * min_allele_freq).ceil() as i32).max(1);
topo.iter()
.copied()
.filter_map(|node_idx| {
let arms: Vec<usize> = nodes[node_idx]
.out_edges
.iter()
.filter(|&&(_, w)| w >= threshold)
.map(|&(to, _)| to)
.collect();
if arms.len() >= 2 {
Some((node_idx, arms))
} else {
None
}
})
.collect()
}
fn find_structural_bubbles(
nodes: &[Node],
topo: &[usize],
n_reads: usize,
cfg: &PoaConfig,
) -> Vec<(usize, Vec<usize>)> {
let threshold = ((n_reads as f64 * cfg.min_allele_freq).ceil() as i32).max(1);
let max_check = cfg.phasing_bubble_min_span.saturating_add(1);
topo.iter()
.copied()
.filter_map(|entry_node| {
let arms: Vec<usize> = nodes[entry_node]
.out_edges
.iter()
.filter(|&&(_, w)| w >= threshold)
.map(|&(to, _)| to)
.collect();
if arms.len() < 2 {
return None;
}
let max_span = arms
.iter()
.map(|&start| {
if nodes[start].in_edges.len() > 1 {
0
} else {
materialize_arm_len(nodes, start, max_check)
}
})
.max()
.unwrap_or(0);
if max_span >= cfg.phasing_bubble_min_span {
Some((entry_node, arms))
} else {
None
}
})
.collect()
}
const ARM_SEQUENCE_CAP: usize = 256;
fn arm_sequence(nodes: &[Node], start: usize) -> Vec<u8> {
if nodes[start].in_edges.len() > 1 {
return vec![];
}
let mut seq = Vec::new();
let mut cur = start;
loop {
if seq.len() >= ARM_SEQUENCE_CAP {
return vec![];
}
seq.push(nodes[cur].base);
match nodes[cur].out_edges.as_slice() {
[(next, _)] => {
if nodes[*next].in_edges.len() > 1 {
break; }
cur = *next;
}
_ => break, }
}
seq
}
fn collect_bubble_sites(
nodes: &[Node],
topo: &[usize],
filtered: &[(usize, u8, i32)],
edge_reads: &HashMap<(usize, usize), Vec<u32>>,
min_allele_freq: f64,
n_reads: usize,
phasing_bubble_min_span: usize,
) -> Vec<BubbleSite> {
let path_pos: HashMap<usize, usize> = filtered
.iter()
.enumerate()
.map(|(i, &(node_idx, _, _))| (node_idx, i))
.collect();
find_bubbles(nodes, topo, n_reads, min_allele_freq)
.into_iter()
.filter_map(|(entry_node, arm_starts)| {
let consensus_pos = *path_pos.get(&entry_node)?;
let arm_read_counts: Vec<u32> = arm_starts
.iter()
.map(|&a| {
edge_reads
.get(&(entry_node, a))
.map_or(0, |v| v.len() as u32)
})
.collect();
let is_structural = arm_starts.iter().any(|&a| {
let len = if nodes[a].in_edges.len() > 1 {
0
} else {
materialize_arm_len(nodes, a, phasing_bubble_min_span)
};
len >= phasing_bubble_min_span
});
let arm_sequences: Vec<Vec<u8>> =
arm_starts.iter().map(|&a| arm_sequence(nodes, a)).collect();
Some(BubbleSite {
consensus_pos,
arm_read_counts,
arm_sequences,
is_structural,
})
})
.collect()
}
fn phasing_groups(
edge_reads: &HashMap<(usize, usize), Vec<u32>>,
bubbles: &[(usize, Vec<usize>)],
n_reads: usize,
min_reads: usize,
) -> Vec<Vec<usize>> {
let n_bubbles = bubbles.len();
let mut sig: Vec<Vec<Option<usize>>> = vec![vec![None; n_bubbles]; n_reads];
for (b, (entry, arm_starts)) in bubbles.iter().enumerate() {
for (arm_idx, &arm_start) in arm_starts.iter().enumerate() {
if let Some(reads) = edge_reads.get(&(*entry, arm_start)) {
for &r in reads {
let r = r as usize;
if r < n_reads {
sig[r][b] = Some(arm_idx);
}
}
}
}
}
let mut assigned: Vec<usize> = Vec::new();
let mut unassigned: Vec<usize> = Vec::new();
for (r, row) in sig.iter().enumerate() {
if row.iter().any(|s| s.is_some()) {
assigned.push(r);
} else {
unassigned.push(r);
}
}
let n = assigned.len();
let mut parent: Vec<usize> = (0..n).collect();
for i in 0..n {
for j in (i + 1)..n {
let ri = assigned[i];
let rj = assigned[j];
let compatible = (0..n_bubbles).all(|b| match (sig[ri][b], sig[rj][b]) {
(Some(a), Some(bv)) => a == bv,
_ => true,
});
if compatible {
let mut pi = i;
while parent[pi] != pi {
pi = parent[pi];
}
let mut pj = j;
while parent[pj] != pj {
pj = parent[pj];
}
if pi != pj {
parent[pj] = pi;
}
}
}
}
for i in 0..n {
let mut root = i;
while parent[root] != root {
root = parent[root];
}
parent[i] = root;
}
let mut group_map: HashMap<usize, Vec<usize>> = HashMap::new();
for (slot, &r) in assigned.iter().enumerate() {
group_map.entry(parent[slot]).or_default().push(r);
}
let mut groups: Vec<Vec<usize>> = group_map.into_values().collect();
groups.sort_unstable_by_key(|g| std::cmp::Reverse(g.len()));
if groups.is_empty() {
groups.push(Vec::new());
}
let mut i = 1;
while i < groups.len() {
if groups[i].len() < min_reads {
let g = groups.remove(i);
groups[0].extend(g);
} else {
i += 1;
}
}
groups[0].extend(unassigned);
groups.retain(|g| !g.is_empty());
groups
}
fn partition_reads_by_bubble(
edge_reads: &HashMap<(usize, usize), Vec<u32>>,
entry_node: usize,
arm_starts: &[usize],
n_reads: usize,
) -> Vec<Vec<usize>> {
let mut arm_order: Vec<usize> = (0..arm_starts.len()).collect();
arm_order.sort_unstable_by(|&a, &b| {
let wa = edge_reads
.get(&(entry_node, arm_starts[a]))
.map_or(0, |v| v.len());
let wb = edge_reads
.get(&(entry_node, arm_starts[b]))
.map_or(0, |v| v.len());
wb.cmp(&wa)
});
let n_arms = arm_order.len().min(2);
let mut groups: Vec<Vec<usize>> = vec![Vec::new(); n_arms];
let mut assigned = vec![false; n_reads];
for (slot, &arm_idx) in arm_order[..n_arms].iter().enumerate() {
let arm_start = arm_starts[arm_idx];
if let Some(reads) = edge_reads.get(&(entry_node, arm_start)) {
for &r in reads {
let r = r as usize;
if r < n_reads && !assigned[r] {
groups[slot].push(r);
assigned[r] = true;
}
}
}
}
let largest = groups
.iter()
.enumerate()
.max_by_key(|(_, g)| g.len())
.map(|(i, _)| i)
.unwrap_or(0);
for (r, &done) in assigned.iter().enumerate() {
if !done {
groups[largest].push(r);
}
}
groups.retain(|g| !g.is_empty());
groups
}
fn choose_seed(group: &[usize], reads: &[Vec<u8>]) -> usize {
if group.len() == 1 {
return 0;
}
let mut lens: Vec<usize> = group.iter().map(|&i| reads[i].len()).collect();
lens.sort_unstable();
let median = lens[lens.len() / 2];
group
.iter()
.enumerate()
.min_by_key(|&(_, &i)| reads[i].len().abs_diff(median))
.map(|(slot, _)| slot)
.unwrap_or(0)
}
impl PoaGraph {
pub fn new(seed: &[u8], config: PoaConfig) -> Result<Self, PoaError> {
if seed.is_empty() {
return Err(PoaError::EmptyInput);
}
let mut nodes: Vec<Node> = Vec::with_capacity(seed.len());
for &base in seed {
let idx = push_node(&mut nodes, base);
nodes[idx].coverage = 1;
}
let n = nodes.len();
let mut edge_reads: HashMap<(usize, usize), Vec<u32>> = HashMap::new();
for i in 0..n.saturating_sub(1) {
add_edge(&mut nodes, i, i + 1);
edge_reads.insert((i, i + 1), vec![0u32]);
}
Ok(PoaGraph {
nodes,
config,
n_reads: 1,
reads: vec![seed.to_vec()],
edge_reads,
warnings: 0,
cached_spine: Vec::new(),
spine_updated_at: 0,
spine_interval: 1,
align_scratch: AlignScratch::new(),
spine_mers: HashMap::new(),
})
}
pub fn add_read(&mut self, read: &[u8]) -> Result<(), PoaError> {
if read.is_empty() {
return Err(PoaError::EmptyInput);
}
let (topo, rank_of) = topological_order(&self.nodes);
let reads_since_update = self.n_reads.saturating_sub(self.spine_updated_at);
if self.cached_spine.is_empty() || reads_since_update >= self.spine_interval {
let new_spine = heaviest_path(&self.nodes, &topo, &rank_of);
let diff = spine_diff(&self.cached_spine, &new_spine);
self.cached_spine = new_spine;
self.spine_updated_at = self.n_reads;
self.spine_mers = build_spine_mers(&self.cached_spine, MINI_K, MINI_W);
if diff <= SPINE_STABLE_THRESHOLD {
self.spine_interval = (self.spine_interval * 2).min(SPINE_MAX_INTERVAL);
} else {
self.spine_interval = 1;
}
}
let read_mers = compute_minimizers(read, MINI_K, MINI_W);
let raw_anchors =
build_anchor_chain(&read_mers, &self.spine_mers, &self.cached_spine, &rank_of);
let anchors: Vec<(usize, usize)> = if raw_anchors.len() >= MINI_MIN_CHAIN {
raw_anchors
} else {
vec![]
};
let ops = align(
&self.nodes,
&topo,
&rank_of,
&self.cached_spine,
read,
&self.config,
&mut self.align_scratch,
&anchors,
)?;
let read_idx = self.n_reads as u32;
add_to_graph(&mut self.nodes, &mut self.edge_reads, read, &ops, read_idx);
self.reads.push(read.to_vec());
self.n_reads += 1;
Ok(())
}
pub fn consensus(&self) -> Result<Consensus, PoaError> {
if self.n_reads < self.config.min_reads {
return Err(PoaError::InsufficientDepth {
got: self.n_reads,
min: self.config.min_reads,
});
}
if self.n_reads == 1 {
let (topo, _) = topological_order(&self.nodes);
let sequence: Vec<u8> = topo.iter().map(|&idx| self.nodes[idx].base).collect();
let coverage: Vec<u32> = topo.iter().map(|_| 1).collect();
let path_weights: Vec<i32> = topo.iter().map(|_| 1).collect();
let mut graph_stats =
compute_stats(&self.nodes, self.config.min_allele_freq, self.n_reads);
graph_stats.median_input_read_len = median_read_len(&self.reads);
return Ok(Consensus {
sequence,
coverage,
path_weights,
n_reads: 1,
graph_stats,
gaps: vec![],
bubble_sites: vec![],
read_indices: vec![],
});
}
let (topo, rank_of) = topological_order(&self.nodes);
let min_cov = self.min_coverage();
let filtered: Vec<(usize, u8, i32)> = match self.config.consensus_mode {
ConsensusMode::HeaviestPath => {
let path = heaviest_path(&self.nodes, &topo, &rank_of);
let meets_floor: Vec<bool> = path
.iter()
.map(|&(node_idx, _, _)| self.nodes[node_idx].coverage >= min_cov)
.collect();
let start = meets_floor.iter().position(|&s| s).unwrap_or(0);
let end = meets_floor
.iter()
.rposition(|&s| s)
.map(|i| i + 1)
.unwrap_or(path.len());
let range = if start < end {
start..end
} else {
0..path.len()
};
const FORK_SEARCH_HOPS: usize = 64;
let range_start = range.start;
range
.filter(|&i| {
let (node_idx, _, _) = path[i];
if meets_floor[i] {
return true;
}
let mut search_idx = i;
let mut hops = 0usize;
let fork_info = loop {
if search_idx == range_start || hops >= FORK_SEARCH_HOPS {
break None;
}
let (cand_idx, _, _) = path[search_idx - 1];
if self.nodes[cand_idx].out_edges.len() >= 2 {
let (arm_idx, _, _) = path[search_idx];
break Some((cand_idx, arm_idx));
}
search_idx -= 1;
hops += 1;
};
let Some((pred_idx, arm_idx)) = fork_info else {
return self.nodes[node_idx].coverage
> self.nodes[node_idx].delete_count;
};
let local_total: i32 = self.nodes[pred_idx]
.out_edges
.iter()
.map(|&(_, ew)| ew)
.sum();
if (local_total.max(0) as u32) < min_cov {
return false;
}
let local_min_cov = coverage_threshold(
local_total.max(0) as usize,
self.config.min_coverage_fraction,
);
if self.nodes[node_idx].coverage as i32 >= local_min_cov as i32 {
return true;
}
if self.nodes[pred_idx].delete_count != 0
|| self.nodes[node_idx].delete_count != 0
{
return false;
}
let arm_weight = self.nodes[pred_idx]
.out_edges
.iter()
.find(|&&(to, _)| to == arm_idx)
.map(|&(_, w)| w)
.unwrap_or(0);
self.nodes[pred_idx]
.out_edges
.iter()
.all(|&(_, w)| w <= arm_weight)
})
.map(|i| path[i])
.collect()
}
ConsensusMode::MajorityFrequency => majority_frequency(&self.nodes, &topo, min_cov),
};
let sequence: Vec<u8> = filtered.iter().map(|&(_, base, _)| base).collect();
let coverage: Vec<u32> = filtered
.iter()
.map(|&(node_idx, _, _)| self.nodes[node_idx].coverage)
.collect();
let path_weights: Vec<i32> = filtered.iter().map(|&(_, _, w)| w).collect();
let mut graph_stats = compute_stats(&self.nodes, self.config.min_allele_freq, self.n_reads);
graph_stats.median_input_read_len = median_read_len(&self.reads);
let gaps = detect_coverage_gaps(&coverage);
let bubble_sites = collect_bubble_sites(
&self.nodes,
&topo,
&filtered,
&self.edge_reads,
self.config.min_allele_freq,
self.n_reads,
self.config.phasing_bubble_min_span,
);
Ok(Consensus {
sequence,
coverage,
path_weights,
n_reads: self.n_reads,
graph_stats,
gaps,
bubble_sites,
read_indices: vec![],
})
}
pub fn stats(&self) -> GraphStats {
let mut s = compute_stats(&self.nodes, self.config.min_allele_freq, self.n_reads);
s.median_input_read_len = median_read_len(&self.reads);
s
}
pub fn warnings_emitted(&self) -> usize {
self.warnings
}
pub fn edge_weights(&self) -> Vec<i32> {
self.nodes
.iter()
.flat_map(|n| n.out_edges.iter().map(|&(_, w)| w))
.collect()
}
pub fn node_coverages(&self) -> Vec<u32> {
let (topo, _) = topological_order(&self.nodes);
topo.iter().map(|&i| self.nodes[i].coverage).collect()
}
pub fn align_read_ops(
&self,
read: &[u8],
) -> Result<(Vec<AlignOp>, usize, Vec<usize>), PoaError> {
let (topo, rank_of) = topological_order(&self.nodes);
let spine = heaviest_path(&self.nodes, &topo, &rank_of);
let ops = align(
&self.nodes,
&topo,
&rank_of,
&spine,
read,
&self.config,
&mut AlignScratch::new(),
&[],
)?;
Ok((ops, 0, rank_of))
}
pub fn align_read_ops_unbanded(
&self,
read: &[u8],
) -> Result<(Vec<AlignOp>, Vec<usize>), PoaError> {
let (topo, rank_of) = topological_order(&self.nodes);
let spine = heaviest_path(&self.nodes, &topo, &rank_of);
let ops = align(
&self.nodes,
&topo,
&rank_of,
&spine,
read,
&self.config,
&mut AlignScratch::new(),
&[],
)?;
Ok((ops, rank_of))
}
pub fn node_count(&self) -> usize {
self.nodes.len()
}
pub fn graph_topology(&self) -> crate::types::GraphTopology {
use crate::types::{GraphEdgeInfo, GraphNodeInfo, GraphTopology};
let (topo, rank_of) = topological_order(&self.nodes);
let spine = heaviest_path(&self.nodes, &topo, &rank_of);
let nodes: Vec<GraphNodeInfo> = topo
.iter()
.enumerate()
.map(|(rank, &node_idx)| {
let n = &self.nodes[node_idx];
GraphNodeInfo {
node_idx,
base: n.base,
coverage: n.coverage,
delete_count: n.delete_count,
topo_rank: rank,
}
})
.collect();
let edges: Vec<GraphEdgeInfo> = topo
.iter()
.enumerate()
.flat_map(|(from_rank, &node_idx)| {
self.nodes[node_idx]
.out_edges
.iter()
.map(move |&(succ_idx, weight)| (from_rank, succ_idx, weight))
.collect::<Vec<_>>()
})
.map(|(from_rank, succ_idx, weight)| GraphEdgeInfo {
from_rank,
to_rank: rank_of[succ_idx],
weight,
})
.collect();
let spine_ranks: Vec<usize> = spine
.iter()
.map(|(node_idx, _, _)| rank_of[*node_idx])
.collect();
GraphTopology {
nodes,
edges,
spine_ranks,
}
}
pub fn node_out_edge_info(&self) -> Vec<(usize, i32, i32)> {
self.nodes
.iter()
.map(|n| {
let count = n.out_edges.len();
let max_w = n.out_edges.iter().map(|&(_, w)| w).max().unwrap_or(0);
let min_w = n.out_edges.iter().map(|&(_, w)| w).min().unwrap_or(0);
(count, max_w, min_w)
})
.collect()
}
pub fn bubble_arm_lengths(
&self,
weight_threshold: i32,
arm_len_threshold: usize,
) -> Vec<(usize, Vec<usize>)> {
let (topo, _) = topological_order(&self.nodes);
let mut result = Vec::new();
for (t, &node_idx) in topo.iter().enumerate() {
let qualifying: Vec<usize> = self.nodes[node_idx]
.out_edges
.iter()
.filter(|&&(_, w)| w >= weight_threshold)
.map(|&(succ, _)| succ)
.collect();
if qualifying.len() >= 2 {
let arm_lens: Vec<usize> = qualifying
.iter()
.map(|&arm_start| materialize_arm_len(&self.nodes, arm_start, 500))
.collect();
let min_len = arm_lens.iter().copied().min().unwrap_or(0);
if min_len >= arm_len_threshold {
result.push((t, arm_lens));
}
}
}
result
}
pub fn consensus_multi(&self) -> Result<Vec<Consensus>, PoaError> {
if self.n_reads < self.config.min_reads {
return Err(PoaError::InsufficientDepth {
got: self.n_reads,
min: self.config.min_reads,
});
}
let (topo, _) = topological_order(&self.nodes);
let structural = find_structural_bubbles(&self.nodes, &topo, self.n_reads, &self.config);
if structural.is_empty() && (self.config.band_width > 0 || self.config.adaptive_band) {
let mut cfg2 = self.config.clone();
cfg2.band_width = 0;
cfg2.adaptive_band = false;
let seed = &self.reads[0];
let mut g2 = PoaGraph::new(seed, cfg2)?;
for read in self.reads.iter().skip(1) {
g2.add_read(read)?;
}
return g2.consensus_multi();
}
let groups = if !structural.is_empty() {
phasing_groups(
&self.edge_reads,
&structural,
self.n_reads,
self.config.min_reads,
)
} else {
let snp_bubbles = find_bubbles(
&self.nodes,
&topo,
self.n_reads,
self.config.min_allele_freq,
);
if !snp_bubbles.is_empty() {
let (entry, arm_starts) = snp_bubbles
.iter()
.max_by_key(|(entry, arms)| {
arms.iter()
.filter_map(|&arm| self.edge_reads.get(&(*entry, arm)))
.map(|v| v.len())
.min()
.unwrap_or(0)
})
.unwrap();
partition_reads_by_bubble(&self.edge_reads, *entry, arm_starts, self.n_reads)
} else {
return Ok(vec![self.consensus()?]);
}
};
if groups.len() < 2 {
return Ok(vec![self.consensus()?]);
}
let mut results = Vec::with_capacity(groups.len());
for group in &groups {
if group.len() < self.config.min_reads {
return Err(PoaError::InsufficientDepth {
got: group.len(),
min: self.config.min_reads,
});
}
let seed_slot = choose_seed(group, &self.reads);
let seed = &self.reads[group[seed_slot]];
let mut sub = PoaGraph::new(seed, self.config.clone())?;
for (slot, &read_idx) in group.iter().enumerate() {
if slot == seed_slot {
continue;
}
sub.add_read(&self.reads[read_idx])?;
}
let mut c = sub.consensus()?;
c.read_indices = group.to_vec();
results.push(c);
}
Ok(results)
}
fn min_coverage(&self) -> u32 {
coverage_threshold(self.n_reads, self.config.min_coverage_fraction)
}
}
fn coverage_threshold(population: usize, min_coverage_fraction: f64) -> u32 {
if min_coverage_fraction > 0.0 {
((population as f64 * min_coverage_fraction).ceil() as u32).max(1)
} else if population <= 1 {
1
} else {
((population / 2 + 1).max(2)) as u32
}
}