HARK/distribution.py
class MarkovProcess(Distribution):
    """
    A representation of a discrete Markov process.

    Parameters
    ----------
    transition_matrix : np.array
        An array of floats representing a probability mass for
        each state transition.
    seed : int
        Seed for random number generator.

    """

    transition_matrix = None

    def __init__(self, transition_matrix, seed=0):
        """
        Initialize a discrete distribution.

        """
        self.transition_matrix = transition_matrix

        # Set up the RNG
        super().__init__(seed)

    def draw(self, state):
        """
        Draw new states fromt the transition matrix.

        Parameters
        ----------
        state : int or nd.array
            The state or states (1-D array) from which to draw new states.

        Returns
        -------
        new_state : int or nd.array
            New states.
        """

        def sample(s):
            return self.RNG.choice(
                self.transition_matrix.shape[1], p=self.transition_matrix[s, :]
            )

        array_sample = np.frompyfunc(sample, 1, 1)

        return array_sample(state)
HARK/distribution.py
class Uniform(Distribution):
    """
    A Uniform distribution.

    Parameters
    ----------
    bot : float or [float]
        One or more bottom values.
        Number of elements T in mu determines number
        of rows of output.
    top : float or [float]
        One or more top values.
        Number of elements T in top determines number of
        rows of output.
    seed : int
        Seed for random number generator.
    """

    bot = None
    top = None

    def __init__(self, bot=0.0, top=1.0, seed=0):
        self.bot = np.array(bot)
        self.top = np.array(top)
        # Set up the RNG
        self.RNG = np.random.RandomState(seed)

    def draw(self, N):
        """
        Generate arrays of uniform draws.  The bot and top inputs can be numbers or
        list-likes.  If a number, output is a length N array of draws from the
        uniform distribution on [bot,top]. If a list, output is a length T list
        whose t-th entry is a length N array with draws from the uniform distribution
        on [bot[t],top[t]].

        Parameters
        ----------
        N : int
            Number of draws in each row.

        Returns
        -------
        draws : np.array or [np.array]
            T-length list of arrays of uniform draws each of size N, or a single
            array of size N (if sigma is a scalar).
        """
        draws = []
        for j in range(self.bot.size):
            draws.append(
                self.bot.item(j)
                + (self.top.item(j) - self.bot.item(j)) * self.RNG.rand(N)
            )
        return draws[0] if len(draws) == 1 else draws

    def approx(self, N):
        """
        Makes a discrete approximation to this uniform distribution.

        Parameters
        ----------
        N : int
            The number of points in the discrete approximation

        Returns
        -------
        d : DiscreteDistribution
            Probability associated with each point in array of discrete
            points for discrete probability mass function.
        """
        pmf = np.ones(N) / float(N)
        center = (self.top + self.bot) / 2.0
        width = (self.top - self.bot) / 2.0
        X = center + width * np.linspace(-(N - 1.0) / 2.0, (N - 1.0) / 2.0, N) / (
            N / 2.0
        )
        return DiscreteDistribution(
            pmf, X, seed=self.RNG.randint(0, 2 ** 31 - 1, dtype="int32")
        )
HARK/distribution.py
class MVNormal(Distribution):
    """
    A Multivariate Normal distribution.

    Parameters
    ----------
    mu : numpy array
        Mean vector.
    Sigma : 2-d numpy array. Each dimension must have length equal to that of
            mu.
        Variance-covariance matrix.
    seed : int
        Seed for random number generator.
    """

    mu = None
    Sigma = None

    def __init__(self, mu=np.array([1, 1]), Sigma=np.array([[1, 0], [0, 1]]), seed=0):
        self.mu = mu
        self.Sigma = Sigma
        self.M = len(self.mu)
        super().__init__(seed)

    def draw(self, N):
        """
        Generate an array of multivariate normal draws.

        Parameters
        ----------
        N : int
            Number of multivariate draws.

        Returns
        -------
        draws : np.array
            Array of dimensions N x M containing the random draws, where M is
            the dimension of the multivariate normal and N is the number of
            draws. Each row represents a draw.
        """
        draws = self.RNG.multivariate_normal(self.mu, self.Sigma, N)

        return draws

    def approx(self, N, equiprobable=False):
        """
        Returns a discrete approximation of this distribution.

        The discretization will have N**M points, where M is the dimension of
        the multivariate normal.

        It uses the fact that:
            - Being positive definite, Sigma can be factorized as Sigma = QVQ',
              with V diagonal. So, letting A=Q*sqrt(V), Sigma = A*A'.
            - If Z is an N-dimensional multivariate standard normal, then
              A*Z ~ N(0,A*A' = Sigma).

        The idea therefore is to construct an equiprobable grid for a standard
        normal and multiply it by matrix A.
        """

        # Start by computing matrix A.
        v, Q = np.linalg.eig(self.Sigma)
        sqrtV = np.diag(np.sqrt(v))
        A = np.matmul(Q, sqrtV)

        # Now find a discretization for a univariate standard normal.
        if equiprobable:
            z_approx = Normal().approx_equiprobable(N)
        else:
            z_approx = Normal().approx(N)

        # Now create the multivariate grid and pmf
        Z = np.array(list(product(*[z_approx.X] * self.M)))
        pmf = np.prod(np.array(list(product(*[z_approx.pmf] * self.M))), axis=1)

        # Apply mean and standard deviation to the Z grid
        X = np.tile(np.reshape(self.mu, (1, self.M)), (N ** self.M, 1)) + np.matmul(
            Z, A.T
        )

        # Construct and return discrete distribution
        return DiscreteDistribution(
            pmf, X, seed=self.RNG.randint(0, 2 ** 31 - 1, dtype="int32")
        )
