feat(alignmentSieve): opt-in --keepTags for shifted BAM output

docs/content/tools/alignmentSieve.rst

@@ -84,3 +84,16 @@ As can be seen, such fragments are considered to be on the ``-`` strand, so nega
 
 .. note::
     If the ``--shift`` or ``--ATACshift`` options are used, then only properly-paired reads will be used.
+
+Shifted BAM output
+^^^^^^^^^^^^^^^^^^
+
+With ``--shift`` or ``--ATACshift`` and BAM output, each kept read is rewritten as a single match (``M``) covering the shifted fragment, positioned at the shifted start. SEQ and QUAL are empty and all auxiliary tags are dropped. Expanding shifts and shifts longer than the read are supported because no reference or read bases are needed.
+
+If a downstream step needs a tag such as a single-cell barcode, copy it with ``--keepTags``:
+
+.. code:: bash
+
+    $ alignmentSieve -b possorted.bam --ATACshift --keepTags CB -o shifted.bam
+
+Only the listed tags are written, with their input values. Each tag must be a valid SAM tag name (a letter followed by a letter or digit). Tags that describe the original alignment (``AS``, ``NM``, ``MD``, ``MC``, ``SA``, ``XA``, ``CC``, ``CP``, ``CG``, ``OA``, ``OC``, ``OP``, ``TS``) or the original SEQ and QUAL (``BQ``, ``MM``, ``ML``) are refused, because the rewritten CIGAR, moved mates and empty SEQ and QUAL would make them wrong. ``--keepTags`` needs ``--shift`` or ``--ATACshift`` and BAM output. Without it, behaviour is unchanged.

galaxy/wrapper/alignmentSieve.xml

@@ -69,6 +69,10 @@
         #elif $ATACshift:
             --ATACshift
         #end if
+        #if ($shift or $ATACshift) and $keepTags:
+            #set tags = " ".join(["'{}'".format(x) for x in $keepTags.split(" ")])
+            --keepTags $tags
+        #end if
         #if $BED:
             --BED
             -o '$outFile'
@@ -98,6 +102,10 @@
                   take strand into consideration for strand-specific
                   protocols. Note that only properly paired reads are considered."/>
         <param argument="--ATACshift" type="boolean" label="Shift fragment ends as appropriate for ATAC-seq"/>
+        <param argument="--keepTags" type="text" label="Auxiliary tags to keep on shifted reads" value="" optional="true" size="30"
+            help="Space-separated two-character SAM tags (for example CB UB) copied from the input onto shifted BAM records. By default shifting drops all tags. Only used together with a shift and BAM output. Tags that describe the original alignment (such as NM, MD or SA) are refused.">
+            <validator type="regex">^[A-Za-z][A-Za-z0-9]( [A-Za-z][A-Za-z0-9])*$</validator>
+        </param>
         <param argument="filterRNAstrand" type="select" optional="true" label="Only include reads originating from fragments from the forward or reverse strand."
             help="By default (the no option), all reads are processed, regardless of the strand they originated from. For RNAseq, it can be useful to separately create bigWig files for the forward or reverse strands.
                   Note that this tools assumes that a dUTP-based method was used, so fragments will be assigned to the reverse strand if the second read in a pair is reverse complemented.">

pydeeptools/deeptools/alignmentSieve2.py

@@ -1,4 +1,5 @@
 import argparse
+import re
 import signal
 import sys
 from importlib.metadata import version
@@ -7,6 +8,12 @@
 from deeptools.hp import r_alignmentsieve
 from deeptools.parserCommon import existingFile
 
+# SAM tags that describe the original alignment (CIGAR, mismatches, mate or
+# supplementary/secondary alignments, junctions) or the original SEQ/QUAL (BQ, MM, ML).
+# Shifting rewrites the CIGAR, moves mates and empties SEQ and QUAL, so these would be
+# wrong on the output.
+ALIGNMENT_TAGS = frozenset(['AS', 'NM', 'MD', 'MC', 'SA', 'XA', 'CC', 'CP', 'CG', 'OA', 'OC', 'OP', 'TS', 'BQ', 'MM', 'ML', 'CM', 'PT'])
+
 
 def parseArguments():
     parser = argparse.ArgumentParser(
@@ -73,6 +80,13 @@ def parseArguments():
                          action='store_true',
                          help='Shift the produced BAM file or BEDPE regions as commonly done for ATAC-seq. This is equivalent to --shift 4 -5 5 -4.')
 
