<commit_before>package org.broadinstitute.hellbender.tools.exome;

import com.google.common.annotations.VisibleForTesting;
import com.google.common.primitives.Doubles;
import org.apache.commons.math3.distribution.BetaDistribution;
import org.apache.commons.math3.stat.descriptive.moment.Mean;
import org.apache.commons.math3.stat.descriptive.moment.Variance;
import org.apache.spark.api.java.JavaSparkContext;
import org.broadinstitute.hellbender.utils.MathUtils;
import org.broadinstitute.hellbender.utils.mcmc.*;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.function.Function;
import java.util.stream.Collectors;
import java.util.stream.IntStream;

/**
 * Represents an ACNV segmented model for copy ratio fit to target-coverage data.
 * The log_2 coverages in each segment are fit by a mixture model with a normal-distribution component
 * and a uniform outlier component.  The variance of the normal-distribution component and the relative
 * contribution of the uniform outlier component in all segments are both assumed to be global parameters.
 * The mean of the normal-distribution component in each segment is taken to be a segment-level parameter.
 * The component (i.e., normal or outlier) that each target coverage is drawn from is determined by a latent
 * target-level indicator parameter.
 *
 * @author <NAME> &lt;<EMAIL>&gt;
 */
public final class ACNVCopyRatioModeller {
    private static final double EPSILON = 1E-10;
    private static final double VARIANCE_MIN = EPSILON;
    private static final double VARIANCE_MAX = 10.;

    private static final double OUTLIER_PROBABILITY_INITIAL = 0.05;
    private static final double OUTLIER_PROBABILITY_PRIOR_ALPHA = 5.;
    private static final double OUTLIER_PROBABILITY_PRIOR_BETA = 95.;

    private final SegmentedModel segmentedModel;

    private final double outlierUniformLogLikelihood;
    private final double varianceSliceSamplingWidth;
    private final double meanSliceSamplingWidth;

    private final ParameterizedModel<CopyRatioState, CopyRatioDataCollection> model;

    private final List<Double> varianceSamples = new ArrayList<>();
    private final List<Double> outlierProbabilitySamples = new ArrayList<>();
    private final List<SegmentMeans> segmentMeansSamples = new ArrayList<>();
    private final List<OutlierIndicators> outlierIndicatorsSamples = new ArrayList<>();

    //Calculates the exponent for a normal distribution; used in log-likelihood calculation below.
    private static double normalTerm(final double quantity, final double mean, final double variance) {
        return (quantity - mean) * (quantity - mean) / (2. * variance);
    }

