<commit_before>'''
Functions for generating simulated data and shocks.
'''

from __future__ import division
import warnings                             # A library for runtime warnings
import numpy as np                          # Numerical Python

def drawMeanOneLognormal(N, sigma=1.0, seed=0):
    '''
    Generate arrays of mean one lognormal draws. The sigma input can be a number
    or list-like.  If a number, output is a length N array of draws from the
    lognormal distribution with standard deviation sigma. If a list, output is
    a length T list whose t-th entry is a length N array of draws from the
    lognormal with standard deviation sigma[t].

    Parameters
    ----------
    N : int
        Number of draws in each row.
    sigma : float or [float]
        One or more standard deviations. Number of elements T in sigma
        determines number of rows of output.
    seed : int
        Seed for random number generator.

    Returns:
    ------------
    draws : np.array or [np.array]
        T-length list of arrays of mean one lognormal draws each of size N, or
        a single array of size N (if sigma is a scalar).
    '''
    # Set up the RNG
    RNG = np.random.RandomState(seed)
    
    if type(sigma) == float: # Return a single array of length N
        mu = -0.5*sigma**2
        draws = RNG.lognormal(mean=mu, sigma=sigma, size=N)
    else: # Set up empty list to populate, then loop and populate list with draws
        draws=[]
        for sig in sigma:            
            mu = -0.5*(sig**2)
            draws.append(RNG.lognormal(mean=mu, sigma=sig, size=N))            
    return draws
    
def drawLognormal(N,mu=0.0,sigma=1.0,seed=0):
    '''
    Generate arrays of mean one lognormal draws. The sigma input can be a number
    or list-like.  If a number, output is a length N array of draws from the
    lognormal distribution with standard deviation sigma. If a list, output is
    a length T list whose t-th entry is a length N array of draws from the
    lognormal with standard deviation sigma[t].

    Parameters
    ----------
    N : int
        Number of draws in each row.
    sigma : float or [float]
        One or more standard deviations. Number of elements T in sigma
        determines number of rows of output.
    seed : int
        Seed for random number generator.

    Returns:
    ------------
    draws : np.array or [np.array]
        T-length list of arrays of mean one lognormal draws each of size N, or
        a single array of size N (if sigma is a scalar).
    '''
    # Set up the RNG
    RNG = np.random.RandomState(seed)
    
    if type(sigma) == float: # Return a single array of length N
        draws = RNG.lognormal(mean=mu, sigma=sigma, size=N)
    else: # Set up empty list to populate, then loop and populate list with draws
        draws=[]
        for j in range(len(sigma)):       
            draws.append(RNG.lognormal(mean=mu[j], sigma=sigma[j], size=N))            
    return draws
    
    
def drawNormal(N, mu=0.0, sigma=1.0, seed=0):
    '''
    Generate arrays of normal draws.  The mu and sigma inputs can be numbers or
    list-likes.  If a number, output is a length N array of draws from the normal
    distribution with mean mu and standard deviation sigma. If a list, output is
    a length T list whose t-th entry is a length N array with draws from the
    normal distribution with mean mu[t] and standard deviation sigma[t].

    Parameters
    ----------
    N : int
        Number of draws in each row.
    mu : float or [float]
        One or more means.  Number of elements T in mu determines number of rows
        of output.
    sigma : float or [float]
        One or more standard deviations. Number of elements T in sigma
        determines number of rows of output.
    seed : int<commit_msg>drawLognormal can handle sigma = 0<commit_after>'''
Functions for generating simulated data and shocks.
'''

from __future__ import division
import warnings                             # A library for runtime warnings
import numpy as np                          # Numerical Python

def drawMeanOneLognormal(N, sigma=1.0, seed=0):
    '''
    Generate arrays of mean one lognormal draws. The sigma input can be a number
    or list-like.  If a number, output is a length N array of draws from the
    lognormal distribution with standard deviation sigma. If a list, output is
    a length T list whose t-th entry is a length N array of draws from the
    lognormal with standard deviation sigma[t].

    Parameters
    ----------
    N : int
        Number of draws in each row.
    sigma : float or [float]
        One or more standard deviations. Number of elements T in sigma
        determines number of rows of output.
    seed : int
        Seed for random number generator.

    Returns:
    ------------
    draws : np.array or [np.array]
        T-length list of arrays of mean one lognormal draws each of size N, or
        a single array of size N (if sigma is a scalar).
    '''
    # Set up the RNG
    RNG = np.random.RandomState(seed)
    
    if type(sigma) == float: # Return a single array of length N
        mu = -0.5*sigma**2
        draws = RNG.lognormal(mean=mu, sigma=sigma, size=N)
    else: # Set up empty list to populate, then loop and populate list with draws
        draws=[]
        for sig in sigma:            
            mu = -0.5*(sig**2)
            draws.append(RNG.lognormal(mean=mu, sigma=sig, size=N))            
    return draws
    
def drawLognormal(N,mu=0.0,sigma=1.0,seed=0):
    '''
    Generate arrays of mean one lognormal draws. The sigma input can be a number
    or list-like.  If a number, output is a length N array of draws from the
    lognormal distribution with standard deviation sigma. If a list, output is
    a length T list whose t-th entry is a length N array of draws from the
    lognormal with standard deviation sigma[t].

    Parameters
    ----------
    N : int
        Number of draws in each row.
    mu : float or [float]
        One or more means.  Number of elements T in mu determines number
        of rows of output.
    sigma : float or [float]
        One or more standard deviations. Number of elements T in sigma
        determines number of rows of output.
    seed : int
        Seed for random number generator.

    Returns:
    ------------
    draws : np.array or [np.array]
        T-length list of arrays of mean one lognormal draws each of size N, or
        a single array of size N (if sigma is a scalar).
    '''
    # Set up the RNG
    RNG = np.random.RandomState(seed)
    
    if type(sigma) == float: # Return a single array of length N
        if sigma == 0:
            draws = np.exp(mu)*np.ones(N)
        else:
            draws = RNG.lognormal(mean=mu, sigma=sigma, size=N)
    else: # Set up empty list to populate, then loop and populate list with draws
        draws=[]
        for j in range(len(sigma)):
            if sigma[j] == 0:
                draws.append(np.exp(mu[j])*np.ones(N))
            else:
                draws.append(RNG.lognormal(mean=mu[j], sigma=sigma[j], size=N))            
    return draws
    
    
def drawNormal(N, mu=0.0, sigma=1.0, seed=0):
    '''
    Generate arrays of normal draws.  The mu and sigma inputs can be numbers or
    list-likes.  If a number, output is a length N array of draws from the normal
    distribution with mean mu and standard deviation sigma. If a list, output is
    a length T list whose t-th entry is a length N array with draws from the
    normal distribution with mean mu[t] and standard deviation sigma[t].

    Parameters
    ----------
    N : int
        Number of draws in each row.
    mu : float or [float]
        One or more means.  Number of elements T in mu determines number of rows
        of output.
    sigma : float or [float]
        One or more standard deviations. Number of elements T in sigma
        determines number of rows of output.
    seed : int
{'fraction_non_alphanumeric': 0.0788152610441767, 'fraction_numerical': 0.005145582329317269, 'mean_word_length': 2.083978328173375, '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': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7011365', 'n_tokens_mistral': 2251, 'n_tokens_neox': 2197, 'n_words': 1040}