+    general.add_argument('--keepTags',
+                         action='extend',
+                         nargs='+',
+                         default=[],
+                         metavar='TAG',
+                         help=f"Copy these two-letter auxiliary tags (for example CB for a cell barcode) from the input onto shifted BAM records. By default shifting drops all tags, SEQ and QUAL. Only has an effect together with --shift or --ATACshift and BAM output. Tags that describe the alignment itself ({', '.join(sorted(ALIGNMENT_TAGS))}) are rejected because shifting rewrites the CIGAR, which makes them wrong.")
+
     output = parser.add_argument_group('Output arguments')
     output.add_argument('--BED',
                         action='store_true',
@@ -168,6 +182,18 @@ def main(args=None):
             print("Warning! The --ATACshift option is used, but a --shift option is provided as well. The latter will be ignored in favor of 4 -5 5 -4.")
         args.shift = [4, -5, 5, -4]
 
+    for tag in args.keepTags:
+        if not re.fullmatch(r'[A-Za-z][A-Za-z0-9]', tag):
+            sys.exit(f"--keepTags expects SAM tags (a letter followed by a letter or digit), got '{tag}'.")
+        if tag in ALIGNMENT_TAGS:
+            sys.exit(f"--keepTags: '{tag}' describes the original alignment and would be wrong after shifting.")
+    if len(set(args.keepTags)) != len(args.keepTags):
+        sys.exit("--keepTags: each tag may be given only once.")
+    if args.keepTags and not args.shift:
+        sys.exit("--keepTags only applies together with --shift or --ATACshift.")
+    if args.keepTags and args.BED:
+        sys.exit("--keepTags only applies to BAM output, not --BED.")
+
     if args.ignoreDuplicates:
         args.samFlagExclude |= 0x400
 
@@ -197,4 +223,5 @@ def main(args=None):
         args.maxFragmentLength,
         args.label,
         args.smartLabels,
+        args.keepTags,
     )

pydeeptools/deeptools/test/test_alignmentsieve.py

@@ -155,3 +155,102 @@ def test_with_cram_output_with_shifts():
     aln_seive.main(args)
 
     _assert_equals(BAMFILE_OUT8, output_bam_file)
+
+
+def _bam_with_cell_barcodes(path: str) -> None:
+    """Copy the fixture, adding a CB tag derived from the read name."""
+    with pysam.AlignmentFile(BAMFILE_IN, "rb") as fin, pysam.AlignmentFile(path, "wb", template=fin) as fout:
+        for read in fin:
+            read.set_tag("CB", f"CELL{hash(read.query_name) % 7}-1")
+            fout.write(read)
+
+
+def test_alsieve_shift_drops_tags_by_default():
+    """Shifting without --keepTags drops every tag, SEQ and QUAL (unchanged behaviour)."""
+    _, inbam = tempfile.mkstemp(suffix=".bam")
+    _, outbam = tempfile.mkstemp(suffix=".bam")
+    _bam_with_cell_barcodes(inbam)
+    aln_seive.main(f"--bam {inbam} -o {outbam} --ATACshift".split())
+    with pysam.AlignmentFile(outbam, "rb") as f:
+        reads = list(f)
+    assert reads
+    assert all(not r.get_tags() for r in reads)
+
+
+def test_alsieve_keeptags_cb_matches_default_alignments():
+    """--keepTags CB keeps only CB, with the input value, and does not change positions."""
+    _, inbam = tempfile.mkstemp(suffix=".bam")
+    _, plain = tempfile.mkstemp(suffix=".bam")
+    _, kept = tempfile.mkstemp(suffix=".bam")
+    _bam_with_cell_barcodes(inbam)
+    aln_seive.main(f"--bam {inbam} -o {plain} --ATACshift".split())
+    aln_seive.main(f"--bam {inbam} -o {kept} --ATACshift --keepTags CB".split())
+    with pysam.AlignmentFile(inbam, "rb") as fi:
+        expected_cb = {(r.query_name, r.flag): r.get_tag("CB") for r in fi}
+    with pysam.AlignmentFile(plain, "rb") as fp, pysam.AlignmentFile(kept, "rb") as fk:
+        plain_reads, kept_reads = list(fp), list(fk)
+    assert len(plain_reads) == len(kept_reads) > 0
+    for p, k in zip(plain_reads, kept_reads):
+        assert (p.query_name, p.flag, p.reference_start, p.cigarstring, p.template_length) == (
+            k.query_name, k.flag, k.reference_start, k.cigarstring, k.template_length)
+        assert [t for t, _ in k.get_tags()] == ["CB"]
+        assert k.get_tag("CB") == expected_cb[(k.query_name, k.flag)]
+
+
+def test_alsieve_keeptags_rejects_alignment_tags():
+    """Tags describing the original alignment are wrong after a CIGAR rewrite and are refused."""
+    _, outbam = tempfile.mkstemp(suffix=".bam")
+    for tag in ("AS", "NM", "MD", "SA", "BQ", "MM", "ML", "CM", "PT"):
+        try:
+            aln_seive.main(f"--bam {BAMFILE_IN} -o {outbam} --ATACshift --keepTags {tag}".split())
+        except SystemExit as e:
+            assert tag in str(e.code)
+        else:
+            raise AssertionError(f"--keepTags {tag} was accepted")
+
+
+def test_alsieve_keeptags_rejects_malformed_tag_names():
+    """Non-SAM tag names fail as a CLI error, not as a Rust panic (e.g. multi-byte characters)."""
+    _, outbam = tempfile.mkstemp(suffix=".bam")
+    for tag in ("éA", "A-", "1A", "C", "CBX"):
+        try:
+            aln_seive.main(["--bam", BAMFILE_IN, "-o", outbam, "--ATACshift", "--keepTags", tag])
+        except SystemExit as e:
+            assert "SAM tags" in str(e.code)
+        else:
+            raise AssertionError(f"--keepTags {tag!r} was accepted")
+
+
+def test_alsieve_keeptags_rejects_duplicates():
+    """A repeated tag would be written twice on the same record."""
+    _, outbam = tempfile.mkstemp(suffix=".bam")
+    for tags in ("--keepTags CB CB", "--keepTags CB --keepTags CB"):
+        try:
+            aln_seive.main(f"--bam {BAMFILE_IN} -o {outbam} --ATACshift {tags}".split())
+        except SystemExit as e:
+            assert "only once" in str(e.code)
+        else:
+            raise AssertionError(f"{tags!r} was accepted")
+
+def test_alsieve_keeptags_repeated_option_accumulates():
+    """--keepTags CB --keepTags UB keeps both tags, not just the last group."""
+    inbam = tempfile.mkstemp(suffix=".bam")[1]
+    outbam = tempfile.mkstemp(suffix=".bam")[1]
+    with pysam.AlignmentFile(BAMFILE_IN) as src, pysam.AlignmentFile(inbam, "wb", template=src) as dst:
+        for r in src:
+            r.set_tag("CB", "AAAC")
+            r.set_tag("UB", "TTTG")
+            dst.write(r)
+    aln_seive.main(f"--bam {inbam} -o {outbam} --ATACshift --keepTags CB --keepTags UB".split())
+    with pysam.AlignmentFile(outbam) as out:
+        for r in out:
+            assert sorted(t for t, _ in r.get_tags()) == ["CB", "UB"]
+
+def test_alsieve_keeptags_needs_shift_and_bam_output():
+    _, outbam = tempfile.mkstemp(suffix=".bam")
+    for extra in ("", "--ATACshift --BED"):
+        try:
+            aln_seive.main(f"--bam {BAMFILE_IN} -o {outbam} --keepTags CB {extra}".split())
+        except SystemExit:
+            continue
+        raise AssertionError(f"--keepTags CB {extra!r} was accepted")

