blob: d10ce786899c903efc449653ed1cec3f19c1f52e [file] [log] [blame]
Guido van Rossume7b146f2000-02-04 15:28:42 +00001"""Random variable generators.
Guido van Rossumff03b1a1994-03-09 12:55:02 +00002
Guido van Rossume7b146f2000-02-04 15:28:42 +00003 distributions on the real line:
4 ------------------------------
5 normal (Gaussian)
6 lognormal
7 negative exponential
8 gamma
9 beta
Guido van Rossumff03b1a1994-03-09 12:55:02 +000010
Guido van Rossume7b146f2000-02-04 15:28:42 +000011 distributions on the circle (angles 0 to 2pi)
12 ---------------------------------------------
13 circular uniform
14 von Mises
15
16Translated from anonymously contributed C/C++ source.
17
18Multi-threading note: the random number generator used here is not
19thread-safe; it is possible that two calls return the same random
20value. See whrandom.py for more info.
21"""
Guido van Rossumd03e1191998-05-29 17:51:31 +000022
Guido van Rossum33d7f1a1998-05-20 16:28:24 +000023import whrandom
Guido van Rossum33575611998-07-31 13:40:05 +000024from whrandom import random, uniform, randint, choice, randrange # For export!
Guido van Rossum95bfcda1994-03-09 14:21:05 +000025from math import log, exp, pi, e, sqrt, acos, cos, sin
Guido van Rossumff03b1a1994-03-09 12:55:02 +000026
Guido van Rossum33d7f1a1998-05-20 16:28:24 +000027# Interfaces to replace remaining needs for importing whrandom
28# XXX TO DO: make the distribution functions below into methods.
29
30def makeseed(a=None):
Tim Peters0c9886d2001-01-15 01:18:21 +000031 """Turn a hashable value into three seed values for whrandom.seed().
Guido van Rossum33d7f1a1998-05-20 16:28:24 +000032
Tim Peters0c9886d2001-01-15 01:18:21 +000033 None or no argument returns (0, 0, 0), to seed from current time.
Guido van Rossum33d7f1a1998-05-20 16:28:24 +000034
Tim Peters0c9886d2001-01-15 01:18:21 +000035 """
36 if a is None:
37 return (0, 0, 0)
38 a = hash(a)
39 a, x = divmod(a, 256)
40 a, y = divmod(a, 256)
41 a, z = divmod(a, 256)
42 x = (x + a) % 256 or 1
43 y = (y + a) % 256 or 1
44 z = (z + a) % 256 or 1
45 return (x, y, z)
Guido van Rossum33d7f1a1998-05-20 16:28:24 +000046
47def seed(a=None):
Tim Peters0c9886d2001-01-15 01:18:21 +000048 """Seed the default generator from any hashable value.
Guido van Rossum33d7f1a1998-05-20 16:28:24 +000049
Tim Peters0c9886d2001-01-15 01:18:21 +000050 None or no argument seeds from current time.
Guido van Rossum33d7f1a1998-05-20 16:28:24 +000051
Tim Peters0c9886d2001-01-15 01:18:21 +000052 """
53 x, y, z = makeseed(a)
54 whrandom.seed(x, y, z)
Guido van Rossum33d7f1a1998-05-20 16:28:24 +000055
56class generator(whrandom.whrandom):
Tim Peters0c9886d2001-01-15 01:18:21 +000057 """Random generator class."""
Guido van Rossum33d7f1a1998-05-20 16:28:24 +000058
Tim Peters0c9886d2001-01-15 01:18:21 +000059 def __init__(self, a=None):
60 """Constructor. Seed from current time or hashable value."""
61 self.seed(a)
Guido van Rossum33d7f1a1998-05-20 16:28:24 +000062
Tim Peters0c9886d2001-01-15 01:18:21 +000063 def seed(self, a=None):
64 """Seed the generator from current time or hashable value."""
65 x, y, z = makeseed(a)
66 whrandom.whrandom.seed(self, x, y, z)
Guido van Rossum33d7f1a1998-05-20 16:28:24 +000067
68def new_generator(a=None):
Tim Peters0c9886d2001-01-15 01:18:21 +000069 """Return a new random generator instance."""
70 return generator(a)
Guido van Rossum33d7f1a1998-05-20 16:28:24 +000071
Guido van Rossumff03b1a1994-03-09 12:55:02 +000072# Housekeeping function to verify that magic constants have been
73# computed correctly
74
75def verify(name, expected):
Tim Peters0c9886d2001-01-15 01:18:21 +000076 computed = eval(name)
77 if abs(computed - expected) > 1e-7:
78 raise ValueError, \
79'computed value for %s deviates too much (computed %g, expected %g)' % \
80(name, computed, expected)
Guido van Rossumff03b1a1994-03-09 12:55:02 +000081
82# -------------------- normal distribution --------------------
83
Guido van Rossumcc32ac91994-03-15 16:10:24 +000084NV_MAGICCONST = 4*exp(-0.5)/sqrt(2.0)
Guido van Rossumff03b1a1994-03-09 12:55:02 +000085verify('NV_MAGICCONST', 1.71552776992141)
86def normalvariate(mu, sigma):
Tim Peters0c9886d2001-01-15 01:18:21 +000087 # mu = mean, sigma = standard deviation
Guido van Rossumff03b1a1994-03-09 12:55:02 +000088
Tim Peters0c9886d2001-01-15 01:18:21 +000089 # Uses Kinderman and Monahan method. Reference: Kinderman,
90 # A.J. and Monahan, J.F., "Computer generation of random
91 # variables using the ratio of uniform deviates", ACM Trans
92 # Math Software, 3, (1977), pp257-260.
Guido van Rossumff03b1a1994-03-09 12:55:02 +000093
Tim Peters0c9886d2001-01-15 01:18:21 +000094 while 1:
95 u1 = random()
96 u2 = random()
97 z = NV_MAGICCONST*(u1-0.5)/u2
98 zz = z*z/4.0
99 if zz <= -log(u2):
100 break
101 return mu+z*sigma
Guido van Rossumff03b1a1994-03-09 12:55:02 +0000102
103# -------------------- lognormal distribution --------------------
104
105def lognormvariate(mu, sigma):
Tim Peters0c9886d2001-01-15 01:18:21 +0000106 return exp(normalvariate(mu, sigma))
Guido van Rossumff03b1a1994-03-09 12:55:02 +0000107
108# -------------------- circular uniform --------------------
109
110def cunifvariate(mean, arc):
Tim Peters0c9886d2001-01-15 01:18:21 +0000111 # mean: mean angle (in radians between 0 and pi)
112 # arc: range of distribution (in radians between 0 and pi)
Guido van Rossumff03b1a1994-03-09 12:55:02 +0000113
Tim Peters0c9886d2001-01-15 01:18:21 +0000114 return (mean + arc * (random() - 0.5)) % pi
Guido van Rossumff03b1a1994-03-09 12:55:02 +0000115
116# -------------------- exponential distribution --------------------
117
118def expovariate(lambd):
Tim Peters0c9886d2001-01-15 01:18:21 +0000119 # lambd: rate lambd = 1/mean
120 # ('lambda' is a Python reserved word)
Guido van Rossumff03b1a1994-03-09 12:55:02 +0000121
Tim Peters0c9886d2001-01-15 01:18:21 +0000122 u = random()
123 while u <= 1e-7:
124 u = random()
125 return -log(u)/lambd
Guido van Rossumff03b1a1994-03-09 12:55:02 +0000126
127# -------------------- von Mises distribution --------------------
128
Guido van Rossumcc32ac91994-03-15 16:10:24 +0000129TWOPI = 2.0*pi
Guido van Rossumff03b1a1994-03-09 12:55:02 +0000130verify('TWOPI', 6.28318530718)
131
132def vonmisesvariate(mu, kappa):
Tim Peters0c9886d2001-01-15 01:18:21 +0000133 # mu: mean angle (in radians between 0 and 2*pi)
134 # kappa: concentration parameter kappa (>= 0)
135 # if kappa = 0 generate uniform random angle
Guido van Rossum58102971998-04-06 14:12:13 +0000136
Tim Peters0c9886d2001-01-15 01:18:21 +0000137 # Based upon an algorithm published in: Fisher, N.I.,
138 # "Statistical Analysis of Circular Data", Cambridge
139 # University Press, 1993.
Guido van Rossum58102971998-04-06 14:12:13 +0000140
Tim Peters0c9886d2001-01-15 01:18:21 +0000141 # Thanks to Magnus Kessler for a correction to the
142 # implementation of step 4.
Guido van Rossum58102971998-04-06 14:12:13 +0000143
Tim Peters0c9886d2001-01-15 01:18:21 +0000144 if kappa <= 1e-6:
145 return TWOPI * random()
Guido van Rossumff03b1a1994-03-09 12:55:02 +0000146
Tim Peters0c9886d2001-01-15 01:18:21 +0000147 a = 1.0 + sqrt(1.0 + 4.0 * kappa * kappa)
148 b = (a - sqrt(2.0 * a))/(2.0 * kappa)
149 r = (1.0 + b * b)/(2.0 * b)
Guido van Rossumff03b1a1994-03-09 12:55:02 +0000150
Tim Peters0c9886d2001-01-15 01:18:21 +0000151 while 1:
152 u1 = random()
Guido van Rossumff03b1a1994-03-09 12:55:02 +0000153
Tim Peters0c9886d2001-01-15 01:18:21 +0000154 z = cos(pi * u1)
155 f = (1.0 + r * z)/(r + z)
156 c = kappa * (r - f)
Guido van Rossumff03b1a1994-03-09 12:55:02 +0000157
Tim Peters0c9886d2001-01-15 01:18:21 +0000158 u2 = random()
Guido van Rossumff03b1a1994-03-09 12:55:02 +0000159
Tim Peters0c9886d2001-01-15 01:18:21 +0000160 if not (u2 >= c * (2.0 - c) and u2 > c * exp(1.0 - c)):
161 break
Guido van Rossumff03b1a1994-03-09 12:55:02 +0000162
Tim Peters0c9886d2001-01-15 01:18:21 +0000163 u3 = random()
164 if u3 > 0.5:
165 theta = (mu % TWOPI) + acos(f)
166 else:
167 theta = (mu % TWOPI) - acos(f)
Guido van Rossumff03b1a1994-03-09 12:55:02 +0000168
Tim Peters0c9886d2001-01-15 01:18:21 +0000169 return theta
Guido van Rossumff03b1a1994-03-09 12:55:02 +0000170
171# -------------------- gamma distribution --------------------
172
Guido van Rossumcc32ac91994-03-15 16:10:24 +0000173LOG4 = log(4.0)
Guido van Rossumff03b1a1994-03-09 12:55:02 +0000174verify('LOG4', 1.38629436111989)
175
176def gammavariate(alpha, beta):
Tim Peters0c9886d2001-01-15 01:18:21 +0000177 # beta times standard gamma
178 ainv = sqrt(2.0 * alpha - 1.0)
179 return beta * stdgamma(alpha, ainv, alpha - LOG4, alpha + ainv)
Guido van Rossumff03b1a1994-03-09 12:55:02 +0000180
Guido van Rossumcc32ac91994-03-15 16:10:24 +0000181SG_MAGICCONST = 1.0 + log(4.5)
Guido van Rossumff03b1a1994-03-09 12:55:02 +0000182verify('SG_MAGICCONST', 2.50407739677627)
183
184def stdgamma(alpha, ainv, bbb, ccc):
Tim Peters0c9886d2001-01-15 01:18:21 +0000185 # ainv = sqrt(2 * alpha - 1)
186 # bbb = alpha - log(4)
187 # ccc = alpha + ainv
Guido van Rossumff03b1a1994-03-09 12:55:02 +0000188
Tim Peters0c9886d2001-01-15 01:18:21 +0000189 if alpha <= 0.0:
190 raise ValueError, 'stdgamma: alpha must be > 0.0'
Guido van Rossumff03b1a1994-03-09 12:55:02 +0000191
Tim Peters0c9886d2001-01-15 01:18:21 +0000192 if alpha > 1.0:
Guido van Rossumff03b1a1994-03-09 12:55:02 +0000193
Tim Peters0c9886d2001-01-15 01:18:21 +0000194 # Uses R.C.H. Cheng, "The generation of Gamma
195 # variables with non-integral shape parameters",
196 # Applied Statistics, (1977), 26, No. 1, p71-74
Guido van Rossumff03b1a1994-03-09 12:55:02 +0000197
Tim Peters0c9886d2001-01-15 01:18:21 +0000198 while 1:
199 u1 = random()
200 u2 = random()
201 v = log(u1/(1.0-u1))/ainv
202 x = alpha*exp(v)
203 z = u1*u1*u2
204 r = bbb+ccc*v-x
205 if r + SG_MAGICCONST - 4.5*z >= 0.0 or r >= log(z):
206 return x
Guido van Rossumff03b1a1994-03-09 12:55:02 +0000207
Tim Peters0c9886d2001-01-15 01:18:21 +0000208 elif alpha == 1.0:
209 # expovariate(1)
210 u = random()
211 while u <= 1e-7:
212 u = random()
213 return -log(u)
Guido van Rossumff03b1a1994-03-09 12:55:02 +0000214
Tim Peters0c9886d2001-01-15 01:18:21 +0000215 else: # alpha is between 0 and 1 (exclusive)
Guido van Rossumff03b1a1994-03-09 12:55:02 +0000216
Tim Peters0c9886d2001-01-15 01:18:21 +0000217 # Uses ALGORITHM GS of Statistical Computing - Kennedy & Gentle
Guido van Rossumff03b1a1994-03-09 12:55:02 +0000218
Tim Peters0c9886d2001-01-15 01:18:21 +0000219 while 1:
220 u = random()
221 b = (e + alpha)/e
222 p = b*u
223 if p <= 1.0:
224 x = pow(p, 1.0/alpha)
225 else:
226 # p > 1
227 x = -log((b-p)/alpha)
228 u1 = random()
229 if not (((p <= 1.0) and (u1 > exp(-x))) or
230 ((p > 1) and (u1 > pow(x, alpha - 1.0)))):
231 break
232 return x
Guido van Rossumff03b1a1994-03-09 12:55:02 +0000233
Guido van Rossum95bfcda1994-03-09 14:21:05 +0000234
235# -------------------- Gauss (faster alternative) --------------------
236
Guido van Rossum95bfcda1994-03-09 14:21:05 +0000237gauss_next = None
238def gauss(mu, sigma):
Guido van Rossumcc32ac91994-03-15 16:10:24 +0000239
Tim Peters0c9886d2001-01-15 01:18:21 +0000240 # When x and y are two variables from [0, 1), uniformly
241 # distributed, then
242 #
243 # cos(2*pi*x)*sqrt(-2*log(1-y))
244 # sin(2*pi*x)*sqrt(-2*log(1-y))
245 #
246 # are two *independent* variables with normal distribution
247 # (mu = 0, sigma = 1).
248 # (Lambert Meertens)
249 # (corrected version; bug discovered by Mike Miller, fixed by LM)
Guido van Rossumcc32ac91994-03-15 16:10:24 +0000250
Tim Peters0c9886d2001-01-15 01:18:21 +0000251 # Multithreading note: When two threads call this function
252 # simultaneously, it is possible that they will receive the
253 # same return value. The window is very small though. To
254 # avoid this, you have to use a lock around all calls. (I
255 # didn't want to slow this down in the serial case by using a
256 # lock here.)
Guido van Rossumd03e1191998-05-29 17:51:31 +0000257
Tim Peters0c9886d2001-01-15 01:18:21 +0000258 global gauss_next
Guido van Rossumcc32ac91994-03-15 16:10:24 +0000259
Tim Peters0c9886d2001-01-15 01:18:21 +0000260 z = gauss_next
261 gauss_next = None
262 if z is None:
263 x2pi = random() * TWOPI
264 g2rad = sqrt(-2.0 * log(1.0 - random()))
265 z = cos(x2pi) * g2rad
266 gauss_next = sin(x2pi) * g2rad
Guido van Rossumcc32ac91994-03-15 16:10:24 +0000267
Tim Peters0c9886d2001-01-15 01:18:21 +0000268 return mu + z*sigma
Guido van Rossum95bfcda1994-03-09 14:21:05 +0000269
270# -------------------- beta --------------------
271
272def betavariate(alpha, beta):
Guido van Rossumcc32ac91994-03-15 16:10:24 +0000273
Tim Peters0c9886d2001-01-15 01:18:21 +0000274 # Discrete Event Simulation in C, pp 87-88.
Guido van Rossumcc32ac91994-03-15 16:10:24 +0000275
Tim Peters0c9886d2001-01-15 01:18:21 +0000276 y = expovariate(alpha)
277 z = expovariate(1.0/beta)
278 return z/(y+z)
Guido van Rossum95bfcda1994-03-09 14:21:05 +0000279
Guido van Rossum5bdea891997-12-09 19:43:18 +0000280# -------------------- Pareto --------------------
Guido van Rossumcf4559a1997-12-02 02:47:39 +0000281
282def paretovariate(alpha):
Tim Peters0c9886d2001-01-15 01:18:21 +0000283 # Jain, pg. 495
Guido van Rossumcf4559a1997-12-02 02:47:39 +0000284
Tim Peters0c9886d2001-01-15 01:18:21 +0000285 u = random()
286 return 1.0 / pow(u, 1.0/alpha)
Guido van Rossumcf4559a1997-12-02 02:47:39 +0000287
Guido van Rossum5bdea891997-12-09 19:43:18 +0000288# -------------------- Weibull --------------------
Guido van Rossumcf4559a1997-12-02 02:47:39 +0000289
290def weibullvariate(alpha, beta):
Tim Peters0c9886d2001-01-15 01:18:21 +0000291 # Jain, pg. 499; bug fix courtesy Bill Arms
Guido van Rossumcf4559a1997-12-02 02:47:39 +0000292
Tim Peters0c9886d2001-01-15 01:18:21 +0000293 u = random()
294 return alpha * pow(-log(u), 1.0/beta)
Guido van Rossumcf4559a1997-12-02 02:47:39 +0000295
Guido van Rossum6c395ba1999-08-18 13:53:28 +0000296# -------------------- shuffle --------------------
297# Not quite a random distribution, but a standard algorithm.
298# This implementation due to Tim Peters.
299
300def shuffle(x, random=random, int=int):
301 """x, random=random.random -> shuffle list x in place; return None.
302
303 Optional arg random is a 0-argument function returning a random
304 float in [0.0, 1.0); by default, the standard random.random.
305
306 Note that for even rather small len(x), the total number of
307 permutations of x is larger than the period of most random number
308 generators; this implies that "most" permutations of a long
309 sequence can never be generated.
310 """
311
312 for i in xrange(len(x)-1, 0, -1):
Tim Peters0c9886d2001-01-15 01:18:21 +0000313 # pick an element in x[:i+1] with which to exchange x[i]
Guido van Rossum6c395ba1999-08-18 13:53:28 +0000314 j = int(random() * (i+1))
315 x[i], x[j] = x[j], x[i]
316
Guido van Rossumff03b1a1994-03-09 12:55:02 +0000317# -------------------- test program --------------------
318
Guido van Rossum2922c6d1994-05-06 14:28:19 +0000319def test(N = 200):
Tim Peters0c9886d2001-01-15 01:18:21 +0000320 print 'TWOPI =', TWOPI
321 print 'LOG4 =', LOG4
322 print 'NV_MAGICCONST =', NV_MAGICCONST
323 print 'SG_MAGICCONST =', SG_MAGICCONST
324 test_generator(N, 'random()')
325 test_generator(N, 'normalvariate(0.0, 1.0)')
326 test_generator(N, 'lognormvariate(0.0, 1.0)')
327 test_generator(N, 'cunifvariate(0.0, 1.0)')
328 test_generator(N, 'expovariate(1.0)')
329 test_generator(N, 'vonmisesvariate(0.0, 1.0)')
330 test_generator(N, 'gammavariate(0.5, 1.0)')
331 test_generator(N, 'gammavariate(0.9, 1.0)')
332 test_generator(N, 'gammavariate(1.0, 1.0)')
333 test_generator(N, 'gammavariate(2.0, 1.0)')
334 test_generator(N, 'gammavariate(20.0, 1.0)')
335 test_generator(N, 'gammavariate(200.0, 1.0)')
336 test_generator(N, 'gauss(0.0, 1.0)')
337 test_generator(N, 'betavariate(3.0, 3.0)')
338 test_generator(N, 'paretovariate(1.0)')
339 test_generator(N, 'weibullvariate(1.0, 1.0)')
Guido van Rossumff03b1a1994-03-09 12:55:02 +0000340
341def test_generator(n, funccall):
Tim Peters0c9886d2001-01-15 01:18:21 +0000342 import time
343 print n, 'times', funccall
344 code = compile(funccall, funccall, 'eval')
345 sum = 0.0
346 sqsum = 0.0
347 smallest = 1e10
348 largest = -1e10
349 t0 = time.time()
350 for i in range(n):
351 x = eval(code)
352 sum = sum + x
353 sqsum = sqsum + x*x
354 smallest = min(x, smallest)
355 largest = max(x, largest)
356 t1 = time.time()
357 print round(t1-t0, 3), 'sec,',
358 avg = sum/n
359 stddev = sqrt(sqsum/n - avg*avg)
360 print 'avg %g, stddev %g, min %g, max %g' % \
361 (avg, stddev, smallest, largest)
Guido van Rossumff03b1a1994-03-09 12:55:02 +0000362
363if __name__ == '__main__':
Tim Peters0c9886d2001-01-15 01:18:21 +0000364 test()