<commit_before>from HARK.utilities import memoize
import math
import numpy as np
from scipy.special import erf, erfc
import scipy.stats as stats
import warnings


class DiscreteDistribution():
    """
    A representation of a discrete probability distribution.

    """
    pmf = None
    X = None

    def __init__(self, pmf, X):
        '''
        Initialize a discrete distribution.

        Parameters
        ----------
        pmf : np.array
            An array of floats representing a probability mass function. 
        X : np.array or [np.array]
            Discrete point values for each probability mass.
            May be multivariate (list of arrays).
        
        '''
        self.pmf = pmf
        self.X = X

        # Very quick and incomplete parameter check:
        # TODO: Check that pmf adn X arrays have same length.

    def dim(self):
        if isinstance(self.X, list):
            return len(self.X)
        else:
            return 1

    def draw_events(self, n, seed=0):
        '''
        Draws N 'events' from the distribution PMF.
        These events are indices into X.
        '''
        # Set up the RNG
        RNG = np.random.RandomState(seed)
        # Generate a cumulative distribution
        base_draws = RNG.uniform(size=n)
        cum_dist = np.cumsum(self.pmf)

        # Convert the basic uniform draws into discrete draws
        indices = cum_dist.searchsorted(base_draws)

        return indices

    def drawDiscrete(self, N,X=None,exact_match=False,seed=0):
        '''
        Simulates N draws from a discrete distribution with probabilities P and outcomes X.

        Parameters
        ----------
        N : int
            Number of draws to simulate.
        X : None, int, or np.array
            If None, then use this distribution's X for point values.
            If an int, then the index of X for the point values.
            If an np.array, use the array for the point values.
        exact_match : boolean
            Whether the draws should "exactly" match the discrete distribution (as
            closely as possible given finite draws).  When True, returned draws are
            a random permutation of the N-length list that best fits the discrete
            distribution.  When False (default), each draw is independent from the
            others and the result could deviate from the input.
        seed : int
            Seed for random number generator.

        Returns
        -------
        draws : np.array
            An array draws from the discrete distribution; each element is a value in X.
        '''
        if X is None:
            X = self.X
        elif isinstance(X,int):
            X = self.X[Xdim]
        else:
            X = X
        
        # Set up the RNG
        RNG = np.random.RandomState(seed)

        if exact_match:
            events = np.arange(self.pmf.size) # just a list of integers
            cutoffs = np.round(np.cumsum(self.pmf)*N).astype(int) # cutoff points between discrete outcomes
            top = 0
            # Make a list of event indices that closely matches the discrete distribution
            event_list        = []
            for j in range(events.size):
                bot = top
                top = cutoffs[j]
                event_list += (top-bot)*[events[j]]
                # Randomly permute the event indices and store the corresponding results
                event_draws = RNG.permutation(event_list)
                draws = X[event_draws]
        else:
            indices = self.draw_events(self, N, seed=seed)
            draws = np.asarray(X)[indices]

        return draws

