<commit_before># -*- coding: utf-8 -*-
import scipy.stats as ss
import numpy.random as npr
from functools import partial

from . import core


def spr_op(distribution, size, input_dict):
    prng = npr.RandomState(0)
    prng.set_state(input_dict['random_state'])
    size = (input_dict['n'],)+tuple(size)
    data = distribution.rvs(*input_dict['data'], size=size, random_state=prng)
    return core.to_output(input_dict, data=data, random_state=prng.get_state())


class ScipyRV(core.RandomStateMixin, core.Operation):
    """
    Allows any distribution inheriting scipy.stats.rv_continuous or
    scipy.stats.rv_discrete. In the latter case methods pdf and logpdf
    are mapped to pmf and logpmf.

    Examples
    --------
    ScipyRV('tau', scipy.stats.norm, 5, size=(2,3))
    """
    def __init__(self, name, distribution, *params, size=(1,)):
        if isinstance(distribution, str):
            distribution = getattr(ss, distribution)
        self.distribution = distribution
        if not isinstance(size, tuple):
            size = (size,)
        op = partial(spr_op, distribution, size)
        super(ScipyRV, self).__init__(name, op, *params)

    @property
    def is_discrete(self):
        return isinstance(self.distribution, ss.rv_discrete)

    def pdf(self, x, *params):
        """
        Probability density function at x of the given RV.
        """
        params = self._get_params(*params)
        if self.is_discrete:
            return self.distribution.pmf(x, *params)
        else:<commit_msg>Allow alias strings<commit_after># -*- coding: utf-8 -*-
import scipy.stats as ss
import numpy.random as npr
from functools import partial

from . import core


def spr_op(distribution, size, input_dict):
    prng = npr.RandomState(0)
    prng.set_state(input_dict['random_state'])
    size = (input_dict['n'],)+tuple(size)
    data = distribution.rvs(*input_dict['data'], size=size, random_state=prng)
    return core.to_output(input_dict, data=data, random_state=prng.get_state())


class ScipyRV(core.RandomStateMixin, core.Operation):
    """
    Allows any distribution inheriting scipy.stats.rv_continuous or
    scipy.stats.rv_discrete. In the latter case methods pdf and logpdf
    are mapped to pmf and logpmf.

    Examples
    --------
    ScipyRV('tau', scipy.stats.norm, 5, size=(2,3))
    """

    # Convert some common names to scipy equivalents
    ALIASES = {'normal': 'norm'}

    def __init__(self, name, distribution, *params, size=(1,)):
        if isinstance(distribution, str):
            distribution = getattr(ss, self.ALIASES.get(distribution, distribution))
        self.distribution = distribution
        if not isinstance(size, tuple):
            size = (size,)
        op = partial(spr_op, distribution, size)
        super(ScipyRV, self).__init__(name, op, *params)

    @property
    def is_discrete(self):
        return isinstance(self.distribution, ss.rv_discrete)

    def pdf(self, x, *params):
        """
        Probability density function at x of the given RV.
        """
        params = self._get_params(*params)
        if self.is_discrete:
            return self.distribution.pmf(x, *params)
        else:
{'fraction_non_alphanumeric': 0.11989956057752668, 'fraction_numerical': 0.003766478342749529, 'mean_word_length': 3.023989898989899, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, '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': '3911174', 'n_tokens_mistral': 1013, 'n_tokens_neox': 978, 'n_words': 250}