diff --git a/modkit-core/src/extract/args.rs b/modkit-core/src/extract/args.rs index 4c2d31f5..77e70dde 100644 --- a/modkit-core/src/extract/args.rs +++ b/modkit-core/src/extract/args.rs @@ -2,6 +2,11 @@ use clap::Args; use std::path::PathBuf; #[derive(Args)] +#[command(group( + clap::ArgGroup::new("motif_or_cpg") + .args(["motif", "cpg"]) + .multiple(true) +))] pub(super) struct InputArgs { /// Path to modBAM file to extract read-level information from, or one of /// `-` or `stdin` to specify a stream from standard input. If a file @@ -113,7 +118,7 @@ pub(super) struct InputArgs { /// column. "." will be used when an aligned position does not match a /// motif. #[clap(help_heading = "Modified Base Options")] - #[arg(long, requires = "motif", default_value_t = false)] + #[arg(long, requires = "motif_or_cpg", default_value_t = false)] pub annotate_motifs: bool, /// Only output counts at CpG motifs. Requires a reference sequence to be /// provided. @@ -125,7 +130,7 @@ pub(super) struct InputArgs { #[arg( long, short = 'k', - requires = "motif", + requires = "motif_or_cpg", default_value_t = false, hide_short_help = true )] diff --git a/modkit-core/src/extract/subcommand.rs b/modkit-core/src/extract/subcommand.rs index 42b4d6ed..39f4079d 100644 --- a/modkit-core/src/extract/subcommand.rs +++ b/modkit-core/src/extract/subcommand.rs @@ -276,7 +276,8 @@ impl EntryExtractFull { ); }); - let with_motifs = self.input_args.motif.is_some(); + let with_motifs = + self.input_args.motif.is_some() || self.input_args.cpg; let output_header = if self.input_args.no_headers { None } else { @@ -716,7 +717,8 @@ impl EntryExtractCalls { // TODO(arand) once I refactor extract, I'll want to keep this around. drop(io_threadpool); - let with_motifs = self.input_args.motif.is_some(); + let with_motifs = + self.input_args.motif.is_some() || self.input_args.cpg; let output_header = if self.input_args.no_headers { None } else { @@ -1194,3 +1196,68 @@ impl EntryReadStats { !(self.filter_threshold.is_empty() && self.mod_thresholds.is_none()) } } + +#[cfg(test)] +mod tests { + use clap::Parser; + + use super::ExtractMods; + + #[derive(Parser)] + #[command(name = "extract")] + struct ExtractCli { + #[command(subcommand)] + command: ExtractMods, + } + + fn parse_extract(subcommand: &str, extra_args: &[&str]) -> bool { + let mut args = vec![ + "extract", + subcommand, + "input.bam", + "output.tsv", + "--reference", + "reference.fa", + ]; + args.extend_from_slice(extra_args); + ExtractCli::try_parse_from(args).is_ok() + } + + #[test] + fn motif_dependent_flags_accept_motif_or_cpg() { + let actual = ["full", "calls"] + .into_iter() + .flat_map(|subcommand| { + ["--annotate-motifs", "--mask"].into_iter().map( + move |dependent_arg| { + ( + subcommand, + dependent_arg, + parse_extract( + subcommand, + &[dependent_arg, "--motif", "CG", "0"], + ), + parse_extract( + subcommand, + &[dependent_arg, "--cpg"], + ), + !parse_extract(subcommand, &[dependent_arg]), + ) + }, + ) + }) + .collect::>(); + let expected = ["full", "calls"] + .into_iter() + .flat_map(|subcommand| { + ["--annotate-motifs", "--mask"].into_iter().map( + move |dependent_arg| { + (subcommand, dependent_arg, true, true, true) + }, + ) + }) + .collect::>(); + + assert_eq!(actual, expected); + } +} diff --git a/modkit-core/src/extract/writer.rs b/modkit-core/src/extract/writer.rs index 237827ba..86f1e01e 100644 --- a/modkit-core/src/extract/writer.rs +++ b/modkit-core/src/extract/writer.rs @@ -67,7 +67,7 @@ impl PositionModCalls { let inferred = self.base_mod_probs.inferred_unmodified; let motif_hits = motif_position_lookup.and_then(|lu| { match (self.ref_position, profile.chrom_id, self.alignment_strand) { - (Some(i), Some(tid), Some(strand)) if i > 0i64 => { + (Some(i), Some(tid), Some(strand)) if i >= 0i64 => { let pos = i as usize; let motif_hits = lu.get_motif_hits( tid, @@ -289,7 +289,7 @@ impl OutwriterWithMemory .and_then(|chrom_id| self.tid_to_name.get(&chrom_id)); let position_calls = PositionModCalls::from_profile(&profile); for call in position_calls { - call.to_row( + if let Some(row) = call.to_row( profile, chrom_name, &self.caller, @@ -298,10 +298,10 @@ impl OutwriterWithMemory false, motif_position_lookup, self.with_motifs, - ) - .map(|s| self.tsv_writer.write(s.as_bytes())) - .transpose()?; - rows_written += 1; + ) { + self.tsv_writer.write(row.as_bytes())?; + rows_written += 1; + } } self.number_of_written_reads += 1; } @@ -450,3 +450,118 @@ impl Ok(()) } } + +#[cfg(test)] +mod tests { + use rustc_hash::FxHashMap; + + use super::*; + use crate::mod_base_code::{DnaBase, ModCodeRepr}; + use crate::motifs::motif_bed::RegexMotif; + use crate::read_ids_to_base_mod_probs::ModProfile; + + fn mod_profile(position: usize, probability: f32) -> ModProfile { + ModProfile::new( + position, + Some(position as i64), + 0, + 0, + 2, + probability, + ModCodeRepr::Code('a'), + 30, + Kmer::new(b"AA", position, 1), + Strand::Positive, + Some(Strand::Positive), + DnaBase::A, + false, + ) + } + + fn motif_lookup_at_zero() -> MotifPositionLookup { + let positions = + FxHashMap::from_iter([((0, 0), vec![(0, Strand::Positive)])]); + let motifs = vec![RegexMotif::parse_string("A", 0).unwrap()]; + MotifPositionLookup::new(positions, motifs) + } + + #[test] + fn motif_annotation_includes_reference_position_zero() { + let profile = mod_profile(0, 1.0); + let motif_lookup = motif_lookup_at_zero(); + let reference_seqs = + HashMap::from([("chr1".to_string(), b"AA".to_vec())]); + let mod_profile_row = profile.to_row( + "read", + "chr1", + Some(0), + Some(0), + Some(2), + &reference_seqs, + 0, + Some(&motif_lookup), + true, + ); + + let read_profile = ReadBaseModProfile::new( + "read".to_string(), + Some(0), + 0, + Some(0), + Some(2), + vec![profile], + ); + let position_call = + PositionModCalls::from_profile(&read_profile).pop().unwrap(); + let position_call_row = position_call + .to_row( + &read_profile, + Some(&"chr1".to_string()), + &MultipleThresholdModCaller::new_passthrough(), + &reference_seqs, + false, + false, + Some(&motif_lookup), + true, + ) + .unwrap(); + + assert_eq!( + ( + mod_profile_row.trim_end().split(TAB).last(), + position_call_row.trim_end().split(TAB).last(), + ), + (Some("A,0"), Some("A,0")) + ); + } + + #[test] + fn rows_written_counts_only_emitted_rows() { + let read_profile = ReadBaseModProfile::new( + "read".to_string(), + Some(0), + 0, + Some(0), + Some(2), + vec![mod_profile(0, 1.0), mod_profile(1, 0.2)], + ); + let item = ReadsBaseModProfile::new(vec![read_profile], 0, 0); + let caller = MultipleThresholdModCaller::new( + HashMap::new(), + HashMap::new(), + 0.9, + ); + let mut writer = TsvWriterWithContigNames::new_with_caller( + TsvWriter::new_null(), + HashMap::new(), + HashMap::new(), + caller, + true, + false, + ) + .unwrap(); + + assert_eq!(writer.write(item, None).unwrap(), 1); + assert_eq!(writer.num_reads(), 1); + } +} diff --git a/modkit-core/src/read_ids_to_base_mod_probs.rs b/modkit-core/src/read_ids_to_base_mod_probs.rs index 39513aa2..7fd751e2 100644 --- a/modkit-core/src/read_ids_to_base_mod_probs.rs +++ b/modkit-core/src/read_ids_to_base_mod_probs.rs @@ -498,7 +498,7 @@ impl ModProfile { let query_kmer = format!("{}", self.query_kmer); let motif_hits = motif_positions_lookup.and_then(|lu| { match (self.ref_position, tid, self.alignment_strand) { - (Some(i), Some(tid), Some(strand)) if i > 0i64 => { + (Some(i), Some(tid), Some(strand)) if i >= 0i64 => { let pos = i as usize; let motif_hits = lu.get_motif_hits( tid, diff --git a/modkit/tests/test_extract.rs b/modkit/tests/test_extract.rs index 729153bb..63f9d354 100644 --- a/modkit/tests/test_extract.rs +++ b/modkit/tests/test_extract.rs @@ -92,6 +92,24 @@ fn check_mod_profiles_same( } } +fn assert_cpg_motif_schema(output_fp: &Path, expected_columns: usize) { + let mut lines = BufReader::new(File::open(output_fp).unwrap()) + .lines() + .map(|line| line.unwrap()); + let header = lines.next().expect("expected extract header"); + let header_fields = header.split('\t').collect::>(); + assert_eq!(header_fields.len(), expected_columns); + assert_eq!(header_fields.last(), Some(&"motifs")); + + let rows = lines.collect::>(); + assert!(!rows.is_empty(), "expected at least one extracted row"); + for row in rows { + let fields = row.split('\t').collect::>(); + assert_eq!(fields.len(), expected_columns); + assert_eq!(fields.last(), Some(&"CG,0")); + } +} + // #[test] // fn test_common_parse_extract_full() { // let mut reader = csv::ReaderBuilder::new() @@ -470,6 +488,8 @@ fn test_extract_implicit_mod_calls() { fn test_extract_cpg_motif() { let extract_tsv = std::env::temp_dir().join("test_extract_cpg_motif_extract.tsv"); + let calls_tsv = + std::env::temp_dir().join("test_extract_cpg_motif_calls.tsv"); let reference_fasta_fp = "../tests/resources/CGI_ladder_3.6kb_ref.fa"; let cpg_positions = parse_bed_file( &Path::new("../tests/resources/CGI_ladder_3.6kb_ref_CG.bed") @@ -502,6 +522,22 @@ fn test_extract_cpg_motif() { ]) .unwrap(); + run_modkit(&[ + "extract", + "calls", + "../tests/resources/2_reads_all_context.bam", + calls_tsv.to_str().unwrap(), + "--cpg", + "--reference", + reference_fasta_fp, + "--no-filtering", + "--force", + ]) + .unwrap(); + + assert_cpg_motif_schema(&extract_tsv, 22); + assert_cpg_motif_schema(&calls_tsv, 24); + let mod_profile = parse_mod_profile(&extract_tsv).unwrap(); for (read, mod_data) in mod_profile { for row in mod_data {