HARK/distribution.py
class MeanOneLogNormal(Lognormal):
    def __init__(self, sigma=1.0, seed=0):
        mu = -0.5 * sigma ** 2
        super().__init__(mu=mu, sigma=sigma, seed=seed)
HARK/distribution.py
class Bernoulli(Distribution):
    """
    A Bernoulli distribution.

    Parameters
    ----------
    p : float or [float]
        Probability or probabilities of the event occurring (True).

    seed : int
        Seed for random number generator.
    """

    p = None

    def __init__(self, p=0.5, seed=0):
        self.p = np.array(p)
        # Set up the RNG
        super().__init__(seed)

    def draw(self, N):
        """
        Generates arrays of booleans drawn from a simple Bernoulli distribution.
        The input p can be a float or a list-like of floats; its length T determines
        the number of entries in the output.  The t-th entry of the output is an
        array of N booleans which are True with probability p[t] and False otherwise.

        Arguments
        ---------
        N : int
            Number of draws in each row.

        Returns
        -------
        draws : np.array or [np.array]
            T-length list of arrays of Bernoulli draws each of size N, or a single
        array of size N (if sigma is a scalar).
        """
        draws = []
        for j in range(self.p.size):
            draws.append(self.RNG.uniform(size=N) < self.p.item(j))
        return draws[0] if len(draws) == 1 else draws
HARK/distribution.py
class Normal(Distribution):
    """
    A Normal distribution.

    Parameters
    ----------
    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.
    """

    mu = None
    sigma = None

    def __init__(self, mu=0.0, sigma=1.0, seed=0):
        self.mu = np.array(mu)
        self.sigma = np.array(sigma)
        super().__init__(seed)

    def draw(self, N):
        """
        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.

        Returns
        -------
        draws : np.array or [np.array]
            T-length list of arrays of normal draws each of size N, or a single array
            of size N (if sigma is a scalar).
        """
        draws = []
        for t in range(self.sigma.size):
            draws.append(self.sigma.item(t) * self.RNG.randn(N) + self.mu.item(t))

        return draws

    def approx(self, N):
        """
        Returns a discrete approximation of this distribution.
        """
        x, w = np.polynomial.hermite.hermgauss(N)
        # normalize w
        pmf = w * np.pi ** -0.5
        # correct x
        X = math.sqrt(2.0) * self.sigma * x + self.mu
        return DiscreteDistribution(
            pmf, X, seed=self.RNG.randint(0, 2 ** 31 - 1, dtype="int32")
        )

    def approx_equiprobable(self, N):
        CDF = np.linspace(0, 1, N + 1)
        lims = stats.norm.ppf(CDF)
        scores = (lims - self.mu) / self.sigma
        pdf = stats.norm.pdf(scores)

        # Find conditional means using Mills's ratio
        pmf = np.diff(CDF)
        X = self.mu - np.diff(pdf) / pmf

        return DiscreteDistribution(
            pmf, X, seed=self.RNG.randint(0, 2 ** 31 - 1, dtype="int32")
        )
