make changes suggested in PR

Author: jdidion

src/main/scala/com/fulcrumgenomics/vcf/DownsampleVcf.scala

@@ -15,9 +15,8 @@ import scala.math.log10
 import scala.util.Random
 
 object DownsampleVcf extends LazyLogging {
-  /** Removes variants that are within a specified distance from a previous variant
-   * The end position of the current variant is compared with the start position of the following variant
-   *
+  /** Removes variants that are within a specified distance from a previous variant.
+   * The end position of the current variant is compared with the start position of the following variant.
    * @param variants   an iterator of the variants to process
    * @param windowSize the interval (exclusive) in which to check for additional variants.
    *                   windowSize considers the distance between the end position of a variant
@@ -32,7 +31,7 @@ object DownsampleVcf extends LazyLogging {
 
       def hasNext: Boolean = iter.hasNext
 
-      def isInOrder(current: Variant, next: Variant, currentIndex: Int, nextIndex: Int): Boolean = {
+      private def isInOrder(current: Variant, next: Variant, currentIndex: Int, nextIndex: Int): Boolean = {
         (currentIndex < nextIndex) || (currentIndex == nextIndex && current.end <= next.pos)
       }
 
@@ -52,32 +51,33 @@ object DownsampleVcf extends LazyLogging {
     }
   }
 
-  /** Downsamples variants using Allele Depths
-   *
+  /** Downsamples variants by randomly sampling the total allele depths at the given proportion.
    * @param oldAds     an indexed seq of the original allele depths
    * @param proportion the proportion to use for downsampling,
    *                   calculated using total base count from the index and a target base count
    * @return a new IndexedSeq of allele depths of the same length as `oldAds`
    */
-  def downsampleADs(oldAds: IndexedSeq[Int], proportion: Double, random: Random): IndexedSeq[Int] = {
+  def downsampleADs(oldAds: IterableOnce[Int], proportion: Double, random: Random): IndexedSeq[Int] = {
     require(proportion <= 1, f"proportion must be less than 1: proportion = ${proportion}")
-    oldAds.map(s => Range(0, s).iterator.map(_ => random.nextDouble()).count(_ < proportion))
+    oldAds.iterator.toIndexedSeq.map(s => Range(0, s).iterator.map(_ => random.nextDouble()).count(_ < proportion))
   }
 
   /**
-   * Does the downsampling on a Variant
-   * @param variant the variant with the genotype to downsample
-   * @param proportions a map of downsampling target proportions for each sample
+   * Re-genotypes a variant for each sample after downsampling the allele counts based on the given 
+   * per-sample proportions.
+   * @param variant the variant to downsample and re-genotype
+   * @param proportions proportion to downsample the allele counts for each sample prior to re-genotyping
    * @param random random number generator for downsampling
-   * @param epsilon the error rate for genotyping
-   * @return a new variant with updated genotypes
+   * @param epsilon the sequencing error rate for genotyping
+   * @return a new variant with updated genotypes, downsampled ADs, and recomputed PLs
    */
-  // Returns a new variant that has downsampled ADs, recomputed PLs and updated genotypes
-  def downsampleAndRegenotype(variant: Variant, proportions: Map[String, Double], random: Random, epsilon: Double = 0.01): Variant = {
+  def downsampleAndRegenotype(variant: Variant,
+                              proportions: Map[String, Double],
+                              random: Random, epsilon: Double=0.01): Variant = {
     try {
-      variant.copy(genotypes = variant.genotypes.map { case (sample, gt) =>
+      variant.copy(genotypes=variant.genotypes.map { case (sample, gt) =>
         val proportion = proportions(sample)
-        sample -> downsampleAndRegenotype(gt = gt, proportion = proportion, random = random, epsilon = epsilon)
+        sample -> downsampleAndRegenotype(gt=gt, proportion=proportion, random=random, epsilon=epsilon)
       })
     } catch {
       case e: MatchError => throw new Exception(
@@ -87,15 +87,15 @@ object DownsampleVcf extends LazyLogging {
   }
 
   /**
-   * Does the downsampling on a Genotype
+   * Re-genotypes a sample after downsampling the allele counts based on the given proportion.
    * @param gt the genotype to downsample
-   * @param proportion the proportion to use for downsampling allele depths
+   * @param proportion proportion to downsample the allele count prior to re-genotyping
    * @param random random number generator for downsampling
-   * @param epsilon the error rate for genotyping
-   * @return a new Genotype with updated allele depths, PLs and genotype
+   * @param epsilon the sequencing error rate for genotyping
+   * @return a new Genotype with updated allele depths, PLs, and genotype
    */
   def downsampleAndRegenotype(gt: Genotype, proportion: Double, random: Random, epsilon: Double): Genotype = {
-    val oldAds = gt[IndexedSeq[Int]]("AD")
+    val oldAds = gt.getOrElse[IndexedSeq[Int]]("AD", throw new Exception(s"AD tag not found for sample ${gt.sample}"))
     val newAds = downsampleADs(oldAds, proportion, random)
     val Seq(aa, ab, bb) = computePls(newAds)
     val Seq(alleleA, alleleB) = gt.alleles.toSeq
@@ -106,20 +106,21 @@ object DownsampleVcf extends LazyLogging {
       else if (bb < ab && bb < aa) IndexedSeq(alleleB, alleleB)
       else IndexedSeq(alleleA, alleleB)
     }
-    gt.copy(attrs = Map("PL" -> IndexedSeq(aa, ab, bb), "AD" -> newAds, "DP" -> newAds.sum), calls = calls)
+    gt.copy(attrs=Map("PL" -> IndexedSeq(aa, ab, bb), "AD" -> newAds, "DP" -> newAds.sum), calls=calls)
   }
 
   /**
-   * Compute the genotype likelihoods given the allele depths.
-   * @param ads The allele depths to generate likelihoods from
-   * @return a list of three likelihoods
+   * Compute the genotype likelihoods given the allele depths, assuming a diploid genotype (i.e. 
+   * two allele depths).
+   * @param ads The input depths for the two alleles A and B.
+   * @return a list of three likelihoods for the alleles AA, AB, and BB.
    */
   def computePls(ads: IndexedSeq[Int]): IndexedSeq[Int] = {
+    require(ads.length == 2, "there must be exactly two allele depths")
     val likelihoods = Likelihoods(ads(0), ads(1))
     IndexedSeq(likelihoods.aa.round.toInt, likelihoods.ab.round.toInt, likelihoods.bb.round.toInt)
   }
 
-
   object Likelihoods {
     /** Computes the likelihoods for each possible genotype.
      *
@@ -128,7 +129,7 @@ object DownsampleVcf extends LazyLogging {
      * @param epsilon      the error rate for genotyping
      * @return a new `Likelihood` that has the likelihoods of AA, AB, and BB
      */
-    def apply(alleleDepthA: Int, alleleDepthB: Int, epsilon: Double = 0.01): Likelihoods = {
+    def apply(alleleDepthA: Int, alleleDepthB: Int, epsilon: Double=0.01): Likelihoods = {
       val aGivenAA = log10(1 - epsilon)
       val aGivenBB = log10(epsilon)
       val aGivenAB = log10((1 - epsilon) / 2)
@@ -154,102 +155,143 @@ object DownsampleVcf extends LazyLogging {
    * @param bb likelihood of BB
    */
   case class Likelihoods(aa: Double, ab: Double, bb: Double) {
+    /**
+      * Returns the likelihoods as a list of phred-scaled integers (i.e, the value of the PL tag).
+      * @return a list of phred-scaled likelihooodS for AA, AB, BB.
+      */
     def pls = IndexedSeq(aa.round.toInt, ab.round.toInt, bb.round.toInt)
   }
 }
 
-  @clp(group = ClpGroups.VcfOrBcf, description =
-    """
-      |DownsampleVcf takes a vcf file and metadata with sequencing info and
-      |1. winnows the vcf to remove variants within a specified distance to each other,
-      |2. downsamples the variants using the provided allele depths and target base count by
-      |   re-computing/downsampling the allele depths for the new target base count
-      |   and re-computing the genotypes based on the new allele depths
-      |and writes a new downsampled vcf file.
-      |For single-sample VCFs, the metadata file can be omitted, and instead you can specify originalBases.
-  """)
-  class DownsampleVcf
-  (@arg(flag = 'i', doc = "The vcf to downsample.") val input: PathToVcf,
-   @arg(flag = 'm', doc = "Index file with bases per sample.") val metadata: Option[FilePath] = None,
-   @arg(flag = 'b', doc = "Original number of bases (for single-sample VCF)") val originalBases: Option[Double] = None,
-   @arg(flag = 'n', doc = "Target number of bases to downsample to.") val downsampleToBases: Double,
-   @arg(flag = 'o', doc = "Output file name.") val output: PathToVcf,
-   @arg(flag = 'w', doc = "Winnowing window size.") val windowSize: Int = 150,
-   @arg(flag = 'e', doc = "Error rate for genotyping.") val epsilon: Double = 0.01,
-   @arg(flag = 'c', doc = "True to write out no-calls.") val writeNoCall: Boolean = false)
-    extends FgBioTool {
-    Io.assertReadable(input)
-    Io.assertReadable(metadata)
-    Io.assertCanWriteFile(output)
-    require(downsampleToBases > 0, "target base count must be greater than zero")
-    require(windowSize >= 0, "window size must be greater than or equal to zero")
-    require(0 <= epsilon && epsilon <= 1, "epsilon/error rate must be between 0 and 1")
-    originalBases match {
-      case Some(x) =>
-        require(x > 0, "originalBases must be greater than zero")
-        require(metadata.isEmpty, "Must pass either originalBases (for single-sample VCF) or metadata, not both")
-      case None =>
-        require(metadata.isDefined, "Must pass either originalBases (for single-sample VCF) or metadata, not both")
+@clp(group=ClpGroups.VcfOrBcf, description =
+  """
+    |Re-genotypes a VCF after downsampling the allele counts.
+    |
+    |The input VCF must have at least one sample.
+    |
+    |If the input VCF contains a single sample, the downsampling target may be specified as a
+    |proportion of the original read depth using `--proportion=(0..1)`, or as the combination of
+    |the original and target _number of sequenced bases_ (`--originalBases` and
+    |`--downsampleToBases`). For multi-sample VCFs, the downsampling target must be specified using
+    |`--downsampleToBases`, and a metadata file with the total number of sequenced bases per sample
+    |is required as well. The metadata file must follow the
+    |[[https://www.internationalgenome.org/category/meta-data/] 1000 Genomes index format], but the
+    |only required columns are `SAMPLE_NAME` and `BASE_COUNT`. A propportion for each sample is
+    |calculated by dividing the _target number of sequenced bases_ by the _original number of
+    |sequenced bases_.
+    |
+    |The tool first (optionally) winnows the VCF file to remove variants within a distance to each
+    |other specified by `--window-size` (the default value of `0` disables winnowing). Next, each
+    |sample at each variant is examined independently. The allele depths per-genotype are randoml
+    |downsampled given the proportion. The downsampled allele depths are then used to re-compute
+    |allele likelhoods and produce a new genotype.
+    |
+    |The tool outputs a downsampled VCF file with the winnowed variants removed, and with the
+    |genotype calls and `DP`, `AD`, and `PL` tags updated for each sample at each retained variant.
+""")
+class DownsampleVcf
+(@arg(flag='i', doc="The vcf to downsample.") val input: PathToVcf,
+  @arg(flag='p', doc="Proportion of bases to retain (for single-sample VCF).") val proportion: Option[Double] = None,
+  @arg(flag='b', doc="Original number of bases (for single-sample VCF).") val originalBases: Option[Double] = None,
+  @arg(flag='m', doc="Index file with bases per sample.") val metadata: Option[FilePath] = None,
+  @arg(flag='n', doc="Target number of bases to downsample to.") val downsampleToBases: Option[Double],
+  @arg(flag='o', doc="Output file name.") val output: PathToVcf,
+  @arg(flag='w', doc="Winnowing window size.") val windowSize: Int = 0,
+  @arg(flag='e', doc="Sequencing Error rate for genotyping.") val epsilon: Double = 0.01,
+  @arg(flag='c', doc="True to write out no-calls.") val writeNoCall: Boolean = false,
+  @arg(flag='s', doc="Random seed value.") val seed: Int = 42,
+  ) extends FgBioTool {
+  Io.assertReadable(input)
+  Io.assertCanWriteFile(output)
+  require(windowSize >= 0, "window size must be greater than or equal to zero")
+  require(0 <= epsilon && epsilon <= 1, "epsilon/error rate must be between 0 and 1")
+  (proportion, originalBases, metadata, downsampleToBases) match {
+    case (Some(x), None, None, None) => 
+      require(x > 0, "proportion must be greater than 0")
+      require(x < 1, "proportion must be less than 1")
+    case (None, Some(original), None, Some(target)) =>
+      require(original > 0, "originalBases must be greater than zero")
+      require(target > 0, "target base count must be greater than zero")
+    case (None, None, Some(metadata), Some(target)) =>
+      Io.assertReadable(metadata)
+      require(target > 0, "target base count must be greater than zero")
+    case (None, _, _, None) =>
+      throw new IllegalArgumentException(
+        "exactly one of proportion or downsampleToBases must be specified"
+      )
+    case _ =>
+      throw new IllegalArgumentException(
+        "exactly one of proportion, originalBases, or metadata must be specified"
+      )
+  }
+  
+  override def execute(): Unit = {
+    val vcf = VcfSource(input)
+    val proportions = (
+      (proportion, originalBases, metadata, downsampleToBases) match {
+        case (Some(x), None, None, None) =>
+          require(vcf.header.samples.length == 1, "--original-bases requires a single-sample VCF")
+          LazyList(vcf.header.samples.head -> x)
+        case (None, Some(original), None, Some(target)) =>
+          require(vcf.header.samples.length == 1, "--original-bases requires a single-sample VCF")
+          LazyList(vcf.header.samples.head -> math.min(target / original, 1.0))
+        case (None, None, Some(metadata), Some(target)) =>
+          Sample.read(metadata)
+            .filter(s => vcf.header.samples.contains(s.SAMPLE_NAME))
+            .map(sample => sample.SAMPLE_NAME -> math.min(target / sample.BASE_COUNT.toDouble, 1.0))
+        case _ =>
+          throw new RuntimeException("unexpected parameter combination")
+      }
+    ).toMap
+    proportions.foreach { case (s, p) => logger.info(f"Downsampling $s with proportion ${p}%.4f") }
+    
+    val inputProgress = ProgressLogger(logger, noun="variants read")
+    val inputVariants = ProgressLogger.ProgressLoggingIterator(vcf.iterator).progress(inputProgress)
+    val winnowed = if (windowSize > 0) {
+      val winnowed = winnowVariants(inputVariants, windowSize=windowSize, dict=vcf.header.dict)
+      val winnowedProgress = ProgressLogger(logger, noun="variants retained")
+      ProgressLogger.ProgressLoggingIterator(winnowed).progress(winnowedProgress)
+    } else {
+      inputVariants
     }
+    val outputVcf = VcfWriter(path=output, header=buildOutputHeader(vcf.header))
 
-    override def execute(): Unit = {
-      val vcf = VcfSource(input)
-      val progress = ProgressLogger(logger, noun = "variants")
-      val proportions = (
-        originalBases match {
-          case Some(x) =>
-            require(vcf.header.samples.length == 1, "--original-bases requires a single-sample VCF")
-            LazyList(vcf.header.samples.head -> math.min(downsampleToBases / x, 1.0))
-          case _ =>
-            Sample.read(metadata.getOrElse(throw new RuntimeException))
-              .filter(s => vcf.header.samples.contains(s.SAMPLE_NAME))
-              .map(sample => sample.SAMPLE_NAME -> math.min(downsampleToBases / sample.BASE_COUNT.toDouble, 1.0))
-        }
-      ).toMap
-      proportions.foreach { case (s, p) => logger.info(f"Downsampling $s with proportion ${p}%.4f") }
-
-      val winnowed = if (windowSize > 0) winnowVariants(vcf.iterator, windowSize = windowSize, dict = vcf.header.dict) else vcf.iterator
-      val outputVcf = VcfWriter(path = output, header = buildOutputHeader(vcf.header))
-
-      val random = new Random(42)
-      winnowed.foreach { v =>
-        val ds = downsampleAndRegenotype(v, proportions = proportions, random = random, epsilon = epsilon)
-        if (writeNoCall) {
-          outputVcf += ds
-          progress.record(ds)
-        }
-        else if (!ds.gts.forall(g => g.isNoCall)) {
-          outputVcf += ds
-          progress.record(ds)
-        }
+    val progress = ProgressLogger(logger, noun="variants written")
+    val random = new Random(seed)
+    winnowed.foreach { v =>
+      val ds = downsampleAndRegenotype(v, proportions=proportions, random=random, epsilon=epsilon)
+      if (writeNoCall || !ds.gts.forall(g => g.isNoCall)) {
+        outputVcf += ds
+        progress.record(ds)
       }
-      
-      progress.logLast()
-      vcf.safelyClose()
-      outputVcf.close()
     }
+    
+    progress.logLast()
+    vcf.safelyClose()
+    outputVcf.close()
+  }
 
-    def buildOutputHeader(in: VcfHeader): VcfHeader = {
-      val fmts = Seq.newBuilder[VcfFormatHeader]
-      fmts ++= in.formats
+  private def buildOutputHeader(in: VcfHeader): VcfHeader = {
+    val fmts = Seq.newBuilder[VcfFormatHeader]
+    fmts ++= in.formats
 
-      if (!in.format.contains("AD")) {
-        fmts += VcfFormatHeader(id="AD", count=VcfCount.OnePerAllele, kind=VcfFieldType.Integer, description="Per allele depths.")
-      }
-
-      if (!in.format.contains("DP")) {
-        fmts += VcfFormatHeader(id="DP", count=VcfCount.Fixed(1), kind=VcfFieldType.Integer, description="Total depth across alleles.")
-      }
+    if (!in.format.contains("AD")) {
+      fmts += VcfFormatHeader(id="AD", count=VcfCount.OnePerAllele, kind=VcfFieldType.Integer, description="Per allele depths.")
+    }
 
-      if (!in.format.contains("PL")) {
-        fmts += VcfFormatHeader(id="PL", count=VcfCount.OnePerGenotype, kind=VcfFieldType.Integer, description="Per genotype phred scaled likelihoods.")
-      }
+    if (!in.format.contains("DP")) {
+      fmts += VcfFormatHeader(id="DP", count=VcfCount.Fixed(1), kind=VcfFieldType.Integer, description="Total depth across alleles.")
+    }
 
-      in.copy(formats = fmts.result())
+    if (!in.format.contains("PL")) {
+      fmts += VcfFormatHeader(id="PL", count=VcfCount.OnePerGenotype, kind=VcfFieldType.Integer, description="Per genotype phred scaled likelihoods.")
     }
+
+    in.copy(formats=fmts.result())
   }
+}
 
-object Sample {
+private object Sample {
   /** Load a set of samples from the 1KG metadata file. */
   def read(path: FilePath): Seq[Sample] = {
     val lines = Io.readLines(path).dropWhile(_.startsWith("##")).map(line => line.dropWhile(_ == '#'))