Skip to main content

salmon_model/
gcbias.rs

1//! Fragment-GC bias model (`GCFragModel`).
2//!
3//! A port of salmon's `GCFragModel` (`include/.../model/GCFragModel.hpp`): a
4//! `condBins × numGCBins` table of fragment counts, where a fragment is keyed by
5//! its **GC fraction** (0–100, → `numGCBins` bins) and a **conditioning context**
6//! (a coarse sequence/context fraction, → `condBins` bins). An *observed* model
7//! (from mapped fragments) and an *expected* model (from the transcriptome) are
8//! each normalized per conditioning row into a distribution, then divided —
9//! [`gc_ratio`] — to give the per-`(context, GC)` bias factor used to weight
10//! fragments in the bias-corrected effective length.
11//!
12//! salmon accumulates the observed model in log space and the expected in linear
13//! space, but both are normalized to linear distributions before the ratio is
14//! taken, so the normalized result is independent of the accumulation space.
15//! We therefore accumulate in **linear** space (summing weights) throughout —
16//! numerically equivalent after [`GcFragModel::normalize`], and simpler.
17
18/// salmon's default number of conditioning (context) bins (`numConditionalGCBins`).
19pub const DEFAULT_COND_BINS: usize = 3;
20/// salmon's default number of fragment-GC bins for accumulation
21/// (`salmon::defaults::numFragGCBins`).
22pub const DEFAULT_GC_BINS: usize = 25;
23/// salmon's `ratio()` clamp (`gcBias = gcCounts.ratio(transcriptGCDist, 1000.0)`).
24pub const GC_MAX_RATIO: f64 = 1000.0;
25/// salmon's per-row normalization prior.
26const NORM_PRIOR: f64 = 0.1;
27
28/// Bin a 0–100 fraction into `n` bins (salmon's `GCDesc::fragBin(n)`/`contextBin(n)`:
29/// `min(n-1, floor(frac / (100/n)))`). With `n == 101` this is the identity.
30#[inline]
31pub fn bin_frac(frac: i32, n: usize) -> usize {
32    if n == 101 {
33        return (frac.clamp(0, 100)) as usize;
34    }
35    let w = 100.0 / n as f64;
36    let b = (frac as f64 / w) as i32;
37    b.clamp(0, n as i32 - 1) as usize
38}
39
40/// `(condBins × numGCBins)` table of fragment-GC counts (then a normalized
41/// distribution / ratio after [`normalize`](Self::normalize) / [`gc_ratio`]).
42#[derive(Debug, Clone)]
43pub struct GcFragModel {
44    cond_bins: usize,
45    gc_bins: usize,
46    /// row-major `counts[ctx * gc_bins + gc]`
47    counts: Vec<f64>,
48    normalized: bool,
49    /// Precomputed `[0..=100] -> bin` maps for the context and GC fractions
50    /// (both always lie in `[0, 100]`). [`bin_frac`] is ~44% of the eff-length
51    /// convolution self-time (two `100.0/n` + `frac/w` divides per [`get`]/[`inc`]
52    /// call, called per fragment); these LUTs replace both divides with a byte
53    /// load. Each entry is exactly `bin_frac(frac, n)`, so the result is
54    /// bit-identical to the divide path (verified in tests).
55    ctx_lut: Vec<u8>,
56    gc_lut: Vec<u8>,
57}
58
59impl GcFragModel {
60    pub fn new(cond_bins: usize, gc_bins: usize) -> Self {
61        let ctx_lut = (0..=100).map(|f| bin_frac(f, cond_bins) as u8).collect();
62        let gc_lut = (0..=100).map(|f| bin_frac(f, gc_bins) as u8).collect();
63        Self {
64            cond_bins,
65            gc_bins,
66            counts: vec![0.0; cond_bins * gc_bins],
67            normalized: false,
68            ctx_lut,
69            gc_lut,
70        }
71    }
72
73    /// salmon's default 3 × 101 model.
74    pub fn default_model() -> Self {
75        Self::new(DEFAULT_COND_BINS, DEFAULT_GC_BINS)
76    }
77
78    #[inline]
79    fn idx(&self, ctx_frac: i32, gc_frac: i32) -> usize {
80        // Clamp-then-LUT is bit-identical to `bin_frac` for every i32 input:
81        // for frac in [0,100] the LUT entry *is* `bin_frac(frac, n)`, and a frac
82        // outside [0,100] clamps to the same boundary bin `bin_frac` would
83        // produce (0 or n-1). Inputs are always lrint(100·ratio) in [0,100].
84        let ctx = if self.cond_bins > 1 {
85            self.ctx_lut[ctx_frac.clamp(0, 100) as usize] as usize
86        } else {
87            0
88        };
89        let gc = self.gc_lut[gc_frac.clamp(0, 100) as usize] as usize;
90        ctx * self.gc_bins + gc
91    }
92
93    /// The flattened `cond_bins × gc_bins` table (`counts[ctx * gc_bins + gc]`,
94    /// row-major), for dumping to the aux bias files.
95    pub fn dump(&self) -> &[f64] {
96        &self.counts
97    }
98
99    /// Accumulate `weight` for a fragment with the given GC and context fractions.
100    pub fn inc(&mut self, gc_frac: i32, ctx_frac: i32, weight: f64) {
101        debug_assert!(!self.normalized, "cannot inc a normalized model");
102        let i = self.idx(ctx_frac, gc_frac);
103        self.counts[i] += weight;
104    }
105
106    /// Value at a `(context, GC)` cell (a normalized density after `normalize`,
107    /// or a clamped ratio for a model produced by [`gc_ratio`]).
108    #[inline]
109    pub fn get(&self, gc_frac: i32, ctx_frac: i32) -> f64 {
110        self.counts[self.idx(ctx_frac, gc_frac)]
111    }
112
113    /// Merge another (compatible) model's counts. Both must be pre-normalization.
114    pub fn combine_counts(&mut self, other: &GcFragModel) {
115        debug_assert!(!self.normalized && !other.normalized);
116        for (a, b) in self.counts.iter_mut().zip(&other.counts) {
117            *a += *b;
118        }
119    }
120
121    /// Normalize each conditioning row into a distribution over GC bins, with a
122    /// pseudocount `prior` (salmon's default 0.1). Idempotent.
123    pub fn normalize(&mut self) {
124        if self.normalized {
125            return;
126        }
127        for r in 0..self.cond_bins {
128            let base = r * self.gc_bins;
129            let row = &mut self.counts[base..base + self.gc_bins];
130            let row_mass: f64 = row.iter().map(|c| NORM_PRIOR + *c).sum();
131            if row_mass > 0.0 {
132                let norm = 1.0 / row_mass;
133                for c in row.iter_mut() {
134                    *c = (NORM_PRIOR + *c) * norm;
135                }
136            }
137        }
138        self.normalized = true;
139    }
140}
141
142/// Per-cell bias ratio `observed / expected`, clamped to `[1/max_ratio, max_ratio]`
143/// (salmon's `GCFragModel::ratio`). Both models are normalized first.
144pub fn gc_ratio(
145    observed: &mut GcFragModel,
146    expected: &mut GcFragModel,
147    max_ratio: f64,
148) -> GcFragModel {
149    observed.normalize();
150    expected.normalize();
151    let min_ratio = 1.0 / max_ratio;
152    let mut out = GcFragModel::new(observed.cond_bins, observed.gc_bins);
153    for i in 0..out.counts.len() {
154        let e = expected.counts[i];
155        let rat = if e != 0.0 {
156            observed.counts[i] / e
157        } else {
158            max_ratio
159        };
160        out.counts[i] = rat.clamp(min_ratio, max_ratio);
161    }
162    out.normalized = true; // a ratio model is used directly, not re-normalized
163    out
164}
165
166// ===========================================================================
167// Fragment-GC content (salmon's `Transcript::GCCount_` / `gcFrac` / `gcDesc`)
168// ===========================================================================
169
170use crate::seqbias::{
171    conditional_cdf, log_bias, revcomp_bytes, SBModel, CONTEXT_LEFT, CONTEXT_LENGTH, MIN_ALPHA,
172    MIN_CDF_MASS,
173};
174
175/// salmon's fragment-length sampling stride for the GC convolution
176/// (`pdfSampFactor` = `biasSpeedSamp`, default 5).
177pub const GC_SAMP_STRIDE: usize = 5;
178
179/// `gcDesc` context-window geometry (salmon `Transcript::gcDesc`):
180/// `outsideContext = 3`, `insideContext = 2`.
181const OUTSIDE_5P: i32 = 4; // outsideContext + 1
182const OUTSIDE_3P: i32 = 3; // outsideContext
183const INSIDE_5P: i32 = 1; // insideContext - 1
184const INSIDE_3P: i32 = 2; // insideContext
185
186/// Round to nearest (ties to even), matching C++ `std::lrint` under the default
187/// rounding mode.
188///
189/// On x86-64 this is a single `cvtsd2si` instruction (which rounds per the
190/// current — default round-to-nearest-even — MXCSR mode), identical to salmon's
191/// `std::lrint`. The portable `f64::round_ties_even() as i32` lowers to a libm
192/// `roundeven` *function call* on this baseline build; since this is called
193/// ~2× per `gc_desc` over tens of billions of fragment evaluations under
194/// `--gcBias`, that call was the dominant per-iteration cost (≈4× the GC-bias
195/// runtime vs C++). Other architectures keep the portable path.
196#[inline]
197fn lrint(x: f64) -> i32 {
198    #[cfg(target_arch = "x86_64")]
199    {
200        // SAFETY: SSE2 is part of the x86-64 baseline, so `_mm_set_sd` /
201        // `_mm_cvtsd_si32` are always available. `cvtsd2si` matches `std::lrint`
202        // (round-to-nearest-even); our arguments are small (0–100), no overflow.
203        #[allow(unsafe_code)]
204        unsafe {
205            core::arch::x86_64::_mm_cvtsd_si32(core::arch::x86_64::_mm_set_sd(x))
206        }
207    }
208    #[cfg(not(target_arch = "x86_64"))]
209    {
210        x.round_ties_even() as i32
211    }
212}
213
214/// Cumulative G+C counts (salmon's `Transcript::GCCount_`): `prefix[p]` is the
215/// number of `G`/`C` bases in `seq[0..=p]`.
216pub fn gc_prefix(seq: &[u8]) -> Vec<u32> {
217    let mut prefix = Vec::with_capacity(seq.len());
218    let mut acc = 0u32;
219    for &b in seq {
220        if matches!(b, b'G' | b'g' | b'C' | b'c') {
221            acc += 1;
222        }
223        prefix.push(acc);
224    }
225    prefix
226}
227
228/// GC count in the closed interval `[a, b]` from a cumulative-count prefix.
229#[inline]
230fn gc_in(prefix: &[u32], a: i32, b: i32) -> i64 {
231    let lo = if a > 0 {
232        prefix[(a - 1) as usize] as i64
233    } else {
234        0
235    };
236    prefix[b as usize] as i64 - lo
237}
238
239/// Fragment GC fraction (0–100) over the closed interval `[s, e]`
240/// (salmon's `Transcript::gcFrac`): `round(100·GC[s,e] / (e−s+1))`.
241#[inline]
242pub fn gc_frac(prefix: &[u32], s: i32, e: i32) -> i32 {
243    lrint(100.0 * gc_in(prefix, s, e) as f64 / (e - s + 1) as f64)
244}
245
246/// A rank-enabled bitvector over the *concatenated* reference sequence marking
247/// G/C positions (salmon's `--reduceGCMemory` representation). Replaces the
248/// per-transcript dense cumulative-GC arrays (`Vec<Vec<u32>>`, 4 bytes/base) with
249/// one bitvector (~1 bit/base + 25% rank overhead) — GC in any interval is a
250/// difference of two O(1) ranks. Build once over the whole transcriptome.
251pub struct GcRank {
252    rank: sux::rank_sel::Rank9,
253}
254
255impl GcRank {
256    /// Build the rank bitvector from the concatenated reference bytes
257    /// (`seq[i]` is G/C ⇒ bit `i` set).
258    pub fn new(concat: &[u8]) -> GcRank {
259        use sux::traits::bit_vec_ops::BitVecOpsMut;
260        let mut bv = sux::bits::BitVec::new(concat.len());
261        for (i, &b) in concat.iter().enumerate() {
262            if matches!(b, b'G' | b'g' | b'C' | b'c') {
263                bv.set(i, true);
264            }
265        }
266        GcRank {
267            rank: sux::rank_sel::Rank9::new(bv),
268        }
269    }
270
271    /// A cumulative-GC view of the transcript occupying `concat[off..off+len]`.
272    #[inline]
273    pub fn view(&self, off: usize, len: usize) -> GcView<'_> {
274        use sux::traits::Rank;
275        GcView::Rank {
276            rank: &self.rank,
277            off,
278            base: self.rank.rank(off),
279            len,
280        }
281    }
282}
283
284/// A per-transcript cumulative-GC accessor: either the dense prefix array
285/// (default) or a window of the shared [`GcRank`] (`--reduceGCMemory`). Both
286/// return the *same* cumulative counts, so results are identical; `cum(p)` is
287/// the number of G/C bases in transcript-local `[0, p]`.
288#[derive(Clone, Copy)]
289pub enum GcView<'a> {
290    Dense(&'a [u32]),
291    Rank {
292        rank: &'a sux::rank_sel::Rank9,
293        off: usize,
294        /// `rank(off)` — precomputed so `cum` is a single rank query.
295        base: usize,
296        len: usize,
297    },
298}
299
300impl GcView<'_> {
301    /// Cumulative G/C count in transcript-local `[0, p]` (inclusive).
302    #[inline]
303    pub fn cum(&self, p: i32) -> i64 {
304        match self {
305            GcView::Dense(prefix) => prefix[p as usize] as i64,
306            GcView::Rank {
307                rank, off, base, ..
308            } => {
309                use sux::traits::Rank;
310                (rank.rank(off + p as usize + 1) - base) as i64
311            }
312        }
313    }
314
315    /// Transcript length in bases.
316    #[inline]
317    pub fn ref_len(&self) -> usize {
318        match self {
319            GcView::Dense(prefix) => prefix.len(),
320            GcView::Rank { len, .. } => *len,
321        }
322    }
323}
324
325/// A whole-transcriptome source of per-transcript [`GcView`]s: either dense
326/// per-transcript prefixes (default) or one shared [`GcRank`] window per
327/// transcript (`--reduceGCMemory`). Lets callers stay representation-agnostic.
328#[derive(Clone, Copy)]
329pub enum GcStore<'a> {
330    Dense(&'a [Vec<u32>]),
331    Rank {
332        rank: &'a GcRank,
333        offsets: &'a [u64],
334    },
335}
336
337impl<'a> GcStore<'a> {
338    /// The cumulative-GC view for transcript `tid`. The view borrows the
339    /// underlying data (lifetime `'a`), not `self`, so it outlives a temporary
340    /// `GcStore`.
341    #[inline]
342    pub fn view(&self, tid: usize) -> GcView<'a> {
343        match *self {
344            GcStore::Dense(prefixes) => GcView::Dense(&prefixes[tid]),
345            GcStore::Rank { rank, offsets } => {
346                let off = offsets[tid] as usize;
347                let len = (offsets[tid + 1] - offsets[tid]) as usize;
348                rank.view(off, len)
349            }
350        }
351    }
352}
353
354/// Fragment GC descriptor `(fragFrac, contextFrac)` for the closed fragment
355/// `[s, e]` (salmon's `Transcript::gcDesc`). The context fraction is the GC
356/// content over a 5-base 5' window `[s−3, s+1]` and a 5-base 3' window
357/// `[e−1, e+3]` (edge-clamped). Returns `None` when the context window is empty
358/// (matching salmon's `valid = false`).
359#[inline]
360pub fn gc_desc(v: &GcView, s: i32, e: i32) -> Option<(i32, i32)> {
361    let last = v.ref_len() as i32 - 1;
362    let cs = if s > 0 { v.cum(s - 1) } else { 0 };
363    let ce = v.cum(e);
364
365    let fs = s - OUTSIDE_5P;
366    let fe = s + INSIDE_5P;
367    let ts = e - INSIDE_3P;
368    let te = e + OUTSIDE_3P;
369
370    let fp_left = fs >= 0;
371    let fp_right = fe <= last;
372    let tp_left = ts >= 0;
373    let tp_right = te <= last;
374
375    let fps = if fp_left { v.cum(fs) } else { 0 };
376    let fpe = if fp_right { v.cum(fe) } else { ce };
377    let tps = if tp_left { v.cum(ts) } else { 0 };
378    let tpe = if tp_right { v.cum(te) } else { ce };
379
380    let fs_c = fs.max(0);
381    let fe_c = fe.min(last);
382    let ts_c = ts.max(0);
383    let te_c = te.min(last);
384    let fp_context_size = if !fp_left { fe_c + 1 } else { fe_c - fs_c };
385    let tp_context_size = if !tp_left { te_c + 1 } else { te_c - ts_c };
386    let context_size = (fp_context_size + tp_context_size) as f64;
387    if context_size == 0.0 {
388        return None;
389    }
390
391    let frag_frac = lrint(100.0 * (ce - cs) as f64 / (e - s + 1) as f64);
392    let context_frac = lrint(100.0 * ((fpe - fps) + (tpe - tps)) as f64 / context_size);
393    Some((frag_frac, context_frac))
394}
395
396/// Per-position context-GC cache for a single transcript — salmon's
397/// `populateContextCounts`. [`gc_desc`] recomputes the 5'/3' context window
398/// geometry (four conditional prefix loads, four edge branches, a division)
399/// for *every* `(fragStart, fragEnd)` pair in the effective-length convolution;
400/// since the 5' term depends only on `fragStart` and the 3' term only on
401/// `fragEnd`, we hoist both to O(ref_len) per-position arrays and reduce the
402/// inner loop to two array loads + one division + one `lrint` — matching the
403/// C++ hot path. [`GcContext::desc`] is bit-identical to [`gc_desc`] for every
404/// fragment the convolution actually visits (`fragStart + INSIDE_5P <= last`,
405/// i.e. `fp_right` always holds — guaranteed since `fragStart <= ref_len-2`).
406pub struct GcContext {
407    cum: Vec<u32>,      // cumulative G/C count in [0, p], indexed by position
408    fp_count: Vec<i64>, // 5' window GC count, indexed by fragStart
409    fp_wlen: Vec<i32>,  // 5' window length, indexed by fragStart
410    tp_count: Vec<i64>, // 3' window GC count, indexed by fragEnd
411    tp_wlen: Vec<i32>,  // 3' window length, indexed by fragEnd
412}
413
414impl GcContext {
415    /// Precompute the per-position 5'/3' context arrays (and a transient copy of
416    /// the cumulative-GC counts) from any [`GcView`]. Each entry replicates
417    /// exactly the corresponding sub-expression of [`gc_desc`] (same
418    /// edge-clamping, same `fp_right`/`tp_right` handling). Materializing `cum`
419    /// here keeps the hot inner loop ([`GcContext::desc`]) free of any rank
420    /// queries or dense lookups beyond cached arrays.
421    pub fn build(v: &GcView) -> GcContext {
422        let n = v.ref_len();
423        let last = n as i32 - 1;
424        let cum: Vec<u32> = (0..n).map(|p| v.cum(p as i32) as u32).collect();
425        let at = |i: i32| cum[i as usize] as i64;
426        let mut fp_count = vec![0i64; n];
427        let mut fp_wlen = vec![0i32; n];
428        let mut tp_count = vec![0i64; n];
429        let mut tp_wlen = vec![0i32; n];
430        for p in 0..n {
431            let pos = p as i32;
432            // 3' window for a fragment *ending* at `pos`.
433            let ts = pos - INSIDE_3P;
434            let te = pos + OUTSIDE_3P;
435            let tp_left = ts >= 0;
436            let tp_right = te <= last;
437            let tps = if tp_left { at(ts) } else { 0 };
438            // `tpe = ce = cum[e]` when the 3' window runs off the end.
439            let tpe = if tp_right { at(te) } else { at(pos) };
440            let ts_c = ts.max(0);
441            let te_c = te.min(last);
442            tp_count[p] = tpe - tps;
443            tp_wlen[p] = if !tp_left { te_c + 1 } else { te_c - ts_c };
444
445            // 5' window for a fragment *starting* at `pos`. Within the
446            // convolution `fp_right` always holds; the `cum[p]` fallback is a
447            // placeholder for the never-visited `pos == last` slot.
448            let fs = pos - OUTSIDE_5P;
449            let fe = pos + INSIDE_5P;
450            let fp_left = fs >= 0;
451            let fp_right = fe <= last;
452            let fps = if fp_left { at(fs) } else { 0 };
453            let fpe = if fp_right { at(fe) } else { at(pos) };
454            let fs_c = fs.max(0);
455            let fe_c = fe.min(last);
456            fp_count[p] = fpe - fps;
457            fp_wlen[p] = if !fp_left { fe_c + 1 } else { fe_c - fs_c };
458        }
459        GcContext {
460            cum,
461            fp_count,
462            fp_wlen,
463            tp_count,
464            tp_wlen,
465        }
466    }
467
468    /// GC descriptor `(fragFrac, contextFrac)` for the closed fragment `[s, e]`,
469    /// bit-identical to [`gc_desc`] but using the cached context arrays. `s` and
470    /// `e` must lie in `0..ref_len` with `s + INSIDE_5P <= ref_len - 1`.
471    #[inline]
472    pub fn desc(&self, s: i32, e: i32) -> Option<(i32, i32)> {
473        let context_size = (self.fp_wlen[s as usize] + self.tp_wlen[e as usize]) as f64;
474        if context_size == 0.0 {
475            return None;
476        }
477        let cs = if s > 0 {
478            self.cum[(s - 1) as usize] as i64
479        } else {
480            0
481        };
482        let ce = self.cum[e as usize] as i64;
483        let count = (self.fp_count[s as usize] + self.tp_count[e as usize]) as f64;
484        let frag_frac = lrint(100.0 * (ce - cs) as f64 / (e - s + 1) as f64);
485        let context_frac = lrint(100.0 * count / context_size);
486        Some((frag_frac, context_frac))
487    }
488}
489
490/// Build the *expected* fragment-GC model (salmon's `transcriptGCDist`): slide
491/// every fragment `[fragStart, fragStart+fl−1]` over each expressed transcript,
492/// weighting each by `(alpha/effLen)·(conditionalCDF(fl) − prevMass)` — the same
493/// abundance-density × conditional-FLD weighting used for the expected
494/// sequence-bias model. `k` is the leading offset salmon excludes from the
495/// fragment-start loop (`9` with `--seqBias`, `1` otherwise); `stride` subsamples
496/// fragment lengths ([`GC_SAMP_STRIDE`]).
497#[allow(clippy::too_many_arguments)]
498pub fn build_expected_gc<'a, FS, FP>(
499    num_targets: usize,
500    seq_of: FS,
501    view_of: FP,
502    alphas: &[f64],
503    eff_lens: &[f64],
504    cdf: &[f64],
505    fld_low: usize,
506    fld_high: usize,
507    cond_bins: usize,
508    gc_bins: usize,
509    k: usize,
510    stride: usize,
511) -> GcFragModel
512where
513    FS: Fn(usize) -> &'a [u8] + Sync,
514    FP: Fn(usize) -> GcView<'a> + Sync,
515{
516    let stride = stride.max(1) as i32;
517    // The expected-GC distribution is a sum of independent per-transcript
518    // contributions, each an O(refLen · fldRange/stride) double loop — billions
519    // of `gc_desc` calls in total. salmon parallelizes this over transcripts;
520    // do the same with rayon (per-thread `GcFragModel` partials, reduced via
521    // `combine_counts`). `seq_of`/`prefix_of` must be `Sync` to share across
522    // threads (they are: closures over `&[Vec<_>]`).
523    use rayon::prelude::*;
524    let per_tid = |tid: usize| -> Option<GcFragModel> {
525        if alphas[tid] < MIN_ALPHA || eff_lens[tid] <= 0.0 {
526            return None;
527        }
528        let seq = seq_of(tid);
529        let ref_len = seq.len();
530        if ref_len <= k {
531            return None;
532        }
533        let cdf_max_arg = (cdf.len() - 1).min(ref_len);
534        let cdf_max_val = cdf[cdf_max_arg];
535        if cdf_max_val < MIN_CDF_MASS {
536            return None;
537        }
538        let view = view_of(tid);
539        let ctx = GcContext::build(&view);
540        let weight = alphas[tid] / eff_lens[tid];
541        let cond = |x: i32| conditional_cdf(cdf, cdf_max_arg, cdf_max_val, x);
542        let sp = if fld_low > 0 { fld_low as i32 - 1 } else { 0 };
543        let mut model = GcFragModel::new(cond_bins, gc_bins);
544        for frag_start in 0..(ref_len - k) {
545            let mut prev = cond(sp);
546            let mut fl = fld_low as i32;
547            while fl <= fld_high as i32 {
548                let frag_end = frag_start as i32 + fl - 1;
549                if (frag_end as usize) < ref_len {
550                    if let Some((ff, cf)) = ctx.desc(frag_start as i32, frag_end) {
551                        model.inc(ff, cf, weight * (cond(fl) - prev));
552                    }
553                    prev = cond(fl);
554                } else {
555                    break;
556                }
557                fl += stride;
558            }
559        }
560        Some(model)
561    };
562    (0..num_targets)
563        .into_par_iter()
564        .fold(
565            || GcFragModel::new(cond_bins, gc_bins),
566            |mut acc, tid| {
567                if let Some(m) = per_tid(tid) {
568                    acc.combine_counts(&m);
569                }
570                acc
571            },
572        )
573        .reduce(
574            || GcFragModel::new(cond_bins, gc_bins),
575            |mut a, b| {
576                a.combine_counts(&b);
577                a
578            },
579        )
580}
581
582/// Bias-corrected effective length including fragment-GC bias (and, when
583/// `seq_models` is provided, sequence-specific bias too). Mirrors salmon's
584/// combined `updateEffectiveLengths` convolution: per-position 5'/3' sequence
585/// factors are multiplied by the per-fragment `gcBias.get({fragFrac, contextFrac})`
586/// ratio, convolved with the conditional FLD, and floored at the lower barrier.
587///
588/// `gc_bias` is the normalized observed/expected ratio model ([`gc_ratio`]).
589/// `seq_models` is `(obs_fw, exp_fw, obs_rc, exp_rc)` when `--seqBias` is also on.
590#[allow(clippy::too_many_arguments)]
591pub fn gc_corrected_effective_length(
592    seq: &[u8],
593    prefix: &[u32],
594    cdf: &[f64],
595    fld_low: usize,
596    fld_high: usize,
597    gc_bias: &GcFragModel,
598    seq_models: Option<(&SBModel, &SBModel, &SBModel, &SBModel)>,
599    elen: f64,
600    stride: usize,
601) -> f64 {
602    let k = if seq_models.is_some() {
603        CONTEXT_LENGTH
604    } else {
605        1
606    };
607    let ref_len = seq.len();
608    let unprocessed = (ref_len as i32 - elen as i32).max(0);
609    let cdf_max_arg = (cdf.len() - 1).min(ref_len);
610    let cdf_max_val = cdf[cdf_max_arg];
611    if ref_len < k || unprocessed <= 0 || cdf_max_val < MIN_CDF_MASS {
612        return elen;
613    }
614    let cond = |x: i32| conditional_cdf(cdf, cdf_max_arg, cdf_max_val, x);
615
616    // Per-position sequence-bias factors (1.0 when not seq-correcting).
617    let mut fw = vec![1.0f64; ref_len];
618    let mut rc = vec![1.0f64; ref_len];
619    if let Some((obs_fw, exp_fw, obs_rc, exp_rc)) = seq_models {
620        let cu = CONTEXT_LEFT;
621        let rc_seq = revcomp_bytes(seq);
622        for frag_start in 0..(ref_len - CONTEXT_LENGTH) {
623            let read_start = frag_start + cu;
624            if read_start < ref_len {
625                fw[read_start] = log_bias(
626                    obs_fw,
627                    exp_fw,
628                    &seq[frag_start..frag_start + CONTEXT_LENGTH],
629                    false,
630                )
631                .exp();
632                rc[read_start] = log_bias(
633                    obs_rc,
634                    exp_rc,
635                    &rc_seq[frag_start..frag_start + CONTEXT_LENGTH],
636                    false,
637                )
638                .exp();
639            }
640        }
641        rc.reverse();
642    }
643
644    // Convolve seq×gc fragment factors with the conditional FLD.
645    let ctx = GcContext::build(&GcView::Dense(prefix));
646    let stride = stride.max(1) as i32;
647    let max_len = (ref_len as i32).min(fld_high as i32 + 1);
648    let mut fl = fld_low as i32;
649    let mut done = fl >= max_len;
650    let sp = if fl > 0 { fl - 1 } else { 0 };
651    let mut prev_mass = cond(sp);
652    let mut eff = 0.0f64;
653    while !done {
654        if fl >= max_len {
655            done = true;
656            fl = max_len - 1;
657        }
658        let fl_weight = cond(fl) - prev_mass;
659        prev_mass = cond(fl);
660        let mut mass = 0.0f64;
661        // Hoist the bound: for kstart in [0, kmax) we have
662        // frag_end = kstart+fl-1 <= ref_len-2 < ref_len, so the old per-iteration
663        // `frag_end < ref_len` guard is always true and is dropped (it kept
664        // `ref_len` live in the inner loop, forcing a spill/reload).
665        let kmax = ref_len as i32 - fl;
666        let mut kstart = 0i32;
667        while kstart < kmax {
668            let frag_start = kstart;
669            let frag_end = kstart + fl - 1;
670            let mut frag_factor = fw[frag_start as usize] * rc[frag_end as usize];
671            if let Some((ff, cf)) = ctx.desc(frag_start, frag_end) {
672                frag_factor *= gc_bias.get(ff, cf);
673            }
674            mass += frag_factor;
675            kstart += 1;
676        }
677        eff += fl_weight * mass;
678        fl += stride;
679    }
680
681    let offset = (unprocessed as f64).max(1.0);
682    eff.max(elen.min(offset))
683}
684
685#[cfg(test)]
686mod tests {
687    use super::*;
688
689    #[test]
690    fn bin_frac_identity_at_101() {
691        assert_eq!(bin_frac(0, 101), 0);
692        assert_eq!(bin_frac(57, 101), 57);
693        assert_eq!(bin_frac(100, 101), 100);
694    }
695
696    #[test]
697    fn bin_frac_coarse() {
698        // 3 bins: width 33.33 -> [0,33),[33,66),[66,100]
699        assert_eq!(bin_frac(0, 3), 0);
700        assert_eq!(bin_frac(30, 3), 0);
701        assert_eq!(bin_frac(40, 3), 1);
702        assert_eq!(bin_frac(70, 3), 2);
703        assert_eq!(bin_frac(100, 3), 2); // clamped to n-1
704    }
705
706    #[test]
707    fn normalize_makes_rows_sum_to_one() {
708        let mut m = GcFragModel::new(2, 5);
709        for gc in 0..5 {
710            m.inc(gc * 25, 10, (gc + 1) as f64); // row 0 (ctx 10 -> bin 0)
711        }
712        m.normalize();
713        let row0: f64 = (0..5).map(|c| m.counts[c]).sum();
714        assert!((row0 - 1.0).abs() < 1e-9, "row0 sum = {row0}");
715    }
716
717    #[test]
718    fn ratio_of_identical_models_is_one() {
719        let mut obs = GcFragModel::new(1, 101);
720        let mut exp = GcFragModel::new(1, 101);
721        for gc in 0..=100 {
722            obs.inc(gc, 0, (gc + 1) as f64);
723            exp.inc(gc, 0, (gc + 1) as f64);
724        }
725        let r = gc_ratio(&mut obs, &mut exp, GC_MAX_RATIO);
726        for gc in 0..=100 {
727            assert!(
728                (r.get(gc, 0) - 1.0).abs() < 1e-9,
729                "gc {gc}: {}",
730                r.get(gc, 0)
731            );
732        }
733    }
734
735    #[test]
736    fn gc_frac_basic() {
737        // 10-base sequence, 5 G/C -> 50%
738        let seq = b"ACGTACGTAC"; // A C G T A C G T A C : GC at 1,2,5,6,9 = 5
739        let p = gc_prefix(seq);
740        assert_eq!(gc_frac(&p, 0, 9), 50);
741        // all-GC window
742        let seq2 = b"GGGGCCCC";
743        let p2 = gc_prefix(seq2);
744        assert_eq!(gc_frac(&p2, 0, 7), 100);
745        // single base
746        assert_eq!(gc_frac(&p, 2, 2), 100); // 'G'
747        assert_eq!(gc_frac(&p, 0, 0), 0); // 'A'
748    }
749
750    #[test]
751    fn gc_desc_matches_frag_and_context() {
752        // 40-base sequence; check an interior fragment
753        let seq: Vec<u8> = b"ACGT".iter().cycle().take(40).copied().collect();
754        let p = gc_prefix(&seq);
755        let (ff, cf) = gc_desc(&GcView::Dense(&p), 10, 29).unwrap();
756        // ACGT repeat is exactly 50% GC everywhere
757        assert_eq!(ff, 50, "fragFrac");
758        assert!((0..=100).contains(&cf), "contextFrac in range: {cf}");
759        // fragFrac must equal gc_frac over the same interval
760        assert_eq!(ff, gc_frac(&p, 10, 29));
761    }
762
763    #[test]
764    fn idx_lut_matches_bin_frac() {
765        // The LUT-based idx must equal the original divide-based formula for
766        // every (ctx, gc) fraction the model can see, including out-of-[0,100].
767        for &(cb, gb) in &[(3usize, 25usize), (1, 101), (3, 101), (4, 50)] {
768            let m = GcFragModel::new(cb, gb);
769            for cf in -5i32..=105 {
770                for ff in -5i32..=105 {
771                    let want_ctx = if cb > 1 { bin_frac(cf, cb) } else { 0 };
772                    let want = want_ctx * gb + bin_frac(ff, gb);
773                    assert_eq!(m.idx(cf, ff), want, "cb={cb} gb={gb} cf={cf} ff={ff}");
774                }
775            }
776        }
777    }
778
779    #[test]
780    fn gc_context_matches_gc_desc() {
781        // Cached context (GcContext::desc) must be bit-identical to the
782        // per-fragment gc_desc for every fragment the convolution can visit
783        // (fragStart in 0..ref_len-fl, fragEnd = fragStart+fl-1).
784        let bases = [b'A', b'C', b'G', b'T', b'C', b'G', b'A', b'T', b'G', b'C'];
785        let seq: Vec<u8> = (0..200).map(|i| bases[(i * 7 + 3) % bases.len()]).collect();
786        let p = gc_prefix(&seq);
787        let v = GcView::Dense(&p);
788        let ctx = GcContext::build(&v);
789        let n = seq.len() as i32;
790        for fl in 1..=120 {
791            let kmax = n - fl;
792            for s in 0..kmax {
793                let e = s + fl - 1;
794                assert_eq!(
795                    ctx.desc(s, e),
796                    gc_desc(&v, s, e),
797                    "mismatch at fl={fl}, s={s}, e={e}"
798                );
799            }
800        }
801    }
802
803    #[test]
804    fn gc_rank_view_matches_dense() {
805        // The rank-bitvector view must return identical cumulative GC (and hence
806        // identical gc_desc) to the dense prefix, for an offset window of a
807        // concatenated buffer.
808        let bases = [b'A', b'C', b'G', b'T', b'G', b'C', b'A', b'T'];
809        let pad: Vec<u8> = (0..37).map(|i| bases[(i * 3) % bases.len()]).collect();
810        let seq: Vec<u8> = (0..150).map(|i| bases[(i * 5 + 1) % bases.len()]).collect();
811        // concat = [pad | seq], so the transcript sits at offset pad.len().
812        let mut concat = pad.clone();
813        concat.extend_from_slice(&seq);
814        let p = gc_prefix(&seq);
815        let rank = GcRank::new(&concat);
816        let dense = GcView::Dense(&p);
817        let rview = rank.view(pad.len(), seq.len());
818        for fl in 1..=100 {
819            let n = seq.len() as i32;
820            for s in 0..(n - fl) {
821                let e = s + fl - 1;
822                assert_eq!(
823                    gc_desc(&dense, s, e),
824                    gc_desc(&rview, s, e),
825                    "fl={fl} s={s}"
826                );
827            }
828        }
829    }
830
831    #[test]
832    fn gc_desc_context_window_geometry() {
833        // Make the 5' context window [s-3, s+1] all GC and the rest AT, so the
834        // contextFrac is dominated by the GC island near the fragment ends.
835        let mut seq = vec![b'A'; 60];
836        for i in 17..=21 {
837            seq[i] = b'G'; // 5' window for s=20 is [17,21]
838        }
839        for i in 39..=43 {
840            seq[i] = b'C'; // 3' window for e=40 is [39,43]
841        }
842        let p = gc_prefix(&seq);
843        let (_ff, cf) = gc_desc(&GcView::Dense(&p), 20, 40).unwrap();
844        // both 5-base context windows are fully G/C -> 100%
845        assert_eq!(cf, 100, "contextFrac");
846    }
847
848    #[test]
849    fn ratio_clamped() {
850        let mut obs = GcFragModel::new(1, 2);
851        let mut exp = GcFragModel::new(1, 2);
852        obs.inc(0, 0, 1e6); // bin 0 hugely enriched in obs
853        exp.inc(75, 0, 1e6); // bin 1 enriched in exp
854        let r = gc_ratio(&mut obs, &mut exp, 1000.0);
855        for v in [r.get(0, 0), r.get(75, 0)] {
856            assert!(
857                (1.0 / 1000.0 - 1e-12..=1000.0 + 1e-6).contains(&v),
858                "ratio {v} out of clamp"
859            );
860        }
861    }
862}