Add methods to fix mate info on non-primaries and templates

fgpyo/sam/__init__.py

@@ -158,10 +158,12 @@
 import enum
 import io
 import sys
+from collections.abc import Collection
 from itertools import chain
 from pathlib import Path
 from typing import IO
 from typing import Any
+from typing import Callable
 from typing import Dict
 from typing import Iterable
 from typing import Iterator
@@ -176,6 +178,7 @@
 from pysam import AlignedSegment
 from pysam import AlignmentFile as SamFile
 from pysam import AlignmentHeader as SamHeader
+from typing_extensions import deprecated
 
 import fgpyo.io
 from fgpyo.collections import PeekableIterator
@@ -607,6 +610,161 @@
         return self.elements[index]
 
 
+@enum.unique
+class PairOrientation(enum.Enum):
+    """Enumerations of read pair orientations."""
+
+    FR = "FR"
+    """A pair orientation for forward-reverse reads ("innie")."""
+
+    RF = "RF"
+    """A pair orientation for reverse-forward reads ("outie")."""
+
+    TANDEM = "TANDEM"
+    """A pair orientation for tandem (forward-forward or reverse-reverse) reads."""
+
+    @classmethod
+    def from_recs(  # noqa: C901  # `from_recs` is too complex (11 > 10)
+        cls, rec1: AlignedSegment, rec2: Optional[AlignedSegment] = None
+    ) -> Optional["PairOrientation"]:
+        """Returns the pair orientation if both reads are mapped to the same reference sequence.
+
+        Args:
+            rec1: The first record in the pair.
+            rec2: The second record in the pair. If None, then mate info on `rec1` will be used.
+
+        See:
+            [`htsjdk.samtools.SamPairUtil.getPairOrientation()`](https://github.com/samtools/htsjdk/blob/c31bc92c24bc4e9552b2a913e52286edf8f8ab96/src/main/java/htsjdk/samtools/SamPairUtil.java#L71-L102)
+        """
+
+        if rec2 is None:
+            rec2_is_unmapped = rec1.mate_is_unmapped
+            rec2_reference_id = rec1.next_reference_id
+        else:
+            rec2_is_unmapped = rec2.is_unmapped
+            rec2_reference_id = rec2.reference_id
+
+        if rec1.is_unmapped or rec2_is_unmapped or rec1.reference_id != rec2_reference_id:
+            return None
+
+        if rec2 is None:
+            rec2_is_forward = rec1.mate_is_forward
+            rec2_reference_start = rec1.next_reference_start
+        else:
+            rec2_is_forward = rec2.is_forward
+            rec2_reference_start = rec2.reference_start
+
+        if rec1.is_forward is rec2_is_forward:
+            return PairOrientation.TANDEM
+        if rec1.is_forward and rec1.reference_start <= rec2_reference_start:
+            return PairOrientation.FR
+        if rec1.is_reverse and rec2_reference_start < rec1.reference_end:
+            return PairOrientation.FR
+        if rec1.is_reverse and rec2_reference_start >= rec1.reference_end:
+            return PairOrientation.RF
+
+        if rec2 is None:
+            if not rec1.has_tag("MC"):
+                raise ValueError('Cannot determine proper pair status without a mate cigar ("MC")!')
+            rec2_cigar = Cigar.from_cigarstring(str(rec1.get_tag("MC")))
+            rec2_reference_end = rec1.next_reference_start + rec2_cigar.length_on_target()
+        else:
+            rec2_reference_end = rec2.reference_end
+
+        if rec1.reference_start < rec2_reference_end:
+            return PairOrientation.FR
+        else:
+            return PairOrientation.RF
+
+
+def isize(rec1: AlignedSegment, rec2: Optional[AlignedSegment] = None) -> int:
+    """Computes the insert size ("template length" or "TLEN") for a pair of records.
+
+    Args:
+        rec1: The first record in the pair.
+        rec2: The second record in the pair. If None, then mate info on `rec1` will be used.
+    """
+    if rec2 is None:
+        rec2_is_unmapped = rec1.mate_is_unmapped
+        rec2_reference_id = rec1.next_reference_id
+    else:
+        rec2_is_unmapped = rec2.is_unmapped
+        rec2_reference_id = rec2.reference_id
+
+    if rec1.is_unmapped or rec2_is_unmapped or rec1.reference_id != rec2_reference_id:
+        return 0
+
+    if rec2 is None:
+        rec2_is_forward = rec1.mate_is_forward
+        rec2_reference_start = rec1.next_reference_start
+    else:
+        rec2_is_forward = rec2.is_forward
+        rec2_reference_start = rec2.reference_start
+
+    if rec1.is_forward and rec2_is_forward:
+        return rec2_reference_start - rec1.reference_start
+    if rec1.is_reverse and rec2_is_forward:
+        return rec2_reference_start - rec1.reference_end
+
+    if rec2 is None:
+        if not rec1.has_tag("MC"):
+            raise ValueError('Cannot determine proper pair status without a mate cigar ("MC")!')
+        rec2_cigar = Cigar.from_cigarstring(str(rec1.get_tag("MC")))
+        rec2_reference_end = rec1.next_reference_start + rec2_cigar.length_on_target()
+    else:
+        rec2_reference_end = rec2.reference_end
+
+    if rec1.is_forward:
+        return rec2_reference_end - rec1.reference_start
+    else:
+        return rec2_reference_end - rec1.reference_end
+
+
+DefaultProperlyPairedOrientations: set[PairOrientation] = {PairOrientation.FR}
+"""The default orientations for properly paired reads."""
+
+
+def is_proper_pair(
+    rec1: AlignedSegment,
+    rec2: Optional[AlignedSegment] = None,
+    max_insert_size: int = 1000,
+    orientations: Collection[PairOrientation] = DefaultProperlyPairedOrientations,
+    isize: Callable[[AlignedSegment, AlignedSegment], int] = isize,
+) -> bool:
+    """Determines if a pair of records are properly paired or not.
+
+    Criteria for records in a proper pair are:
+        - Both records are aligned
+        - Both records are aligned to the same reference sequence
+        - The pair orientation of the records is one of the valid pair orientations (default "FR")
+        - The inferred insert size is not more than a maximum length (default 1000)
+
+    Args:
+        rec1: The first record in the pair.
+        rec2: The second record in the pair. If None, then mate info on `rec1` will be used.
+        max_insert_size: The maximum insert size to consider a pair "proper".
+        orientations: The valid set of orientations to consider a pair "proper".
+        isize: A function that takes the two alignments and calculates their isize.
+
+    See:
+        [`htsjdk.samtools.SamPairUtil.isProperPair()`](https://github.com/samtools/htsjdk/blob/c31bc92c24bc4e9552b2a913e52286edf8f8ab96/src/main/java/htsjdk/samtools/SamPairUtil.java#L106-L125)
+    """
+    if rec2 is None:
+        rec2_is_mapped = rec1.mate_is_mapped
+        rec2_reference_id = rec1.next_reference_id
+    else:
+        rec2_is_mapped = rec2.is_mapped
+        rec2_reference_id = rec2.reference_id
+
+    return (
+        rec1.is_mapped
+        and rec2_is_mapped
+        and rec1.reference_id == rec2_reference_id
+        and PairOrientation.from_recs(rec1=rec1, rec2=rec2) in orientations
+        and 0 < abs(isize(rec1, rec2)) <= max_insert_size
+    )
+
+
 @attr.s(frozen=True, auto_attribs=True)
 class SupplementaryAlignment:
     """Stores a supplementary alignment record produced by BWA and stored in the SA SAM tag.
@@ -688,48 +846,145 @@
             return []
 
 
-def isize(r1: AlignedSegment, r2: AlignedSegment) -> int:
-    """Computes the insert size for a pair of records."""
-    if r1.is_unmapped or r2.is_unmapped or r1.reference_id != r2.reference_id:
-        return 0
-    else:
-        r1_pos = r1.reference_end if r1.is_reverse else r1.reference_start
-        r2_pos = r2.reference_end if r2.is_reverse else r2.reference_start
-        return r2_pos - r1_pos
+def sum_of_base_qualities(rec: AlignedSegment, min_quality_score: int = 15) -> int:
+    """Calculate the sum of base qualities score for an alignment record.
+
+    This function is useful for calculating the "mate score" as implemented in samtools fixmate.
+
+    Args:
+        rec: The alignment record to calculate the sum of base qualities from.
+        min_quality_score: The minimum base quality score to use for summation.
+
+    See:
+        [`calc_sum_of_base_qualities()`](https://github.com/samtools/samtools/blob/4f3a7397a1f841020074c0048c503a01a52d5fa2/bam_mate.c#L227-L238)
+        [`MD_MIN_QUALITY`](https://github.com/samtools/samtools/blob/4f3a7397a1f841020074c0048c503a01a52d5fa2/bam_mate.c#L42)
+    """
+    score: int = sum(qual for qual in rec.query_qualities if qual >= min_quality_score)
+    return score
+
+
+def _set_common_mate_fields(dest: AlignedSegment, source: AlignedSegment) -> None:
+    """Set common mate info an alignment to its mate's primary alignment.
+
+    Args:
+        dest: The alignment to set the mate info upon.
+        source: The primary alignment to use as a mate reference.
+    """
+    if source.is_read1 is dest.is_read1 or source.is_secondary or source.is_supplementary:
+        raise ValueError("Mate info must be set from a primary of the next ordinal!")
+    if source.query_name != dest.query_name:
+        raise ValueError("Cannot set mate info on alignments with different query names!")
+
+    dest.next_reference_id = source.reference_id
+    dest.next_reference_name = source.reference_name
+    dest.next_reference_start = source.reference_start
+    dest.mate_is_forward = source.is_forward
+    dest.mate_is_mapped = source.is_mapped
+    dest.set_tag("MC", source.cigarstring)
+    dest.set_tag("MQ", source.mapping_quality)
+    dest.set_tag("ms", sum_of_base_qualities(source))
+
+
+def set_mate_info(
+    r1: AlignedSegment,
+    r2: AlignedSegment,
+    is_proper_pair: Callable[[AlignedSegment, AlignedSegment], bool] = is_proper_pair,
+) -> None:
+    """Resets mate pair information between reads in a pair.
+
+    Args:
+        r1: Read 1 (first read in the template).
+        r2: Read 2 with the same query name as r1 (second read in the template).
+        is_proper_pair: A function that takes the two alignments and determines proper pair status.
+    """
+    if r1.query_name != r2.query_name:
+        raise ValueError("Cannot set mate info on alignments with different query names!")
+
+    for dest, source in [(r1, r2), (r2, r1)]:
+        _set_common_mate_fields(dest=dest, source=source)
+
+    template_length = isize(r1, r2)
+    r1.template_length = template_length
+    r2.template_length = -template_length
+
+    proper_pair = is_proper_pair(r1, r2)
+    r1.is_proper_pair = proper_pair
+    r2.is_proper_pair = proper_pair
+
+
+def set_mate_info_on_secondary(
+    secondary: AlignedSegment,
+    mate_primary: AlignedSegment,
+    is_proper_pair: Callable[[AlignedSegment, AlignedSegment], bool] = is_proper_pair,
+) -> None:
+    """Set mate info on a secondary alignment to its mate's primary alignment.
+
+    Args:
+        secondary: The secondary alignment to set mate information upon.
+        mate_primary: The primary alignment of the secondary's mate.
+        is_proper_pair: A function that takes the two alignments and determines proper pair status.
+
+    Raises:
+        ValueError: If primary and secondary are of the same read ordinal.
+        ValueError: If the primary is marked as either secondary or supplementary.
+        ValueError: If the secondary is not marked as secondary.
+    """
+    if not secondary.is_secondary:
+        raise ValueError("Cannot set mate info on an alignment not marked as secondary!")
+
+    _set_common_mate_fields(dest=secondary, source=mate_primary)
+
+    # NB: calculate isize and proper pair as if this secondary alignment was the primary alignment.
+    secondary.is_proper_pair = is_proper_pair(mate_primary, secondary)
+    secondary.template_length = isize(mate_primary, secondary)
+
+
+def set_mate_info_on_supplementary(supp: AlignedSegment, mate_primary: AlignedSegment) -> None:
+    """Set mate info on a supplementary alignment to its mate's primary alignment.
+
+    Args:
+        supp: The supplementary alignment to set mate information upon.
+        mate_primary: The primary alignment of the supplementary's mate.
+
+    Raises:
+        ValueError: If primary and secondary are of the same read ordinal.
+        ValueError: If the primary is marked as either secondary or supplementary.
+        ValueError: If the secondary is not marked as secondary.
+    """
+    if not supp.is_supplementary:
+        raise ValueError("Cannot set mate info on an alignment not marked as supplementary!")
+
+    _set_common_mate_fields(dest=supp, source=mate_primary)
 
+    # NB: for a non-secondary supplemental alignment, set the following to the same as the primary.
+    if not supp.is_secondary:
+        supp.is_proper_pair = mate_primary.is_proper_pair
+        supp.template_length = -mate_primary.template_length
 
+
+@deprecated("Use `set_mate_info()` instead. Deprecated after fgpyo 0.8.0.")
 def set_pair_info(r1: AlignedSegment, r2: AlignedSegment, proper_pair: bool = True) -> None:
-    """Resets mate pair information between reads in a pair. Requires that both r1
-    and r2 are mapped.  Can be handed reads that already have pairing flags setup or
-    independent R1 and R2 records that are currently flagged as SE reads.
+    """Resets mate pair information between reads in a pair.
+
+    Requires that both r1 and r2 are mapped. Can be handed reads that already have pairing flags
+    setup or independent R1 and R2 records that are currently flagged as SE reads.
 
     Args:
-        r1: read 1
-        r2: read 2 with the same queryname as r1
+        r1: Read 1 (first read in the template).
+        r2: Read 2 with the same query name as r1 (second read in the template).
+        proper_pair: whether the pair is proper or not.
     """
-    assert not r1.is_unmapped, f"Cannot process unmapped mate {r1.query_name}/1"
-    assert not r2.is_unmapped, f"Cannot process unmapped mate {r2.query_name}/2"
-    assert r1.query_name == r2.query_name, "Attempting to pair reads with different qnames."
+    if r1.query_name != r2.query_name:
+        raise ValueError("Cannot pair reads with different query names!")
 
     for r in [r1, r2]:
         r.is_paired = True
-        r.is_proper_pair = proper_pair
 
     r1.is_read1 = True
     r1.is_read2 = False
     r2.is_read2 = True
     r2.is_read1 = False
-
-    for src, dest in [(r1, r2), (r2, r1)]:
-        dest.next_reference_id = src.reference_id
-        dest.next_reference_start = src.reference_start
-        dest.mate_is_reverse = src.is_reverse
-        dest.mate_is_unmapped = False
-        dest.set_tag("MC", src.cigarstring)
-
-    insert_size = isize(r1, r2)
-    r1.template_length = insert_size
-    r2.template_length = -insert_size
+    set_mate_info(r1=r1, r2=r2, is_proper_pair=lambda a, b: proper_pair)
 
 
 @attr.s(frozen=True, auto_attribs=True)
@@ -971,6 +1226,10 @@
             for rec in recs:
                 yield rec
 
+    def set_mate_info(self) -> "Template":
+        """Reset all mate information on every record in a template."""
+        return set_mate_info_for_template(self)
+
     def write_to(
         self,
         writer: SamFile,
@@ -1029,6 +1288,31 @@
         return Template.build(recs, validate=False)
 
 
+def set_mate_info_for_template(
+    template: Template,
+    is_proper_pair: Callable[[AlignedSegment, AlignedSegment], bool] = is_proper_pair,
+) -> Template:
+    """Reset all mate information on every record in a template.
+
+    Args:
+        template: The template of alignments to reset all mate information on.
+        is_proper_pair: A function that takes two alignments and determines proper pair status.
+    """
+    if template.r1 is not None and template.r2 is not None:
+        set_mate_info(template.r1, template.r2, is_proper_pair=is_proper_pair)
+    if template.r1 is not None:
+        for rec in template.r2_secondaries:
+            set_mate_info_on_secondary(rec, template.r1, is_proper_pair=is_proper_pair)
+        for rec in template.r2_supplementals:
+            set_mate_info_on_supplementary(rec, template.r1)
+    if template.r2 is not None:
+        for rec in template.r1_secondaries:
+            set_mate_info_on_secondary(rec, template.r2, is_proper_pair=is_proper_pair)
+        for rec in template.r1_supplementals:
+            set_mate_info_on_supplementary(rec, template.r2)
+    return template
+
+
 class SamOrder(enum.Enum):
     """
     Enumerations of possible sort orders for a SAM file.