    /**
     * Constructs a copy-ratio model given a SegmentedModel with segments and a Genome.  Initial point estimates of
     * parameters are set to empirical estimates where available.
     * @param segmentedModel    SegmentedModel with segments and a Genome
     */
    public ACNVCopyRatioModeller(final SegmentedModel segmentedModel) {
        this.segmentedModel = segmentedModel;

        //load segmented coverages from SegmentedModel into CopyRatioDataCollection
        final CopyRatioDataCollection data = new CopyRatioDataCollection(segmentedModel);

        //the uniform log-likelihood for outliers is determined by the minimum and maximum coverages in the dataset;
        //the outlier-probability parameter should be interpreted accordingly
        outlierUniformLogLikelihood = -Math.log(data.coverageRange);

        //set widths for slice sampling of variance and segment-mean posteriors using empirical variance estimate.
        //variance posterior is inverse chi-squared, segment-mean posteriors are Gaussian; the below expressions
        //approximate the standard deviations of these distributions.
        final double varianceEstimate = data.estimateVariance();
        varianceSliceSamplingWidth = Math.sqrt(2. * varianceEstimate / data.numTargets);
        meanSliceSamplingWidth = Math.sqrt(varianceEstimate * data.numSegments / data.numTargets);

        //use empirical segment means and empirical average variance across segments to initialize CopyRatioState
        final CopyRatioState initialState = new CopyRatioState(data);

        //samples log conditional posterior for the variance parameter, assuming uniform prior; this is given by
        //the product of Gaussian likelihoods for each non-outlier target t:
        //  log[product_{non-outlier t} variance^(-1/2) * exp(-(coverage_t - mean_t)^2 / (2 * variance))] + constant
        //where mean_t is identical for all targets in a segment
        final Sampler<Double, CopyRatioState, CopyRatioDataCollection> varianceSampler =
                (rng, state, dataCollection) -> {
            final Function<Double, Double> logConditionalPDF = newVariance -> {
                double ll = 0.;
                int numNotOutliers = 0;
                for (int segment = 0; segment < dataCollection.numSegments; segment++) {
                    final int numTargetsInSegment = dataCollection.getNumTargetsInSegment(segment);
                    final int targetStartInSegment = dataCollection.getStartTargetInSegment(segment);
                    for (int target = 0; target < numTargetsInSegment; target++) {
                        final double coverage = dataCollection.getCoveragesInSegment(segment).get(target);
                        if (!state.getOutlierIndicator(targetStartInSegment + target)) {
                            ll -= normalTerm(coverage, state.getMeanInSegment(segment), newVariance);
                            numNotOutliers++;
                        }
                    }
                }
                ll -= 0.5 * Math.log(newVariance) * numNotOutliers;
                return ll;
            };
            return new SliceSampler(rng, logConditionalPDF, state.variance(),
                    VARIANCE_MIN, VARIANCE_MAX, varianceSliceSamplingWidth).sample();
        };

        //samples log conditional posterior for the outlier-probability parameter, assuming Beta(alpha, beta) prior;
        //this is given by:
        //  log Beta(alpha + number of outlier targets, beta + number of non-outlier targets) + constant
        final Sampler<Double, CopyRatioState, CopyRatioDataCollection> outlierProbabilitySampler =
                (rng, state, dataCollection) -> {
            final int numOutliers = (int) IntStream.range(0, dataCollection.numTargets).filter(state::getOutlierIndicator).count();
            return new BetaDistribution(
                    rng,
                    OUTLIER_PROBABILITY_PRIOR_ALPHA + numOutliers,
                    OUTLIER_PROBABILITY_PRIOR_BETA + dataCollection.numTargets - numOutliers).sample();
        };

        //samples log conditional posteriors for the segment-mean parameters, assuming uniform priors;
        //for each segment s, this is given by the product of Gaussian likelihoods for each non-outlier target t:
        //  log[product_{non-outlier t in s} exp(-(coverage_t - mean_s)^2 / (2 * variance))] + constant
        final Sampler<SegmentMeans, CopyRatioState, CopyRatioDataCollection> segmentMeansSampler =
                (rng, state, dataCollection) -> {
            final List<Double> means = new ArrayList<>(dataCollection.numSegments);
            for (int segment = 0; segment < dataCollection.numSegments; segment++) {
                final int numTargetsInSegment = dataCollection.getNumTargetsInSegment(segment);
                if (numTargetsInSegment == 0) {
                    means.add(Double.NaN);
                } else {
                    //get coverages at non-outlier targets in the segment and construct log conditional PDF
                    final int targetStartInSegment = dataCollection.getStartTargetInSegment(segment);
                    final List<Double> allCoveragesInSegment = dataCollection.getCoveragesInSegment(segment);
                    final List<Double> nonOutlierCoveragesInSegment = IntStream.range(0, numTargetsInSegment).boxed()
                            .filter(target -> !state.getOutlierIndicator(targetStartInSegment + target))
                            .map(allCoveragesInSegment::get)
                            .collect(Collectors.toList());
                    final Function<Double, Double> logConditionalPDF = newMean -> nonOutlierCoveragesInSegment.stream()
                            .mapToDouble(coverage -> -normalTerm(coverage, newMean, state.variance()))
                            .sum();
                    //slice sample within range given by minimum and maximum coverages
                    final double minimumCoverageInSegment = Collections.min(allCoveragesInSegment);
                    final double maximumCoverageInSegment = Collections.max(allCoveragesInSegment);
                    final SliceSampler sampler =
                            new SliceSampler(rng, logConditionalPDF, state.getMeanInSegment(segment),
                                    minimumCoverageInSegment, maximumCoverageInSegment, meanSliceSamplingWidth);
                    means.add(sampler.sample());
                }
            }
            return new SegmentMeans(means);
        };

        //samples log conditional posteriors for the outlier-indicator parameters; for each target t, this is given by:
        //          z_t * [log outlier_prob + outlierUniformLogLikelihood]
        //  + (1 - z_t) * [log(1 - outlier_prob) - log(2 * pi * variance)/2 - (coverage_t - mean_t)^2 / (2 * variance)]
        //  + const
        //where z_t is the indicator for target t, and outlier_prob is the outlier probability.
        //note that we compute the normalizing constant, so that we can sample a new indicator value by simply sampling
        //uniformly in [0, 1] and checking whether the resulting value is less than the probability of being an outlier
        //(corresponding to the first line in the unnormalized expression above)
        final Sampler<OutlierIndicators, CopyRatioState, CopyRatioDataCollection> outlierIndicatorsSampler =<commit_msg>Revert "Fixed edge case for infinite segment-mean bug."<commit_after>package org.broadinstitute.hellbender.tools.exome;

import com.google.common.annotations.VisibleForTesting;
import com.google.common.primitives.Doubles;
import org.apache.commons.math3.distribution.BetaDistribution;
import org.apache.commons.math3.stat.descriptive.moment.Mean;
import org.apache.commons.math3.stat.descriptive.moment.Variance;
import org.apache.spark.api.java.JavaSparkContext;
import org.broadinstitute.hellbender.utils.MathUtils;
import org.broadinstitute.hellbender.utils.mcmc.*;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.function.Function;
import java.util.stream.Collectors;
import java.util.stream.IntStream;

/**
 * Represents an ACNV segmented model for copy ratio fit to target-coverage data.
 * The log_2 coverages in each segment are fit by a mixture model with a normal-distribution component
 * and a uniform outlier component.  The variance of the normal-distribution component and the relative
 * contribution of the uniform outlier component in all segments are both assumed to be global parameters.
 * The mean of the normal-distribution component in each segment is taken to be a segment-level parameter.
 * The component (i.e., normal or outlier) that each target coverage is drawn from is determined by a latent
 * target-level indicator parameter.
 *
 * @author <NAME> &lt;<EMAIL>&gt;
 */
public final class ACNVCopyRatioModeller {
    private static final double EPSILON = 1E-10;
    private static final double VARIANCE_MIN = EPSILON;
    private static final double VARIANCE_MAX = 10.;

    private static final double OUTLIER_PROBABILITY_INITIAL = 0.05;
    private static final double OUTLIER_PROBABILITY_PRIOR_ALPHA = 5.;
    private static final double OUTLIER_PROBABILITY_PRIOR_BETA = 95.;

    private final SegmentedModel segmentedModel;

    private final double outlierUniformLogLikelihood;
    private final double varianceSliceSamplingWidth;
    private final double meanSliceSamplingWidth;

    private final ParameterizedModel<CopyRatioState, CopyRatioDataCollection> model;

    private final List<Double> varianceSamples = new ArrayList<>();
    private final List<Double> outlierProbabilitySamples = new ArrayList<>();
    private final List<SegmentMeans> segmentMeansSamples = new ArrayList<>();
    private final List<OutlierIndicators> outlierIndicatorsSamples = new ArrayList<>();

    //Calculates the exponent for a normal distribution; used in log-likelihood calculation below.
    private static double normalTerm(final double quantity, final double mean, final double variance) {
        return (quantity - mean) * (quantity - mean) / (2. * variance);
    }

    /**
     * Constructs a copy-ratio model given a SegmentedModel with segments and a Genome.  Initial point estimates of
     * parameters are set to empirical estimates where available.
     * @param segmentedModel    SegmentedModel with segments and a Genome
     */
    public ACNVCopyRatioModeller(final SegmentedModel segmentedModel) {
        this.segmentedModel = segmentedModel;

        //load segmented coverages from SegmentedModel into CopyRatioDataCollection
        final CopyRatioDataCollection data = new CopyRatioDataCollection(segmentedModel);

        //the uniform log-likelihood for outliers is determined by the minimum and maximum coverages in the dataset;
        //the outlier-probability parameter should be interpreted accordingly
        outlierUniformLogLikelihood = -Math.log(data.coverageRange);

        //set widths for slice sampling of variance and segment-mean posteriors using empirical variance estimate.
        //variance posterior is inverse chi-squared, segment-mean posteriors are Gaussian; the below expressions
        //approximate the standard deviations of these distributions.
        final double varianceEstimate = data.estimateVariance();
        varianceSliceSamplingWidth = Math.sqrt(2. * varianceEstimate / data.numTargets);
        meanSliceSamplingWidth = Math.sqrt(varianceEstimate * data.numSegments / data.numTargets);

        //use empirical segment means and empirical average variance across segments to initialize CopyRatioState
        final CopyRatioState initialState = new CopyRatioState(data);

        //samples log conditional posterior for the variance parameter, assuming uniform prior; this is given by
        //the product of Gaussian likelihoods for each non-outlier target t:
        //  log[product_{non-outlier t} variance^(-1/2) * exp(-(coverage_t - mean_t)^2 / (2 * variance))] + constant
        //where mean_t is identical for all targets in a segment
        final Sampler<Double, CopyRatioState, CopyRatioDataCollection> varianceSampler =
                (rng, state, dataCollection) -> {
            final Function<Double, Double> logConditionalPDF = newVariance -> {
                double ll = 0.;
                int numNotOutliers = 0;
                for (int segment = 0; segment < dataCollection.numSegments; segment++) {
                    final int numTargetsInSegment = dataCollection.getNumTargetsInSegment(segment);
                    final int targetStartInSegment = dataCollection.getStartTargetInSegment(segment);
                    for (int target = 0; target < numTargetsInSegment; target++) {
                        final double coverage = dataCollection.getCoveragesInSegment(segment).get(target);
                        if (!state.getOutlierIndicator(targetStartInSegment + target)) {
                            ll -= normalTerm(coverage, state.getMeanInSegment(segment), newVariance);
                            numNotOutliers++;
                        }
                    }
                }
                ll -= 0.5 * Math.log(newVariance) * numNotOutliers;
                return ll;
            };
            return new SliceSampler(rng, logConditionalPDF, state.variance(),
                    VARIANCE_MIN, VARIANCE_MAX, varianceSliceSamplingWidth).sample();
        };

        //samples log conditional posterior for the outlier-probability parameter, assuming Beta(alpha, beta) prior;
        //this is given by:
        //  log Beta(alpha + number of outlier targets, beta + number of non-outlier targets) + constant
        final Sampler<Double, CopyRatioState, CopyRatioDataCollection> outlierProbabilitySampler =
                (rng, state, dataCollection) -> {
            final int numOutliers = (int) IntStream.range(0, dataCollection.numTargets).filter(state::getOutlierIndicator).count();
            return new BetaDistribution(
                    rng,
                    OUTLIER_PROBABILITY_PRIOR_ALPHA + numOutliers,
                    OUTLIER_PROBABILITY_PRIOR_BETA + dataCollection.numTargets - numOutliers).sample();
        };

        //samples log conditional posteriors for the segment-mean parameters, assuming uniform priors;
        //for each segment s, this is given by the product of Gaussian likelihoods for each non-outlier target t:
        //  log[product_{non-outlier t in s} exp(-(coverage_t - mean_s)^2 / (2 * variance))] + constant
        final Sampler<SegmentMeans, CopyRatioState, CopyRatioDataCollection> segmentMeansSampler =
                (rng, state, dataCollection) -> {
            final List<Double> means = new ArrayList<>(dataCollection.numSegments);
            for (int segment = 0; segment < dataCollection.numSegments; segment++) {
                final int numTargetsInSegment = dataCollection.getNumTargetsInSegment(segment);
                if (numTargetsInSegment == 0) {
                    means.add(Double.NaN);
                } else {
                    //get coverages at non-outlier targets in the segment and construct log conditional PDF
                    final int targetStartInSegment = dataCollection.getStartTargetInSegment(segment);
                    final List<Double> allCoveragesInSegment = dataCollection.getCoveragesInSegment(segment);
                    final List<Double> nonOutlierCoveragesInSegment = IntStream.range(0, numTargetsInSegment).boxed()
                            .filter(target -> !state.getOutlierIndicator(targetStartInSegment + target))
                            .map(allCoveragesInSegment::get)
                            .collect(Collectors.toList());
                    final Function<Double, Double> logConditionalPDF = newMean -> nonOutlierCoveragesInSegment.stream()
                            .mapToDouble(coverage -> -normalTerm(coverage, newMean, state.variance()))
                            .sum();
                    //slice sample within range given by minimum and maximum non-outlier coverages
                    final double minimumCoverageInSegment = Collections.min(nonOutlierCoveragesInSegment);
                    final double maximumCoverageInSegment = Collections.max(nonOutlierCoveragesInSegment);
                    final SliceSampler sampler =
                            new SliceSampler(rng, logConditionalPDF, state.getMeanInSegment(segment),
                                    minimumCoverageInSegment, maximumCoverageInSegment, meanSliceSamplingWidth);
                    means.add(sampler.sample());
                }
            }
            return new SegmentMeans(means);
        };

        //samples log conditional posteriors for the outlier-indicator parameters; for each target t, this is given by:
        //          z_t * [log outlier_prob + outlierUniformLogLikelihood]
        //  + (1 - z_t) * [log(1 - outlier_prob) - log(2 * pi * variance)/2 - (coverage_t - mean_t)^2 / (2 * variance)]
        //  + const
        //where z_t is the indicator for target t, and outlier_prob is the outlier probability.
        //note that we compute the normalizing constant, so that we can sample a new indicator value by simply sampling
        //uniformly in [0, 1] and checking whether the resulting value is less than the probability of being an outlier
        //(corresponding to the first line in the unnormalized expression above)
        final Sampler<OutlierIndicators, CopyRatioState, CopyRatioDataCollection> outlierIndicatorsSampler =
{'fraction_non_alphanumeric': 0.07350049164208457, 'fraction_numerical': 0.004031465093411996, 'mean_word_length': 2.806324850299401, 'pattern_counts': {'":': 0, '<': 51, '<?xml version=': 0, '>': 59, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6747362', 'n_tokens_mistral': 5287, 'n_tokens_neox': 4939, 'n_words': 1578}