Przeglądaj źródła

uniqueness update with RC

Thomas 1 miesiąc temu
rodzic
commit
be50e652d3
2 zmienionych plików z 354 dodań i 111 usunięć
  1. 300 74
      src/uniqueness/genome_index.rs
  2. 54 37
      src/uniqueness/mod.rs

+ 300 - 74
src/uniqueness/genome_index.rs

@@ -1,6 +1,6 @@
 // src/uniqueness/genome_index.rs
 
-use log::{info, warn};
+use log::info;
 use std::collections::HashMap;
 use std::path::Path;
 
@@ -8,26 +8,38 @@ use anyhow::Context;
 
 use super::{AnchorBias, CacheKey, UniqueInterval, UniquenessIndex};
 use crate::{
-    io::fasta::{FaiEntry, read_fai, read_single_contig_fasta, split_fasta},
-    uniqueness::{UniqueContext, cache::load_from_cache},
+    io::fasta::{read_fai, read_single_contig_fasta, split_fasta, FaiEntry},
+    uniqueness::UniqueContext,
 }; // ton module existant
 
+const CONTIG_SEPARATOR: u8 = 0;
+
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct GenomicPos {
     pub contig: String,
     pub pos: usize, // 0-based
 }
 
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct ContigSpan {
+    pub name: String,
+    pub start: usize,
+    pub end: usize,
+}
+
 pub struct GenomeIndex {
-    contigs: HashMap<String, UniquenessIndex>,
+    index: UniquenessIndex,
+    contigs: Vec<ContigSpan>,
+    by_name: HashMap<String, usize>,
 }
 
 impl GenomeIndex {
     /// Build or load from cache.
     ///
     /// Uses your existing `split_fasta` to stream contig by contig into
-    /// a temp dir, builds/loads a UniquenessIndex per contig, then drops
-    /// the split files.
+    /// a temp dir, concatenates selected contigs with non-FASTA separator
+    /// bytes, builds/loads one global [`UniquenessIndex`], then drops the
+    /// split files.
     ///
     /// `filter`: if Some, only process these contig names.
     pub fn build(
@@ -43,103 +55,317 @@ impl GenomeIndex {
         let fai = read_fai(&fai_path)
             .with_context(|| format!("cannot read FAI {}", fai_path.display()))?;
 
-        // ── Premier passage : quels contigs ont besoin d'être (re)buildés ? ──
-        let mut need_build: Vec<&FaiEntry> = Vec::new();
-        let mut cached: HashMap<String, UniquenessIndex> = HashMap::new();
-
-        for entry in &fai {
-            if let Some(allowed) = filter {
-                if !allowed.contains(&entry.name.as_str()) {
-                    continue;
-                }
-            }
-
-            let key = CacheKey::from_fai(&entry.name, entry.len, fasta_path)?;
-            let fwd_path = cache_dir.join(key.file_name());
-            let rev_path = cache_dir.join(format!("rev.{}", key.file_name()));
-
-            if fwd_path.exists() && rev_path.exists() {
-                match (
-                    load_from_cache(&fwd_path, entry.len),
-                    load_from_cache(&rev_path, entry.len),
-                ) {
-                    (Ok(min_unique_len), Ok(rev_min_unique_len)) => {
-                        info!("[genome_index] cache hit  : {}", entry.name);
-                        cached.insert(
-                            entry.name.clone(),
-                            UniquenessIndex::from_raw(min_unique_len, rev_min_unique_len),
-                        );
-                    }
-                    (fwd, rev) => {
-                        warn!(
-                            "[genome_index] cache miss : {} (fwd={:?}, rev={:?})",
-                            entry.name,
-                            fwd.err(),
-                            rev.err()
-                        );
-                        need_build.push(entry);
-                    }
-                }
-            } else {
-                warn!("[genome_index] cache miss : {}", entry.name);
-                need_build.push(entry);
-            }
-        }
-
-        if need_build.is_empty() {
-            info!("[genome_index] all contigs loaded from cache, FASTA not read.");
-            return Ok(Self { contigs: cached });
-        }
+        let selected: Vec<&FaiEntry> = fai
+            .iter()
+            .filter(|entry| {
+                filter
+                    .map(|allowed| allowed.contains(&entry.name.as_str()))
+                    .unwrap_or(true)
+            })
+            .collect();
+        anyhow::ensure!(!selected.is_empty(), "no FASTA contigs selected");
 
-        let need_names: Vec<&str> = need_build.iter().map(|e| e.name.as_str()).collect();
+        let selected_names: Vec<&str> = selected.iter().map(|e| e.name.as_str()).collect();
         info!(
-            "[genome_index] building {} contig(s): {:?}",
-            need_names.len(),
-            need_names
+            "[genome_index] building global uniqueness over {} contig(s): {:?}",
+            selected_names.len(),
+            selected_names
         );
 
         let split_dir = cache_dir.join("_split_tmp");
         let contig_fastas = split_fasta(fasta_path, &split_dir)?;
+        let contig_fastas_by_name: HashMap<&str, &Path> = contig_fastas
+            .iter()
+            .map(|cf| (cf.name.as_str(), cf.fasta_path.as_path()))
+            .collect();
+
+        let mut genome = Vec::new();
+        let mut contigs = Vec::with_capacity(selected.len());
+        let mut by_name = HashMap::with_capacity(selected.len());
 
-        for cf in &contig_fastas {
-            if !need_names.contains(&cf.name.as_str()) {
-                continue;
+        for entry in selected {
+            if !genome.is_empty() {
+                genome.push(CONTIG_SEPARATOR);
             }
 
-            let entry = need_build.iter().find(|e| e.name == cf.name).unwrap();
-            let seq = read_single_contig_fasta(&cf.fasta_path)?;
+            let fasta_path = contig_fastas_by_name
+                .get(entry.name.as_str())
+                .with_context(|| format!("split FASTA missing selected contig {}", entry.name))?;
+            let seq = read_single_contig_fasta(fasta_path)?;
+            anyhow::ensure!(
+                seq.len() == entry.len,
+                "FASTA/FAI length mismatch for {}: sequence has {}, FAI has {}",
+                entry.name,
+                seq.len(),
+                entry.len
+            );
+            anyhow::ensure!(
+                !seq.contains(&CONTIG_SEPARATOR),
+                "contig {} contains reserved separator byte",
+                entry.name
+            );
 
             info!(
-                "[genome_index] building SA : {} ({:.1} Mbp)",
-                cf.name,
+                "[genome_index] adding contig: {} ({:.1} Mbp)",
+                entry.name,
                 seq.len() as f64 / 1e6
             );
 
-            let key = CacheKey::from_fai(&cf.name, entry.len, fasta_path)?;
-            let idx = UniquenessIndex::build_with_cache(&seq, &key, cache_dir)
-                .with_context(|| format!("build_with_cache failed for {}", cf.name))?;
-            cached.insert(cf.name.clone(), idx);
+            let start = genome.len();
+            genome.extend_from_slice(&seq);
+            let end = genome.len();
+            by_name.insert(entry.name.clone(), contigs.len());
+            contigs.push(ContigSpan {
+                name: entry.name.clone(),
+                start,
+                end,
+            });
         }
 
+        let key = CacheKey::from_genome("genome", &genome);
+        let index = UniquenessIndex::build_with_cache(&genome, &key, cache_dir)
+            .context("build_with_cache failed for concatenated genome")?;
+
         let _ = std::fs::remove_dir_all(&split_dir);
-        Ok(Self { contigs: cached })
+        Ok(Self {
+            index,
+            contigs,
+            by_name,
+        })
     }
 
     pub fn contig_names(&self) -> impl Iterator<Item = &str> {
-        self.contigs.keys().map(String::as_str)
+        self.contigs.iter().map(|span| span.name.as_str())
     }
 
     pub fn contig_len(&self, contig: &str) -> Option<usize> {
-        self.contigs.get(contig).map(|idx| idx.genome_len())
+        let span = self.contig_span(contig)?;
+        Some(span.end - span.start)
     }
 
     pub fn unique_context_at(&self, pos: &GenomicPos) -> Option<UniqueContext> {
-        Some(self.contigs.get(&pos.contig)?.unique_context_at(pos.pos))
+        let span = self.contig_span(&pos.contig)?;
+        let global_pos = self.global_pos(pos)?;
+        let context = self.index.unique_context_at(global_pos);
+
+        Some(UniqueContext {
+            before: context.before.filter(|&k| global_pos >= span.start + k),
+            after: context.after.filter(|&k| global_pos + k <= span.end),
+        })
     }
 
     pub fn query(&self, pos: &GenomicPos, bias: AnchorBias) -> Option<UniqueInterval> {
+        let span = self.contig_span(&pos.contig)?;
+        let global_pos = self.global_pos(pos)?;
+        let global = self.query_global_in_span(global_pos, span, bias)?;
+        Some(UniqueInterval {
+            start: global.start - span.start,
+            end: global.end - span.start,
+        })
+    }
+
+    pub fn global_pos(&self, pos: &GenomicPos) -> Option<usize> {
+        let span = self.contig_span(&pos.contig)?;
+        if pos.pos < span.end - span.start {
+            Some(span.start + pos.pos)
+        } else {
+            None
+        }
+    }
+
+    pub fn global_interval_to_local(
+        &self,
+        interval: &UniqueInterval,
+    ) -> Option<(String, UniqueInterval)> {
+        let span = self.contig_span_for_global_interval(interval)?;
+        Some((
+            span.name.clone(),
+            UniqueInterval {
+                start: interval.start - span.start,
+                end: interval.end - span.start,
+            },
+        ))
+    }
+
+    fn contig_span(&self, contig: &str) -> Option<&ContigSpan> {
+        self.by_name.get(contig).map(|&idx| &self.contigs[idx])
+    }
+
+    fn contig_span_for_global_interval(&self, interval: &UniqueInterval) -> Option<&ContigSpan> {
+        if interval.start >= interval.end {
+            return None;
+        }
+
         self.contigs
-            .get(&pos.contig)?
-            .minimal_unique_interval_containing(pos.pos, bias)
+            .iter()
+            .find(|span| span.start <= interval.start && interval.end <= span.end)
+    }
+
+    fn query_global_in_span(
+        &self,
+        global_pos: usize,
+        span: &ContigSpan,
+        bias: AnchorBias,
+    ) -> Option<UniqueInterval> {
+        debug_assert!(span.start <= global_pos && global_pos < span.end);
+
+        match bias {
+            AnchorBias::Left => {
+                let k = self.index.min_unique_len[global_pos];
+                self.interval_in_span(global_pos, k, span)
+            }
+            AnchorBias::Right => {
+                for a in (span.start..=global_pos).rev() {
+                    let k = self.index.min_unique_len[a];
+                    if let Some(iv) = self.interval_in_span(a, k, span) {
+                        if iv.end > global_pos {
+                            return Some(iv);
+                        }
+                    }
+                }
+                None
+            }
+            AnchorBias::Auto => {
+                let window = 500;
+                let mut best: Option<UniqueInterval> = None;
+                let start = global_pos.saturating_sub(window).max(span.start);
+
+                for a in start..=global_pos {
+                    let k = self.index.min_unique_len[a];
+                    let Some(candidate) = self.interval_in_span(a, k, span) else {
+                        continue;
+                    };
+                    if candidate.end <= global_pos {
+                        continue;
+                    }
+
+                    best = Some(match best {
+                        None => candidate,
+                        Some(current) => {
+                            if candidate.len() < current.len() {
+                                candidate
+                            } else {
+                                current
+                            }
+                        }
+                    });
+                }
+
+                best
+            }
+        }
+    }
+
+    fn interval_in_span(
+        &self,
+        start: usize,
+        len: usize,
+        span: &ContigSpan,
+    ) -> Option<UniqueInterval> {
+        if len == usize::MAX {
+            return None;
+        }
+
+        let end = start.checked_add(len)?;
+        if span.start <= start && start < span.end && end <= span.end {
+            Some(UniqueInterval { start, end })
+        } else {
+            None
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::{collections::HashMap, fs, path::Path};
+
+    use super::*;
+
+    fn write_fasta_and_fai(dir: &Path, records: &[(&str, &[u8])]) -> std::path::PathBuf {
+        let fasta = dir.join("test.fa");
+        let mut fasta_body = String::new();
+        let mut fai_body = String::new();
+
+        for (name, seq) in records {
+            fasta_body.push('>');
+            fasta_body.push_str(name);
+            fasta_body.push('\n');
+            fasta_body.push_str(std::str::from_utf8(seq).unwrap());
+            fasta_body.push('\n');
+
+            fai_body.push_str(name);
+            fai_body.push('\t');
+            fai_body.push_str(&seq.len().to_string());
+            fai_body.push_str("\t0\t0\t0\n");
+        }
+
+        fs::write(&fasta, fasta_body).unwrap();
+        fs::write(dir.join("test.fa.fai"), fai_body).unwrap();
+        fasta
+    }
+
+    #[test]
+    fn repeated_sequence_on_another_contig_is_not_unique() {
+        let dir =
+            std::env::temp_dir().join(format!("pandora_uniqueness_test_{}", std::process::id()));
+        let _ = fs::remove_dir_all(&dir);
+        fs::create_dir_all(&dir).unwrap();
+
+        let fasta = write_fasta_and_fai(&dir, &[("chr1", b"ACGT"), ("chr2", b"ACGT")]);
+        let index = GenomeIndex::build(&fasta, dir.join("cache"), None).unwrap();
+
+        for bias in [AnchorBias::Left, AnchorBias::Right, AnchorBias::Auto] {
+            assert_eq!(
+                index.query(
+                    &GenomicPos {
+                        contig: "chr1".to_owned(),
+                        pos: 0,
+                    },
+                    bias,
+                ),
+                None,
+                "{bias:?} should not use a separator-spanning interval to make chr1:0 unique"
+            );
+        }
+
+        let _ = fs::remove_dir_all(&dir);
+    }
+
+    #[test]
+    fn maps_between_local_and_global_coordinates() {
+        let contigs = vec![
+            ContigSpan {
+                name: "chr1".to_owned(),
+                start: 0,
+                end: 3,
+            },
+            ContigSpan {
+                name: "chr2".to_owned(),
+                start: 4,
+                end: 6,
+            },
+        ];
+        let by_name = HashMap::from([("chr1".to_owned(), 0), ("chr2".to_owned(), 1)]);
+        let index = GenomeIndex {
+            index: UniquenessIndex::build(b"AAA\0TT"),
+            contigs,
+            by_name,
+        };
+
+        assert_eq!(
+            index.global_pos(&GenomicPos {
+                contig: "chr2".to_owned(),
+                pos: 1,
+            }),
+            Some(5)
+        );
+        assert_eq!(
+            index.global_interval_to_local(&UniqueInterval { start: 4, end: 6 }),
+            Some(("chr2".to_owned(), UniqueInterval { start: 0, end: 2 }))
+        );
+        assert_eq!(
+            index.global_interval_to_local(&UniqueInterval { start: 2, end: 5 }),
+            None
+        );
     }
 }

+ 54 - 37
src/uniqueness/mod.rs

@@ -68,6 +68,7 @@
 //! - [`GenomicPos`] — genomic coordinate used by [`GenomeIndex`].
 
 use log::warn;
+use rayon::prelude::*;
 use std::io::{self};
 use std::path::Path;
 
@@ -92,15 +93,20 @@ fn build_min_unique_len(genome: &[u8]) -> Vec<usize> {
         rank[s] = i;
     }
 
-    let mut mul = vec![usize::MAX; n];
-    for p in 0..n {
-        let r = rank[p];
-        let left = if r > 0 { lcp[r] } else { 0 };
-        let right = if r + 1 < n { lcp[r + 1] } else { 0 };
-        let k = left.max(right) + 1;
-        mul[p] = if p + k <= n { k } else { usize::MAX };
-    }
-    mul
+    (0..n)
+        .into_par_iter()
+        .map(|p| {
+            let r = rank[p];
+            let left = if r > 0 { lcp[r] } else { 0 };
+            let right = if r + 1 < n { lcp[r + 1] } else { 0 };
+            let k = left.max(right) + 1;
+            if p + k <= n {
+                k
+            } else {
+                usize::MAX
+            }
+        })
+        .collect()
 }
 
 pub struct UniqueContext {
@@ -146,34 +152,42 @@ impl UniquenessIndex {
         let fwd_path = cache_dir.join(key.file_name());
         let rev_path = cache_dir.join(format!("rev.{}", key.file_name()));
 
-        let min_unique_len = match load_from_cache(&fwd_path, genome.len()) {
-            Ok(mul) => mul,
-            Err(e) => {
-                if fwd_path.exists() {
-                    warn!("[cache] invalid forward ({e}), rebuilding...");
+        let (min_unique_len, rev_min_unique_len) = rayon::join(
+            || -> io::Result<Vec<usize>> {
+                match load_from_cache(&fwd_path, genome.len()) {
+                    Ok(mul) => Ok(mul),
+                    Err(e) => {
+                        if fwd_path.exists() {
+                            warn!("[cache] invalid forward ({e}), rebuilding...");
+                        }
+
+                        let mul = build_min_unique_len(genome);
+                        save_to_cache(&fwd_path, &mul)?;
+                        Ok(mul)
+                    }
                 }
+            },
+            || -> io::Result<Vec<usize>> {
+                match load_from_cache(&rev_path, genome.len()) {
+                    Ok(mul) => Ok(mul),
+                    Err(e) => {
+                        if rev_path.exists() {
+                            warn!("[cache] invalid reverse ({e}), rebuilding...");
+                        }
 
-                let mul = build_min_unique_len(genome);
-                save_to_cache(&fwd_path, &mul)?;
-                mul
-            }
-        };
+                        let mut rev = genome.to_vec();
+                        rev.reverse();
 
-        let rev_min_unique_len = match load_from_cache(&rev_path, genome.len()) {
-            Ok(mul) => mul,
-            Err(e) => {
-                if rev_path.exists() {
-                    warn!("[cache] invalid reverse ({e}), rebuilding...");
+                        let mul = build_min_unique_len(&rev);
+                        save_to_cache(&rev_path, &mul)?;
+                        Ok(mul)
+                    }
                 }
+            },
+        );
 
-                let mut rev = genome.to_vec();
-                rev.reverse();
-
-                let mul = build_min_unique_len(&rev);
-                save_to_cache(&rev_path, &mul)?;
-                mul
-            }
-        };
+        let min_unique_len = min_unique_len?;
+        let rev_min_unique_len = rev_min_unique_len?;
 
         Ok(Self::from_raw(min_unique_len, rev_min_unique_len))
     }
@@ -194,11 +208,14 @@ impl UniquenessIndex {
 
     /// Build without cache (e.g. for tests or one-shot use).
     pub fn build(genome: &[u8]) -> Self {
-        let min_unique_len = build_min_unique_len(genome);
-
-        let mut rev = genome.to_vec();
-        rev.reverse();
-        let rev_min_unique_len = build_min_unique_len(&rev);
+        let (min_unique_len, rev_min_unique_len) = rayon::join(
+            || build_min_unique_len(genome),
+            || {
+                let mut rev = genome.to_vec();
+                rev.reverse();
+                build_min_unique_len(&rev)
+            },
+        );
 
         Self::from_raw(min_unique_len, rev_min_unique_len)
     }