src/alignmentsieve.rs

@@ -29,6 +29,7 @@ pub fn r_alignmentsieve(
     maxfraglen: u32,              // maximum fragment length.
     label: &str,                  // user defined label
     smartlabels: bool,            // derive label from filename
+    keep_tags: Vec<String>,       // 2-letter aux tags to copy onto shifted records.
 ) -> PyResult<()> {
     // Open input BAM once and stream through in order
     let mut bam = Reader::from_path(bamifile)
@@ -47,6 +48,13 @@ pub fn r_alignmentsieve(
 
     let write_filters = filtered_out_readsfile != "None";
     let readshift: Vec<i32> = shift.extract(py).expect("Failed to extract shift");
+    let keep_tags: Vec<[u8; 2]> = keep_tags
+        .iter()
+        .map(|t| match t.as_bytes() {
+            [a, b] => [*a, *b],
+            _ => panic!("--keepTags expects 2-letter tags, got '{}'", t),
+        })
+        .collect();
     // shift is of length 0, 2, or 4.
 
     // Build chrom_names and chrom_sizes from BAM header — indexed by tid for O(1) lookup, zero allocations per record.
@@ -194,7 +202,7 @@ pub fn r_alignmentsieve(
             }
 
             let tid = record.tid() as usize;
-            if let Some(shifted) = apply_shift(&record, &readshift, chrom_sizes[tid]) {
+            if let Some(shifted) = apply_shift(&record, &readshift, chrom_sizes[tid], &keep_tags) {
                 if let Some(ref mut w) = obam {
                     w.write(&shifted).unwrap_or_else(|e| {
                         panic!(
@@ -489,7 +497,12 @@ fn query_alignment_end(record: &bam::Record) -> i64 {
     read_pos
 }
 
-fn apply_shift(record: &bam::Record, shift: &[i32], chrom_len: u64) -> Option<bam::Record> {
+fn apply_shift(
+    record: &bam::Record,
+    shift: &[i32],
+    chrom_len: u64,
+    keep_tags: &[[u8; 2]],
+) -> Option<bam::Record> {
     if !record.is_proper_pair() {
         return None;
     }
@@ -554,14 +567,21 @@ fn apply_shift(record: &bam::Record, shift: &[i32], chrom_len: u64) -> Option<ba
 
     newrec.set(record.qname(), Some(&new_cigar), &[], &[]);
 
-    // Remove all auxiliary tags
+    // Remove all auxiliary tags, then copy back only the ones the user asked for.
+    // The CIGAR above no longer matches the original alignment, so tags that
+    // describe it (NM, MD, MC, SA, ...) are rejected by the CLI, not handled here.
     let tags: Vec<Vec<u8>> = newrec
         .aux_iter()
         .filter_map(|r| r.ok().map(|(t, _)| t.to_vec()))
         .collect();
     for t in &tags {
         let _ = newrec.remove_aux(t);
     }
+    for t in keep_tags {
+        if let Ok(v) = record.aux(t) {
+            let _ = newrec.push_aux(t, v);
+        }
+    }
 
     newrec.set_pos(start);
     newrec.set_insert_size(new_tlen);