com.fulcrumgenomics.umi.GroupReadsByUmi.scala Maven / Gradle / Ivy
The newest version!
/**
* Copyright (c) 2016, Fulcrum Genomics LLC
* All rights reserved.
*
* Redistribution and use in source and binary forms, with or without
* modification, are permitted provided that the following conditions are met:
*
* 1. Redistributions of source code must retain the above copyright notice,
* this list of conditions and the following disclaimer.
*
* 2. Redistributions in binary form must reproduce the above copyright notice,
* this list of conditions and the following disclaimer in the documentation
* and/or other materials provided with the distribution.
*
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
* ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
* LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
* CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
* SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
* INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
* CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
* ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
* POSSIBILITY OF SUCH DAMAGE.
*/
package com.fulcrumgenomics.umi
import com.fulcrumgenomics.FgBioDef._
import com.fulcrumgenomics.bam.api.SamOrder.TemplateCoordinate
import com.fulcrumgenomics.bam.api.{SamOrder, SamRecord, SamSource, SamWriter}
import com.fulcrumgenomics.bam.{Bams, Template}
import com.fulcrumgenomics.cmdline.{ClpGroups, FgBioTool}
import com.fulcrumgenomics.commons.util.{LazyLogging, NumericCounter, SimpleCounter}
import com.fulcrumgenomics.sopt.{arg, clp}
import com.fulcrumgenomics.umi.GroupReadsByUmi._
import com.fulcrumgenomics.util.Metric.{Count, Proportion}
import com.fulcrumgenomics.util.Sequences.countMismatches
import com.fulcrumgenomics.util._
import enumeratum.EnumEntry
import htsjdk.samtools.DuplicateScoringStrategy.ScoringStrategy
import htsjdk.samtools._
import htsjdk.samtools.util.SequenceUtil
import java.util.concurrent.Executors
import java.util.concurrent.atomic.AtomicLong
import scala.collection.immutable.IndexedSeq
import scala.collection.mutable.ListBuffer
import scala.collection.parallel.ExecutionContextTaskSupport
import scala.collection.{BufferedIterator, Iterator, mutable}
import scala.concurrent.ExecutionContext
object GroupReadsByUmi {
/** A type representing an actual UMI that is a string of DNA sequence. */
type Umi = String
/** A type to represent a unique ID assigned to a molecule. */
type MoleculeId = String
/** The suffix appended to molecular identifier reads from the "top strand" of the duplex molecule. These reads have
* 5' coordinate for read 1 less than or equal to the 5' coordinate for read 2. */
val TopStrandDuplex = "/A"
/** The suffix appended to molecular identifier reads from the "bottom strand" of the duplex molecule. These reads have
* 5' coordinate for read 1 greater than the 5' coordinate for read 2. */
val BottomStrandDuplex = "/B"
private val ReadInfoTempAttributeName = "__GRBU_ReadInfo"
/** A case class to represent all the information we need to order reads for duplicate marking / grouping. */
case class ReadInfo(refIndex1: Int,
start1: Int,
strand1: Byte,
refIndex2: Int,
start2: Int,
strand2: Byte,
library: String)
object ReadInfo {
// Use the Max possible value for ref/pos/strand when we have unmapped reads to mirror what's done
// in the TemplateCoordinate sort order where we sort by the _lower_ of the two mates' positions
// and want to use the mapped position when one mate is unmapped
private val UnknownRef = Int.MaxValue
private val UnknownPos = Int.MaxValue
private val UnknownStrand = Byte.MaxValue
/** Looks in all the places the library name can be hiding. Returns the library name
* if one is found, otherwise returns "unknown".
*/
private def library(rec: SamRecord): String = {
val rg = rec.readGroup
if (rg != null && rg.getLibrary != null) rg.getLibrary else "unknown"
}
/** Extract a ReadInfo from a SamRecord; mate cigar must be present if a mate exists and is mapped. */
def apply(rec: SamRecord) : ReadInfo =
rec.transientAttrs.getOrElseUpdate[ReadInfo](GroupReadsByUmi.ReadInfoTempAttributeName, {
val lib = library(rec)
val (ref1, start1, strand1) = positionOf(rec)
val (ref2, start2, strand2) = if (rec.unpaired || rec.mateUnmapped) (UnknownRef, UnknownPos, UnknownStrand) else {
val start = (if (rec.matePositiveStrand) rec.mateUnclippedStart else rec.mateUnclippedEnd)
.getOrElse(throw new IllegalStateException(s"Missing mate cigar tag (MC) for read $rec."))
(rec.mateRefIndex, start, strandToByte(rec.matePositiveStrand))
}
val r1Earlier =
(ref1 < ref2) ||
(ref1 == ref2 && start1 < start2) ||
(ref1 == ref2 && start1 == start2 && strand1 < strand2)
if (r1Earlier) new ReadInfo(ref1, start1, strand1, ref2, start2, strand2, lib)
else new ReadInfo(ref2, start2, strand2, ref1, start1, strand1, lib)
})
/** Extract a ReadInfo from a Template object. R1 primary must be present. */
def apply(t: Template): ReadInfo = {
val r1 = t.r1.getOrElse(throw new IllegalStateException(s"${t.name} did not have a primary R1 record."))
val r2 = t.r2
r1.transientAttrs.getOrElseUpdate[ReadInfo](GroupReadsByUmi.ReadInfoTempAttributeName, {
val lib = library(r1)
val (ref1, start1, strand1) = positionOf(r1)
val (ref2, start2, strand2) = r2.map(positionOf).getOrElse((UnknownRef, UnknownPos, UnknownStrand))
val r1Earlier =
(ref1 < ref2) ||
(ref1 == ref2 && start1 < start2) ||
(ref1 == ref2 && start1 == start2 && strand1 < strand2)
if (r1Earlier) new ReadInfo(ref1, start1, strand1, ref2, start2, strand2, lib)
else new ReadInfo(ref2, start2, strand2, ref1, start1, strand1, lib)
})
}
/** Extracts the refIndex, unclipped 5' end and strand of the read, or Unknown values for unmapped reads. */
private def positionOf(r: SamRecord): (Int, Int, Byte) = {
if (r.unmapped) (UnknownRef, UnknownPos, UnknownStrand) else {
val start = if (r.positiveStrand) r.unclippedStart else r.unclippedEnd
(r.refIndex, start, strandToByte(r.positiveStrand))
}
}
// Encodes a known strand into a byte value
@inline private def strandToByte(positive: Boolean): Byte = if (positive) 0 else 1
}
/** Trait that can be implemented to provide a UMI assignment strategy. */
private[umi] sealed trait UmiAssigner {
private val counter = new AtomicLong()
/** Take in a sequence of UMIs and assign each UMI to a unique UMI group ID. */
def assign(rawUmis: Seq[Umi]) : Map[Umi, MoleculeId]
/** Returns true if the two UMIs are the same. */
def isSameUmi(a: Umi, b: Umi): Boolean = a == b
/** Returns a canonical form of the UMI that is the same for all reads with the same UMI. */
def canonicalize(u: Umi): Umi = u
/** Default implementation of a method to retrieve the next ID based on a counter. */
protected def nextId: MoleculeId = this.counter.getAndIncrement().toString
/** Takes in a map of UMI to "sentinel" UMI, and outputs a map of UMI -> Molecule ID. */
protected def assignIds(assignments: Map[Umi,Umi]): Map[Umi,MoleculeId] = {
val idMap = assignments.values.toSet.map((umi:Umi) => umi -> nextId).toMap
assignments.map { case (k,v) => (k, idMap(v)) }
}
}
/**
* Assigns UMIs based solely on sequence identity.
*/
private[umi] class IdentityUmiAssigner extends UmiAssigner {
override def assign(rawUmis: Seq[Umi]): Map[Umi, MoleculeId] = assignIds(rawUmis.map(umi => umi -> umi.toUpperCase).toMap)
}
/**
* Assigns UMIs based solely on sequence identity.
*/
private[umi] class SimpleErrorUmiAssigner(val maxMismatches: Int) extends UmiAssigner {
override def assign(rawUmis: Seq[Umi]): Map[Umi, MoleculeId] = {
type UmiSet = mutable.SortedSet[Umi]
val umiSets = ListBuffer[UmiSet]()
for (umi <- rawUmis) {
val matchedSets = mutable.Set[UmiSet]()
for (set <- umiSets) {
if (set.exists(other => countMismatches(other, umi) <= maxMismatches)) {
set.add(umi)
matchedSets.add(set)
}
}
matchedSets.size match {
case 0 => umiSets += mutable.SortedSet(umi)
case 1 => matchedSets.head.add(umi)
case _ =>
// Merge the sets and remove all but one from umiSets
val pick = matchedSets.head
matchedSets.tail.foreach(set => {
umiSets -= set
pick ++= set
})
}
}
assignIds(umiSets.flatMap(set => set.map(umi => umi -> set.head)).toMap)
}
}
/**
* Class that implements the directed adjacency graph method from umi_tools.
* See: https://github.com/CGATOxford/UMI-tools
*/
private[umi] class AdjacencyUmiAssigner(final val maxMismatches: Int, val threads: Int = 1) extends UmiAssigner {
private val taskSupport = if (threads < 2) None else {
val ctx = ExecutionContext.fromExecutorService(Executors.newFixedThreadPool(threads))
Some(new ExecutionContextTaskSupport(ctx))
}
/** Represents a node in the adjacency graph; equality is just by UMI sequence. */
final class Node(val umi: Umi, val count: Int, val children: mutable.Buffer[Node] = mutable.Buffer(), var assigned: Boolean = false) {
/** Gets the full set of descendants from this node. */
def descendants: List[Node] = {
val buffer = ListBuffer[Node]()
addDescendants(buffer)
buffer.toList
}
private def addDescendants(buffer: mutable.Buffer[Node]): Unit = children.foreach(child => {
buffer += child
child.addDescendants(buffer)
})
override def equals(other: scala.Any): Boolean = other.isInstanceOf[Node] && this.umi == other.asInstanceOf[Node].umi
override def toString: String = s"$umi [$count]"
}
/** Returns the count of each raw UMI that was observed. */
protected def count(umis: Seq[Umi]): Iterator[(Umi,Long)] = SimpleCounter(umis).iterator
/** Returns whether or not a pair of UMIs match closely enough to be considered adjacent in the graph. */
protected def matches(lhs: Umi, rhs: Umi): Boolean = {
val len = lhs.length
var idx = 0
var mismatches = 0
val tooManyMismatches = this.maxMismatches + 1
while (mismatches < tooManyMismatches && idx < len) {
if (lhs.charAt(idx) != rhs.charAt(idx)) mismatches += 1
idx += 1
}
mismatches <= maxMismatches
}
/** Assigns IDs to each UMI based on the root to which is it mapped. */
protected def assignIdsToNodes(roots: Seq[Node]): Map[Umi, MoleculeId] = {
val mappings = Map.newBuilder[Umi,MoleculeId]
roots.foreach(root => {
val id = nextId
mappings += ((root.umi, id))
root.descendants.foreach(child => mappings += ((child.umi, id)))
})
mappings.result()
}
override def assign(rawUmis: Seq[Umi]): Map[Umi, MoleculeId] = {
// Make a list of counts of all UMIs in order from most to least abundant; we'll consume from this buffer
val orderedNodes = count(rawUmis).map{ case(umi,count) => new Node(umi, count.toInt) }.toIndexedSeq.sortBy((n:Node) => -n.count)
if (orderedNodes.length == 1) {
orderedNodes.head.assigned = true
assignIdsToNodes(orderedNodes)
}
else {
val umiLength = orderedNodes.head.umi.length
require(orderedNodes.forall(_.umi.length == umiLength), f"Multiple UMI lengths: ${orderedNodes.map(_.umi).mkString(", ")}")
val lookup = countIndexLookup(orderedNodes) // Seq of (count, firstIdx) pairs
// A list of all the root UMIs/Nodes that we find
val roots = Seq.newBuilder[Node]
// Now build one or more graphs starting with the most abundant remaining umi
val working = mutable.Queue[Node]()
forloop (from=0, until=orderedNodes.length) { rootIdx =>
val nextRoot = orderedNodes(rootIdx)
if (!nextRoot.assigned) {
roots += nextRoot
working.enqueue(nextRoot)
while (working.nonEmpty) {
val root = working.dequeue()
root.assigned = true
val maxChildCountPlusOne = (root.count / 2 + 1) + 1
val searchFromIdx = lookup
.find { case (count, _) => count < maxChildCountPlusOne }
.map { case (_, idx) => idx }
.getOrElse(-1)
if (searchFromIdx >= 0) {
val hits = taskSupport match {
case None =>
orderedNodes
.drop(searchFromIdx)
.filter(other => !other.assigned && matches(root.umi, other.umi))
case Some(ts) =>
orderedNodes
.drop(searchFromIdx)
.parWith(ts)
.filter(other => !other.assigned && matches(root.umi, other.umi))
.seq
}
root.children ++= hits
working.enqueueAll(hits)
hits.foreach(_.assigned = true)
}
}
}
}
assignIdsToNodes(roots.result())
}
}
/**
* Generates an indexed seq to enable fast identification of the first index in `nodes` where
* a given UMI count is observed. Assumes that the input is sorted from most abundant to least
* abundant. The generated Seq will contain one entry for every unique count seen; the second value
* in the tuple is the first index in the list of input nodes with that count.
*
* E.g. given nodes with counts [10, 10, 10, 9, 3, 3, 3, 3, 2, 1], the output would be:
* [(10, 0), (9, 3), (3, 4) (2, 8), (1, 9)]
*/
private def countIndexLookup(nodes: IndexedSeq[Node]): IndexedSeq[(Int, Int)] = {
val builder = IndexedSeq.newBuilder[(Int, Int)]
val iter = nodes.iterator.zipWithIndex.bufferBetter
while (iter.hasNext) {
val (currNode, currIdx) = iter.next
builder += ((currNode.count, currIdx))
iter.dropWhile { case (node, _) => node.count == currNode.count}
}
builder.result()
}
}
/**
* Version of the adjacency assigner that works for paired UMIs stored as a single tag of
* the form A-B where reads with A-B and B-A are related but not identical.
*
* @param maxMismatches the maximum number of mismatches between UMIs
*/
class PairedUmiAssigner(maxMismatches: Int, threads: Int = 1) extends AdjacencyUmiAssigner(maxMismatches, threads) {
/** String that is prefixed onto the UMI from the read with that maps to a lower coordinate in the genome.. */
private[umi] val lowerReadUmiPrefix: String = ("a" * (maxMismatches+1)) + ":"
/** String that is prefixed onto the UMI from the read with that maps to a higher coordinate in the genome.. */
private[umi] val higherReadUmiPrefix: String = ("b" * (maxMismatches+1)) + ":"
/** Returns true if the two UMIs are the same. */
final override def isSameUmi(a: Umi, b: Umi): Boolean = {
if (a == b) true else {
// same as `a == reverse(b)` but more efficient than creating new strings
val (a1, a2) = split(a)
val (b1, b2) = split(b)
a1 == b2 && a2 == b1
}
}
/** Returns the UMI with the lexically lower half first. */
override def canonicalize(u: Umi): Umi = {
val (a, b) = split(u)
if (a < b) u else s"${b}-${a}"
}
/** Splits the paired UMI into its two parts. */
@inline private def split(umi: Umi): (Umi, Umi) = {
val index = umi.indexOf('-')
if (index == -1) throw new IllegalStateException(s"UMI $umi is not a paired UMI.")
val first = umi.substring(0, index)
val second = umi.substring(index+1, umi.length)
(first, second)
}
/** Takes a UMI of the form "A-B" and returns "B-A". */
def reverse(umi: Umi): Umi = {
val (first, second) = split(umi)
s"${second}-${first}"
}
/** Turns each UMI into the lexically earlier of A-B or B-A and then counts them. */
override protected def count(umis: Seq[Umi]): Iterator[(Umi, Long)] = {
val counter = new SimpleCounter[Umi]()
umis.foreach { umi =>
val reversed = reverse(umi)
val lower = if (umi.compare(reversed) < 0) umi else reversed
counter.count(lower)
}
counter.iterator
}
/** Returns whether or not a paired-UMI matches within mismatch tolerance. */
override protected def matches(lhs: Umi, rhs: Umi): Boolean = super.matches(lhs, rhs) || super.matches(reverse(lhs), rhs)
/** Takes in a map of UMI to "sentinel" UMI, and outputs a map of UMI -> Molecule ID. */
override protected def assignIdsToNodes(roots: Seq[Node]): Map[Umi, MoleculeId] = {
val mappings = mutable.Buffer[(Umi,MoleculeId)]()
roots.foreach(root => {
val id = nextId
val ab = id + GroupReadsByUmi.TopStrandDuplex
val ba = id + GroupReadsByUmi.BottomStrandDuplex
mappings.append((root.umi, ab))
mappings.append((reverse(root.umi), ba))
root.descendants.foreach { child =>
val childUmi = child.umi
val childUmiRev = reverse(child.umi)
// If the root UMI and child UMI are more similar then presumably they originate
// from the same pairing of UMIs, otherwise if the root UMI is more similar to the
// reversed child UMI, they are paired with each other's inverse
if (countMismatches(root.umi, childUmi) < countMismatches(root.umi, childUmiRev)) {
mappings.append((childUmi, ab))
mappings.append((childUmiRev, ba))
}
else {
mappings.append((childUmi, ba))
mappings.append((childUmiRev, ab))
}
}
})
mappings.toMap
}
/** Since we generate mappings for A-B and B-A whether we've seen both or not, we override
* the main assignment method to filter out the ones that shouldn't be in the Map.
*/
override def assign(rawUmis: Seq[Umi]): Map[Umi, MoleculeId] = {
super.assign(rawUmis).filter { case (umi, _) => rawUmis.contains(umi) }
}
}
}
/**
* Metrics produced by `GroupReadsByUmi` to describe the distribution of tag family sizes
* observed during grouping.
*
* @param family_size The family size, or number of templates/read-pairs belonging to the family.
* @param count The number of families (or source molecules) observed with `family_size` observations.
* @param fraction The fraction of all families of all sizes that have this specific `family_size`.
* @param fraction_gt_or_eq_family_size The fraction of all families that have `>= family_size`.
*/
case class TagFamilySizeMetric(family_size: Int,
count: Count,
var fraction: Proportion = 0,
var fraction_gt_or_eq_family_size: Proportion = 0) extends Metric
/** The strategies implemented by [[GroupReadsByUmi]] to identify reads from the same source molecule.*/
sealed trait Strategy extends EnumEntry {
def newStrategy(edits: Int, threads: Int): UmiAssigner
}
object Strategy extends FgBioEnum[Strategy] {
def values: IndexedSeq[Strategy] = findValues
/** Strategy to only reads with identical UMI sequences are grouped together. */
case object Identity extends Strategy {
def newStrategy(edits: Int = 0, threads: Int = 0): UmiAssigner = {
require(edits == 0, "Edits should be zero when using the identity UMI assigner.")
new IdentityUmiAssigner
}
}
/** Strategy to cluster reads into groups based on mismatches between reads in clusters. */
case object Edit extends Strategy { def newStrategy(edits: Int, threads: Int = 0): UmiAssigner = new SimpleErrorUmiAssigner(edits) }
/** Strategy based on the directed adjacency method described in [umi_tools](http://dx.doi.org/10.1101/051755)
* that allows for errors between UMIs but only when there is a count gradient.
*/
case object Adjacency extends Strategy { def newStrategy(edits: Int, threads: Int = 1): UmiAssigner = new AdjacencyUmiAssigner(edits, threads) }
/** Strategy similar to the [[Adjacency]] strategy similar to adjacency but for methods that produce template with a
* pair of UMIs such that a read with A-B is related to but not identical to a read with B-A.
*/
case object Paired extends Strategy { def newStrategy(edits: Int, threads: Int = 1): UmiAssigner = new PairedUmiAssigner(edits, threads)}
}
@clp(group=ClpGroups.Umi, description =
"""
|Groups reads together that appear to have come from the same original molecule. Reads
|are grouped by template, and then templates are sorted by the 5' mapping positions of
|the reads from the template, used from earliest mapping position to latest. Reads that
|have the same end positions are then sub-grouped by UMI sequence.
|
|Accepts reads in any order (including `unsorted`) and outputs reads sorted by:
|
| 1. The lower genome coordinate of the two outer ends of the templates
| 2. The sequencing library
| 3. The assigned UMI tag
| 4. Read Name
|
|If the input is not template-coordinate sorted (i.e. `SO:unsorted GO:query SS:unsorted:template-coordinate`), then
|this tool will re-sort the input. The output will be written in template-coordinate order.
|
|During grouping, reads and templates are filtered out as follows:
|
|1. Templates are filtered if all reads for the template are unmapped
|2. Templates are filtered if any non-secondary, non-supplementary read has mapping quality < `min-map-q`
|3. Templates are filtered if R1 and R2 are mapped to different chromosomes and `--allow-inter-contig` is false
|4. Templates are filtered if any UMI sequence contains one or more `N` bases
|5. Templates are filtered if `--min-umi-length` is specified and the UMI does not meet the length requirement
|6. Reads are filtered out if flagged as secondary and `--include-secondary` is false
|7. Reads are filtered out if flagged as supplementary and `--include-supplementary` is false
|
|Grouping of UMIs is performed by one of four strategies:
|
|1. **identity**: only reads with identical UMI sequences are grouped together. This strategy
| may be useful for evaluating data, but should generally be avoided as it will
| generate multiple UMI groups per original molecule in the presence of errors.
|2. **edit**: reads are clustered into groups such that each read within a group has at least
| one other read in the group with <= edits differences and there are inter-group
| pairings with <= edits differences. Effective when there are small numbers of
| reads per UMI, but breaks down at very high coverage of UMIs.
|3. **adjacency**: a version of the directed adjacency method described in [umi_tools](http://dx.doi.org/10.1101/051755)
| that allows for errors between UMIs but only when there is a count gradient.
|4. **paired**: similar to adjacency but for methods that produce template such that a read with A-B is related
| to but not identical to a read with B-A. Expects the UMI sequences to be stored in a single SAM
| tag separated by a hyphen (e.g. `ACGT-CCGG`) and allows for one of the two UMIs to be absent
| (e.g. `ACGT-` or `-ACGT`). The molecular IDs produced have more structure than for single
| UMI strategies and are of the form `{base}/{A|B}`. E.g. two UMI pairs would be mapped as
| follows AAAA-GGGG -> 1/A, GGGG-AAAA -> 1/B.
|
|Strategies `edit`, `adjacency`, and `paired` make use of the `--edits` parameter to control the matching of
|non-identical UMIs.
|
|By default, all UMIs must be the same length. If `--min-umi-length=len` is specified then reads that have a UMI
|shorter than `len` will be discarded, and when comparing UMIs of different lengths, the first len bases will be
|compared, where `len` is the length of the shortest UMI. The UMI length is the number of [ACGT] bases in the UMI
|(i.e. does not count dashes and other non-ACGT characters). This option is not implemented for reads with UMI pairs
|(i.e. using the paired assigner).
|
|If the `--mark-duplicates` option is given, reads will also have their duplicate flag set in the BAM file.
|Each tag-family is treated separately, and a single template within the tag family is chosen to be the "unique"
|template and marked as non-duplicate, while all other templates in the tag family are then marked as duplicate.
|One limitation of duplicate-marking mode, vs. e.g. Picard MarkDuplicates, is that read pairs with one unmapped read
|are duplicate-marked independently from read pairs with both reads mapped.
|
|Several parameters have different defaults depending on whether duplicates are being marked or not (all are
|directly settable on the command line):
|
| 1. `--min-map-q` defaults to 0 in duplicate marking mode and 1 otherwise
| 2. `--include-secondary` defaults to true in duplicate marking mode and false otherwise
| 3. `--include-supplementary` defaults to true in duplicate marking mode and false otherwise
|
|Multi-threaded operation is supported via the `--threads/-@` option. This only applies to the Adjacency and Paired
|strategies. Additionally the only operation that is multi-threaded is the comparisons of UMIs at the same genomic
|position. Running with e.g. `--threads 8` can provide a _substantial_ reduction in runtime when there are many
|UMIs observed at the same genomic location, such as can occur in amplicon sequencing or ultra-deep coverage data.
"""
)
class GroupReadsByUmi
(@arg(flag='i', doc="The input BAM file.") val input: PathToBam = Io.StdIn,
@arg(flag='o', doc="The output BAM file.") val output: PathToBam = Io.StdOut,
@arg(flag='f', doc="Optional output of tag family size counts.") val familySizeHistogram: Option[FilePath] = None,
@arg(flag='t', doc="The tag containing the raw UMI.") val rawTag: String = ConsensusTags.UmiBases,
@arg(flag='T', doc="The output tag for UMI grouping.") val assignTag: String = ConsensusTags.MolecularId,
@arg(flag='d', doc="Turn on duplicate marking mode.") val markDuplicates: Boolean = false,
@arg(flag='S', doc="Include secondary reads.") val includeSecondary: Option[Boolean] = None,
@arg(flag='U', doc="Include supplementary reads.") val includeSupplementary: Option[Boolean] = None,
@arg(flag='m', doc="Minimum mapping quality for mapped reads.") val minMapQ: Option[Int] = None,
@arg(flag='n', doc="Include non-PF reads.") val includeNonPfReads: Boolean = false,
@arg(flag='s', doc="The UMI assignment strategy.") val strategy: Strategy,
@arg(flag='e', doc="The allowable number of edits between UMIs.") val edits: Int = 1,
@arg(flag='l', doc= """The minimum UMI length. If not specified then all UMIs must have the same length,
|otherwise discard reads with UMIs shorter than this length and allow for differing UMI lengths.
|""")
val minUmiLength: Option[Int] = None,
@arg(flag='x', doc= """
|DEPRECATED: this option will be removed in future versions and inter-contig reads will be
|automatically processed.""")
@deprecated val allowInterContig: Boolean = true,
@arg(flag='@', doc="Number of threads to use when comparing UMIs. Only recommended for amplicon or similar data.") val threads: Int = 1,
)extends FgBioTool with LazyLogging {
import GroupReadsByUmi._
require(this.minUmiLength.forall(_ => this.strategy != Strategy.Paired), "Paired strategy cannot be used with --min-umi-length")
private val assigner = strategy.newStrategy(this.edits, this.threads)
// Give values to unset parameters that are different in duplicate marking mode
private val _minMapQ = this.minMapQ.getOrElse(if (this.markDuplicates) 0 else 1)
private val _includeSecondaryReads = this.includeSecondary.getOrElse(this.markDuplicates)
private val _includeSupplementaryReads = this.includeSupplementary.getOrElse(this.markDuplicates)
/** Checks that the read's mapq is over a minimum, and if the read is paired, that the mate mapq is also over the min. */
private def mapqOk(rec: SamRecord, minMapQ: Int): Boolean = {
if (rec.mapped && rec.mapq < minMapQ) {
false
}
else if (rec.unpaired || rec.mateUnmapped) {
true
}
else {
rec.get[Int](SAMTag.MQ.name()) match {
case None => fail(s"Mate mapping quality (MQ) tag not present on read ${rec.name}.")
case Some(mq) => mq >= minMapQ
}
}
}
// A handful of counters for tracking reads
private var (filteredNonPf, filteredPoorAlignment, filteredNsInUmi, filterUmisTooShort, kept) = (0L, 0L, 0L, 0L, 0L)
override def execute(): Unit = {
Io.assertReadable(input)
Io.assertCanWriteFile(output)
this.familySizeHistogram.foreach(f => Io.assertCanWriteFile(f))
val in = SamSource(input)
val header = in.header
val skipSorting = SamOrder(header).contains(TemplateCoordinate)
// True if the input will be sorted and no differences in UMIs are tolerated and the Molecular ID tag is MI,
// false otherwise. Pre-sorted (TemplateCoordinate sorted) input does not enable this optimization as we rely on
// copying the `RX` tag into the `MI` tag so that same raw UMIs are grouped together in the sorted stream of records.
//
// True here enables an optimization where, when bringing groups of reads into memory, we can _also_ group by UMI
// thus reducing the number of reads in memory. This is helpful since edits=0 is often used for data that has
// high numbers of reads with the same start/stop coordinates.
// We do this by setting the MI tag to the canonicalized (optionally truncated) UMI prior to sorting, so that
// reads with the same UMI are grouped together in the sorted stream of records.
val canTakeNextGroupByUmi = {
!skipSorting &&
(this.assignTag == ConsensusTags.MolecularId) &&
(this.edits == 0 || this.strategy == Strategy.Identity)
}
// Filter and sort the input BAM file
logger.info("Filtering the input.")
val filteredIterator = in.iterator
.filter(r => this._includeSecondaryReads || !r.secondary )
.filter(r => this._includeSupplementaryReads || !r.supplementary )
.filter(r => (includeNonPfReads || r.pf) || { filteredNonPf += 1; false })
.filter(r => (r.mapped || (r.paired && r.mateMapped)) || { filteredPoorAlignment += 1; false })
.filter(r => (allowInterContig || r.unpaired || r.refIndex == r.mateRefIndex) || { filteredPoorAlignment += 1; false })
.filter(r => mapqOk(r, this._minMapQ) || { filteredPoorAlignment += 1; false })
.filter(r => !r.get[String](rawTag).exists(_.contains('N')) || { filteredNsInUmi += 1; false })
.filter { r =>
this.minUmiLength.forall { l =>
r.get[String](this.rawTag).forall { umi =>
l <= umi.toUpperCase.count(c => SequenceUtil.isUpperACGTN(c.toByte))
}
} || { filterUmisTooShort += 1; false}
}
.tapEach { r =>
// If we're able to also group by the UMI because edits aren't allowed, push the trimmed, canonicalized UMI
// into the assign tag (which must be MI if canTakeNextGroupByUmi is true), since that is used by the
// SamOrder to sort the reads _and_ we'll overwrite it on the way out!
// Note that here we trim UMIs (if enabled) to the minimum UMI length for sorting, but that when doing the
// actual grouping later we go back to the raw tag (RX) and use as much of the UMI as possible.
if (canTakeNextGroupByUmi) {
val umi = this.assigner.canonicalize(r[String](rawTag).toUpperCase)
val truncated = this.minUmiLength match {
case None => umi
case Some(n) => umi.substring(0, n)
}
r(this.assignTag) = truncated
}
kept += 1
r
}
val tagFamilySizeCounter = new NumericCounter[Int]()
val outHeader = header.clone()
SamOrder.TemplateCoordinate.applyTo(outHeader)
val out = SamWriter(output, outHeader)
val templateCoordinateIterator = {
// Skip sorting if the input is already template coordinate!
val iter = if (skipSorting) filteredIterator else {
logger.info("Sorting the input to TemplateCoordinate order.")
val sorter = Bams.sorter(SamOrder.TemplateCoordinate, header)
val sortProgress = ProgressLogger(logger, verb="Sorted")
filteredIterator.foreach { rec =>
sortProgress.record(rec)
sorter += rec
}
sortProgress.logLast()
sorter.iterator
}
// Note: this should never have to sort, just groups into templates
Bams.templateIterator(iter, out.header, Bams.MaxInMemory, Io.tmpDir)
}
// If we sorted then log stats now, since the `sorter.iterator` above has consumed all the records fed to it,
// and so the stats should be ready
if (!skipSorting) logStats()
// Output the reads in the new ordering
logger.info("Assigning reads to UMIs and outputting.")
while (templateCoordinateIterator.hasNext) {
// Take the next set of templates by position and assign UMIs
val templates = takeNextGroup(templateCoordinateIterator, canTakeNextGroupByUmi=canTakeNextGroupByUmi)
assignUmiGroups(templates)
// Then group the records in the right order (assigned tag, read name, r1, r2)
val templatesByMi = templates.groupBy { t => t.r1.get.apply[String](this.assignTag) }
// If marking duplicates, assign bitflag to all duplicate reads
if (this.markDuplicates) {
templatesByMi.values.foreach(t => setDuplicateFlags(t))
}
// Then output the records in the right order (assigned tag, read name, r1, r2)
templatesByMi.keys.toSeq.sortBy(id => (id.length, id)).foreach(tag => {
templatesByMi(tag).sortBy(t => t.name).flatMap(t => t.primaryReads).foreach(rec => {
out += rec
})
})
// Count up the family sizes
templatesByMi.values.foreach(ps => tagFamilySizeCounter.count(ps.size))
}
out.close()
// If we skipped sorting, log stats after we have consumed all the records
if (skipSorting) logStats()
// Write out the family size histogram
this.familySizeHistogram match {
case None => ()
case Some(p) =>
val ms = tagFamilySizeCounter.toSeq.map { case (size, count) => TagFamilySizeMetric(size, count)}
val total = ms.map(_.count.toDouble).sum
ms.foreach(m => m.fraction = m.count / total)
ms.tails.foreach { tail => tail.headOption.foreach(m => m.fraction_gt_or_eq_family_size = tail.map(_.fraction).sum) }
Metric.write(p, ms)
}
}
private def logStats(): Unit = {
logger.info(f"Accepted $kept%,d reads for grouping.")
if (filteredNonPf > 0) logger.info(f"Filtered out $filteredNonPf%,d non-PF reads.")
logger.info(f"Filtered out $filteredPoorAlignment%,d reads due to mapping issues.")
logger.info(f"Filtered out $filteredNsInUmi%,d reads that contained one or more Ns in their UMIs.")
this.minUmiLength.foreach { _ => logger.info(f"Filtered out $filterUmisTooShort%,d reads that contained UMIs that were too short.") }
}
/** Consumes the next group of templates with all matching end positions and returns them. */
def takeNextGroup(iterator: BufferedIterator[Template], canTakeNextGroupByUmi: Boolean) : Seq[Template] = {
val first = iterator.next()
val firstEnds = ReadInfo(first)
val firstUmi = first.r1.get.apply[String](this.assignTag)
val builder = IndexedSeq.newBuilder[Template]
builder += first
while (
iterator.hasNext &&
firstEnds == ReadInfo(iterator.head) &&
// This last condition only works because we put a canonicalized UMI into rec(assignTag) if canTakeNextGroupByUmi
(!canTakeNextGroupByUmi || firstUmi == iterator.head.r1.get.apply[String](this.assignTag))
) {
builder += iterator.next()
}
builder.result()
}
/**
* Takes in templates and writes an attribute (specified by `assignTag`) denoting
* sub-grouping into UMI groups by original molecule.
*/
def assignUmiGroups(templates: Seq[Template]): Unit = {
val startMs = System.currentTimeMillis
val umis = truncateUmis(templates.map { t => umiForRead(t) })
val rawToId = this.assigner.assign(umis)
templates.iterator.zip(umis.iterator).foreach { case (template, umi) =>
val id = rawToId(umi)
template.primaryReads.foreach(r => r(this.assignTag) = id)
}
val endMs = System.currentTimeMillis()
val durationMs = endMs - startMs
if (durationMs >= 2500) {
logger.debug(f"Grouped ${rawToId.size}%,d UMIs from ${templates.size}%,d templates in ${durationMs}%,d ms." )
}
}
/** When a minimum UMI length is specified, truncates all the UMIs to the length of the shortest UMI. For the paired
* assigner, truncates the first UMI and second UMI separately.*/
private def truncateUmis(umis: Seq[Umi]): Seq[Umi] = this.minUmiLength match {
case None => umis
case Some(length) =>
this.assigner match {
case _: PairedUmiAssigner =>
throw new IllegalStateException("Cannot used the paired umi assigner when min-umi-length is defined.")
case _ =>
val minLength = umis.map(_.length).min
require(length <= minLength, s"Bug: UMI found that had shorter length than expected ($minLength < $length)")
umis.map(_.substring(0, minLength))
}
}
/**
* Retrieves the UMI to use for a template. In the case of single-umi strategies this is just
* the upper-case UMI sequence.
*
* For Paired the UMI is extended to encode which "side" of the template the read came from - the earlier
* or later read on the genome. This is necessary to ensure that, when the two paired UMIs are the same
* or highly similar, that the A vs. B groups are constructed correctly.
*/
private def umiForRead(t: Template): Umi = {
// Check that all the primary reads have the UMI defined
t.primaryReads.foreach { rec =>
val umi = rec.get[String](this.rawTag)
if (!umi.exists(_.nonEmpty)) fail(s"Record '$rec' was missing the raw UMI tag '${this.rawTag}'")
}
val umi = t.r1.getOrElse(fail(s"R1 must be present for ${t.name}")).apply[String](this.rawTag)
(t.r1, t.r2, this.assigner) match {
case (Some(r1), Some(r2), paired: PairedUmiAssigner) =>
if (!this.allowInterContig) require(r1.refIndex == r2.refIndex, s"Mates on different references not supported: ${r1.name}")
val pos1 = if (r1.positiveStrand) r1.unclippedStart else r1.unclippedEnd
val pos2 = if (r2.positiveStrand) r2.unclippedStart else r2.unclippedEnd
val r1Lower = r1.refIndex < r2.refIndex || (r1.refIndex == r2.refIndex && (pos1 < pos2 || (pos1 == pos2 && r1.positiveStrand)))
val umis = umi.split("-", -1) // Split and ensure we return empty strings for missing UMIs.
require(umis.length == 2, s"Paired strategy used but umi did not contain 2 segments delimited by a '-': $umi")
if (r1Lower) paired.lowerReadUmiPrefix + ":" + umis(0) + "-" + paired.higherReadUmiPrefix + ":" + umis(1)
else paired.higherReadUmiPrefix + ":" + umis(0) + "-" + paired.lowerReadUmiPrefix + ":" + umis(1)
case (_, _, _: PairedUmiAssigner) =>
fail(s"Template ${t.name} has only one read, paired-reads required for paired strategy.")
case (Some(r1), _, _) =>
r1[String](this.rawTag)
}
}
/** Sets the duplicate flags on all reads within all templates. */
private def setDuplicateFlags(group: Seq[Template]): Unit = {
val nonDuplicateTemplate = group.maxBy { template =>
template.primaryReads.sumBy { r =>
r.mapq + DuplicateScoringStrategy.computeDuplicateScore(r.asSam, ScoringStrategy.SUM_OF_BASE_QUALITIES)
}
}
group.foreach { template =>
val flag = template ne nonDuplicateTemplate
template.allReads.foreach(_.duplicate = flag)
}
}
}
© 2015 - 2025 Weber Informatics LLC | Privacy Policy