HARK/distribution.py
class Weibull(Distribution):
    """
    A Weibull distribution.

    Parameters
    ----------
    scale : float or [float]
        One or more scales.  Number of elements T in scale
        determines number of
        rows of output.
    shape : float or [float]
        One or more shape parameters. Number of elements T in scale
        determines number of rows of output.
    seed : int
        Seed for random number generator.
    """

    scale = None
    shape = None

    def __init__(self, scale=1.0, shape=1.0, seed=0):
        self.scale = np.array(scale)
        self.shape = np.array(shape)
        # Set up the RNG
        super().__init__(seed)

    def draw(self, N):
        """
        Generate arrays of Weibull draws.  The scale and shape inputs can be
        numbers or list-likes.  If a number, output is a length N array of draws from
        the Weibull distribution with the given scale and shape. If a list, output
        is a length T list whose t-th entry is a length N array with draws from the
        Weibull distribution with scale scale[t] and shape shape[t].

        Note: When shape=1, the Weibull distribution is simply the exponential dist.

        Mean: scale*Gamma(1 + 1/shape)

        Parameters
        ----------
        N : int
            Number of draws in each row.

        Returns:
        ------------
        draws : np.array or [np.array]
            T-length list of arrays of Weibull draws each of size N, or a single
            array of size N (if sigma is a scalar).
        """
        draws = []
        for j in range(self.scale.size):
            draws.append(
                self.scale.item(j)
                * (-np.log(1.0 - self.RNG.rand(N))) ** (1.0 / self.shape.item(j))
            )
        return draws[0] if len(draws) == 1 else draws
