Correct extract motif metadata and output accounting

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
     )]

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::<Vec<_>>();
+        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::<Vec<_>>();
+
+        assert_eq!(actual, expected);
+    }
+}

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<W: Write> OutwriterWithMemory<ReadsBaseModProfile>
                 .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<W: Write> OutwriterWithMemory<ReadsBaseModProfile>
                     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<T: Write, const SIZE: usize>
         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);
+    }
+}

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,

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::<Vec<_>>();
+    assert_eq!(header_fields.len(), expected_columns);
+    assert_eq!(header_fields.last(), Some(&"motifs"));
+
+    let rows = lines.collect::<Vec<_>>();
+    assert!(!rows.is_empty(), "expected at least one extracted row");
+    for row in rows {
+        let fields = row.split('\t').collect::<Vec<_>>();
+        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 {