Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- def choice(a, size=1, replace=True, p=None):
- # Format and Verify input
- if isinstance(a, int):
- if a > 0:
- pop_size = a #population size
- else:
- raise ValueError("a must be greater than 0")
- else:
- a = np.array(a, ndmin=1, copy=0)
- if a.ndim != 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.array(p, dtype=np.double, ndmin=1, copy=0)
- if p.ndim != 1:
- raise ValueError("p must be 1-dimensional")
- if p.size != pop_size:
- raise ValueError("a and p must have same size")
- if np.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:
- cdf = p.cumsum()
- cdf /= cdf[-1]
- uniform_samples = np.random.random(size)
- idx = cdf.searchsorted(uniform_samples, side='right')
- 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 = self.rand(size - n_uniq)
- if n_uniq > 0:
- p[found[0:n_uniq]] = 0
- cdf = np.cumsum(p)
- cdf /= cdf[-1]
- new = cdf.searchsorted(x, side='right')
- 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)
Advertisement
Add Comment
Please, Sign In to add comment