Source code for quantecon._discrete_rv

"""
Generates an array of draws from a discrete random variable with a
specified vector of probabilities.

"""

import numpy as np
from .util import check_random_state


[docs]class DiscreteRV: """ Generates an array of draws from a discrete random variable with vector of probabilities given by q. Parameters ---------- q : array_like(float) Nonnegative numbers that sum to 1. Attributes ---------- q : see Parameters. Q : array_like(float) The cumulative sum of q. """ def __init__(self, q): self._q = np.asarray(q) self.Q = np.cumsum(q) def __repr__(self): return "DiscreteRV with {n} elements".format(n=self._q.size) def __str__(self): return self.__repr__() @property def q(self): """ Getter method for q. """ return self._q @q.setter def q(self, val): """ Setter method for q. """ self._q = np.asarray(val) self.Q = np.cumsum(val)
[docs] def draw(self, k=1, random_state=None): """ Returns k draws from q. For each such draw, the value i is returned with probability q[i]. Parameters ---------- k : scalar(int), optional Number of draws to be returned random_state : int or np.random.RandomState/Generator, optional Random seed (integer) or np.random.RandomState or Generator instance to set the initial state of the random number generator for reproducibility. If None, a randomly initialized RandomState is used. Returns ------- array_like(int) An array of k independent draws from q """ random_state = check_random_state(random_state) return self.Q.searchsorted(random_state.uniform(0, 1, size=k), side='right')