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
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
|
"""
Implementation of `BinomialFunction`
"""
import numpy as np
from brian2.core.base import Nameable
from brian2.core.functions import DEFAULT_FUNCTIONS, Function
from brian2.units.fundamentalunits import check_units
from brian2.utils.stringtools import replace
__all__ = ["BinomialFunction"]
def _pre_calc_constants_approximated(n, p):
loc = n * p
scale = np.sqrt(n * p * (1 - p))
return loc, scale
def _pre_calc_constants(n, p):
reverse = p > 0.5
if reverse:
P = 1.0 - p
else:
P = p
q = 1.0 - P
qn = np.exp(n * np.log(q))
bound = min(n, n * P + 10.0 * np.sqrt(n * P * q + 1))
return reverse, q, P, qn, bound
def _generate_cython_code(n, p, use_normal, name):
# Cython implementation
# Inversion transform sampling
if use_normal:
loc, scale = _pre_calc_constants_approximated(n, p)
cython_code = """
cdef float %NAME%(const int vectorisation_idx):
return _randn(vectorisation_idx) * %SCALE% + %LOC%
"""
cython_code = replace(
cython_code,
{"%SCALE%": f"{scale:.15f}", "%LOC%": f"{loc:.15f}", "%NAME%": name},
)
dependencies = {"_randn": DEFAULT_FUNCTIONS["randn"]}
else:
reverse, q, P, qn, bound = _pre_calc_constants(n, p)
# The following code is an almost exact copy of numpy's
# rk_binomial_inversion function
# (numpy/random/mtrand/distributions.c)
cython_code = """
cdef long %NAME%(const int vectorisation_idx):
cdef long X = 0
cdef double px = %QN%
cdef double U = _rand(vectorisation_idx)
while U > px:
X += 1
if X > %BOUND%:
X = 0
px = %QN%
U = _rand(vectorisation_idx)
else:
U -= px
px = ((%N%-X+1) * %P% * px)/(X*%Q%)
return %RETURN_VALUE%
"""
cython_code = replace(
cython_code,
{
"%N%": f"{int(n)}",
"%P%": f"{p:.15f}",
"%Q%": f"{q:.15f}",
"%QN%": f"{qn:.15f}",
"%BOUND%": f"{bound:.15f}",
"%RETURN_VALUE%": f"{int(n)}-X" if reverse else "X",
"%NAME%": name,
},
)
dependencies = {"_rand": DEFAULT_FUNCTIONS["rand"]}
return cython_code, dependencies
def _generate_cpp_code(n, p, use_normal, name):
# C++ implementation
# Inversion transform sampling
if use_normal:
loc, scale = _pre_calc_constants_approximated(n, p)
cpp_code = """
float %NAME%(const int vectorisation_idx)
{
return _randn(vectorisation_idx) * %SCALE% + %LOC%;
}
"""
cpp_code = replace(
cpp_code,
{"%SCALE%": f"{scale:.15f}", "%LOC%": f"{loc:.15f}", "%NAME%": name},
)
dependencies = {"_randn": DEFAULT_FUNCTIONS["randn"]}
else:
reverse, q, P, qn, bound = _pre_calc_constants(n, p)
# The following code is an almost exact copy of numpy's
# rk_binomial_inversion function
# (numpy/random/mtrand/distributions.c)
cpp_code = """
long %NAME%(const int vectorisation_idx)
{
long X = 0;
double px = %QN%;
double U = _rand(vectorisation_idx);
while (U > px)
{
X++;
if (X > %BOUND%)
{
X = 0;
px = %QN%;
U = _rand(vectorisation_idx);
} else
{
U -= px;
px = ((%N%-X+1) * %P% * px)/(X*%Q%);
}
}
return %RETURN_VALUE%;
}
"""
cpp_code = replace(
cpp_code,
{
"%N%": f"{int(n)}",
"%P%": f"{P:.15f}",
"%Q%": f"{q:.15f}",
"%QN%": f"{qn:.15f}",
"%BOUND%": f"{bound:.15f}",
"%RETURN_VALUE%": f"{int(n)}-X" if reverse else "X",
"%NAME%": name,
},
)
dependencies = {"_rand": DEFAULT_FUNCTIONS["rand"]}
return {"support_code": cpp_code}, dependencies
class BinomialFunction(Function, Nameable):
r"""
BinomialFunction(n, p, approximate=True, name='_binomial*')
A function that generates samples from a binomial distribution.
Parameters
----------
n : int
Number of samples
p : float
Probablility
approximate : bool, optional
Whether to approximate the binomial with a normal distribution if
:math:`n p > 5 \wedge n (1 - p) > 5`. Defaults to ``True``.
"""
#: Container for implementing functions for different targets
#: This container can be extended by other codegeneration targets/devices
#: The key has to be the name of the target, the value a function
#: that takes three parameters (n, p, use_normal) and returns a tuple of
#: (code, dependencies)
implementations = {"cpp": _generate_cpp_code, "cython": _generate_cython_code}
@check_units(n=1, p=1)
def __init__(self, n, p, approximate=True, name="_binomial*"):
Nameable.__init__(self, name)
# Python implementation
use_normal = approximate and (n * p > 5) and n * (1 - p) > 5
if use_normal:
loc = n * p
scale = np.sqrt(n * p * (1 - p))
def sample_function(vectorisation_idx):
try:
N = len(vectorisation_idx)
except TypeError:
N = int(vectorisation_idx)
return np.random.normal(loc, scale, size=N)
else:
def sample_function(vectorisation_idx):
try:
N = len(vectorisation_idx)
except TypeError:
N = int(vectorisation_idx)
return np.random.binomial(n, p, size=N)
Function.__init__(
self,
pyfunc=lambda: sample_function(1),
arg_units=[],
return_unit=1,
stateless=False,
auto_vectorise=True,
)
self.implementations.add_implementation("numpy", sample_function)
for target, func in BinomialFunction.implementations.items():
code, dependencies = func(n=n, p=p, use_normal=use_normal, name=self.name)
self.implementations.add_implementation(
target, code, dependencies=dependencies, name=self.name
)
|
