Pythonで3次元単位球面上に一様な乱数を生成する

お題の通り. 単純に

v = numpy.random.uniform(-1, 1, size=3)
v /= numpy.linalg.norm(v)

等とすると一様でなく偏った分布になってしまう. この場合, 例えば, 一行目で生成したvの長さが1以下の場合に限って正規化して返し, それ以上の場合は生成しなおすようにすれば良い(下のコードのcall_sprnd1を参照).

それ以外にも色々な生成法があるのだが, 実際どの程度速度に違いが出るのか, 適当に書いたコードで試してみた. ただし, 面倒だったので乱数区間の端っこの扱い(開区間か閉区間か)が若干いい加減なので注意.

import time
import numpy
import random as rng

PI_2 = numpy.pi * 2

def call_wrong_sprnd():
  v = numpy.array([rng.random(), rng.random(), rng.random()]) * 2 - 1
  return v / numpy.linalg.norm(v)

def call_sprnd1():
  while True:
    v = numpy.array([rng.random(), rng.random(), rng.random()]) * 2 - 1
    r = numpy.linalg.norm(v)
    if r <= 1:
      return v / r

def call_sprnd2():
  phi, z = rng.random() * PI_2, rng.random() * 2 - 1
  a = numpy.sqrt(1 - z * z)
  x = numpy.cos(phi)
  y = numpy.sin(phi)
  return numpy.array([a * x, a * y, z])

def call_sprnd3():
  v = numpy.array([
      rng.normalvariate(0, 1),
      rng.normalvariate(0, 1),
      rng.normalvariate(0, 1)])
  return v / numpy.linalg.norm(v)

def call_sprnd4():
  a, b, r2 = 0, 0, 1
  while r2 > 0.25:
    a = rng.random() - 0.5
    b = rng.random() - 0.5
    r2 = a * a + b * b
  z = 8 * r2 - 1
  scale = 8 * numpy.sqrt(0.25 - r2)
  return numpy.array([a * scale, b * scale, z])

def call_sprnd5():
  while True:
    X, Y = rng.random() * 2 - 1, rng.random() * 2 - 1
    S = X * X + Y * Y
    if S <= 1:
      break
  Z = 2 * S - 1
  S = numpy.sqrt((1 - Z * Z) / S)
  return numpy.array([X * S, Y * S, Z])

def count_acc_time(f, n=100000, m=5):
  retval = []
  for i in xrange(m):
    tstart = time.time()
    for i in xrange(n):
      f()
    tend = time.time()
    retval.append(tend - tstart)

  retval = numpy.array(retval)
  print '%s : %.3f \pm %.3f' % (f.func_name, numpy.average(retval), numpy.std(retval))
  return retval

def dump_sprnd(f, filename='test.png'):
  from matplotlib import pylab
  from mpl_toolkits.mplot3d import Axes3D

  retval = numpy.array([f() for i in xrange(10000)])
  fig = pylab.figure()
  ax = Axes3D(fig)
  ax.scatter3D(retval[: , 0], retval[: , 1], retval[: , 2], color=(1, 0, 0), marker=".")
  pylab.savefig(filename)
  pylab.show()

if __name__ == '__main__':
  count_acc_time(call_sprnd1)
  count_acc_time(call_sprnd2)
  count_acc_time(call_sprnd3)
  count_acc_time(call_sprnd4)
  count_acc_time(call_sprnd5)

  # dump_sprnd(call_wrong_sprnd)

アルゴリズムについては, 以下のページに書かれていたものを参考にした.

http://d.hatena.ne.jp/nakt/20100116
http://hep.planet-koo.com/index.php?g=root&sid=rt_g_grandom_sphere

結果は以下の通り.

call_sprnd1 : 3.952 \pm 0.009
call_sprnd2 : 1.124 \pm 0.000
call_sprnd3 : 1.919 \pm 0.001
call_sprnd4 : 1.075 \pm 0.001
call_sprnd5 : 1.023 \pm 0.001

結果としては予想通り, sprnd4とsprnd5が速い. sprnd4とsprnd5は実質的には同じもので, gsl_ran_dir_3dで実装されているアルゴリズムです.

http://www.gnu.org/software/gsl/manual/html_node/Spherical-Vector-Distributions.html
http://dl.acm.org/citation.cfm?id=362377

sprnd2のsin, cosとかモジュールの選び方でもっと効率化できると思うが, ざっと書いた感じではこの程度の違いです. これを見るとループのない方法としてsprnd2でも十分速いですね.