<commit_before>""" Utilities for computing the probability mass function. """
import scipy.stats
import scipy.interpolate
import numpy
from fgivenx.parallel import parallel_apply
from fgivenx.io import CacheException, Cache


def PMF(samples, y):
    """ Compute the probability mass function.

        The set of samples defines a probability density P(y),
        which is computed using a kernel density estimator.

        From P(y) we define:

                    /
        PMF(p) =    | P(t) dt
                    /
                P(t) < p

        This is the cumulative distribution function expressed as a
        function of the probability

        We aim to compute M(y), which indicates the amount of
        probability contained outside the iso-probability contour
        passing through y.


         ^ P(y)                   ...
         |                     | .   .
         |                     |       .
        p|- - - - - - - - - - .+- - - - . - - - - - - - - - - -
         |                   .#|        #.
         |                  .##|        ##.
         |                  .##|        ##.
         |                 .###|        ###.     M(p)
         |                 .###|        ###.     is the
         |                 .###|        ###.     shaded area
         |                .####|        ####.
         |                .####|        ####.
         |              ..#####|        #####..
         |          ....#######|        #######....
         |         .###########|        ###########.
         +---------------------+-------------------------------> y
                              t

         ^ M(p)                        ^ M(y)
         |                             |
        1|                +++         1|         +
         |               +             |        + +
         |       ++++++++              |       +   +
         |     ++                      |     ++     ++
         |   ++                        |   ++         ++
         |+++                          |+++             +++
         +---------------------> p     +---------------------> y
        0                   1

        Parameters
        ----------
        samples: array-like
            Array of samples from a probability density P(y).

        y: array-like (optional)
            Array to evaluate the PDF at

        Returns
        -------
        PMF evaluated at each y value

    """
    # Remove any nans from the samples
    samples = numpy.array(samples)
    samples = samples[~numpy.isnan(samples)]
    try:
        # Compute the kernel density estimate
        kernel = scipy.stats.gaussian_kde(samples)

        # Add two more samples definitely outside the range and sort them
        mn = min(samples) - 10*numpy.sqrt(kernel.covariance[0,0])
        mx = max(samples) + 10*numpy.sqrt(kernel.covariance[0,0]) 
        y_ = numpy.linspace(mn, mx, len(y)*10)

        # Compute the probabilities at each of the extended samples
        ps_ = kernel(y_)

        # Compute the masses
        ms = []
        for yi in y:
            # compute the probability at this y value
            p = kernel(yi)
            if p <= min(ps_) or yi < mn or yi > mx:
                m = 0.
            else:
                # Find out which samples have greater probability than P(y)
                bools = ps_>p

                # Compute indices where to start and stop the integration
                stops = numpy.where(numpy.logical_and(~bools[:-1], bools[1:]))[0]
                starts = numpy.where(numpy.logical_and(bools[:-1], ~bools[1:]))[0]

                # Compute locations
                starts =  [mn] + [y_[i] if numpy.isclose(kernel(y_[i]),p) else y_[i+1] if numpy.isclose(kernel(y_[i+1]),p) else scipy.optimize.brentq(lambda u: kernel(u)-p,y_[i], y_[i+1]) for i in starts]
                stops = [y_[i] if numpy.isclose(kernel(y_[i]),p) else y_[i+1] if numpy.isclose(kernel(y_[i+1]),p) else scipy.optimize.brentq(lambda u: kernel(u)-p,y_[i], y_[i+1]) for i in stops] + [mx]

                # Sum up the masses
                m = sum(kernel.integrate_box_1d(a, b) for a, b in zip(starts, stops))
            ms.append(m)
        return numpy.array(ms)

    except numpy.linalg.LinAlgError:
        return numpy.zeros_like(y)