HARK/distribution.py
class Lognormal(Distribution):
    """
    A Lognormal distribution

    Parameters
    ----------
    mu : float or [float]
        One or more means of underlying normal distribution.
        Number of elements T in mu determines number of rows of output.
    sigma : float or [float]
        One or more standard deviations of underlying normal distribution.
        Number of elements T in sigma determines number of rows of output.
    seed : int
        Seed for random number generator.
    """

    mu = None
    sigma = None

    def __init__(self, mu=0.0, sigma=1.0, seed=0):
        self.mu = np.array(mu)
        self.sigma = np.array(sigma)
        # Set up the RNG
        super().__init__(seed)

        if self.mu.size != self.sigma.size:
            raise Exception(
                "mu and sigma must be of same size, are %s, %s"
                % ((self.mu.size), (self.sigma.size))
            )

    def draw(self, N):
        """
        Generate arrays of lognormal draws. The sigma parameter 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.

        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).
        """

        draws = []
        for j in range(self.mu.size):
            draws.append(
                self.RNG.lognormal(
                    mean=self.mu.item(j), sigma=self.sigma.item(j), size=N
                )
            )
        # TODO: change return type to np.array?
        return draws[0] if len(draws) == 1 else draws

    def approx(self, N, tail_N=0, tail_bound=None, 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.
        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.

        Returns
        -------
        d : DiscreteDistribution
            Probability associated with each point in array of discrete
            points for discrete probability mass function.
        """
        tail_bound = tail_bound if tail_bound is not None else [0.02, 0.98]
        # Find the CDF boundaries of each segment
        if self.sigma > 0.0:
            if tail_N > 0:
                lo_cut = tail_bound[0]
                hi_cut = tail_bound[1]
            else:
                lo_cut = 0.0
                hi_cut = 1.0
            inner_size = hi_cut - lo_cut
            inner_CDF_vals = [
                lo_cut + x * N ** (-1.0) * inner_size for x in range(1, N)
            ]
            if inner_size < 1.0:
                scale = 1.0 / tail_order
                mag = (1.0 - scale ** tail_N) / (1.0 - scale)
            lower_CDF_vals = [0.0]
            if lo_cut > 0.0:
                for x in range(tail_N - 1, -1, -1):
                    lower_CDF_vals.append(
                        lower_CDF_vals[-1] + lo_cut * scale ** x / mag
                    )
            upper_CDF_vals = [hi_cut]
            if hi_cut < 1.0:
                for x in range(tail_N):
                    upper_CDF_vals.append(
                        upper_CDF_vals[-1] + (1.0 - hi_cut) * scale ** x / mag
                    )
            CDF_vals = lower_CDF_vals + inner_CDF_vals + upper_CDF_vals
            temp_cutoffs = list(
                stats.lognorm.ppf(
                    CDF_vals[1:-1], s=self.sigma, loc=0, scale=np.exp(self.mu)
                )
            )
            cutoffs = [0] + temp_cutoffs + [np.inf]
            CDF_vals = np.array(CDF_vals)

            K = CDF_vals.size - 1  # number of points in approximation
            pmf = CDF_vals[1 : (K + 1)] - CDF_vals[0:K]
            X = np.zeros(K)
            for i in range(K):
                zBot = cutoffs[i]
                zTop = cutoffs[i + 1]
                # Manual check to avoid the RuntimeWarning generated by "divide by zero"
                # with np.log(zBot).
                if zBot == 0:
                    tempBot = np.inf
                else:
                    tempBot = (self.mu + self.sigma ** 2 - np.log(zBot)) / (
                        np.sqrt(2) * self.sigma
                    )
                tempTop = (self.mu + self.sigma ** 2 - np.log(zTop)) / (
                    np.sqrt(2) * self.sigma
                )
                if tempBot <= 4:
                    X[i] = (
                        -0.5
                        * np.exp(self.mu + (self.sigma ** 2) * 0.5)
                        * (erf(tempTop) - erf(tempBot))
                        / pmf[i]
                    )
                else:
                    X[i] = (
                        -0.5
                        * np.exp(self.mu + (self.sigma ** 2) * 0.5)
                        * (erfc(tempBot) - erfc(tempTop))
                        / pmf[i]
                    )

        else:
            pmf = np.ones(N) / N
            X = np.exp(self.mu) * np.ones(N)
        return DiscreteDistribution(
            pmf, X, seed=self.RNG.randint(0, 2 ** 31 - 1, dtype="int32")
        )

    @classmethod
    def from_mean_std(cls, mean, std, seed=0):
        """
        Construct a LogNormal distribution from its
        mean and standard deviation.

        This is unlike the normal constructor for the distribution,
        which takes the mu and sigma for the normal distribution
        that is the logarithm of the Log Normal distribution.

        Parameters
        ----------
        mean : float or [float]
            One or more means.  Number of elements T in mu determines number
            of rows of output.
        std : 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
        ---------
        LogNormal

        """
        mean_squared = mean ** 2
        variance = std ** 2
        mu = np.log(mean / (np.sqrt(1.0 + variance / mean_squared)))
        sigma = np.sqrt(np.log(1.0 + variance / mean_squared))

        return cls(mu=mu, sigma=sigma, seed=seed)
HARK/distribution.py
def combine_indep_dstns(*distributions, seed=0):
    """
    Given n lists (or tuples) whose elements represent n independent, discrete
    probability spaces (probabilities and values), construct a joint pmf over
    all combinations of these independent points.  Can take multivariate discrete
    distributions as inputs.

    Parameters
    ----------
    distributions : [np.array]
        Arbitrary number of distributions (pmfs).  Each pmf is a list or tuple.
        For each pmf, the first vector is probabilities and all subsequent vectors
        are values.  For each pmf, this should be true:
        len(X_pmf[0]) == len(X_pmf[j]) for j in range(1,len(distributions))

    Returns
    -------
    A DiscreteDistribution, consisting of:

    P_out: np.array
        Probability associated with each point in X_out.

    X_out: np.array (as many as in *distributions)
        Discrete points for the joint discrete probability mass function.
    """
    # Get information on the distributions
    dist_lengths = ()
    dist_dims = ()
    for dist in distributions:
        dist_lengths += (len(dist.pmf),)
        dist_dims += (dist.dim(),)
    number_of_distributions = len(distributions)

    # Initialize lists we will use
    X_out = []
    P_temp = []

    # Now loop through the distributions, tiling and flattening as necessary.
    for dd, dist in enumerate(distributions):

        # The shape we want before we tile
        dist_newshape = (
            (1,) * dd + (len(dist.pmf),) + (1,) * (number_of_distributions - dd)
        )

        # The tiling we want to do
        dist_tiles = dist_lengths[:dd] + (1,) + dist_lengths[dd + 1 :]

        # Now we are ready to tile.
        # We don't use the np.meshgrid commands, because they do not
        # easily support non-symmetric grids.

        # First deal with probabilities
        Pmesh = np.tile(dist.pmf.reshape(dist_newshape), dist_tiles)  # Tiling
        flatP = Pmesh.ravel()  # Flatten the tiled arrays
        P_temp += [
            flatP,
        ]  # Add the flattened arrays to the output lists

        # Then loop through each value variable
        for n in range(dist_dims[dd]):
            if dist.dim() > 1:
                Xmesh = np.tile(dist.X[n].reshape(dist_newshape), dist_tiles)
            else:
                Xmesh = np.tile(dist.X.reshape(dist_newshape), dist_tiles)
            flatX = Xmesh.ravel()
            X_out += [
                flatX,
            ]

    # We're done getting the flattened X_out arrays we wanted.
    # However, we have a bunch of flattened P_temp arrays, and just want one
    # probability array. So get the probability array, P_out, here.
    P_out = np.prod(np.array(P_temp), axis=0)

    assert np.isclose(np.sum(P_out), 1), "Probabilities do not sum to 1!"
    return DiscreteDistribution(P_out, X_out, seed=seed)
HARK/distribution.py
def calc_expectation(dstn, func=lambda x: x, *args):
    """
    Expectation of a function, given an array of configurations of its inputs
    along with a DiscreteDistribution object that specifies the probability
    of each configuration.

    Parameters
    ----------
    dstn : DiscreteDistribution
        The N-valued distribution over which the function is to be evaluated.
    func : function
        The function to be evaluated.
        This function should take an array of size N x M.
        It may also take other arguments *args
        Please see numpy.apply_along_axis() for guidance on
        design of func.
        Defaults to identity function.
    *args : scalar or np.array
        One or more constants or arrays of input values for func,
        representing the non-stochastic arguments.
        The arrays must all have the same shape, and the expectation is computed
        at f(dstn, args[0], args[1],...,args[M]).

    Returns
    -------
    f_exp : np.array or scalar
        The expectation of the function at the queried values.
        Scalar if only one value.
    """
    N = dstn.dim()

    dstn_array = np.column_stack(dstn.X)

    if N > 1:
        # numpy is weird about 1-D arrays.
        dstn_array = dstn_array.T

    f_query = np.apply_along_axis(func, 0, dstn_array, *args)

    # Compute expectations over the values
    f_exp = np.dot(f_query, np.vstack(dstn.pmf))

    # a hack.
    if f_exp.size == 1:
        f_exp = f_exp.flat[0]
    elif f_exp.shape[0] == f_exp.size:
        f_exp = f_exp.flatten()

    return f_exp
HARK/distribution.py
class IndexDistribution(Distribution):
    """
    This class provides a way to define a distribution that
    is conditional on an index.

    The current implementation combines a defined distribution
    class (such as Bernoulli, LogNormal, etc.) with information
    about the conditions on the parameters of the distribution.

    For example, an IndexDistribution can be defined as
    a Bernoulli distribution whose parameter p is a function of
    a different inpute parameter.

    Parameters
    ----------

    engine : Distribution class
        A Distribution subclass.

    conditional: dict
        Information about the conditional variation
        on the input parameters of the engine distribution.
        Keys should match the arguments to the engine class
        constructor.

    seed : int
        Seed for random number generator.
    """

    conditional = None
    engine = None

    def __init__(self, engine, conditional, RNG = None, seed=0):
        
        if RNG is None:
            # Set up the RNG
            super().__init__(seed)
        else:
            # If an RNG is received, use it in whatever state it is in.
            self.RNG = RNG
            # The seed will still be set, even if it is not used for the RNG,
            # for whenever self.reset() is called.
            # Note that self.reset() will stop using the RNG that was passed
            # and create a new one.
            self.seed = seed

        self.conditional = conditional
        self.engine = engine
        
        
        self.dstns = []
        
        # Test one item to determine case handling
        item0 = list(self.conditional.values())[0]
        
        if type(item0) is list:
            # Create and store all the conditional distributions
            for y in range(len(item0)):
                cond = {key: val[y] for (key, val) in self.conditional.items()}
                self.dstns.append(self.engine(seed=self.RNG.randint(0, 2 ** 31 - 1), **cond))
                
        elif type(item0) is float:
            
            self.dstns = [self.engine(seed=self.RNG.randint(0, 2 ** 31 - 1), **conditional)]

        else:
            raise (
                Exception(
                    f"IndexDistribution: Unhandled case for __getitem__ access. y: {y}; conditional: {self.conditional}"
                )
            )
            
    def __getitem__(self, y):
        
        return self.dstns[y]

    def approx(self, N, **kwds):
        """
        Approximation of the distribution.

        Parameters
        ----------
        N : init
            Number of discrete points to approximate
            continuous distribution into.

        kwds: dict
            Other keyword arguments passed to engine
            distribution approx() method.

        Returns:
        ------------
        dists : [DiscreteDistribution]
            A list of DiscreteDistributions that are the
            approximation of engine distribution under each condition.

            TODO: It would be better if there were a conditional discrete
            distribution representation. But that integrates with the
            solution code. This implementation will return the list of
            distributions representations expected by the solution code.
        """

        # test one item to determine case handling
        item0 = list(self.conditional.values())[0]

        if type(item0) is float:
            # degenerate case. Treat the parameterization as constant.
            return self.dstns[0].approx(N, **kwds)

        if type(item0) is list:
            return TimeVaryingDiscreteDistribution(
                [self[i].approx(N, **kwds) for i, _ in enumerate(item0)]
            )

    def draw(self, condition):
        """
        Generate arrays of draws.
        The input is an array containing the conditions.
        The output is an array of the same length (axis 1 dimension)
        as the conditions containing random draws of the conditional
        distribution.

        Parameters
        ----------
        condition : np.array
            The input conditions to the distribution.

        Returns:
        ------------
        draws : np.array
        """
        # for now, assume that all the conditionals
        # are of the same type.
        # this matches the HARK 'time-varying' model architecture.

        # test one item to determine case handling
        item0 = list(self.conditional.values())[0]

        if type(item0) is float:
            # degenerate case. Treat the parameterization as constant.
            N = condition.size

            return self.engine(
                seed=self.RNG.randint(0, 2 ** 31 - 1), **self.conditional
            ).draw(N)

        if type(item0) is list:
            # conditions are indices into list
            # somewhat convoluted sampling strategy retained
            # for test backwards compatibility

            draws = np.zeros(condition.size)

            for c in np.unique(condition):
                these = c == condition
                N = np.sum(these)

                cond = {key: val[c] for (key, val) in self.conditional.items()}
                draws[these] = self[c].draw(N)

            return draws
HARK/distribution.py
def calc_normal_style_pars_from_lognormal_pars(avg_lognormal, std_lognormal):
    varLognormal = std_lognormal ** 2
    varNormal = math.log(1 + varLognormal / avg_lognormal ** 2)
    avgNormal = math.log(avg_lognormal) - varNormal * 0.5
    std_normal = math.sqrt(varNormal)
    return avgNormal, std_normal
HARK/distribution.py
def calc_lognormal_style_pars_from_normal_pars(mu_normal, std_normal):
    varNormal = std_normal ** 2
    avg_lognormal = math.exp(mu_normal + varNormal * 0.5)
    varLognormal = (math.exp(varNormal) - 1) * avg_lognormal ** 2
    std_lognormal = math.sqrt(varLognormal)
    return avg_lognormal, std_lognormal
HARK/distribution.py
class DiscreteDistribution(Distribution):
    """
    A representation of a discrete probability 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).
    seed : int
        Seed for random number generator.
    """

    pmf = None
    X = None

    def __init__(self, pmf, X, seed=0):
        self.pmf = pmf
        self.X = X
        # Set up the RNG
        super().__init__(seed)

        # Very quick and incomplete parameter check:
        # TODO: Check that pmf and 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):
        """
        Draws N 'events' from the distribution PMF.
        These events are indices into X.
        """
        # Generate a cumulative distribution
        base_draws = self.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 draw(self, N, X=None, exact_match=False):
        """
        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.

        Returns
        -------
        draws : np.array
            An array of draws from the discrete distribution; each element is a value in X.
        """
        if X is None:
            X = self.X
            J = self.dim()
        elif isinstance(X, int):
            X = self.X[X]
            J = 1
        else:
            X = X
            J = 1

        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
            indices = self.RNG.permutation(event_list)

        # Draw event indices randomly from the discrete distribution
        else:
            indices = self.draw_events(N)

        # Create and fill in the output array of draws based on the output of event indices
        if J > 1:
            draws = np.zeros((J, N))
            for j in range(J):
                draws[j, :] = X[j][indices]
        else:
            draws = np.asarray(X)[indices]

        return draws


Question: import unittest
import numpy as np
from HARK.distribution import (
    Bernoulli,
    calc_expectation,
    calc_lognormal_style_pars_from_normal_pars,
    calc_normal_style_pars_from_lognormal_pars,
    combine_indep_dstns,
    DiscreteDistribution,
    IndexDistribution,
    Lognormal,
    MarkovProcess,
    MeanOneLogNormal,
    MVNormal,
    Normal,
    Uniform,
    Weibull,
)




class DiscreteDistributionTests(unittest.TestCase):
    """
    Tests for distribution.py sampling distributions
    with default seed.
    """

    def test_draw(self):
        self.assertEqual(

Answer:             DiscreteDistribution(np.ones(1), np.zeros(1)).draw(1)[0],