<commit_before>package is.hail.stats

import breeze.linalg.{DenseMatrix, DenseVector, SparseVector}
import is.hail.expr._
import is.hail.utils._
import is.hail.variant.{Genotype, VariantDataset}

import scala.collection.mutable

object RegressionUtils {
  def toDouble(t: Type, code: String): Any => Double = t match {
    case TInt => _.asInstanceOf[Int].toDouble
    case TLong => _.asInstanceOf[Long].toDouble
    case TFloat => _.asInstanceOf[Float].toDouble
    case TDouble => _.asInstanceOf[Double]
    case TBoolean => _.asInstanceOf[Boolean].toDouble
    case _ => fatal(s"Sample annotation `$code' must be numeric or Boolean, got $t")
  }

  def getPhenoCovCompleteSamples(
    vds: VariantDataset,
    ySA: String,
    covSA: Array[String]): (DenseVector[Double], DenseMatrix[Double], IndexedSeq[String]) = {

    val symTab = Map(
      "s" -> (0, TSample),
      "sa" -> (1, vds.saSignature))

    val ec = EvalContext(symTab)

    val (yT, yQ) = Parser.parseExpr(ySA, ec)
    val yToDouble = toDouble(yT, ySA)
    val yIS = vds.sampleIdsAndAnnotations.map { case (s, sa) =>
      ec.setAll(s, sa)
      Option(yQ()).map(yToDouble)
    }

    val (covT, covQ) = Parser.parseExprs(covSA.mkString(","), ec)
    val covToDouble = (covT, covSA).zipped.map(toDouble)
    val covIS = vds.sampleIdsAndAnnotations.map { case (s, sa) =>
      ec.setAll(s, sa)
      (covQ().map(Option(_)), covToDouble).zipped.map(_.map(_))
    }

    val (yForCompleteSamples, covForCompleteSamples, completeSamples) =
      (yIS, covIS, vds.sampleIds)
        .zipped
        .filter((y, c, s) => y.isDefined && c.forall(_.isDefined))

    val n = completeSamples.size
    if (n == 0)
      fatal("No complete samples: each sample is missing its phenotype or some covariate")

    val yArray = yForCompleteSamples.map(_.get).toArray
    if (yArray.toSet.size == 1)<commit_msg>better parsing of covariate list (#1502)<commit_after>package is.hail.stats

import breeze.linalg.{DenseMatrix, DenseVector, SparseVector}
import is.hail.expr._
import is.hail.utils._
import is.hail.variant.{Genotype, VariantDataset}

import scala.collection.mutable

object RegressionUtils {
  def toDouble(t: Type, code: String): Any => Double = t match {
    case TInt => _.asInstanceOf[Int].toDouble
    case TLong => _.asInstanceOf[Long].toDouble
    case TFloat => _.asInstanceOf[Float].toDouble
    case TDouble => _.asInstanceOf[Double]
    case TBoolean => _.asInstanceOf[Boolean].toDouble
    case _ => fatal(s"Sample annotation `$code' must be numeric or Boolean, got $t")
  }

  def getPhenoCovCompleteSamples(
    vds: VariantDataset,
    ySA: String,
    covSA: Array[String]): (DenseVector[Double], DenseMatrix[Double], IndexedSeq[String]) = {

    val symTab = Map(
      "s" -> (0, TSample),
      "sa" -> (1, vds.saSignature))

    val ec = EvalContext(symTab)

    val (yT, yQ) = Parser.parseExpr(ySA, ec)
    val yToDouble = toDouble(yT, ySA)
    val yIS = vds.sampleIdsAndAnnotations.map { case (s, sa) =>
      ec.setAll(s, sa)
      Option(yQ()).map(yToDouble)
    }

    val (covT, covQ0) = covSA.map(Parser.parseExpr(_, ec)).unzip
    val covQ = () => covQ0.map(_.apply())
    val covToDouble = (covT, covSA).zipped.map(toDouble)
    val covIS = vds.sampleIdsAndAnnotations.map { case (s, sa) =>
      ec.setAll(s, sa)
      (covQ().map(Option(_)), covToDouble).zipped.map(_.map(_))
    }

    val (yForCompleteSamples, covForCompleteSamples, completeSamples) =
      (yIS, covIS, vds.sampleIds)
        .zipped
        .filter((y, c, s) => y.isDefined && c.forall(_.isDefined))

    val n = completeSamples.size
    if (n == 0)
      fatal("No complete samples: each sample is missing its phenotype or some covariate")

    val yArray = yForCompleteSamples.map(_.get).toArray
    if (yArray.toSet.size == 1)
{'fraction_non_alphanumeric': 0.13849394418114797, 'fraction_numerical': 0.003686150605581885, 'mean_word_length': 3.4536928487690504, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 28, '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': '4933733', 'n_tokens_mistral': 1369, 'n_tokens_neox': 1325, 'n_words': 305}