<commit_before>package org.bdgenomics.guacamole

import org.bdgenomics.guacamole.Common._
import org.apache.spark.rdd._
import org.apache.spark.SparkContext._
import org.apache.spark.broadcast.Broadcast
import org.apache.commons.math3
import org.apache.spark._
import scala.reflect.ClassTag
import org.bdgenomics.guacamole.Common.Arguments.{ Loci, Base }
import org.kohsuke.args4j.{ Option => Opt }
import org.bdgenomics.guacamole.pileup.Pileup
import scala.Some
import org.apache.spark.serializer.JavaSerializer
import scala.collection.mutable.{ HashMap => MutableHashMap }

object DistributedUtil extends Logging {
  trait Arguments extends Base with Loci {
    @Opt(name = "-parallelism", usage = "Num variant calling tasks. Set to 0 (default) to use the number of Spark partitions.")
    var parallelism: Int = 0

    @Opt(name = "-partition-accuracy",
      usage = "Num micro partitions to use per task in loci partitioning. Set to 0 to partition loci uniformly. Default: 250.")
    var partitioningAccuracy: Int = 250
  }

  /**
   * Partition a LociSet among tasks according to the strategy specified in args.
   */
  def partitionLociAccordingToArgs(args: Arguments, loci: LociSet, readsRDDs: RDD[MappedRead]*): LociMap[Long] = {
    val tasks = if (args.parallelism > 0) args.parallelism else readsRDDs(0).partitions.length
    if (args.partitioningAccuracy == 0)
      partitionLociUniformly(tasks, loci)
    else
      partitionLociByApproximateReadDepth(tasks, loci, args.partitioningAccuracy, readsRDDs: _*)
  }

  /**
   * Assign loci from a LociSet to partitions. Contiguous intervals of loci will tend to get assigned to the same
   * partition.
   *
   * This implementation assigns loci uniformly, i.e. each task gets about the same number of loci. A smarter
   * implementation would know about the reads (depth of coverage), and try to assign each task loci corresponding to
   * about the same number of reads.
   *
   * @param tasks number of partitions
   * @param loci loci to partition
   * @return LociMap of locus -> task assignments
   */
  def partitionLociUniformly(tasks: Long, loci: LociSet): LociMap[Long] = {
    assume(tasks >= 1)
    val lociPerTask = math.max(1, loci.count.toDouble / tasks.toDouble)
    progress("Splitting loci evenly among %,d tasks = ~%,.0f loci per task".format(tasks, lociPerTask))
    val builder = LociMap.newBuilder[Long]
    var lociAssigned = 0L
    var task = 0L
    def remainingForThisTask = math.round(((task + 1) * lociPerTask) - lociAssigned).toLong
    for (contig <- loci.contigs; range <- loci.onContig(contig).ranges) {
      var start = range.start
      val end = range.end
      while (start < end) {
        val length: Long = math.min(remainingForThisTask, end - start)
        builder.put(contig, start, start + length, task)
        start += length
        lociAssigned += length
        if (remainingForThisTask == 0) task += 1
      }
    }
    val result = builder.result
    assert(lociAssigned == loci.count)
    assert(result.count == loci.count)
    result
  }

