diff --git a/cli/Cargo.lock b/cli/Cargo.lock index e0e0326..f3ac425 100644 --- a/cli/Cargo.lock +++ b/cli/Cargo.lock @@ -105,13 +105,14 @@ checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" [[package]] name = "biosynth" -version = "0.1.29" +version = "0.1.30" dependencies = [ "anyhow", "chrono", "clap", "csv", "dashmap", + "flate2", "glob", "indicatif", "noodles", diff --git a/cli/Cargo.toml b/cli/Cargo.toml index 671c031..d0c87e6 100644 --- a/cli/Cargo.toml +++ b/cli/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "biosynth" -version = "0.1.29" +version = "0.1.30" edition = "2021" rust-version = "1.91" authors = ["Madhava Jay "] @@ -31,5 +31,6 @@ reqwest = { version = "0.12", default-features = false, features = ["blocking", noodles = { version = "0.104.0", features = ["vcf", "tabix", "bgzf", "core"] } glob = "0.3" dashmap = "6" +flate2 = "1.1" [dev-dependencies] diff --git a/cli/src/commands/cohort_bed.rs b/cli/src/commands/cohort_bed.rs index 869408d..16c609d 100644 --- a/cli/src/commands/cohort_bed.rs +++ b/cli/src/commands/cohort_bed.rs @@ -9,7 +9,7 @@ use anyhow::{bail, Context, Result}; use dashmap::DashMap; use rayon::prelude::*; -use crate::genotype_reader::{detect_delimiter, RowOutcome, RowParser}; +use crate::genotype_reader::{detect_delimiter, open_text_reader, RowOutcome, RowParser}; use crate::CohortBedArgs; const LOOKAHEAD_LINES: usize = 2048; @@ -43,12 +43,20 @@ pub fn run_cohort_bed(args: CohortBedArgs) -> Result<()> { // id). After a worker sees the panel once, subsequent samples resolve rsids // from this local map with zero shared-map access — so the ~1e9 lookups stop // contending on the DashMap shards (which all samples hit identically). + // A bad file must not abort the cohort: log it and treat it as an empty + // sample (all-missing column) so the rest of the pipeline keeps running. let per_sample: Vec> = samples .par_iter() - .map_init(HashMap::::new, |local, (_sid, path)| { - parse_sample(path, &interner, local) + .map_init(HashMap::::new, |local, (sid, path)| { + parse_sample(path, &interner, local).unwrap_or_else(|e| { + eprintln!( + "⚠️ cohort-bed: skipping unreadable sample {sid} ({}): {e:#}", + path.display() + ); + Vec::new() + }) }) - .collect::>>()?; + .collect::>(); eprintln!( "🧬 cohort-bed: parsed {} samples in {:.2}s", samples.len(), @@ -323,9 +331,8 @@ fn parse_sample( interner: &Interner, local: &mut HashMap, ) -> Result> { - use std::io::{BufRead, BufReader}; - let file = File::open(path).with_context(|| format!("Open {:?}", path))?; - let mut reader = BufReader::new(file); + use std::io::BufRead; + let mut reader = open_text_reader(path)?; let mut buffered: Vec = Vec::new(); let mut buf = String::new(); while buffered.len() < LOOKAHEAD_LINES { @@ -344,7 +351,12 @@ fn parse_sample( let mut out: Vec<(u32, u8, u8)> = Vec::new(); let mut handle = |parser: &mut RowParser, line: &str| -> Result<()> { - if let RowOutcome::Parsed(row) = parser.consume_line(line)? { + // A bad row must not break the file: skip it and keep going. + let outcome = match parser.consume_line(line) { + Ok(o) => o, + Err(_) => return Ok(()), + }; + if let RowOutcome::Parsed(row) = outcome { // variant_id = rsid, else chrom:pos (matches read_pipeline_genotypes) let rsid = if row.rsid.is_empty() { format!("{}:{}", row.chrom, row.pos) @@ -414,7 +426,7 @@ fn discover_samples(data_dir: &Path) -> Result> { .with_context(|| format!("Read dir {:?}", dir))? .filter_map(|e| e.ok()) .map(|e| e.path()) - .filter(|p| p.is_file() && p.extension().and_then(|x| x.to_str()) == Some("txt")) + .filter(|p| p.is_file() && is_genotype_text_file(p)) .collect(); txts.sort(); if let Some(first) = txts.into_iter().next() { @@ -428,3 +440,15 @@ fn discover_samples(data_dir: &Path) -> Result> { } Ok(samples) } + +fn is_genotype_text_file(path: &Path) -> bool { + match path.extension().and_then(|x| x.to_str()) { + Some(ext) if ext.eq_ignore_ascii_case("txt") => true, + Some(ext) if ext.eq_ignore_ascii_case("gz") => path + .file_stem() + .and_then(|stem| Path::new(stem).extension()) + .and_then(|x| x.to_str()) + .is_some_and(|inner| inner.eq_ignore_ascii_case("txt")), + _ => false, + } +} diff --git a/cli/src/commands/fast_allele_freq.rs b/cli/src/commands/fast_allele_freq.rs index 9420101..d27fef6 100644 --- a/cli/src/commands/fast_allele_freq.rs +++ b/cli/src/commands/fast_allele_freq.rs @@ -325,7 +325,7 @@ fn write_allele_freq(path: &Path, loci: &LociMap) -> Result<()> { Ok(()) } -fn collect_input_files(inputs: &[PathBuf]) -> Result> { +pub(crate) fn collect_input_files(inputs: &[PathBuf]) -> Result> { let mut files = Vec::new(); for input in inputs { if input.is_file() { diff --git a/cli/src/commands/genotype_to_vcf.rs b/cli/src/commands/genotype_to_vcf.rs index 3661033..05ba1e5 100644 --- a/cli/src/commands/genotype_to_vcf.rs +++ b/cli/src/commands/genotype_to_vcf.rs @@ -1,6 +1,6 @@ use std::collections::HashMap; use std::fs::File; -use std::io::{BufRead, BufReader, BufWriter, Write}; +use std::io::{BufRead, BufWriter, Write}; use std::path::{Path, PathBuf}; use std::process::Command; @@ -9,10 +9,10 @@ use rayon::prelude::*; use rayon::ThreadPoolBuilder; use crate::download::ensure_reference_db; -use crate::genotype_reader::{detect_delimiter, RowOutcome, RowParser}; +use crate::genotype_reader::{detect_delimiter, open_text_reader, RowOutcome, RowParser}; use crate::rsid_cache::{default_cache_path, normalize_rsid, RsidCache}; use crate::stats::{ReferenceVariant, StatsStore}; -use crate::util::collect_input_files; +use crate::util::{collect_input_files, genotype_file_stem}; use crate::GenotypeToVcfArgs; const LOOKAHEAD_LINES: usize = 2048; @@ -133,8 +133,7 @@ fn convert_file( missing_log_path: &Path, include_metrics: bool, ) -> Result { - let input_file = File::open(input).with_context(|| format!("Open {:?}", input))?; - let mut reader = BufReader::new(input_file); + let mut reader = open_text_reader(input)?; let mut buffered_lines = Vec::new(); let mut buffer = String::new(); @@ -566,10 +565,7 @@ fn resolve_output_paths(args: &GenotypeToVcfArgs, input: &Path) -> Result, gzip: bool) -> PathBuf { - let base_name = input - .file_stem() - .and_then(|stem| stem.to_str()) - .unwrap_or("output"); + let base_name = genotype_file_stem(input).unwrap_or_else(|| "output".to_string()); let filename = if gzip { format!("{base_name}.vcf.gz") } else { @@ -601,9 +597,7 @@ fn default_sample_name(input: &Path) -> String { return prefix.to_string(); } } - input - .file_stem() - .and_then(|stem| stem.to_str()) + genotype_file_stem(input) .map(|stem| stem.replace(' ', "_")) .filter(|stem| !stem.is_empty()) .unwrap_or_else(|| "SAMPLE".to_string()) diff --git a/cli/src/commands/het_signal.rs b/cli/src/commands/het_signal.rs new file mode 100644 index 0000000..c5010d1 --- /dev/null +++ b/cli/src/commands/het_signal.rs @@ -0,0 +1,333 @@ +use std::collections::HashSet; +use std::fs::{self, File}; +use std::io::{BufRead, BufWriter, Write}; +use std::path::{Path, PathBuf}; +use std::time::Instant; + +use anyhow::{bail, Context, Result}; +use rayon::prelude::*; + +use crate::genotype_reader::{detect_delimiter, open_text_reader, RowOutcome, RowParser}; +use crate::HetSignalArgs; + +const LOOKAHEAD_LINES: usize = 2048; +const SUBSAMPLE_STEP: u32 = 14; // autosomes only (matches sba.SUBSAMPLE_STEP) + +/// Per-sample het-signal cache for the sex-biased admixture NMF, reproducing +/// `fast_sex_biased_admixture._signals_for_region`. For each sample it emits, per +/// region (autosomes 1-22 subsampled by rsid_num%14, and X), the kept +/// `(rsid_num, signal)` arrays where signal = 1 (het) / 0 (hom); no-calls dropped. +/// The NMF (sklearn, seed 42) then runs on these in Python -> identical results. +/// +/// Output `.bvhs` binary per sample (replaces the per-sample `.npz` compact_cache): +/// magic "BVHS", version u8=1, +/// auto_n u32, auto_ids[u32], auto_sig[i8], +/// x_n u32, x_ids[u32], x_sig[i8] (all little-endian) +pub fn run_het_signal(args: HetSignalArgs) -> Result<()> { + let overall = Instant::now(); + let samples = discover_samples(&args.input)?; + if samples.is_empty() { + bail!( + "No sample subdirectories with .txt files found under {:?}", + args.input + ); + } + fs::create_dir_all(&args.out_dir).with_context(|| format!("Create {:?}", args.out_dir))?; + eprintln!("▶️ het-signal: {} samples", samples.len()); + + let out_dir = args.out_dir.clone(); + // A bad file must not abort the run: log it and skip its cache. + let written: usize = samples + .par_iter() + .map(|(sid, path)| { + match parse_sample(path) + .and_then(|sig| write_bvhs(&out_dir.join(format!("{sid}.bvhs")), &sig)) + { + Ok(()) => 1usize, + Err(e) => { + eprintln!( + "⚠️ het-signal: skipping unreadable sample {sid} ({}): {e:#}", + path.display() + ); + 0 + } + } + }) + .sum(); + + eprintln!( + "✅ het-signal: wrote {} caches in {:.2}s", + written, + overall.elapsed().as_secs_f64() + ); + println!( + "✅ Wrote {} .bvhs caches to {}", + written, + args.out_dir.display() + ); + Ok(()) +} + +#[derive(Default)] +struct SampleSignals { + auto_ids: Vec, + auto_sig: Vec, + x_ids: Vec, + x_sig: Vec, + // Per-sample stats (over ALL rows, matching _compact_stats) for the results table. + auto_snps: u64, + x_snps: u64, + auto_het: f64, + x_het: f64, + auto_lrr_mean: f64, + x_lrr_mean: f64, +} + +#[derive(Default, Clone, Copy)] +struct ChromAcc { + count: u64, + het: u64, + lrr_sum: f64, + lrr_cnt: u64, +} + +/// is_het per `read_txt`: 2-char, alleles differ, no '0'. (Note: only '0' is +/// excluded here, unlike the signal's '0'/'-' check — matches the Python.) +fn is_het(gt: &str) -> bool { + let c: Vec = gt.chars().collect(); + c.len() == 2 && c[0] != c[1] && !c.contains(&'0') +} + +/// mean of finite values; NaN if none (matches `_finite_mean`). +fn finite_mean(v: &[f64]) -> f64 { + let f: Vec = v.iter().copied().filter(|x| x.is_finite()).collect(); + if f.is_empty() { + f64::NAN + } else { + f.iter().sum::() / f.len() as f64 + } +} + +/// rs -> Some(number); anything else None. Matches `_rsid_num` + the +/// `^rs\d+$` filter (full match: "rs" then all digits). +fn rsid_num(rsid: &str) -> Option { + let rest = rsid.strip_prefix("rs")?; + if rest.is_empty() || !rest.bytes().all(|b| b.is_ascii_digit()) { + return None; + } + rest.parse::().ok().filter(|&n| n > 0) +} + +/// `_genotype_signal`: 2-char, no '0'/'-', -> 1 (het) / 0 (hom); else None (no-call). +fn genotype_signal(gt: &str) -> Option { + let g = gt.trim().to_ascii_uppercase(); + let chars: Vec = g.chars().collect(); + if chars.len() != 2 || chars.contains(&'0') || chars.contains(&'-') { + return None; + } + Some(if chars[0] != chars[1] { 1 } else { 0 }) +} + +fn parse_sample(path: &Path) -> Result { + let mut reader = open_text_reader(path)?; + let mut buffered: Vec = Vec::new(); + let mut buf = String::new(); + while buffered.len() < LOOKAHEAD_LINES { + buf.clear(); + if reader.read_line(&mut buf)? == 0 { + break; + } + buffered.push(buf.clone()); + } + if buffered.is_empty() { + return Ok(SampleSignals::default()); + } + let delim = detect_delimiter(&buffered); + let mut parser = RowParser::new(delim); + let mut out = SampleSignals::default(); + // dedup rsid keep-first, per region (regions are disjoint by chrom). + let mut seen_auto: HashSet = HashSet::new(); + let mut seen_x: HashSet = HashSet::new(); + // Per-chrom stats over ALL rows (autosomes 1-22 + X), matching _compact_stats. + let mut chrom_acc: std::collections::HashMap = std::collections::HashMap::new(); + let mut x_acc = ChromAcc::default(); + + let mut handle = |parser: &mut RowParser, line: &str| -> Result<()> { + // A bad row must not break the file: skip it and keep going. + let outcome = match parser.consume_line(line) { + Ok(o) => o, + Err(_) => return Ok(()), + }; + if let RowOutcome::Parsed(row) = outcome { + let chrom_int = row.chrom.parse::().ok(); + let is_auto = chrom_int.is_some_and(|c| (1..=22).contains(&c)); + let is_x = row.chrom == "X"; + // Stats accumulation over ALL rows of autosomes/X (before any filter). + if is_auto || is_x { + let acc = if is_auto { + chrom_acc.entry(chrom_int.unwrap()).or_default() + } else { + &mut x_acc + }; + acc.count += 1; + if let Some(g) = &row.genotype { + if is_het(&g.to_ascii_uppercase()) { + acc.het += 1; + } + } + if let Some(lrr) = row + .lrr + .as_deref() + .and_then(|s| s.trim().parse::().ok()) + { + if lrr.is_finite() { + acc.lrr_sum += lrr; + acc.lrr_cnt += 1; + } + } + } + let Some(num) = rsid_num(&row.rsid) else { + return Ok(()); + }; + if !is_auto && !is_x { + return Ok(()); + } + // dedup by rsid keep-first within region + if is_auto { + if !seen_auto.insert(num) { + return Ok(()); + } + if num % SUBSAMPLE_STEP != 0 { + return Ok(()); + } + } else if !seen_x.insert(num) { + return Ok(()); + } + let gt = match &row.genotype { + Some(g) => g, + None => return Ok(()), + }; + let Some(sig) = genotype_signal(gt) else { + return Ok(()); // no-call dropped + }; + if is_auto { + out.auto_ids.push(num); + out.auto_sig.push(sig); + } else { + out.x_ids.push(num); + out.x_sig.push(sig); + } + } + Ok(()) + }; + + for line in &buffered { + handle(&mut parser, line)?; + } + loop { + buf.clear(); + if reader.read_line(&mut buf)? == 0 { + break; + } + handle(&mut parser, &buf)?; + } + // Per-chrom means over autosomes; auto_het/auto_lrr = mean of per-chrom means. + let mut auto_het_means = Vec::new(); + let mut auto_lrr_means = Vec::new(); + for c in 1..=22u32 { + if let Some(acc) = chrom_acc.get(&c) { + if acc.count > 0 { + out.auto_snps += acc.count; + auto_het_means.push(acc.het as f64 / acc.count as f64); + auto_lrr_means.push(if acc.lrr_cnt > 0 { + acc.lrr_sum / acc.lrr_cnt as f64 + } else { + f64::NAN + }); + } + } + } + out.auto_het = finite_mean(&auto_het_means); + out.auto_lrr_mean = finite_mean(&auto_lrr_means); + out.x_snps = x_acc.count; + out.x_het = if x_acc.count > 0 { + x_acc.het as f64 / x_acc.count as f64 + } else { + f64::NAN + }; + out.x_lrr_mean = if x_acc.lrr_cnt > 0 { + x_acc.lrr_sum / x_acc.lrr_cnt as f64 + } else { + f64::NAN + }; + Ok(out) +} + +fn write_bvhs(path: &Path, sig: &SampleSignals) -> Result<()> { + let mut w = BufWriter::new(File::create(path).with_context(|| format!("Create {:?}", path))?); + w.write_all(b"BVHS")?; + w.write_all(&[2u8])?; // v2: signal arrays + per-sample stats + write_arrays(&mut w, &sig.auto_ids, &sig.auto_sig)?; + write_arrays(&mut w, &sig.x_ids, &sig.x_sig)?; + // stats: auto_snps u64, x_snps u64, then auto_het/x_het/auto_lrr_mean/x_lrr_mean f64 + w.write_all(&sig.auto_snps.to_le_bytes())?; + w.write_all(&sig.x_snps.to_le_bytes())?; + for v in [sig.auto_het, sig.x_het, sig.auto_lrr_mean, sig.x_lrr_mean] { + w.write_all(&v.to_le_bytes())?; + } + w.flush()?; + Ok(()) +} + +fn write_arrays(w: &mut W, ids: &[u32], sigs: &[i8]) -> Result<()> { + w.write_all(&(ids.len() as u32).to_le_bytes())?; + for &id in ids { + w.write_all(&id.to_le_bytes())?; + } + for &s in sigs { + w.write_all(&[s as u8])?; + } + Ok(()) +} + +/// Discover samples: sorted subdirs, first sorted *.txt in each (matches the flows). +fn discover_samples(data_dir: &Path) -> Result> { + let mut dirs: Vec = fs::read_dir(data_dir) + .with_context(|| format!("Read dir {:?}", data_dir))? + .filter_map(|e| e.ok()) + .map(|e| e.path()) + .filter(|p| p.is_dir()) + .collect(); + dirs.sort(); + let mut samples = Vec::new(); + for dir in dirs { + let mut txts: Vec = fs::read_dir(&dir) + .with_context(|| format!("Read dir {:?}", dir))? + .filter_map(|e| e.ok()) + .map(|e| e.path()) + .filter(|p| p.is_file() && is_genotype_text_file(p)) + .collect(); + txts.sort(); + if let Some(first) = txts.into_iter().next() { + let sid = dir + .file_name() + .and_then(|n| n.to_str()) + .unwrap_or_default() + .to_string(); + samples.push((sid, first)); + } + } + Ok(samples) +} + +fn is_genotype_text_file(path: &Path) -> bool { + match path.extension().and_then(|x| x.to_str()) { + Some(ext) if ext.eq_ignore_ascii_case("txt") => true, + Some(ext) if ext.eq_ignore_ascii_case("gz") => path + .file_stem() + .and_then(|stem| Path::new(stem).extension()) + .and_then(|x| x.to_str()) + .is_some_and(|inner| inner.eq_ignore_ascii_case("txt")), + _ => false, + } +} diff --git a/cli/src/commands/long_emit.rs b/cli/src/commands/long_emit.rs index 10b80b8..369adb8 100644 --- a/cli/src/commands/long_emit.rs +++ b/cli/src/commands/long_emit.rs @@ -8,13 +8,14 @@ use std::time::Instant; use anyhow::{bail, Context, Result}; use crate::download::ensure_reference_db; -use crate::genotype_reader::{detect_delimiter, RowOutcome, RowParser}; +use crate::genotype_reader::{detect_delimiter, open_text_reader, RowOutcome, RowParser}; use crate::long_rows::{ is_snp, locus_key, normalize_sequence, parse_alternates, parse_genotype_dosages, parse_vcf_gt_dosages, vcf_gt_from_sample, vcf_sample_name, LongRow, LongRowWriter, }; use crate::rsid_cache::normalize_rsid; use crate::stats::{ReferenceVariant, StatsStore}; +use crate::util::genotype_file_stem; use crate::{EmitLongArgs, WarnDetail}; use rusqlite::OptionalExtension; @@ -250,8 +251,7 @@ fn parse_genotype_file_to_pending( resolver: &mut ReferenceResolver, missing_logger: &mut MissingRefLogger, ) -> Result<(Vec, EmitStats)> { - let input_file = File::open(input).with_context(|| format!("Open {:?}", input))?; - let mut reader = BufReader::new(input_file); + let mut reader = open_text_reader(input)?; let mut buffered_lines: Vec = Vec::new(); let mut buffer = String::new(); @@ -643,10 +643,7 @@ fn resolve_output_path(input: &Path, output: Option<&PathBuf>) -> Result) -> Result String { - input - .file_stem() - .and_then(|stem| stem.to_str()) + genotype_file_stem(input) .map(|stem| stem.replace(' ', "_")) .filter(|stem| !stem.is_empty()) .unwrap_or_else(|| "SAMPLE".to_string()) diff --git a/cli/src/commands/project_bed.rs b/cli/src/commands/project_bed.rs index 1e95e49..8dbb039 100644 --- a/cli/src/commands/project_bed.rs +++ b/cli/src/commands/project_bed.rs @@ -7,7 +7,7 @@ use std::time::Instant; use anyhow::{bail, Context, Result}; use rayon::prelude::*; -use crate::genotype_reader::{detect_delimiter, RowOutcome, RowParser}; +use crate::genotype_reader::{detect_delimiter, open_text_reader, RowOutcome, RowParser}; use crate::ProjectBedArgs; const LOOKAHEAD_LINES: usize = 2048; @@ -41,10 +41,20 @@ pub fn run_project_bed(args: ProjectBedArgs) -> Result<()> { // interner contention (each sample matches against the fixed 76k panel). let min_gs = args.min_gs; let parse_start = Instant::now(); + // A bad file must not abort the run: log it and treat it as 0 matches (the + // overlap filter then drops it), so the rest of the cohort still projects. let parsed: Vec = samples .par_iter() - .map(|(_sid, path)| parse_sample(path, &panel, min_gs)) - .collect::>>()?; + .map(|(sid, path)| { + parse_sample(path, &panel, min_gs).unwrap_or_else(|e| { + eprintln!( + "⚠️ project-bed: skipping unreadable sample {sid} ({}): {e:#}", + path.display() + ); + SampleHits { rows: Vec::new() } + }) + }) + .collect::>(); eprintln!( "🧬 project-bed: parsed {} samples in {:.2}s", samples.len(), @@ -198,8 +208,7 @@ fn load_panel(path: &Path) -> Result { /// Parse one sample: apply fast_convert's projection filters and match to the /// panel by variant key. Dedup repeated keys keep-first. Returns matched rows. fn parse_sample(path: &Path, panel: &Panel, min_gs: f32) -> Result { - let file = File::open(path).with_context(|| format!("Open {:?}", path))?; - let mut reader = BufReader::new(file); + let mut reader = open_text_reader(path)?; let mut buffered: Vec = Vec::new(); let mut buf = String::new(); while buffered.len() < LOOKAHEAD_LINES { @@ -218,7 +227,12 @@ fn parse_sample(path: &Path, panel: &Panel, min_gs: f32) -> Result { let mut rows: Vec<(u32, u8, u8)> = Vec::new(); let mut handle = |parser: &mut RowParser, line: &str| -> Result<()> { - if let RowOutcome::Parsed(row) = parser.consume_line(line)? { + // A bad row must not break the file: skip it and keep going. + let outcome = match parser.consume_line(line) { + Ok(o) => o, + Err(_) => return Ok(()), + }; + if let RowOutcome::Parsed(row) = outcome { // autosome 1-22 (chrom already cleaned by RowParser) let chrom_int: i64 = match row.chrom.parse() { Ok(c) if (1..=22).contains(&c) => c, @@ -412,7 +426,7 @@ fn discover_samples(data_dir: &Path) -> Result> { .with_context(|| format!("Read dir {:?}", dir))? .filter_map(|e| e.ok()) .map(|e| e.path()) - .filter(|p| p.is_file() && p.extension().and_then(|x| x.to_str()) == Some("txt")) + .filter(|p| p.is_file() && is_genotype_text_file(p)) .collect(); txts.sort(); if let Some(first) = txts.into_iter().next() { @@ -426,3 +440,15 @@ fn discover_samples(data_dir: &Path) -> Result> { } Ok(samples) } + +fn is_genotype_text_file(path: &Path) -> bool { + match path.extension().and_then(|x| x.to_str()) { + Some(ext) if ext.eq_ignore_ascii_case("txt") => true, + Some(ext) if ext.eq_ignore_ascii_case("gz") => path + .file_stem() + .and_then(|stem| Path::new(stem).extension()) + .and_then(|x| x.to_str()) + .is_some_and(|inner| inner.eq_ignore_ascii_case("txt")), + _ => false, + } +} diff --git a/cli/src/commands/target_study_dosage.rs b/cli/src/commands/target_study_dosage.rs new file mode 100644 index 0000000..c9bed69 --- /dev/null +++ b/cli/src/commands/target_study_dosage.rs @@ -0,0 +1,509 @@ +use std::collections::HashMap; +use std::collections::HashSet; +use std::fs::File; +use std::io::{BufWriter, Write}; +use std::path::{Path, PathBuf}; +use std::sync::Arc; +use std::time::Instant; + +use anyhow::{bail, Context, Result}; +use rayon::prelude::*; +use rayon::ThreadPoolBuilder; + +use crate::commands::fast_allele_freq::collect_input_files; +use crate::commands::long_emit::{ + collect_merged_long_rows, default_participant, summarize_warning_counts, MissingRefLogger, + ReferenceResolver, SharedReference, +}; +use crate::download::ensure_reference_db; +use crate::long_rows::LongRow; +use crate::stats::StatsStore; +use crate::TargetStudyDosageArgs; + +const MISSING: u8 = 255; + +#[derive(Clone)] +struct Variant { + chrom: String, + pos: i64, + rsid: String, + reference: String, + alternate: String, + locus_key: String, +} + +struct SampleResult { + row_idx: usize, + hits: Vec<(usize, u8)>, + counts: std::collections::BTreeMap, +} + +struct OutputState<'a> { + n_variants: usize, + matrix: &'a mut [u8], + allele_count: &'a mut [i64], + n_obs: &'a mut [i64], + num_homo: &'a mut [i64], + num_hetero: &'a mut [i64], + warning_counts: &'a mut std::collections::BTreeMap, +} + +pub fn run_target_study_dosage(args: TargetStudyDosageArgs) -> Result<()> { + let overall_start = Instant::now(); + + let sqlite_path = ensure_reference_db(Some(&args.sqlite), args.force_download)?; + let store = StatsStore::connect_read_only(&sqlite_path)?; + let preload_start = Instant::now(); + let shared = Arc::new(SharedReference::load(&store)?); + eprintln!( + "📚 target-study-dosage: reference preloaded in {:.2}s", + preload_start.elapsed().as_secs_f64() + ); + + let variants = read_variants(&args.variants_tsv)?; + if variants.is_empty() { + bail!("No variants found in {}", args.variants_tsv.display()); + } + let key_to_col: HashMap = variants + .iter() + .enumerate() + .map(|(idx, variant)| (variant.locus_key.clone(), idx)) + .collect(); + + let files = collect_input_files(&args.inputs)?; + if files.is_empty() { + bail!("No input genotype files found under the provided --input paths"); + } + let mut tasks: Vec<(String, PathBuf)> = files + .into_iter() + .map(|path| (dosage_participant(&path), path)) + .collect(); + tasks.sort_by(|a, b| a.0.cmp(&b.0)); + let tasks: Vec<(usize, String, PathBuf)> = tasks + .into_iter() + .enumerate() + .map(|(idx, (pid, path))| (idx, pid, path)) + .collect(); + + let force_single = args.missing_ref_log.is_some(); + let threads = if force_single { + 1 + } else { + resolve_threads(args.threads) + }; + let batch_samples = if args.batch_samples > 0 { + args.batch_samples + } else { + threads.max(1) + }; + eprintln!( + "▶️ target-study-dosage: {} files, {} variants, threads={}, batch_samples={}", + tasks.len(), + variants.len(), + threads, + batch_samples + ); + + let n_samples = tasks.len(); + let n_variants = variants.len(); + let mut matrix = vec![MISSING; n_samples * n_variants]; + let mut allele_count = vec![0_i64; n_variants]; + let mut n_obs = vec![0_i64; n_variants]; + let mut num_homo = vec![0_i64; n_variants]; + let mut num_hetero = vec![0_i64; n_variants]; + let mut warning_counts = std::collections::BTreeMap::new(); + + if threads <= 1 { + let mut resolver = ReferenceResolver::shared(shared.clone()); + let mut logger = MissingRefLogger::new(args.missing_ref_log.as_deref(), args.warn_detail)?; + for (batch_idx, batch) in tasks.chunks(batch_samples).enumerate() { + for (row_idx, participant, path) in batch { + match parse_sample( + *row_idx, + participant, + path, + &key_to_col, + &mut resolver, + &mut logger, + ) { + Ok(mut result) => { + // The sequential logger is cumulative and may also write + // per-row details, so merge its counts once after the loop. + result.counts.clear(); + let mut output = OutputState { + n_variants, + matrix: &mut matrix, + allele_count: &mut allele_count, + n_obs: &mut n_obs, + num_homo: &mut num_homo, + num_hetero: &mut num_hetero, + warning_counts: &mut warning_counts, + }; + apply_sample_result(result, &mut output) + } + Err(err) => eprintln!("⚠️ skipping {}: {err:#}", path.display()), + } + } + eprintln!( + "🧬 target-study-dosage: parsed {}/{} samples", + ((batch_idx + 1) * batch_samples).min(n_samples), + n_samples + ); + } + merge_counts(&mut warning_counts, logger.counts().clone()); + } else { + let pool = ThreadPoolBuilder::new() + .num_threads(threads) + .build() + .context("build target-study-dosage thread pool")?; + for (batch_idx, batch) in tasks.chunks(batch_samples).enumerate() { + let results: Vec> = pool.install(|| { + batch + .par_iter() + .map(|(row_idx, participant, path)| { + let mut resolver = ReferenceResolver::shared(shared.clone()); + let mut logger = MissingRefLogger::new(None, crate::WarnDetail::None) + .expect("count-only logger has no file handle and cannot fail"); + parse_sample( + *row_idx, + participant, + path, + &key_to_col, + &mut resolver, + &mut logger, + ) + }) + .collect() + }); + for (task, result) in batch.iter().zip(results) { + match result { + Ok(result) => { + let mut output = OutputState { + n_variants, + matrix: &mut matrix, + allele_count: &mut allele_count, + n_obs: &mut n_obs, + num_homo: &mut num_homo, + num_hetero: &mut num_hetero, + warning_counts: &mut warning_counts, + }; + apply_sample_result(result, &mut output) + } + Err(err) => eprintln!("⚠️ skipping {}: {err:#}", task.2.display()), + } + } + eprintln!( + "🧬 target-study-dosage: parsed {}/{} samples", + ((batch_idx + 1) * batch_samples).min(n_samples), + n_samples + ); + } + } + + write_npy_u8_2d(&args.dosage_npy, n_samples, n_variants, &matrix)?; + write_samples(&args.samples_tsv, &tasks)?; + write_variants(&args.study_variants_tsv, &variants)?; + write_allele_freq( + &args.allele_freq_tsv, + &variants, + &allele_count, + &n_obs, + &num_homo, + &num_hetero, + )?; + + if args.warn_detail != crate::WarnDetail::None { + let summary = summarize_warning_counts(&warning_counts); + if !summary.is_empty() { + eprintln!("⚠️ target-study-dosage warnings: {summary}"); + } + } + eprintln!( + "✅ target-study-dosage: {} samples x {} variants, {:.2}s", + n_samples, + n_variants, + overall_start.elapsed().as_secs_f64() + ); + println!("✅ Wrote dosage matrix to {}", args.dosage_npy.display()); + Ok(()) +} + +fn parse_sample( + row_idx: usize, + participant: &str, + path: &Path, + key_to_col: &HashMap, + resolver: &mut ReferenceResolver, + logger: &mut MissingRefLogger, +) -> Result { + let (rows, _stats) = collect_merged_long_rows(path, participant, resolver, logger)?; + let mut non_target_alt_present: HashSet = HashSet::new(); + for row in &rows { + if row.dosage > 0 && !key_to_col.contains_key(&row.locus_key) { + if let Some(base_key) = locus_base_key(&row.locus_key) { + non_target_alt_present.insert(base_key); + } + } + } + let mut hits = Vec::new(); + for row in rows { + push_hit(&row, key_to_col, &non_target_alt_present, &mut hits); + } + Ok(SampleResult { + row_idx, + hits, + counts: logger.counts().clone(), + }) +} + +fn push_hit( + row: &LongRow, + key_to_col: &HashMap, + non_target_alt_present: &HashSet, + hits: &mut Vec<(usize, u8)>, +) { + if row.dosage < 0 { + return; + } + if let Some(base_key) = locus_base_key(&row.locus_key) { + if non_target_alt_present.contains(&base_key) { + return; + } + } + if let Some(&col) = key_to_col.get(&row.locus_key) { + hits.push((col, row.dosage as u8)); + } +} + +fn locus_base_key(locus_key: &str) -> Option { + let (base, _alt) = locus_key.rsplit_once('-')?; + Some(base.to_string()) +} + +fn apply_sample_result(result: SampleResult, output: &mut OutputState<'_>) { + merge_counts(output.warning_counts, result.counts); + let row_offset = result.row_idx * output.n_variants; + for (col, dosage) in result.hits { + let cell = &mut output.matrix[row_offset + col]; + if *cell != MISSING { + continue; + } + *cell = dosage; + output.allele_count[col] += dosage as i64; + output.n_obs[col] += 1; + if dosage == 2 { + output.num_homo[col] += 1; + } else if dosage == 1 { + output.num_hetero[col] += 1; + } + } +} + +fn read_variants(path: &Path) -> Result> { + let mut reader = csv::ReaderBuilder::new() + .delimiter(b'\t') + .from_path(path) + .with_context(|| format!("Open variants TSV {}", path.display()))?; + let headers = reader.headers()?.clone(); + let col = |name: &str| -> Result { + headers + .iter() + .position(|h| h == name) + .with_context(|| format!("{}: missing {name} column", path.display())) + }; + let chrom_col = col("chrom")?; + let pos_col = col("pos")?; + let rsid_col = col("rsid")?; + let ref_col = col("ref")?; + let alt_col = col("alt")?; + + let mut variants = Vec::new(); + for record in reader.records() { + let record = record?; + let chrom = normalize_chrom(record.get(chrom_col).unwrap_or("")); + let pos: i64 = record + .get(pos_col) + .unwrap_or("") + .parse() + .with_context(|| format!("{}: invalid pos in variants TSV", path.display()))?; + let rsid = record.get(rsid_col).unwrap_or("").to_string(); + let reference = record + .get(ref_col) + .unwrap_or("") + .trim() + .to_ascii_uppercase(); + let alternate = record + .get(alt_col) + .unwrap_or("") + .trim() + .to_ascii_uppercase(); + if reference.len() != 1 || alternate.len() != 1 { + continue; + } + let locus_key = format!("{chrom}-{pos}-{reference}-{alternate}"); + variants.push(Variant { + chrom, + pos, + rsid, + reference, + alternate, + locus_key, + }); + } + Ok(variants) +} + +fn normalize_chrom(value: &str) -> String { + let text = value.trim(); + let text = text + .strip_prefix("chr") + .or_else(|| text.strip_prefix("CHR")) + .unwrap_or(text); + if text == "M" { + "MT".to_string() + } else { + text.to_ascii_uppercase() + } +} + +fn dosage_participant(path: &Path) -> String { + if let Some(parent) = path + .parent() + .and_then(|p| p.file_name()) + .and_then(|p| p.to_str()) + { + if !matches!(parent, "" | "." | "geno" | "input" | "participants") { + return parent.replace(' ', "_"); + } + } + default_participant(path) +} + +fn write_npy_u8_2d(path: &Path, rows: usize, cols: usize, data: &[u8]) -> Result<()> { + if data.len() != rows * cols { + bail!("matrix data length does not match requested shape"); + } + if let Some(parent) = path.parent() { + if !parent.as_os_str().is_empty() { + std::fs::create_dir_all(parent) + .with_context(|| format!("Create directory {:?}", parent))?; + } + } + let mut writer = + BufWriter::new(File::create(path).with_context(|| format!("Create {:?}", path))?); + let mut header = + format!("{{'descr': '|u1', 'fortran_order': False, 'shape': ({rows}, {cols}), }}"); + let preamble_len = 10; + let padding = (16 - ((preamble_len + header.len() + 1) % 16)) % 16; + header.push_str(&" ".repeat(padding)); + header.push('\n'); + if header.len() > u16::MAX as usize { + bail!("NPY v1 header is too large"); + } + writer.write_all(b"\x93NUMPY")?; + writer.write_all(&[1, 0])?; + writer.write_all(&(header.len() as u16).to_le_bytes())?; + writer.write_all(header.as_bytes())?; + writer.write_all(data)?; + writer.flush()?; + Ok(()) +} + +fn write_samples(path: &Path, tasks: &[(usize, String, PathBuf)]) -> Result<()> { + if let Some(parent) = path.parent() { + if !parent.as_os_str().is_empty() { + std::fs::create_dir_all(parent) + .with_context(|| format!("Create directory {:?}", parent))?; + } + } + let mut writer = + BufWriter::new(File::create(path).with_context(|| format!("Create {:?}", path))?); + writeln!(writer, "sample_id\tfile")?; + for (_idx, sample_id, file) in tasks { + writeln!(writer, "{sample_id}\t{}", file.display())?; + } + writer.flush()?; + Ok(()) +} + +fn write_variants(path: &Path, variants: &[Variant]) -> Result<()> { + if let Some(parent) = path.parent() { + if !parent.as_os_str().is_empty() { + std::fs::create_dir_all(parent) + .with_context(|| format!("Create directory {:?}", parent))?; + } + } + let mut writer = + BufWriter::new(File::create(path).with_context(|| format!("Create {:?}", path))?); + writeln!(writer, "chrom\tpos\trsid\tref\talt\tlocus_key")?; + for variant in variants { + writeln!( + writer, + "{}\t{}\t{}\t{}\t{}\t{}", + variant.chrom, + variant.pos, + variant.rsid, + variant.reference, + variant.alternate, + variant.locus_key + )?; + } + writer.flush()?; + Ok(()) +} + +fn write_allele_freq( + path: &Path, + variants: &[Variant], + allele_count: &[i64], + n_obs: &[i64], + num_homo: &[i64], + num_hetero: &[i64], +) -> Result<()> { + if let Some(parent) = path.parent() { + if !parent.as_os_str().is_empty() { + std::fs::create_dir_all(parent) + .with_context(|| format!("Create directory {:?}", parent))?; + } + } + let mut writer = + BufWriter::new(File::create(path).with_context(|| format!("Create {:?}", path))?); + writeln!( + writer, + "locus_key\tallele_count\tallele_number\tnum_homo\tnum_hetero\tallele_freq\trsid" + )?; + for (idx, variant) in variants.iter().enumerate() { + let allele_number = 2 * n_obs[idx]; + let allele_freq = if allele_number > 0 { + allele_count[idx] as f64 / allele_number as f64 + } else { + 0.0 + }; + writeln!( + writer, + "{}\t{}\t{allele_number}\t{}\t{}\t{allele_freq:.6}\t{}", + variant.locus_key, allele_count[idx], num_homo[idx], num_hetero[idx], variant.rsid + )?; + } + writer.flush()?; + Ok(()) +} + +fn merge_counts( + into: &mut std::collections::BTreeMap, + from: std::collections::BTreeMap, +) { + for (code, n) in from { + *into.entry(code).or_insert(0) += n; + } +} + +fn resolve_threads(requested: usize) -> usize { + if requested > 0 { + return requested; + } + std::thread::available_parallelism() + .map(|n| n.get()) + .unwrap_or(1) + .max(1) +} diff --git a/cli/src/genotype.rs b/cli/src/genotype.rs index 954fc3a..28a4170 100644 --- a/cli/src/genotype.rs +++ b/cli/src/genotype.rs @@ -1,10 +1,11 @@ use std::collections::{BTreeSet, HashMap}; -use std::fs::File; -use std::io::{BufRead, BufReader}; +use std::io::BufRead; use std::path::Path; use anyhow::{bail, Context, Result}; +use crate::genotype_reader::open_text_reader; + const LOOKAHEAD_LINES: usize = 2048; const COMMENT_PREFIXES: [&str; 2] = ["#", "//"]; const RSID_ALIASES: &[&str] = &["rsid", "name", "snp", "marker", "id"]; @@ -63,8 +64,8 @@ pub fn process_file(path: &Path, mut on_variant: F) -> Result where F: FnMut(&VariantRecord, &FileMetadata) -> Result<()>, { - let file = File::open(path).with_context(|| format!("Failed to open {:?}", path))?; - let mut reader = BufReader::new(file); + let mut reader = + open_text_reader(path).with_context(|| format!("Failed to open {:?}", path))?; let mut buffered_lines: Vec = Vec::new(); let mut buffer = String::new(); diff --git a/cli/src/genotype_reader.rs b/cli/src/genotype_reader.rs index 483f897..f9e6224 100644 --- a/cli/src/genotype_reader.rs +++ b/cli/src/genotype_reader.rs @@ -1,6 +1,10 @@ use std::collections::{BTreeSet, HashMap}; +use std::fs::File; +use std::io::{BufRead, BufReader}; +use std::path::Path; -use anyhow::Result; +use anyhow::{Context, Result}; +use flate2::read::MultiGzDecoder; const COMMENT_PREFIXES: [&str; 2] = ["#", "//"]; @@ -48,6 +52,21 @@ pub enum RowOutcome { Ignored, } +pub fn open_text_reader(path: &Path) -> Result> { + let file = File::open(path).with_context(|| format!("Open {:?}", path))?; + if is_gzip_path(path) { + Ok(Box::new(BufReader::new(MultiGzDecoder::new(file)))) + } else { + Ok(Box::new(BufReader::new(file))) + } +} + +pub fn is_gzip_path(path: &Path) -> bool { + path.extension() + .and_then(|ext| ext.to_str()) + .is_some_and(|ext| ext.eq_ignore_ascii_case("gz")) +} + #[derive(Debug, Clone, Copy)] pub enum Delimiter { Tab, @@ -174,12 +193,28 @@ impl RowParser { } let header = self.header.as_ref().expect("header must be set"); + let genotype_aliases = self.alias_map.get("genotype").cloned(); let mut row_map: HashMap = HashMap::new(); + let mut vendor_missing_gt = false; for (idx, value) in fields.into_iter().enumerate() { if idx >= header.len() { continue; } - row_map.insert(normalize_name(&header[idx]), strip_inline_comment(&value)); + let norm = normalize_name(&header[idx]); + // Vendor missing genotype (#N/A) -> drop the whole row, matching + // genotype_normalizer. Detect on the RAW value before + // strip_inline_comment turns "#N/A" into "" (its leading '#'). + if matches!(value.trim(), "#N/A" | "#n/a") + && genotype_aliases + .as_ref() + .is_some_and(|a| a.contains(norm.as_str())) + { + vendor_missing_gt = true; + } + row_map.insert(norm, strip_inline_comment(&value)); + } + if vendor_missing_gt { + return Ok(RowOutcome::Skipped); } if self.carigenetics || looks_like_illumina_row(&row_map) { diff --git a/cli/src/main.rs b/cli/src/main.rs index 8b20d2e..e238197 100644 --- a/cli/src/main.rs +++ b/cli/src/main.rs @@ -17,6 +17,7 @@ use crate::commands::fast_allele_freq::run_fast_allele_freq; use crate::commands::fst::run_fst; use crate::commands::genostats::run_genostats; use crate::commands::genotype_to_vcf::run_genotype_to_vcf; +use crate::commands::het_signal::run_het_signal; use crate::commands::list_missing_cache::run_list_missing_cache; use crate::commands::list_missing_rsids::run_list_missing_rsids; use crate::commands::load_non_rsids::run_load_non_rsids; @@ -29,6 +30,7 @@ use crate::commands::reference_load::run_reference_load; use crate::commands::resolve_rsids::run_resolve_rsids; use crate::commands::sync_rsid_cache::run_sync_rsid_cache; use crate::commands::synthetic::run_synthetic; +use crate::commands::target_study_dosage::run_target_study_dosage; use crate::commands::update::run_update; mod commands { @@ -38,6 +40,7 @@ mod commands { pub mod fst; pub mod genostats; pub mod genotype_to_vcf; + pub mod het_signal; pub mod list_missing_cache; pub mod list_missing_rsids; pub mod load_non_rsids; @@ -50,6 +53,7 @@ mod commands { pub mod resolve_rsids; pub mod sync_rsid_cache; pub mod synthetic; + pub mod target_study_dosage; pub mod update; } @@ -74,12 +78,17 @@ enum Commands { AggregateLong(AggregateLongArgs), /// Fused parse+aggregate: genotype files -> allele frequency TSV, no .bvlr. FastAlleleFreq(FastAlleleFreqArgs), + /// Build a target-panel-aligned study dosage matrix and allele frequency TSV. + #[command(visible_alias = "hgp1k-study-dosage")] + TargetStudyDosage(TargetStudyDosageArgs), /// Pairwise Weir & Cockerham 1984 FST matrix from merged allele freq/number. Fst(FstArgs), /// Build a cohort PLINK .bed/.bim/.fam from genotype files (for PCA/QC). CohortBed(CohortBedArgs), /// Build a gnomAD-loadings-oriented PLINK bed for PCA projection. ProjectBed(ProjectBedArgs), + /// Per-sample het-signal caches for sex-biased admixture NMF. + HetSignal(HetSignalArgs), /// Inner-join per-population allele frequency TSVs into merged matrices. MergeAlleleFreq(MergeAlleleFreqArgs), /// Convert a .bvlr file into TSV for debugging. @@ -283,6 +292,46 @@ pub struct FastAlleleFreqArgs { pub max_ram_gb: f64, } +#[derive(Args, Clone)] +pub struct TargetStudyDosageArgs { + /// Input genotype file(s) or directories (scanned recursively). + #[arg(short = 'i', long = "input")] + pub inputs: Vec, + /// Target variants TSV with chrom, pos, rsid, ref, alt columns. + #[arg(long)] + pub variants_tsv: PathBuf, + /// Path to the SQLite database containing rsid_reference data. + #[arg(long, default_value = "data/genostats.sqlite")] + pub sqlite: PathBuf, + /// Force re-download of the reference database. + #[arg(long, action = ArgAction::SetTrue)] + pub force_download: bool, + /// Output uint8 NumPy .npy matrix path. Shape is samples x variants. + #[arg(long)] + pub dosage_npy: PathBuf, + /// Output sample metadata TSV path. + #[arg(long)] + pub samples_tsv: PathBuf, + /// Output variant metadata TSV path, aligned to matrix columns. + #[arg(long)] + pub study_variants_tsv: PathBuf, + /// Output allele frequency TSV path. + #[arg(long)] + pub allele_freq_tsv: PathBuf, + /// Optional log file for missing reference rows (per-row TSV; forces single thread). + #[arg(long)] + pub missing_ref_log: Option, + /// Warning verbosity: none, summary (default), or full (per-row). + #[arg(long, value_enum, default_value_t = WarnDetail::Summary)] + pub warn_detail: WarnDetail, + /// Worker threads for parsing (0 = auto/all cores). + #[arg(long, default_value = "0")] + pub threads: usize, + /// Number of samples to parse in each parallel batch. 0 = threads. + #[arg(long, default_value = "0")] + pub batch_samples: usize, +} + #[derive(Args, Clone)] pub struct CohortBedArgs { /// Data dir containing one subdirectory per sample, each with a genotype .txt. @@ -318,6 +367,16 @@ pub struct ProjectBedArgs { pub min_loadings_ratio: f64, } +#[derive(Args, Clone)] +pub struct HetSignalArgs { + /// Data dir with one subdir per sample, each with a genotype .txt. + #[arg(short = 'i', long = "input")] + pub input: PathBuf, + /// Output dir for per-sample .bvhs caches. + #[arg(long)] + pub out_dir: PathBuf, +} + #[derive(Args, Clone)] pub struct FstArgs { /// Merged allele-frequency matrix TSV (locus_key + one column per population). @@ -547,9 +606,11 @@ fn main() -> Result<()> { Commands::EmitLong(args) => run_long_emit(args), Commands::AggregateLong(args) => run_long_aggregate(args), Commands::FastAlleleFreq(args) => run_fast_allele_freq(args), + Commands::TargetStudyDosage(args) => run_target_study_dosage(args), Commands::Fst(args) => run_fst(args), Commands::CohortBed(args) => run_cohort_bed(args), Commands::ProjectBed(args) => run_project_bed(args), + Commands::HetSignal(args) => run_het_signal(args), Commands::MergeAlleleFreq(args) => run_merge_allele_freq(args), Commands::DumpLong(args) => run_long_dump(args), Commands::ListMissingCache(args) => run_list_missing_cache(args), diff --git a/cli/src/stats.rs b/cli/src/stats.rs index c058dd6..87a9063 100644 --- a/cli/src/stats.rs +++ b/cli/src/stats.rs @@ -441,6 +441,12 @@ fn init_schema(conn: &Connection) -> Result<()> { migrate_legacy_schema(conn)?; conn.execute_batch( r#" + CREATE TABLE IF NOT EXISTS reference_sources ( + source_id INTEGER PRIMARY KEY CHECK(source_id BETWEEN 0 AND 255), + code TEXT NOT NULL UNIQUE, + label TEXT NOT NULL, + description TEXT + ); CREATE TABLE IF NOT EXISTS rsid_reference ( rsid INTEGER PRIMARY KEY, chromosome TEXT NOT NULL, @@ -454,7 +460,9 @@ fn init_schema(conn: &Connection) -> Result<()> { position INTEGER NOT NULL, reference TEXT NOT NULL, alternates TEXT NOT NULL, - source TEXT + source TEXT, + source_id INTEGER CHECK(source_id BETWEEN 0 AND 255), + FOREIGN KEY(source_id) REFERENCES reference_sources(source_id) ); CREATE TABLE IF NOT EXISTS grch38_non_rsids ( snp_name TEXT PRIMARY KEY, @@ -464,7 +472,9 @@ fn init_schema(conn: &Connection) -> Result<()> { reference TEXT NOT NULL, alternates TEXT NOT NULL, source TEXT NOT NULL DEFAULT 'ensembl_grch38', - note TEXT + note TEXT, + source_id INTEGER CHECK(source_id BETWEEN 0 AND 255), + FOREIGN KEY(source_id) REFERENCES reference_sources(source_id) ); CREATE TABLE IF NOT EXISTS illumina_probe_manifest ( row_order INTEGER PRIMARY KEY, @@ -479,16 +489,38 @@ fn init_schema(conn: &Connection) -> Result<()> { ); CREATE INDEX IF NOT EXISTS idx_rsid_reference_pos ON rsid_reference(chromosome, position); + CREATE INDEX IF NOT EXISTS idx_rsid_reference_user_source_id + ON rsid_reference_user(source_id); CREATE INDEX IF NOT EXISTS idx_grch38_non_rsids_pos ON grch38_non_rsids(chromosome, position); CREATE INDEX IF NOT EXISTS idx_grch38_non_rsids_rsid ON grch38_non_rsids(rsid); + CREATE INDEX IF NOT EXISTS idx_grch38_non_rsids_source_id + ON grch38_non_rsids(source_id); CREATE INDEX IF NOT EXISTS idx_illumina_probe_manifest_rsid ON illumina_probe_manifest(rsid); CREATE INDEX IF NOT EXISTS idx_illumina_probe_manifest_pos ON illumina_probe_manifest(chromosome, position); "#, )?; + if !has_column(conn, "rsid_reference_user", "source_id")? { + conn.execute_batch( + "ALTER TABLE rsid_reference_user + ADD COLUMN source_id INTEGER CHECK(source_id BETWEEN 0 AND 255);", + )?; + } + if !has_column(conn, "grch38_non_rsids", "source_id")? { + conn.execute_batch( + "ALTER TABLE grch38_non_rsids + ADD COLUMN source_id INTEGER CHECK(source_id BETWEEN 0 AND 255);", + )?; + } + conn.execute_batch( + "CREATE INDEX IF NOT EXISTS idx_rsid_reference_user_source_id + ON rsid_reference_user(source_id); + CREATE INDEX IF NOT EXISTS idx_grch38_non_rsids_source_id + ON grch38_non_rsids(source_id);", + )?; Ok(()) } diff --git a/cli/src/util.rs b/cli/src/util.rs index b27b1c1..0614b3a 100644 --- a/cli/src/util.rs +++ b/cli/src/util.rs @@ -37,6 +37,24 @@ pub fn collect_input_files(inputs: &[PathBuf]) -> Result> { Ok(files) } +pub fn genotype_file_stem(path: &Path) -> Option { + let mut name = path.file_name()?.to_str()?.to_string(); + if name.to_ascii_lowercase().ends_with(".gz") { + name.truncate(name.len() - 3); + } + for ext in [".txt", ".tsv", ".csv"] { + if name.to_ascii_lowercase().ends_with(ext) { + name.truncate(name.len() - ext.len()); + break; + } + } + if name.is_empty() { + None + } else { + Some(name) + } +} + fn canonicalize_path>(path: P) -> Result { let path_ref = path.as_ref(); match fs::canonicalize(path_ref) { @@ -47,7 +65,16 @@ fn canonicalize_path>(path: P) -> Result { fn is_candidate_file(path: &Path) -> bool { if let Some(ext) = path.extension().and_then(|e| e.to_str()) { - let ext_lower = ext.to_lowercase(); + let ext_lower = ext.to_ascii_lowercase(); + if ext_lower == "gz" { + return path + .file_stem() + .and_then(|stem| Path::new(stem).extension()) + .and_then(|e| e.to_str()) + .is_some_and(|inner| { + matches!(inner.to_ascii_lowercase().as_str(), "txt" | "tsv" | "csv") + }); + } return matches!(ext_lower.as_str(), "txt" | "tsv" | "csv"); } true diff --git a/data/genostats.sqlite b/data/genostats.sqlite index f0bc303..4b8e003 100644 Binary files a/data/genostats.sqlite and b/data/genostats.sqlite differ