/
walker_alias.py
107 lines (89 loc) · 3.35 KB
/
walker_alias.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
import numpy
from chainer import cuda
class WalkerAlias(object):
"""Implementation of Walker's alias method.
This method generates a random sample from given probabilities
:math:`p_1, \\dots, p_n` in :math:`O(1)` time.
It is more efficient than :func:`~numpy.random.choice`.
This class works on both CPU and GPU.
Args:
probs (float list): Probabilities of entries. They are normalized with
`sum(probs)`.
See: `Wikipedia article <https://en.wikipedia.org/wiki/Alias_method>`_
"""
def __init__(self, probs):
prob = numpy.array(probs, numpy.float32)
prob /= numpy.sum(prob)
threshold = numpy.ndarray(len(probs), numpy.float32)
values = numpy.ndarray(len(probs) * 2, numpy.int32)
il, ir = 0, 0
pairs = list(zip(prob, range(len(probs))))
pairs.sort()
for prob, i in pairs:
p = prob * len(probs)
while p > 1 and ir < il:
values[ir * 2 + 1] = i
p -= 1.0 - threshold[ir]
ir += 1
threshold[il] = p
values[il * 2] = i
il += 1
# fill the rest
for i in range(ir, len(probs)):
values[i * 2 + 1] = 0
assert((values < len(threshold)).all())
self.threshold = threshold
self.values = values
self.use_gpu = False
def to_gpu(self):
"""Make a sampler GPU mode.
"""
if not self.use_gpu:
self.threshold = cuda.to_gpu(self.threshold)
self.values = cuda.to_gpu(self.values)
self.use_gpu = True
def to_cpu(self):
"""Make a sampler CPU mode.
"""
if self.use_gpu:
self.threshold = cuda.to_cpu(self.threshold)
self.values = cuda.to_cpu(self.values)
self.use_gpu = False
def sample(self, shape):
"""Generates a random sample based on given probabilities.
Args:
shape (tuple of int): Shape of a return value.
Returns:
Returns a generated array with the given shape. If a sampler is in
CPU mode the return value is a :class:`numpy.ndarray` object, and
if it is in GPU mode the return value is a :class:`cupy.ndarray`
object.
"""
if self.use_gpu:
return self.sample_gpu(shape)
else:
return self.sample_cpu(shape)
def sample_cpu(self, shape):
ps = numpy.random.uniform(0, 1, shape)
pb = ps * len(self.threshold)
index = pb.astype(numpy.int32)
left_right = (self.threshold[index] < pb - index).astype(numpy.int32)
return self.values[index * 2 + left_right]
def sample_gpu(self, shape):
ps = cuda.cupy.random.uniform(size=shape, dtype=numpy.float32)
vs = cuda.elementwise(
'T ps, raw T threshold , raw S values, int32 b',
'int32 vs',
'''
T pb = ps * b;
int index = __float2int_rd(pb);
// fill_uniform sometimes returns 1.0, so we need to check index
if (index >= b) {
index = 0;
}
int lr = threshold[index] < pb - index;
vs = values[index * 2 + lr];
''',
'walker_alias_sample'
)(ps, self.threshold, self.values, len(self.threshold))
return vs