  /**
   * Assign loci from a LociSet to partitions, where each partition overlaps about the same number of reads.
   *
   * The approach we take is:
   *
   *  (1) chop up the loci uniformly into many genomic "micro partitions."
   *
   *  (2) for each micro partition, calculate the number of reads that overlap it.
   *
   *  (3) using these counts, assign loci to real ("macro") partitions, making the approximation of uniform depth within
   *      each micro partition.
   *
   *  Some advantages of this approach are:
   *
   *  - Stages (1) and (3), which are done locally by the Spark master, are constant time with respect to the number
   *    of reads.
   *
   *  - Stage (2), where runtime does depend on the number of reads, is done in parallel with Spark.
   *
   *  - We can tune the accuracy vs. performance trade off by setting `microTasks`.
   *
   *  - Does not require a distributed sort.
   *
   * @param tasks number of partitions
   * @param loci loci to partition
   * @param accuracy integer >= 1. Higher values of this will result in a more exact but also more expensive computation.
   *                 Specifically, this is the number of micro partitions to use per task to estimate the read depth.
   *                 In the extreme case, setting this to greater than the number of loci per task will result in an
   *                 exact calculation.
   * @param readRDDs: reads RDD 1, reads RDD 2, ...
   *                Any number RDD[MappedRead] arguments giving the reads to base the partitioning on.
   * @return LociMap of locus -> task assignments.
   */
  def partitionLociByApproximateReadDepth(tasks: Int, loci: LociSet, accuracy: Int, readRDDs: RDD[MappedRead]*): LociMap[Long] = {
    // Step (1). Split loci uniformly into micro partitions.
    assume(tasks >= 1)
    assume(loci.count > 0)
    assume(readRDDs.length > 0)
    val numMicroPartitions: Int = if (accuracy * tasks < loci.count) accuracy * tasks else loci.count.toInt
    progress("Splitting loci by read depth among %,d tasks using %,d micro partitions.".format(tasks, numMicroPartitions))
    val microPartitions = partitionLociUniformly(numMicroPartitions, loci)
    progress("Done calculating micro partitions.")
    val sc = readRDDs(0).sparkContext
    val broadcastMicroPartitions = sc.broadcast(microPartitions)

    // Step (2)
    // Total up reads overlapping each micro partition. We keep the totals as an array of Longs.
    def addArray(first: Array[Long], second: Array[Long]): Array[Long] = {
      assert(first.length == second.length)
      val result = new Array[Long](first.length)
      for (i <- 0 until first.length) {
        result(i) = first(i) + second(i)
      }
      result
    }

    var num = 1
    val counts = readRDDs.map(reads => {
      progress("Collecting read counts for RDD %d of %d.".format(num, readRDDs.length))
      num += 1
      reads.mapPartitions(readIterator => {
        val microPartitions = broadcastMicroPartitions.value
        assert(microPartitions.count > 0)
        val counts = new Array[Long](numMicroPartitions)
        for (read <- readIterator) {
          val contigMap = microPartitions.onContig(read.referenceContig)<commit_msg>minor optimization: change for -> while in DistributedUtil<commit_after>package org.bdgenomics.guacamole

import org.bdgenomics.guacamole.Common._
import org.apache.spark.rdd._
import org.apache.spark.SparkContext._
import org.apache.spark.broadcast.Broadcast
import org.apache.commons.math3
import org.apache.spark._
import scala.reflect.ClassTag
import org.bdgenomics.guacamole.Common.Arguments.{ Loci, Base }
import org.kohsuke.args4j.{ Option => Opt }
import org.bdgenomics.guacamole.pileup.Pileup
import scala.Some
import org.apache.spark.serializer.JavaSerializer
import scala.collection.mutable.{ HashMap => MutableHashMap }

object DistributedUtil extends Logging {
  trait Arguments extends Base with Loci {
    @Opt(name = "-parallelism", usage = "Num variant calling tasks. Set to 0 (default) to use the number of Spark partitions.")
    var parallelism: Int = 0

    @Opt(name = "-partition-accuracy",
      usage = "Num micro partitions to use per task in loci partitioning. Set to 0 to partition loci uniformly. Default: 250.")
    var partitioningAccuracy: Int = 250
  }

  /**
   * Partition a LociSet among tasks according to the strategy specified in args.
   */
  def partitionLociAccordingToArgs(args: Arguments, loci: LociSet, readsRDDs: RDD[MappedRead]*): LociMap[Long] = {
    val tasks = if (args.parallelism > 0) args.parallelism else readsRDDs(0).partitions.length
    if (args.partitioningAccuracy == 0)
      partitionLociUniformly(tasks, loci)
    else
      partitionLociByApproximateReadDepth(tasks, loci, args.partitioningAccuracy, readsRDDs: _*)
  }

  /**
   * Assign loci from a LociSet to partitions. Contiguous intervals of loci will tend to get assigned to the same
   * partition.
   *
   * This implementation assigns loci uniformly, i.e. each task gets about the same number of loci. A smarter
   * implementation would know about the reads (depth of coverage), and try to assign each task loci corresponding to
   * about the same number of reads.
   *
   * @param tasks number of partitions
   * @param loci loci to partition
   * @return LociMap of locus -> task assignments
   */
  def partitionLociUniformly(tasks: Long, loci: LociSet): LociMap[Long] = {
    assume(tasks >= 1)
    val lociPerTask = math.max(1, loci.count.toDouble / tasks.toDouble)
    progress("Splitting loci evenly among %,d tasks = ~%,.0f loci per task".format(tasks, lociPerTask))
    val builder = LociMap.newBuilder[Long]
    var lociAssigned = 0L
    var task = 0L
    def remainingForThisTask = math.round(((task + 1) * lociPerTask) - lociAssigned).toLong
    for (contig <- loci.contigs; range <- loci.onContig(contig).ranges) {
      var start = range.start
      val end = range.end
      while (start < end) {
        val length: Long = math.min(remainingForThisTask, end - start)
        builder.put(contig, start, start + length, task)
        start += length
        lociAssigned += length
        if (remainingForThisTask == 0) task += 1
      }
    }
    val result = builder.result
    assert(lociAssigned == loci.count)
    assert(result.count == loci.count)
    result
  }

  /**
   * Assign loci from a LociSet to partitions, where each partition overlaps about the same number of reads.
   *
   * The approach we take is:
   *
   *  (1) chop up the loci uniformly into many genomic "micro partitions."
   *
   *  (2) for each micro partition, calculate the number of reads that overlap it.
   *
   *  (3) using these counts, assign loci to real ("macro") partitions, making the approximation of uniform depth within
   *      each micro partition.
   *
   *  Some advantages of this approach are:
   *
   *  - Stages (1) and (3), which are done locally by the Spark master, are constant time with respect to the number
   *    of reads.
   *
   *  - Stage (2), where runtime does depend on the number of reads, is done in parallel with Spark.
   *
   *  - We can tune the accuracy vs. performance trade off by setting `microTasks`.
   *
   *  - Does not require a distributed sort.
   *
   * @param tasks number of partitions
   * @param loci loci to partition
   * @param accuracy integer >= 1. Higher values of this will result in a more exact but also more expensive computation.
   *                 Specifically, this is the number of micro partitions to use per task to estimate the read depth.
   *                 In the extreme case, setting this to greater than the number of loci per task will result in an
   *                 exact calculation.
   * @param readRDDs: reads RDD 1, reads RDD 2, ...
   *                Any number RDD[MappedRead] arguments giving the reads to base the partitioning on.
   * @return LociMap of locus -> task assignments.
   */
  def partitionLociByApproximateReadDepth(tasks: Int, loci: LociSet, accuracy: Int, readRDDs: RDD[MappedRead]*): LociMap[Long] = {
    // Step (1). Split loci uniformly into micro partitions.
    assume(tasks >= 1)
    assume(loci.count > 0)
    assume(readRDDs.length > 0)
    val numMicroPartitions: Int = if (accuracy * tasks < loci.count) accuracy * tasks else loci.count.toInt
    progress("Splitting loci by read depth among %,d tasks using %,d micro partitions.".format(tasks, numMicroPartitions))
    val microPartitions = partitionLociUniformly(numMicroPartitions, loci)
    progress("Done calculating micro partitions.")
    val sc = readRDDs(0).sparkContext
    val broadcastMicroPartitions = sc.broadcast(microPartitions)

    // Step (2)
    // Total up reads overlapping each micro partition. We keep the totals as an array of Longs.
    def addArray(first: Array[Long], second: Array[Long]): Array[Long] = {
      assert(first.length == second.length)
      val result = new Array[Long](first.length)
      var i = 0
      while (i < first.length) {
        result(i) = first(i) + second(i)
        i += 1
      }
      result
    }

    var num = 1
    val counts = readRDDs.map(reads => {
      progress("Collecting read counts for RDD %d of %d.".format(num, readRDDs.length))
      num += 1
      reads.mapPartitions(readIterator => {
        val microPartitions = broadcastMicroPartitions.value
        assert(microPartitions.count > 0)
        val counts = new Array[Long](numMicroPartitions)
        for (read <- readIterator) {
          val contigMap = microPartitions.onContig(read.referenceContig)
{'fraction_non_alphanumeric': 0.09095241877829852, 'fraction_numerical': 0.006593057431090635, 'mean_word_length': 3.121112929623568, 'pattern_counts': {'":': 0, '<': 15, '<?xml version=': 0, '>': 30, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5042962', 'n_tokens_mistral': 3885, 'n_tokens_neox': 3617, 'n_words': 1296}