def compute_pmf(fsamps, y, **kwargs):
    """ Compute the pmf defined by fsamps at each x for each y.

    Parameters
    ----------<commit_msg>Bug fix to mass function. Now much more robust<commit_after>""" Utilities for computing the probability mass function. """
import scipy.stats
import scipy.interpolate
import numpy
from fgivenx.parallel import parallel_apply
from fgivenx.io import CacheException, Cache


def PMF(samples, y):
    """ Compute the probability mass function.

        The set of samples defines a probability density P(y),
        which is computed using a kernel density estimator.

        From P(y) we define:

                    /
        PMF(p) =    | P(t) dt
                    /
                P(t) < p

        This is the cumulative distribution function expressed as a
        function of the probability

        We aim to compute M(y), which indicates the amount of
        probability contained outside the iso-probability contour
        passing through y.


         ^ P(y)                   ...
         |                     | .   .
         |                     |       .
        p|- - - - - - - - - - .+- - - - . - - - - - - - - - - -
         |                   .#|        #.
         |                  .##|        ##.
         |                  .##|        ##.
         |                 .###|        ###.     M(p)
         |                 .###|        ###.     is the
         |                 .###|        ###.     shaded area
         |                .####|        ####.
         |                .####|        ####.
         |              ..#####|        #####..
         |          ....#######|        #######....
         |         .###########|        ###########.
         +---------------------+-------------------------------> y
                              t

         ^ M(p)                        ^ M(y)
         |                             |
        1|                +++         1|         +
         |               +             |        + +
         |       ++++++++              |       +   +
         |     ++                      |     ++     ++
         |   ++                        |   ++         ++
         |+++                          |+++             +++
         +---------------------> p     +---------------------> y
        0                   1

        Parameters
        ----------
        samples: array-like
            Array of samples from a probability density P(y).

        y: array-like (optional)
            Array to evaluate the PDF at

        Returns
        -------
        PMF evaluated at each y value

    """
    # Remove any nans from the samples
    samples = numpy.array(samples)
    samples = samples[~numpy.isnan(samples)]
    try:
        # Compute the kernel density estimate
        kernel = scipy.stats.gaussian_kde(samples)

        # Add two more samples definitely outside the range and sort them
        mn = min(samples) - 10*numpy.sqrt(kernel.covariance[0,0])
        mx = max(samples) + 10*numpy.sqrt(kernel.covariance[0,0]) 
        y_ = numpy.linspace(mn, mx, len(y)*10)

        # Compute the probabilities at each of the extended samples
        ps_ = kernel(y_)

        # Compute the masses
        ms = []
        for yi in y:
            # compute the probability at this y value
            p = kernel(yi)
            if p <= max(ps_)*1e-5:
                m = 0.
            else:
                # Find out which samples have greater probability than P(y)
                bools = ps_>p

                # Compute indices where to start and stop the integration
                stops = numpy.where(numpy.logical_and(~bools[:-1], bools[1:]))[0]
                starts = numpy.where(numpy.logical_and(bools[:-1], ~bools[1:]))[0]

                # Compute locations
                starts = [-numpy.inf] + [scipy.optimize.brentq(lambda u: kernel(u)-p,y_[i], y_[i+1]) for i in starts]
                stops = [scipy.optimize.brentq(lambda u: kernel(u)-p,y_[i], y_[i+1]) for i in stops] + [numpy.inf]

                # Sum up the masses
                m = sum(kernel.integrate_box_1d(a, b) for a, b in zip(starts, stops))
            ms.append(m)
        return numpy.array(ms)

    except numpy.linalg.LinAlgError:
        return numpy.zeros_like(y)

def compute_pmf(fsamps, y, **kwargs):
    """ Compute the pmf defined by fsamps at each x for each y.

    Parameters
    ----------
{'fraction_non_alphanumeric': 0.14918521918751435, 'fraction_numerical': 0.006196924489327519, 'mean_word_length': 1.0338389731621938, 'pattern_counts': {'":': 0, '<': 8, '<?xml version=': 0, '>': 12, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 34, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4176972', 'n_tokens_mistral': 2452, 'n_tokens_neox': 2305, 'n_words': 612}