Lib/random.py

   1 #       R A N D O M   V A R I A B L E   G E N E R A T O R S
   2 #
   3 #       distributions on the real line:
   4 #       ------------------------------
   5 #              normal (Gaussian)
   6 #              lognormal
   7 #              negative exponential
   8 #              gamma
   9 #              beta
  10 #
  11 #       distributions on the circle (angles 0 to 2pi)
  12 #       ---------------------------------------------
  13 #              circular uniform
  14 #              von Mises
  15
  16 # Translated from anonymously contributed C/C++ source.
  17
  18 from whrandom import random, uniform, randint, choice # Also for export!
  19 from math import log, exp, pi, e, sqrt, acos, cos, sin
  20
  21 # Housekeeping function to verify that magic constants have been
  22 # computed correctly
  23
  24 def verify(name, expected):
  25         computed = eval(name)
  26         if abs(computed - expected) > 1e-7:
  27                 raise ValueError, \
  28   'computed value for %s deviates too much (computed %g, expected %g)' % \
  29   (name, computed, expected)
  30
  31 # -------------------- normal distribution --------------------
  32
  33 NV_MAGICCONST = 4*exp(-0.5)/sqrt(2.0)
  34 verify('NV_MAGICCONST', 1.71552776992141)
  35 def normalvariate(mu, sigma):
  36         # mu = mean, sigma = standard deviation
  37
  38         # Uses Kinderman and Monahan method. Reference: Kinderman,
  39         # A.J. and Monahan, J.F., "Computer generation of random
  40         # variables using the ratio of uniform deviates", ACM Trans
  41         # Math Software, 3, (1977), pp257-260.
  42
  43         while 1:
  44                 u1 = random()
  45                 u2 = random()
  46                 z = NV_MAGICCONST*(u1-0.5)/u2
  47                 zz = z*z/4.0
  48                 if zz <= -log(u2):
  49                         break
  50         return mu+z*sigma
  51
  52 # -------------------- lognormal distribution --------------------
  53
  54 def lognormvariate(mu, sigma):
  55         return exp(normalvariate(mu, sigma))
  56
  57 # -------------------- circular uniform --------------------
  58
  59 def cunifvariate(mean, arc):
  60         # mean: mean angle (in radians between 0 and pi)
  61         # arc:  range of distribution (in radians between 0 and pi)
  62
  63         return (mean + arc * (random() - 0.5)) % pi
  64
  65 # -------------------- exponential distribution --------------------
  66
  67 def expovariate(lambd):
  68         # lambd: rate lambd = 1/mean
  69         # ('lambda' is a Python reserved word)
  70
  71         u = random()
  72         while u <= 1e-7:
  73                 u = random()
  74         return -log(u)/lambd
  75
  76 # -------------------- von Mises distribution --------------------
  77
  78 TWOPI = 2.0*pi
  79 verify('TWOPI', 6.28318530718)
  80
  81 def vonmisesvariate(mu, kappa):
  82         # mu:    mean angle (in radians between 0 and 180 degrees)
  83         # kappa: concentration parameter kappa (>= 0)
  84
  85         # if kappa = 0 generate uniform random angle
  86         if kappa <= 1e-6:
  87                 return TWOPI * random()
  88
  89         a = 1.0 + sqrt(1.0 + 4.0 * kappa * kappa)
  90         b = (a - sqrt(2.0 * a))/(2.0 * kappa)
  91         r = (1.0 + b * b)/(2.0 * b)
  92
  93         while 1:
  94                 u1 = random()
  95
  96                 z = cos(pi * u1)
  97                 f = (1.0 + r * z)/(r + z)
  98                 c = kappa * (r - f)
  99
 100                 u2 = random()
 101
 102                 if not (u2 >= c * (2.0 - c) and u2 > c * exp(1.0 - c)):
 103                         break
 104
 105         u3 = random()
 106         if u3 > 0.5:
 107                 theta = mu + 0.5*acos(f)
 108         else:
 109                 theta = mu - 0.5*acos(f)
 110
 111         return theta % pi
 112
 113 # -------------------- gamma distribution --------------------
 114
 115 LOG4 = log(4.0)
 116 verify('LOG4', 1.38629436111989)
 117
 118 def gammavariate(alpha, beta):
 119         # beta times standard gamma
 120         ainv = sqrt(2.0 * alpha - 1.0)
 121         return beta * stdgamma(alpha, ainv, alpha - LOG4, alpha + ainv)
 122
 123 SG_MAGICCONST = 1.0 + log(4.5)
 124 verify('SG_MAGICCONST', 2.50407739677627)
 125
 126 def stdgamma(alpha, ainv, bbb, ccc):
 127         # ainv = sqrt(2 * alpha - 1)
 128         # bbb = alpha - log(4)
 129         # ccc = alpha + ainv
 130
 131         if alpha <= 0.0:
 132                 raise ValueError, 'stdgamma: alpha must be > 0.0'
 133
 134         if alpha > 1.0:
 135
 136                 # Uses R.C.H. Cheng, "The generation of Gamma
 137                 # variables with non-integral shape parameters",
 138                 # Applied Statistics, (1977), 26, No. 1, p71-74
 139
 140                 while 1:
 141                         u1 = random()
 142                         u2 = random()
 143                         v = log(u1/(1.0-u1))/ainv
 144                         x = alpha*exp(v)
 145                         z = u1*u1*u2
 146                         r = bbb+ccc*v-x
 147                         if r + SG_MAGICCONST - 4.5*z >= 0.0 or r >= log(z):
 148                                 return x
 149
 150         elif alpha == 1.0:
 151                 # expovariate(1)
 152                 u = random()
 153                 while u <= 1e-7:
 154                         u = random()
 155                 return -log(u)
 156
 157         else:   # alpha is between 0 and 1 (exclusive)
 158
 159                 # Uses ALGORITHM GS of Statistical Computing - Kennedy & Gentle
 160
 161                 while 1:
 162                         u = random()
 163                         b = (e + alpha)/e
 164                         p = b*u
 165                         if p <= 1.0:
 166                                 x = pow(p, 1.0/alpha)
 167                         else:
 168                                 # p > 1
 169                                 x = -log((b-p)/alpha)
 170                         u1 = random()
 171                         if not (((p <= 1.0) and (u1 > exp(-x))) or
 172                                   ((p > 1)  and  (u1 > pow(x, alpha - 1.0)))):
 173                                 break
 174                 return x
 175
 176
 177 # -------------------- Gauss (faster alternative) --------------------
 178
 179 gauss_next = None
 180 def gauss(mu, sigma):
 181
 182         # When x and y are two variables from [0, 1), uniformly
 183         # distributed, then
 184         #
 185         #    cos(2*pi*x)*log(1-y)
 186         #    sin(2*pi*x)*log(1-y)
 187         #
 188         # are two *independent* variables with normal distribution
 189         # (mu = 0, sigma = 1).
 190         # (Lambert Meertens)
 191
 192         global gauss_next
 193
 194         if gauss_next != None:
 195                 z = gauss_next
 196                 gauss_next = None
 197         else:
 198                 x2pi = random() * TWOPI
 199                 log1_y = log(1.0 - random())
 200                 z = cos(x2pi) * log1_y
 201                 gauss_next = sin(x2pi) * log1_y
 202
 203         return mu + z*sigma
 204
 205 # -------------------- beta --------------------
 206
 207 def betavariate(alpha, beta):
 208
 209         # Discrete Event Simulation in C, pp 87-88.
 210
 211         y = expovariate(alpha)
 212         z = expovariate(1.0/beta)
 213         return z/(y+z)
 214
 215 # -------------------- test program --------------------
 216
 217 def test(N = 200):
 218         print 'TWOPI         =', TWOPI
 219         print 'LOG4          =', LOG4
 220         print 'NV_MAGICCONST =', NV_MAGICCONST
 221         print 'SG_MAGICCONST =', SG_MAGICCONST
 222         test_generator(N, 'random()')
 223         test_generator(N, 'normalvariate(0.0, 1.0)')
 224         test_generator(N, 'lognormvariate(0.0, 1.0)')
 225         test_generator(N, 'cunifvariate(0.0, 1.0)')
 226         test_generator(N, 'expovariate(1.0)')
 227         test_generator(N, 'vonmisesvariate(0.0, 1.0)')
 228         test_generator(N, 'gammavariate(0.5, 1.0)')
 229         test_generator(N, 'gammavariate(0.9, 1.0)')
 230         test_generator(N, 'gammavariate(1.0, 1.0)')
 231         test_generator(N, 'gammavariate(2.0, 1.0)')
 232         test_generator(N, 'gammavariate(20.0, 1.0)')
 233         test_generator(N, 'gammavariate(200.0, 1.0)')
 234         test_generator(N, 'gauss(0.0, 1.0)')
 235         test_generator(N, 'betavariate(3.0, 3.0)')
 236
 237 def test_generator(n, funccall):
 238         import time
 239         print n, 'times', funccall
 240         code = compile(funccall, funccall, 'eval')
 241         sum = 0.0
 242         sqsum = 0.0
 243         smallest = 1e10
 244         largest = -1e10
 245         t0 = time.time()
 246         for i in range(n):
 247                 x = eval(code)
 248                 sum = sum + x
 249                 sqsum = sqsum + x*x
 250                 smallest = min(x, smallest)
 251                 largest = max(x, largest)
 252         t1 = time.time()
 253         print round(t1-t0, 3), 'sec,',
 254         avg = sum/n
 255         stddev = sqrt(sqsum/n - avg*avg)
 256         print 'avg %g, stddev %g, min %g, max %g' % \
 257                   (avg, stddev, smallest, largest)
 258
 259 if __name__ == '__main__':
 260         test()