def approxLognormal(N, mu=0.0, sigma=1.0, tail_N=0, tail_bound=[0.02,0.98], tail_order=np.e):
    '''
    Construct a discrete approximation to a lognormal distribution with underlying
    normal distribution N(mu,sigma).  Makes an equiprobable distribution by
    default, but user can optionally request augmented tails with exponentially
    sized point masses.  This can improve solution accuracy in some models.

    Parameters
    ----------
    N: int
        Number of discrete points in the "main part" of the approximation.
    mu: float
        Mean of underlying normal distribution.
    sigma: float
        Standard deviation of underlying normal distribution.
    tail_N: int
        Number of points in each "tail part" of the approximation; 0 = no tail.
    tail_bound: [float]
        CDF boundaries of the tails vs main portion; tail_bound[0] is the lower
        tail bound, tail_bound[1] is the upper tail bound.  Inoperative when
        tail_N = 0.  Can make "one tailed" approximations with 0.0 or 1.0.
    tail_order: float
        Factor by which consecutive point masses in a "tail part" differ in
        probability.  Should be >= 1 for sensible spacing.
<commit_msg>fix bug in draw_events<commit_after>from HARK.utilities import memoize
import math
import numpy as np
from scipy.special import erf, erfc
import scipy.stats as stats
import warnings


class DiscreteDistribution():
    """
    A representation of a discrete probability distribution.

    """
    pmf = None
    X = None

    def __init__(self, pmf, X):
        '''
        Initialize a discrete distribution.

        Parameters
        ----------
        pmf : np.array
            An array of floats representing a probability mass function. 
        X : np.array or [np.array]
            Discrete point values for each probability mass.
            May be multivariate (list of arrays).
        
        '''
        self.pmf = pmf
        self.X = X

        # Very quick and incomplete parameter check:
        # TODO: Check that pmf adn X arrays have same length.

    def dim(self):
        if isinstance(self.X, list):
            return len(self.X)
        else:
            return 1

    def draw_events(self, n, seed=0):
        '''
        Draws N 'events' from the distribution PMF.
        These events are indices into X.
        '''
        # Set up the RNG
        RNG = np.random.RandomState(seed)
        # Generate a cumulative distribution
        base_draws = RNG.uniform(size=n)
        cum_dist = np.cumsum(self.pmf)

        # Convert the basic uniform draws into discrete draws
        indices = cum_dist.searchsorted(base_draws)

        return indices

    def drawDiscrete(self, N,X=None,exact_match=False,seed=0):
        '''
        Simulates N draws from a discrete distribution with probabilities P and outcomes X.

        Parameters
        ----------
        N : int
            Number of draws to simulate.
        X : None, int, or np.array
            If None, then use this distribution's X for point values.
            If an int, then the index of X for the point values.
            If an np.array, use the array for the point values.
        exact_match : boolean
            Whether the draws should "exactly" match the discrete distribution (as
            closely as possible given finite draws).  When True, returned draws are
            a random permutation of the N-length list that best fits the discrete
            distribution.  When False (default), each draw is independent from the
            others and the result could deviate from the input.
        seed : int
            Seed for random number generator.

        Returns
        -------
        draws : np.array
            An array draws from the discrete distribution; each element is a value in X.
        '''
        if X is None:
            X = self.X
        elif isinstance(X,int):
            X = self.X[Xdim]
        else:
            X = X
        
        # Set up the RNG
        RNG = np.random.RandomState(seed)

        if exact_match:
            events = np.arange(self.pmf.size) # just a list of integers
            cutoffs = np.round(np.cumsum(self.pmf)*N).astype(int) # cutoff points between discrete outcomes
            top = 0
            # Make a list of event indices that closely matches the discrete distribution
            event_list        = []
            for j in range(events.size):
                bot = top
                top = cutoffs[j]
                event_list += (top-bot)*[events[j]]
                # Randomly permute the event indices and store the corresponding results
                event_draws = RNG.permutation(event_list)
                draws = X[event_draws]
        else:
            indices = self.draw_events(N, seed=seed)
            draws = np.asarray(X)[indices]

        return draws

def approxLognormal(N, mu=0.0, sigma=1.0, tail_N=0, tail_bound=[0.02,0.98], tail_order=np.e):
    '''
    Construct a discrete approximation to a lognormal distribution with underlying
    normal distribution N(mu,sigma).  Makes an equiprobable distribution by
    default, but user can optionally request augmented tails with exponentially
    sized point masses.  This can improve solution accuracy in some models.

    Parameters
    ----------
    N: int
        Number of discrete points in the "main part" of the approximation.
    mu: float
        Mean of underlying normal distribution.
    sigma: float
        Standard deviation of underlying normal distribution.
    tail_N: int
        Number of points in each "tail part" of the approximation; 0 = no tail.
    tail_bound: [float]
        CDF boundaries of the tails vs main portion; tail_bound[0] is the lower
        tail bound, tail_bound[1] is the upper tail bound.  Inoperative when
        tail_N = 0.  Can make "one tailed" approximations with 0.0 or 1.0.
    tail_order: float
        Factor by which consecutive point masses in a "tail part" differ in
        probability.  Should be >= 1 for sensible spacing.

{'fraction_non_alphanumeric': 0.07068413992364049, 'fraction_numerical': 0.004953049220926633, 'mean_word_length': 1.9009278659084106, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, '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': '877163', 'n_tokens_mistral': 2582, 'n_tokens_neox': 2475, 'n_words': 1026}