Skip to main content

salmon_model/
fld.rs

1//! Fragment-length distribution.
2//!
3//! Direct port of salmon's `FragmentLengthDistribution`
4//! (`src/model/FragmentLengthDistribution.cpp`): a log-space histogram seeded
5//! with a Gaussian (or uniform) prior, updated by adding a binomial smoothing
6//! kernel around each observed length. All masses and probabilities are in log
7//! space. Updates are lock-free so worker threads can call [`add_val`] with a
8//! shared reference, matching the C++ design.
9//!
10//! [`add_val`]: FragmentLengthDistribution::add_val
11
12use salmon_core::atomic::AtomicF64;
13use salmon_core::math::{log_add, LOG_0, LOG_EPSILON};
14use salmon_core::{LibraryFormat, ReadOrientation, ReadStrandedness, ReadType};
15use statrs::distribution::{Binomial, ContinuousCDF, Discrete, Normal};
16use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
17use std::sync::{Arc, RwLock};
18
19/// Tracks the observed distribution of fragment lengths.
20#[derive(Debug)]
21pub struct FragmentLengthDistribution {
22    /// logged binomial smoothing kernel
23    kernel: Vec<f64>,
24    /// logged observed mass per length bin
25    hist: Vec<AtomicF64>,
26    /// logged total observed mass (including pseudo-counts)
27    tot_mass: AtomicF64,
28    /// logged sum of length*mass, for fast mean computation
29    sum: AtomicF64,
30    /// minimum observed length (bin units)
31    min: AtomicUsize,
32    /// internal bin size
33    bin_size: usize,
34
35    /// cached normalized PMF, valid once [`cache`](Self::cache) is called
36    cached_pmf: Vec<f64>,
37    /// cached CMF
38    cached_cmf: Vec<f64>,
39    have_cache: bool,
40
41    /// Periodically-refreshed snapshot of the (un-normalized) log-PMF used during
42    /// the *online* phase, indexed by raw length. Reading the live `hist`/`tot_mass`
43    /// directly (two separate atomic loads on a concurrently-updated distribution)
44    /// returns slightly different values for the same length across calls, which
45    /// breaks the weight symmetry of exact-duplicate transcripts and is then
46    /// amplified by the VBEM `α<1` prior. Mirroring C++ salmon (`cachedPMF_` +
47    /// `LogCMFCache`), worker threads instead capture an immutable snapshot of this
48    /// once per fragment ([`online_snapshot`](Self::online_snapshot)) so every
49    /// transcript of a given length in that fragment gets an identical value;
50    /// [`refresh_online`](Self::refresh_online) rebuilds it at mini-batch
51    /// boundaries.
52    online_pmf: RwLock<Arc<Vec<f64>>>,
53
54    /// Periodically-refreshed snapshot of the (normalized) log-CMF, the
55    /// cumulative companion to [`online_pmf`](Self::online_pmf). Used for the
56    /// ambiguous (orphan / single-end) fragment-length probability and for the
57    /// `pmf(flen) − cmf(txpLen)` length-conditioning of proper pairs, both of
58    /// which need cumulative mass. Rebuilt alongside `online_pmf` in
59    /// [`refresh_online`](Self::refresh_online); mirrors C++ salmon's
60    /// `LogCMFCache`.
61    online_cmf: RwLock<Arc<Vec<f64>>>,
62}
63
64impl FragmentLengthDistribution {
65    /// Construct a distribution.
66    ///
67    /// * `alpha` – total pseudo-count mass (linear space).
68    /// * `max_val` – maximum representable length.
69    /// * `prior_mu` – Gaussian prior mean; if `<= 0`, a uniform prior is used.
70    /// * `prior_sigma` – Gaussian prior standard deviation.
71    /// * `kernel_n` – binomial kernel trials; must be even (after binning).
72    /// * `kernel_p` – binomial kernel success probability.
73    /// * `bin_size` – internal length binning (use 1 for no binning).
74    pub fn new(
75        alpha: f64,
76        max_val: usize,
77        prior_mu: f64,
78        prior_sigma: f64,
79        kernel_n: usize,
80        kernel_p: f64,
81        bin_size: usize,
82    ) -> Self {
83        assert!(bin_size >= 1, "bin_size must be >= 1");
84        let max_val = max_val / bin_size;
85        let kernel_n = kernel_n / bin_size;
86        assert!(
87            kernel_n.is_multiple_of(2),
88            "kernel_n must be even after binning"
89        );
90
91        let tot = alpha.ln();
92        let hist: Vec<AtomicF64>;
93        let mut sum = LOG_0;
94        let mut tot_mass;
95
96        if prior_mu > 0.0 {
97            let norm = Normal::new(
98                prior_mu / bin_size as f64,
99                prior_sigma / (bin_size * bin_size) as f64,
100            )
101            .expect("valid normal prior");
102            hist = (0..=max_val).map(|_| AtomicF64::new(LOG_0)).collect();
103            tot_mass = LOG_0;
104            for (i, slot) in hist.iter().enumerate() {
105                let norm_mass = norm.cdf(i as f64 + 0.5) - norm.cdf(i as f64 - 0.5);
106                let mass = if norm_mass != 0.0 {
107                    tot + norm_mass.ln()
108                } else {
109                    LOG_EPSILON
110                };
111                slot.store(mass);
112                sum = log_add(sum, (i as f64).ln() + mass);
113                tot_mass = log_add(tot_mass, mass);
114            }
115        } else {
116            // uniform prior
117            let per = tot - (max_val as f64).ln();
118            hist = (0..=max_val).map(|_| AtomicF64::new(per)).collect();
119            hist[0].store(LOG_0);
120            let h1 = hist.get(1).map(|a| a.load()).unwrap_or(per);
121            sum = h1 + ((max_val * (max_val + 1)) as f64).ln() - 2.0_f64.ln();
122            tot_mass = tot;
123        }
124
125        // binomial smoothing kernel
126        let binom = Binomial::new(kernel_p, kernel_n as u64).expect("valid binomial kernel");
127        let kernel: Vec<f64> = (0..=kernel_n).map(|i| binom.pmf(i as u64).ln()).collect();
128
129        Self {
130            kernel,
131            hist,
132            tot_mass: AtomicF64::new(tot_mass),
133            sum: AtomicF64::new(sum),
134            min: AtomicUsize::new(max_val),
135            bin_size,
136            cached_pmf: Vec::new(),
137            cached_cmf: Vec::new(),
138            have_cache: false,
139            online_pmf: RwLock::new(Arc::new(Vec::new())),
140            online_cmf: RwLock::new(Arc::new(Vec::new())),
141        }
142    }
143
144    /// salmon's default fragment-length distribution: pseudo-count 1.0, max
145    /// length 1000, no Gaussian prior (uniform), kernel `n=4, p=0.5`.
146    pub fn default_for_paired() -> Self {
147        Self::new(1.0, 1000, 0.0, 0.0, 4, 0.5, 1)
148    }
149
150    pub fn max_val(&self) -> usize {
151        (self.hist.len() - 1) * self.bin_size
152    }
153
154    pub fn min_val(&self) -> usize {
155        let m = self.min.load(Ordering::Relaxed);
156        if m == self.hist.len() - 1 {
157            1
158        } else {
159            m
160        }
161    }
162
163    /// Add `mass` (log space) for an observed fragment of length `len`,
164    /// spreading it over the smoothing kernel. Lock-free; safe to call from
165    /// multiple threads. (Must not race with [`cache`](Self::cache).)
166    pub fn add_val(&self, len: usize, mass: f64) {
167        let mut len = len / self.bin_size;
168        let max_v = self.max_val() / self.bin_size;
169        if len > max_v {
170            len = max_v;
171        }
172        self.min.fetch_min(len, Ordering::Relaxed);
173
174        let half = self.kernel.len() / 2;
175        // offset can go negative conceptually; use isize math then bound-check.
176        let mut offset = len as isize - half as isize;
177        for &k in &self.kernel {
178            if offset > 0 && (offset as usize) < self.hist.len() {
179                let o = offset as usize;
180                let k_mass = mass + k;
181                self.hist[o].log_add_assign(k_mass);
182                self.sum.log_add_assign((o as f64).ln() + k_mass);
183                self.tot_mass.log_add_assign(k_mass);
184            }
185            offset += 1;
186        }
187    }
188
189    /// Logged probability of observing a fragment of length `len`.
190    pub fn pmf(&self, len: usize) -> f64 {
191        if self.have_cache {
192            return *self
193                .cached_pmf
194                .get(len)
195                .unwrap_or_else(|| self.cached_pmf.last().unwrap());
196        }
197        let mut l = len / self.bin_size;
198        let max_v = self.max_val() / self.bin_size;
199        if l > max_v {
200            l = max_v;
201        }
202        self.hist[l].load() - self.tot_mass.load()
203    }
204
205    /// Rebuild the online log-PMF snapshot from the current histogram (one pass
206    /// over the length bins, with a single `tot_mass` read so the snapshot is
207    /// internally consistent). Call at mini-batch boundaries during the online
208    /// phase; no-op once the final [`cache`](Self::cache) has been taken. Cheap
209    /// relative to mapping a batch, and decouples per-fragment reads from the
210    /// concurrent `add_val` writes so identical lengths read identical values.
211    pub fn refresh_online(&self) {
212        if self.have_cache {
213            return;
214        }
215        let max_raw = self.max_val();
216        let max_v = max_raw / self.bin_size;
217        let tot = self.tot_mass.load();
218        // Per-bin cumulative mass (matches `cmf()`), so the snapshot CMF at raw
219        // index `raw` equals `cmf(raw)`. Built first, then both the PMF and CMF
220        // snapshots are expanded over raw indices from the same `tot` read so
221        // they are mutually consistent.
222        let mut bin_cum = Vec::with_capacity(max_v + 1);
223        let mut cum = LOG_0;
224        for b in 0..=max_v {
225            cum = log_add(cum, self.hist[b].load() - tot);
226            bin_cum.push(cum);
227        }
228        let mut v = Vec::with_capacity(max_raw + 1);
229        let mut c = Vec::with_capacity(max_raw + 1);
230        for raw in 0..=max_raw {
231            let l = (raw / self.bin_size).min(max_v);
232            v.push(self.hist[l].load() - tot);
233            c.push(bin_cum[l]);
234        }
235        *self.online_pmf.write().unwrap() = Arc::new(v);
236        *self.online_cmf.write().unwrap() = Arc::new(c);
237    }
238
239    /// Cheap (one `Arc` clone) immutable handle to the current online log-PMF
240    /// snapshot. Capture once per fragment and index by raw length: every
241    /// transcript of a given length then reads an identical value even if another
242    /// thread refreshes the shared snapshot meanwhile. Empty until the first
243    /// [`refresh_online`](Self::refresh_online) (the pre-burn-in window, where this
244    /// term is not folded into the eq-class weight anyway).
245    pub fn online_snapshot(&self) -> Arc<Vec<f64>> {
246        self.online_pmf.read().unwrap().clone()
247    }
248
249    /// Cheap (one `Arc` clone) immutable handle to the current online log-CMF
250    /// snapshot, the cumulative companion to [`online_snapshot`](Self::online_snapshot).
251    /// Capture once per fragment for the ambiguous (orphan / single-end)
252    /// fragment-length probability and the proper-pair length-conditioning.
253    /// Empty until the first [`refresh_online`](Self::refresh_online).
254    pub fn online_cmf_snapshot(&self) -> Arc<Vec<f64>> {
255        self.online_cmf.read().unwrap().clone()
256    }
257
258    /// Logged cumulative mass up to and including `len`.
259    pub fn cmf(&self, len: usize) -> f64 {
260        if self.have_cache {
261            return *self
262                .cached_cmf
263                .get(len)
264                .unwrap_or_else(|| self.cached_cmf.last().unwrap());
265        }
266        let mut l = len / self.bin_size;
267        let max_v = self.max_val() / self.bin_size;
268        if l > max_v {
269            l = max_v;
270        }
271        let mut cum = LOG_0;
272        for i in 0..=l {
273            cum = log_add(cum, self.hist[i].load());
274        }
275        cum - self.tot_mass.load()
276    }
277
278    /// Total observed mass (log space).
279    pub fn tot_mass(&self) -> f64 {
280        self.tot_mass.load()
281    }
282
283    /// Mean observed length.
284    pub fn mean(&self) -> f64 {
285        (self.sum.load() - self.tot_mass.load()).exp()
286    }
287
288    /// Standard deviation of the observed length distribution, computed from the
289    /// cached normalized PMF (call after [`cache`](Self::cache)).
290    pub fn sd(&self) -> f64 {
291        let lp = self.log_pmf();
292        if lp.is_empty() {
293            return 0.0;
294        }
295        let mut mean = 0.0;
296        for (l, &p) in lp.iter().enumerate() {
297            mean += (l as f64) * p.exp();
298        }
299        let mut var = 0.0;
300        for (l, &p) in lp.iter().enumerate() {
301            let d = l as f64 - mean;
302            var += d * d * p.exp();
303        }
304        var.max(0.0).sqrt()
305    }
306
307    /// Freeze the distribution and precompute normalized PMF/CMF for fast,
308    /// allocation-free lookup. Call once after updates have stopped.
309    pub fn cache(&mut self) {
310        if self.have_cache {
311            return;
312        }
313        let max_v = self.max_val();
314        // normalized PMF over [0, max_v]
315        let mut pmf = Vec::with_capacity(max_v + 1);
316        let mut tot = LOG_0;
317        for i in 0..=max_v {
318            let p = self.pmf(i);
319            pmf.push(p);
320            tot = log_add(tot, p);
321        }
322        for p in &mut pmf {
323            *p -= tot;
324        }
325        // CMF from the normalized PMF
326        let mut cmf = Vec::with_capacity(pmf.len());
327        let mut cum = LOG_0;
328        for &p in &pmf {
329            cum = log_add(cum, p);
330            cmf.push(cum);
331        }
332        self.cached_pmf = pmf;
333        self.cached_cmf = cmf;
334        self.have_cache = true;
335    }
336
337    /// Reconstruct a *cached* distribution directly from a (log-space) PMF,
338    /// e.g. one serialized into a RAD header during a previous run. `log_pmf` is
339    /// indexed by raw length over `[0, log_pmf.len())`. The CMF, conditional
340    /// means, mean and sd are all re-derived from it via [`cache`](Self::cache),
341    /// so a reconstructed distribution is interchangeable with the original for
342    /// every read-side use. The masses need not be pre-normalized — `cache`
343    /// normalizes them — but a normalized PMF round-trips exactly.
344    pub fn from_log_pmf(log_pmf: &[f64]) -> Self {
345        let max_val = log_pmf.len().saturating_sub(1);
346        let mut d = Self::new(1.0, max_val, 0.0, 1.0, 4, 0.5, 1);
347        // Replace the prior histogram with the supplied masses and recompute the
348        // aggregate statistics (so `mean`/`sd` are consistent), then cache.
349        let mut tot = LOG_0;
350        let mut sm = LOG_0;
351        for (i, &p) in log_pmf.iter().enumerate() {
352            d.hist[i].store(p);
353            tot = log_add(tot, p);
354            if i > 0 {
355                sm = log_add(sm, (i as f64).ln() + p);
356            }
357        }
358        d.tot_mass.store(tot);
359        d.sum.store(sm);
360        d.min.store(0, Ordering::Relaxed);
361        d.cache();
362        d
363    }
364
365    /// The cached, normalized log-PMF over `[0, max_val]`. Requires [`cache`](Self::cache).
366    pub fn log_pmf(&self) -> &[f64] {
367        debug_assert!(self.have_cache, "call cache() before log_pmf()");
368        &self.cached_pmf
369    }
370
371    /// Cumulative conditional means `E[L | L ≤ i]` over `[0, max_val]`, i.e.
372    /// salmon's `correctionFactorsFromMass` (`DistributionUtils.cpp`):
373    /// `cm[i] = (Σ_{l≤i} l·pmf[l]) / (Σ_{l≤i} pmf[l])`.
374    ///
375    /// These are the per-length correction factors `computeSmoothedEffectiveLengths`
376    /// subtracts from the reference length to get the base effective length. The
377    /// ratio is invariant to the PMF normalization, so the cached (normalized) PMF
378    /// gives the same values as salmon's `100·exp(logPMF)` mass. Requires
379    /// [`cache`](Self::cache).
380    pub fn conditional_means(&self) -> Vec<f64> {
381        debug_assert!(self.have_cache, "call cache() before conditional_means()");
382        let n = self.cached_pmf.len();
383        let mut cms = vec![0.0f64; n];
384        let mut vals = 0.0; // Σ l·pmf[l]
385        let mut mult = 0.0; // Σ pmf[l]
386        for i in 0..n {
387            let p = self.cached_pmf[i].exp();
388            vals += (i as f64) * p;
389            mult += p;
390            cms[i] = if mult > 0.0 { vals / mult } else { 0.0 };
391        }
392        cms
393    }
394}
395
396/// Index a length into a (log) CMF snapshot, clamping out-of-range lengths to
397/// the last bin (which holds the total mass). Returns [`LOG_0`] for an empty
398/// snapshot.
399#[inline]
400fn cmf_at(cmf: &[f64], len: i32) -> f64 {
401    if cmf.is_empty() {
402        return LOG_0;
403    }
404    let i = (len.max(0) as usize).min(cmf.len() - 1);
405    cmf[i]
406}
407
408/// Logged ambiguous-fragment-length probability for an orphan / single-end
409/// read, given a (log) CMF snapshot. Direct port of C++ salmon's
410/// `LogCMFCache::getAmbigFragLengthProb` (`DistributionUtils.cpp`).
411///
412/// The mapped mate bounds the maximum possible fragment length: a forward read
413/// at `pos` can extend downstream to the transcript 3' end (`txp_len − pos`); a
414/// reverse read's outer (5') end sits at `pos + read_len`, bounding the upstream
415/// extent toward the 5' end. The weight is the FLD mass up to that bound,
416/// *conditioned* on the mass up to the full transcript length — i.e.
417/// `cmf(maxFragLen) − cmf(txpLen)` — so orphan weights sit on the same
418/// length-conditioned scale as proper pairs. Returns [`LOG_EPSILON`] when the
419/// transcript admits no representable fragment mass, and `LOG_1` (= 0) when no
420/// snapshot is available yet (pre-burn-in), leaving the weight unmodelled.
421pub fn ambig_frag_log_prob(cmf: &[f64], fwd: bool, pos: i32, read_len: i32, txp_len: i32) -> f64 {
422    if cmf.is_empty() {
423        return 0.0; // LOG_1: no model yet
424    }
425    let stxp = txp_len.max(0);
426    let max_frag_len = if fwd {
427        stxp - pos.clamp(0, stxp)
428    } else {
429        (pos + read_len).clamp(0, stxp)
430    };
431    let ref_cm = cmf_at(cmf, stxp);
432    if ref_cm <= LOG_0 {
433        return LOG_EPSILON;
434    }
435    cmf_at(cmf, max_frag_len) - ref_cm
436}
437
438/// salmon's base effective length (`computeSmoothedEffectiveLengths`):
439/// `effLen = refLen − E[L | L ≤ refLen]`, clamped back to `refLen` if it would
440/// fall below 1. `cond_means` is [`FragmentLengthDistribution::conditional_means`].
441///
442/// This replaces the truncated-PMF `Σ pmf(l)·(refLen−l+1)` estimate (which falls
443/// back to the raw `refLen` for any transcript shorter than the FLD mean), matching
444/// salmon's behaviour exactly.
445pub fn smoothed_effective_length(cond_means: &[f64], ref_len: usize) -> f64 {
446    if cond_means.is_empty() {
447        return ref_len as f64;
448    }
449    let max_len = cond_means.len();
450    let cf = if ref_len >= max_len {
451        cond_means[max_len - 1]
452    } else {
453        cond_means[ref_len]
454    };
455    let eff = ref_len as f64 - cf;
456    if eff < 1.0 {
457        ref_len as f64
458    } else {
459        eff
460    }
461}
462
463/// Order-independent accumulator for *deriving* a fragment-length distribution
464/// (and library format) from uniquely-mapped proper pairs.
465///
466/// Unlike [`FragmentLengthDistribution`] — which accumulates in **log space** and
467/// is required for the online phase (it folds in per-fragment forgetting mass and
468/// is read per fragment) — this stores plain **integer counts** per length,
469/// bucketed by orientation. Integer increments are commutative, so the tallies
470/// are independent of thread / chunk order; the FLD is then built **once**,
471/// deterministically (fixed length order), in [`finish`](Self::finish). Keeping
472/// it a separate type means the online FLD's float/log-space accumulation is
473/// never disturbed, and no runtime dispatch is needed.
474#[derive(Debug)]
475pub struct DiscreteFld {
476    /// per-length counts for opposite-strand (inward/outward) proper pairs
477    opp: Vec<AtomicU64>,
478    /// per-length counts for same-strand proper pairs
479    same: Vec<AtomicU64>,
480    n_opp: AtomicU64,
481    n_same: AtomicU64,
482    /// observed-format tally (indexed by [`LibraryFormat::format_id`]) for
483    /// order-independent `-l A` auto-detection
484    fmt_counts: [AtomicU64; 12],
485    max_len: usize,
486}
487
488impl DiscreteFld {
489    /// Create an accumulator covering raw fragment lengths `[0, fld_max]`.
490    pub fn new(fld_max: usize) -> Self {
491        Self {
492            opp: (0..=fld_max).map(|_| AtomicU64::new(0)).collect(),
493            same: (0..=fld_max).map(|_| AtomicU64::new(0)).collect(),
494            n_opp: AtomicU64::new(0),
495            n_same: AtomicU64::new(0),
496            fmt_counts: std::array::from_fn(|_| AtomicU64::new(0)),
497            max_len: fld_max,
498        }
499    }
500
501    /// Record one uniquely-mapped proper pair from its mate strands: fragment
502    /// length, orientation bucket, and observed format (derived from `is_fw` /
503    /// `mate_fw`, mirroring the RAD reader's `rad_frag_format` so the mapping pass
504    /// and the RAD-derive path agree). Works in sketch mode, where no precomputed
505    /// `LibraryFormat` is available. Thread-safe and order-independent.
506    pub fn add(&self, len: usize, is_fw: bool, mate_fw: bool) {
507        let l = len.min(self.max_len);
508        // opposite-strand (inward/outward) vs same-strand, exactly as the RAD
509        // reader classifies placements.
510        let (orientation, strandedness) = if is_fw != mate_fw {
511            let s = if is_fw {
512                ReadStrandedness::SA
513            } else {
514                ReadStrandedness::AS
515            };
516            (ReadOrientation::Toward, s)
517        } else {
518            let s = if is_fw {
519                ReadStrandedness::S
520            } else {
521                ReadStrandedness::A
522            };
523            (ReadOrientation::Same, s)
524        };
525        if orientation == ReadOrientation::Same {
526            self.same[l].fetch_add(1, Ordering::Relaxed);
527            self.n_same.fetch_add(1, Ordering::Relaxed);
528        } else {
529            self.opp[l].fetch_add(1, Ordering::Relaxed);
530            self.n_opp.fetch_add(1, Ordering::Relaxed);
531        }
532        let fmt = LibraryFormat::new(ReadType::PairedEnd, orientation, strandedness);
533        self.fmt_counts[fmt.format_id() as usize].fetch_add(1, Ordering::Relaxed);
534    }
535
536    /// Total unique proper pairs seen across both orientation buckets.
537    pub fn count(&self) -> u64 {
538        self.n_opp.load(Ordering::Relaxed) + self.n_same.load(Ordering::Relaxed)
539    }
540
541    /// Build the FLD from the majority-orientation bucket (deterministically, in
542    /// fixed length order) and infer the library format from the format tally.
543    /// `fld_mean`/`fld_sd` seed the prior. Returns the cached FLD and the detected
544    /// format (`None` when no proper pairs were seen).
545    pub fn finish(
546        &self,
547        fld_mean: f64,
548        fld_sd: f64,
549    ) -> (FragmentLengthDistribution, Option<LibraryFormat>) {
550        let n_opp = self.n_opp.load(Ordering::Relaxed);
551        let n_same = self.n_same.load(Ordering::Relaxed);
552        let chosen = if n_opp >= n_same {
553            &self.opp
554        } else {
555            &self.same
556        };
557        let mut fld =
558            FragmentLengthDistribution::new(1.0, self.max_len, fld_mean, fld_sd, 4, 0.5, 1);
559        // `add_val(len, ln count)` equals `count` unit `add_val(len, 0)` calls in
560        // log space, so this reproduces the per-fragment FLD — but deterministically.
561        for (len, c) in chosen.iter().enumerate() {
562            let c = c.load(Ordering::Relaxed);
563            if c > 0 {
564                fld.add_val(len, (c as f64).ln());
565            }
566        }
567        fld.cache();
568        let tally: Vec<u64> = self
569            .fmt_counts
570            .iter()
571            .map(|a| a.load(Ordering::Relaxed))
572            .collect();
573        let detected = if tally.iter().sum::<u64>() > 0 {
574            Some(crate::infer_format_from_counts(&tally, ReadType::PairedEnd))
575        } else {
576            None
577        };
578        (fld, detected)
579    }
580}
581
582#[cfg(test)]
583mod tests {
584    use super::*;
585
586    #[test]
587    fn uniform_prior_pmf_normalizes() {
588        let mut fld = FragmentLengthDistribution::new(1.0, 200, 0.0, 0.0, 4, 0.5, 1);
589        fld.cache();
590        let total: f64 = fld.log_pmf().iter().map(|p| p.exp()).sum();
591        assert!((total - 1.0).abs() < 1e-9, "pmf sums to {total}");
592    }
593
594    #[test]
595    fn gaussian_prior_mean_is_near_mu() {
596        let fld = FragmentLengthDistribution::new(1000.0, 1000, 250.0, 25.0, 4, 0.5, 1);
597        let m = fld.mean();
598        assert!((m - 250.0).abs() < 5.0, "mean {m} not near 250");
599    }
600
601    #[test]
602    fn observations_shift_the_distribution() {
603        let mut fld = FragmentLengthDistribution::new(1.0, 1000, 250.0, 25.0, 4, 0.5, 1);
604        // pile observations around 400
605        for _ in 0..100_000 {
606            fld.add_val(400, 0.0); // mass = log(1) = 0
607        }
608        let m = fld.mean();
609        assert!(m > 300.0, "mean {m} did not move toward 400");
610        fld.cache();
611        // length 400 should be among the most probable
612        let p400 = fld.pmf(400);
613        let p250 = fld.pmf(250);
614        assert!(p400 > p250, "p(400)={p400} not > p(250)={p250}");
615    }
616
617    #[test]
618    fn smoothed_efflen_shrinks_short_transcripts() {
619        // Gaussian prior mean 250: a transcript far shorter than the mean should
620        // get a heavily shrunk effective length (NOT the raw refLen the old
621        // truncated-PMF estimate fell back to).
622        let mut fld = FragmentLengthDistribution::new(1000.0, 1000, 250.0, 25.0, 4, 0.5, 1);
623        fld.cache();
624        let cm = fld.conditional_means();
625        // conditional means are non-decreasing
626        for w in cm.windows(2) {
627            assert!(
628                w[1] >= w[0] - 1e-9,
629                "cond means not monotonic: {} < {}",
630                w[1],
631                w[0]
632            );
633        }
634        let short = smoothed_effective_length(&cm, 201);
635        assert!(
636            short < 201.0 && short > 1.0,
637            "short effLen {short} not shrunk"
638        );
639        // a long transcript keeps most of its length
640        let long = smoothed_effective_length(&cm, 5000);
641        assert!(long > 4000.0, "long effLen {long} shrunk too much");
642        // below the 1.0 barrier the raw length is returned
643        let tiny = smoothed_effective_length(&cm, 2);
644        assert_eq!(tiny, 2.0, "tiny transcript should fall back to refLen");
645    }
646
647    #[test]
648    fn ambig_frag_prob_bounds_and_orientation() {
649        let mut fld = FragmentLengthDistribution::new(1000.0, 1000, 250.0, 25.0, 4, 0.5, 1);
650        fld.cache();
651        // Use the cached (frozen) CMF as a stand-in for the online snapshot.
652        let cmf = fld.cached_cmf.clone();
653        let txp_len = 2000i32;
654        // A forward read with ample downstream space (mate fits at the typical
655        // insert) should be near log(1) ≈ 0, since cmf(maxFrag) ≈ cmf(txpLen).
656        let ample = ambig_frag_log_prob(&cmf, true, 100, 75, txp_len);
657        assert!(ample > -0.01, "ample-space orphan logProb {ample} not ~0");
658        // A forward read crammed against the 3' end (little downstream space)
659        // implies an implausibly short fragment -> much smaller probability.
660        let crammed = ambig_frag_log_prob(&cmf, true, txp_len - 50, 75, txp_len);
661        assert!(
662            crammed < ample - 1.0,
663            "crammed orphan {crammed} not << ample {ample}"
664        );
665        // Reverse-strand orientation uses pos + read_len for the upstream bound:
666        // a reverse read whose outer end is near the 5' start is likewise crammed.
667        let rc_crammed = ambig_frag_log_prob(&cmf, false, 0, 50, txp_len);
668        assert!(
669            rc_crammed < ample - 1.0,
670            "rc crammed orphan {rc_crammed} not << ample {ample}"
671        );
672        // Empty snapshot -> unmodelled (LOG_1 = 0).
673        assert_eq!(ambig_frag_log_prob(&[], true, 100, 75, txp_len), 0.0);
674    }
675
676    #[test]
677    fn cmf_is_monotonic() {
678        let mut fld = FragmentLengthDistribution::new(1000.0, 500, 200.0, 30.0, 4, 0.5, 1);
679        fld.cache();
680        let mut prev = f64::NEG_INFINITY;
681        for l in 0..=500 {
682            let c = fld.cmf(l);
683            assert!(c >= prev - 1e-9, "cmf decreased at {l}: {c} < {prev}");
684            prev = c;
685        }
686        assert!((prev - 0.0).abs() < 1e-6, "cmf endpoint {prev} != log(1)");
687    }
688}