Skip to content

Instantly share code, notes, and snippets.

# ntamas/walker.py Last active Sep 26, 2015

Weighted sampling with replacement using Walker's alias method - NumPy version
 #!/usr/bin/env python from numpy import arange, array, bincount, ndarray, ones, where from numpy.random import seed, random, randint __author__ = "Tamas Nepusz, Denis Bzowy" __version__ = "27jul2011" class WalkerRandomSampling(object): """Walker's alias method for random objects with different probablities. Based on the implementation of Denis Bzowy at the following URL: http://code.activestate.com/recipes/576564-walkers-alias-method-for-random-objects-with-diffe/ """ def __init__(self, weights, keys=None): """Builds the Walker tables ``prob`` and ``inx`` for calls to `random()`. The weights (a list or tuple or iterable) can be in any order and they do not even have to sum to 1.""" n = self.n = len(weights) if keys is None: self.keys = keys else: self.keys = array(keys) if isinstance(weights, (list, tuple)): weights = array(weights, dtype=float) elif isinstance(weights, ndarray): if weights.dtype != float: weights = weights.astype(float) else: weights = array(list(weights), dtype=float) if weights.ndim != 1: raise ValueError("weights must be a vector") weights = weights * n / weights.sum() inx = -ones(n, dtype=int) short = where(weights < 1).tolist() long = where(weights > 1).tolist() while short and long: j = short.pop() k = long[-1] inx[j] = k weights[k] -= (1 - weights[j]) if weights[k] < 1: short.append( k ) long.pop() self.prob = weights self.inx = inx def random(self, count=None): """Returns a given number of random integers or keys, with probabilities being proportional to the weights supplied in the constructor. When `count` is ``None``, returns a single integer or key, otherwise returns a NumPy array with a length given in `count`. """ if count is None: u = random() j = randint(self.n) k = j if u <= self.prob[j] else self.inx[j] return self.keys[k] if self.keys is not None else k u = random(count) j = randint(self.n, size=count) k = where(u <= self.prob[j], j, self.inx[j]) return self.keys[k] if self.keys is not None else k if __name__ == "__main__": # little examples, self-contained -- N = 5 Nrand = 1000 randomseed = 1 if randomseed: seed(randomseed) print Nrand, "Walker random sampling with weights .1 .2 .3 .4:" wrand = WalkerRandomSampling(arange(1, N)) nrand = bincount(wrand.random(Nrand)).tolist() s = str(nrand) print s if N==5 and Nrand==1000 and randomseed==1: assert s == "[97, 207, 316, 380]" print Nrand, "Walker random sampling, strings with weights .1 .2 .3 .4:" abcd = dict(A=1, D=4, C=3, B=2) wrand = WalkerRandomSampling(abcd.values(), abcd.keys()) nrand = defaultdict(int) for sample in wrand.random(Nrand): nrand[sample] += 1 s = str(sorted(nrand.iteritems())) print s if N==5 and Nrand==1000 and randomseed==1: assert s == "[('A', 85), ('B', 199), ('C', 343), ('D', 373)]"
to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.