<commit_before>package edu.berkeley.cs.amplab.mlmatrix

import java.util.concurrent.ThreadLocalRandom
import scala.collection.mutable.ArrayBuffer

import breeze.linalg._

import edu.berkeley.cs.amplab.mlmatrix.util.QRUtils
import edu.berkeley.cs.amplab.mlmatrix.util.Utils

import org.apache.spark.rdd.RDD
import org.apache.spark.SparkConf
import org.apache.spark.SparkContext
import org.apache.spark.Accumulator
import org.apache.spark.SparkContext._

class TSQR extends RowPartitionedSolver with Logging with Serializable {

  /**
   * Returns only the R factor of a QR decomposition of the input matrix. This
   * operation is not reproducible since the R factors inside the TSQR tree can
   * be combined in any order.
   */
  def qrR(mat: RowPartitionedMatrix): DenseMatrix[Double] = {
    val localQR = mat.rdd.context.accumulator(0.0, "Time taken for Local QR")

    val qrTree = mat.rdd.map { part =>
      if (part.mat.rows < part.mat.cols) {
        part.mat
      } else {
        val begin = System.nanoTime
        val r = QRUtils.qrR(part.mat)
        localQR += ((System.nanoTime - begin) / 1000000)
        r
      }
    }
    val depth = math.ceil(math.log(mat.rdd.partitions.size)/math.log(2)).toInt
    Utils.treeReduce(qrTree, reduceQR(localQR, _ : DenseMatrix[Double], _ : DenseMatrix[Double]), depth=depth)
  }

  /**
   * Concatenates two matrices in a commutative reduce operation
   * and then returns the R factor of the QR factorization of the concatenated matrix
   */
  private def reduceQR(
      acc: Accumulator[Double],
      a: DenseMatrix[Double],
      b: DenseMatrix[Double]): DenseMatrix[Double] = {
    val begin = System.nanoTime
    val out = QRUtils.qrR(DenseMatrix.vertcat(a, b), false)
    acc += ((System.nanoTime - begin) / 1e6)
    out
  }

  /**
   * Returns both the Q and R factor of a QR decomposiion of the input matrix.
   * The TSQR tree used is always assumed to be binary. The order in which the R
   * factors are combined inside the TSQR tree must be preserved so that the Q
   * can be accurately reconstructed. Thus, the reduce operation used cannot be
   * commutative.
   */

  def qrQR(mat: RowPartitionedMatrix): (RowPartitionedMatrix, DenseMatrix[Double]) = {
    // First step run TSQR, get YTR tree
    val (qrTree, r) = qrYTR(mat)

    var curTreeIdx = qrTree.size - 1

    // Now construct Q by starting at the root of the YTR tree and applying
    // the appropriate Q factors to the identity matrix.
    var qrRevTree = qrTree(curTreeIdx)._2.map { part =>
      val yPart = part._2._1
      val tPart = part._2._2
      val qIn = new DenseMatrix[Double](yPart.rows, yPart.cols)
      for (i <- 0 until yPart.cols) {
        qIn(i, i) =  1.0
      }
      (part._1, QRUtils.applyQ(yPart, tPart, qIn, transpose=false))
    }.flatMap { x =>
      val nrows = x._2.rows
      Iterator((x._1 * 2, x._2),<commit_msg>Change the function comments / docstrings to describe the API more and leave out technical implementation details<commit_after>package edu.berkeley.cs.amplab.mlmatrix

import java.util.concurrent.ThreadLocalRandom
import scala.collection.mutable.ArrayBuffer

import breeze.linalg._

import edu.berkeley.cs.amplab.mlmatrix.util.QRUtils
import edu.berkeley.cs.amplab.mlmatrix.util.Utils

import org.apache.spark.rdd.RDD
import org.apache.spark.SparkConf
import org.apache.spark.SparkContext
import org.apache.spark.Accumulator
import org.apache.spark.SparkContext._

class TSQR extends RowPartitionedSolver with Logging with Serializable {

  /**
   * Returns only the R factor of a QR decomposition of the input matrix.
   * The value of R is not reproducible across multiple invocations of the function
   */
  def qrR(mat: RowPartitionedMatrix): DenseMatrix[Double] = {
    val localQR = mat.rdd.context.accumulator(0.0, "Time taken for Local QR")

    val qrTree = mat.rdd.map { part =>
      if (part.mat.rows < part.mat.cols) {
        part.mat
      } else {
        val begin = System.nanoTime
        val r = QRUtils.qrR(part.mat)
        localQR += ((System.nanoTime - begin) / 1000000)
        r
      }
    }
    val depth = math.ceil(math.log(mat.rdd.partitions.size)/math.log(2)).toInt
    Utils.treeReduce(qrTree, reduceQR(localQR, _ : DenseMatrix[Double], _ : DenseMatrix[Double]), depth=depth)
  }

  /**
   * Concatenates two matrices in a commutative reduce operation
   * and then returns the R factor of the QR factorization of the concatenated matrix
   */
  private def reduceQR(
      acc: Accumulator[Double],
      a: DenseMatrix[Double],
      b: DenseMatrix[Double]): DenseMatrix[Double] = {
    val begin = System.nanoTime
    val out = QRUtils.qrR(DenseMatrix.vertcat(a, b), false)
    acc += ((System.nanoTime - begin) / 1e6)
    out
  }

  /**
   * Returns both the Q and R factor of a QR decomposiion of the input matrix.
   * The returned values of Q and R are reproducible, i.e. they should be the same
   * across multiple invocations of the function
   */

  def qrQR(mat: RowPartitionedMatrix): (RowPartitionedMatrix, DenseMatrix[Double]) = {
    // First step run TSQR, get YTR tree
    val (qrTree, r) = qrYTR(mat)

    var curTreeIdx = qrTree.size - 1

    // Now construct Q by starting at the root of the YTR tree and applying
    // the appropriate Q factors to the identity matrix.
    var qrRevTree = qrTree(curTreeIdx)._2.map { part =>
      val yPart = part._2._1
      val tPart = part._2._2
      val qIn = new DenseMatrix[Double](yPart.rows, yPart.cols)
      for (i <- 0 until yPart.cols) {
        qIn(i, i) =  1.0
      }
      (part._1, QRUtils.applyQ(yPart, tPart, qIn, transpose=false))
    }.flatMap { x =>
      val nrows = x._2.rows
      Iterator((x._1 * 2, x._2),
{'fraction_non_alphanumeric': 0.10098176718092566, 'fraction_numerical': 0.0091164095371669, 'mean_word_length': 2.986722571628232, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 9, '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': '825096', 'n_tokens_mistral': 1964, 'n_tokens_neox': 1831, 'n_words': 552}