Guest User

Untitled

a guest
Jun 5th, 2012
368
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Python 2.29 KB | None | 0 0
  1. def choice(a, size=1, replace=True, p=None):
  2.     # Format and Verify input
  3.     if isinstance(a, int):
  4.         if a > 0:
  5.             pop_size = a #population size
  6.         else:
  7.             raise ValueError("a must be greater than 0")
  8.     else:
  9.         a = np.array(a, ndmin=1, copy=0)
  10.         if a.ndim != 1:
  11.             raise ValueError("a must be 1-dimensional")
  12.         pop_size = a.size
  13.         if pop_size is 0:
  14.             raise ValueError("a must be non-empty")
  15.    
  16.     if None != p:
  17.         p = np.array(p, dtype=np.double, ndmin=1, copy=0)
  18.         if p.ndim != 1:
  19.             raise ValueError("p must be 1-dimensional")
  20.         if p.size != pop_size:
  21.             raise ValueError("a and p must have same size")
  22.         if np.any(p < 0):
  23.             raise ValueError("probabilities are not non-negative")
  24.         if not np.allclose(p.sum(), 1):
  25.             raise ValueError("probabilities do not sum to 1")
  26.    
  27.     # Actual sampling
  28.     if replace:
  29.         if None != p:
  30.             cdf = p.cumsum()
  31.             cdf /= cdf[-1]
  32.             uniform_samples = np.random.random(size)
  33.             idx = cdf.searchsorted(uniform_samples, side='right')
  34.         else:
  35.             idx = self.randint(0, pop_size, size=size)
  36.     else:
  37.         if size > pop_size:
  38.             raise ValueError(''.join(["Cannot take a larger sample than ",
  39.                                       "population when 'replace=False'"]))
  40.        
  41.         if None != p:
  42.             if np.sum(p > 0) < size:
  43.                 raise ValueError("Fewer non-zero entries in p than size")
  44.             n_uniq = 0
  45.             p = p.copy()
  46.             found = np.zeros(size, dtype=np.int)
  47.             while n_uniq < size:
  48.                 x = self.rand(size - n_uniq)
  49.                 if n_uniq > 0:
  50.                     p[found[0:n_uniq]] = 0
  51.                 cdf = np.cumsum(p)
  52.                 cdf /= cdf[-1]
  53.                 new = cdf.searchsorted(x, side='right')
  54.                 new = np.unique(new)
  55.                 found[n_uniq:n_uniq + new.size] = new
  56.                 n_uniq += new.size
  57.             idx = found
  58.         else:
  59.             idx = self.permutation(pop_size)[:size]
  60.    
  61.     #Use samples as indices for a if a is array-like
  62.     if isinstance(a, int):
  63.         return idx
  64.     else:
  65.         return a.take(idx)
Advertisement
Add Comment
Please, Sign In to add comment