diff options
Diffstat (limited to 'numpy/random/mtrand/mtrand.pyx')
| -rw-r--r-- | numpy/random/mtrand/mtrand.pyx | 144 |
1 files changed, 143 insertions, 1 deletions
diff --git a/numpy/random/mtrand/mtrand.pyx b/numpy/random/mtrand/mtrand.pyx index 13c743cd3..3772fa07a 100644 --- a/numpy/random/mtrand/mtrand.pyx +++ b/numpy/random/mtrand/mtrand.pyx @@ -913,6 +913,147 @@ cdef class RandomState: rk_fill(bytes, length, self.internal_state) return bytestring + + def sample(self, a, size, replace=True, p=None): + """ + sample(a, size[, replace, p]) + + Generates a random sample from a given 1-D array + + Parameters + ----------- + a : 1-D array-like or int + If an ndarray, a random sample is generated from its elements. + If an int, the random sample is generated as if a was np.arange(n) + size : int + Positive integer, the size of the sample. + replace : boolean, optional + Whether the sample is with or without replacement + p : 1-D array-like, optional + The probabilities associated with each entry in a. + If not given the sample assumes a uniform distribtion over all + entries in a. + + Returns + -------- + samples : 1-D ndarray, shape (size,) + The generated random samples + + Raises + ------- + ValueError + If a is an int and less than zero, if a or p are not 1-dimensional, + if a is an array-like of size 0, if p is not a vector of + probabilities, if a and p have different lengths, or if + replace=False and the sample size is greater than the population + size + + See Also + --------- + randint, shuffle, permutation + + Examples + --------- + Generate a uniform random sample from np.arange(5) of size 3: + + >>> np.random.sample(5, 3) + array([0, 3, 4]) + >>> #This is equivalent to np.random.randint(0,5,3) + + Generate a non-uniform random sample from np.arange(5) of size 3: + + >>> np.random.sample(5, 3, p=[0.1, 0, 0.3, 0.6, 0]) + array([3, 3, 0]) + + Generate a uniform random sample from np.arange(5) of size 3 without + replacement: + + >>> np.random.sample(5, 3, replace=False) + array([3,1,0]) + >>> #This is equivalent to np.random.shuffle(np.arange(5))[:3] + + Generate a non-uniform random sample from np.arange(5) of size + 3 without replacement: + + >>> np.random.sample(5, 3, replace=False, p=[0.1, 0, 0.3, 0.6, 0]) + array([2, 3, 0]) + + Any of the above can be repeated with an arbitrary array-like + instead of just integers. For instance: + + >>> aa_milne_arr = ['pooh', 'rabbit', 'piglet', 'Christopher'] + >>> np.random.sample(aa_milne_arr, 5, p=[0.5, 0.1, 0.1, 0.3]) + array(['pooh', 'pooh', 'pooh', 'Christopher', 'piglet'], + dtype='|S11') + + """ + + # Format and Verify input + if isinstance(a, int): + if a > 0: + pop_size = a + else: + raise ValueError("a must be greater than 0") + else: + a = np.asarray(a) + if len(a.shape) != 1: + raise ValueError("a must be 1-dimensional") + pop_size = a.size + if pop_size is 0: + raise ValueError("a must be non-empty") + + if None != p: + p = np.asarray(p) + if len(p.shape) != 1: + raise ValueError("p must be 1-dimensional") + if p.size != pop_size: + raise ValueError("a and p must have same size") + if any(p<0): + raise ValueError("probabilities are not non-negative") + if not np.allclose(p.sum(), 1): + raise ValueError("probabilities do not sum to 1") + + # Actual sampling + if replace: + if None != p: + x = np.arange(pop_size) + number_each = np.random.multinomial(size, p) + idx = np.repeat(x, number_each) + self.shuffle(idx) + else: + idx = self.randint(0, pop_size, size=size) + else: + if size > pop_size: + raise ValueError(''.join(["Cannot take a larger sample than ", + "population when 'replace=False'"])) + + if None != p: + if np.sum(p>0) < size: + raise ValueError("Fewer non-zero entries in p than size") + n_uniq = 0 + p = p.copy() + found = np.zeros(size, dtype=np.int) + while n_uniq < size: + x = np.random.rand(size-n_uniq) + if n_uniq > 0: + p[found[0:n_uniq]] = 0 + p = p/p.sum() + cdf = np.cumsum(p) + new = np.searchsorted(cdf, x) + new = np.unique(new) + found[n_uniq:n_uniq+new.size] = new + n_uniq += new.size + idx = found + else: + idx = self.permutation(pop_size)[:size] + + #Use samples as indices for a if a is array-like + if isinstance(a, int): + return idx + else: + return a.take(idx) + + def uniform(self, low=0.0, high=1.0, size=None): """ uniform(low=0.0, high=1.0, size=1) @@ -1028,7 +1169,7 @@ cdef class RandomState: ----- This is a convenience function. If you want an interface that takes a shape-tuple as the first argument, refer to - `random`. + np.random.random_sample . Examples -------- @@ -4331,6 +4472,7 @@ seed = _rand.seed get_state = _rand.get_state set_state = _rand.set_state random_sample = _rand.random_sample +sample = _rand.sample randint = _rand.randint bytes = _rand.bytes uniform = _rand.uniform |