Source code for quantecon.random.utilities

"""
Utilities to Support Random Operations and Generating Vectors and Matrices

"""

import numpy as np
from numba import guvectorize, generated_jit, types

from ..util import check_random_state, searchsorted


# Generating Arrays and Vectors #

[docs]def probvec(m, k, random_state=None, parallel=True): """ Return m randomly sampled probability vectors of dimension k. Parameters ---------- m : scalar(int) Number of probability vectors. k : scalar(int) Dimension of each probability vectors. random_state : int or np.random.RandomState, optional Random seed (integer) or np.random.RandomState instance to set the initial state of the random number generator for reproducibility. If None, a randomly initialized RandomState is used. parallel : bool(default=True) Whether to use multi-core CPU (parallel=True) or single-threaded CPU (parallel=False). (Internally the code is executed through Numba.guvectorize.) Returns ------- x : ndarray(float, ndim=2) Array of shape (m, k) containing probability vectors as rows. Examples -------- >>> qe.random.probvec(2, 3, random_state=1234) array([[ 0.19151945, 0.43058932, 0.37789123], [ 0.43772774, 0.34763084, 0.21464142]]) """ if k == 1: return np.ones((m, k)) # if k >= 2 random_state = check_random_state(random_state) r = random_state.random_sample(size=(m, k-1)) x = np.empty((m, k)) # Parse Parallel Option # if parallel: _probvec_parallel(r, x) else: _probvec_cpu(r, x) return x
def _probvec(r, out): """ Fill `out` with randomly sampled probability vectors as rows. To be complied as a ufunc by guvectorize of Numba. The inputs must have the same shape except the last axis; the length of the last axis of `r` must be that of `out` minus 1, i.e., if out.shape[-1] is k, then r.shape[-1] must be k-1. Parameters ---------- r : ndarray(float) Array containing random values in [0, 1). out : ndarray(float) Output array. """ n = r.shape[0] r.sort() out[0] = r[0] for i in range(1, n): out[i] = r[i] - r[i-1] out[n] = 1 - r[n-1] _probvec_parallel = guvectorize( ['(f8[:], f8[:])'], '(n), (k)', nopython=True, target='parallel', cache=True )(_probvec) _probvec_cpu = guvectorize( ['(f8[:], f8[:])'], '(n), (k)', nopython=True, target='cpu', cache=True )(_probvec)
[docs]def sample_without_replacement(n, k, num_trials=None, random_state=None): """ Randomly choose k integers without replacement from 0, ..., n-1. Parameters ---------- n : scalar(int) Number of integers, 0, ..., n-1, to sample from. k : scalar(int) Number of integers to sample. num_trials : scalar(int), optional(default=None) Number of trials. random_state : int or np.random.RandomState, optional Random seed (integer) or np.random.RandomState instance to set the initial state of the random number generator for reproducibility. If None, a randomly initialized RandomState is used. Returns ------- result : ndarray(int, ndim=1 or 2) Array of shape (k,) if num_trials is None, or of shape (num_trials, k) otherwise, (each row of) which contains k unique random elements chosen from 0, ..., n-1. Examples -------- >>> qe.random.sample_without_replacement(5, 3, random_state=1234) array([0, 2, 1]) >>> qe.random.sample_without_replacement(5, 3, num_trials=4, ... random_state=1234) array([[0, 2, 1], [3, 4, 0], [1, 3, 2], [4, 1, 3]]) """ if n <= 0: raise ValueError('n must be greater than 0') if k > n: raise ValueError('k must be smaller than or equal to n') size = k if num_trials is None else (num_trials, k) random_state = check_random_state(random_state) r = random_state.random_sample(size=size) result = _sample_without_replacement(n, r) return result
@guvectorize(['(i8, f8[:], i8[:])'], '(),(k)->(k)', nopython=True, cache=True) def _sample_without_replacement(n, r, out): """ Main body of `sample_without_replacement`. To be complied as a ufunc by guvectorize of Numba. """ k = r.shape[0] # Logic taken from random.sample in the standard library pool = np.arange(n) for j in range(k): idx = np.intp(np.floor(r[j] * (n-j))) # np.floor returns a float out[j] = pool[idx] pool[idx] = pool[n-j-1]
[docs]@generated_jit(nopython=True) def draw(cdf, size=None): """ Generate a random sample according to the cumulative distribution given by `cdf`. Jit-complied by Numba in nopython mode. Parameters ---------- cdf : array_like(float, ndim=1) Array containing the cumulative distribution. size : scalar(int), optional(default=None) Size of the sample. If an integer is supplied, an ndarray of `size` independent draws is returned; otherwise, a single draw is returned as a scalar. Returns ------- scalar(int) or ndarray(int, ndim=1) Examples -------- >>> cdf = np.cumsum([0.4, 0.6]) >>> qe.random.draw(cdf) 1 >>> qe.random.draw(cdf, 10) array([1, 0, 1, 0, 1, 0, 0, 0, 1, 0]) """ if isinstance(size, types.Integer): def draw_impl(cdf, size): rs = np.random.random(size) out = np.empty(size, dtype=np.int_) for i in range(size): out[i] = searchsorted(cdf, rs[i]) return out else: def draw_impl(cdf, size): r = np.random.random() return searchsorted(cdf, r) return draw_impl