'''
Generate multivariate von Mises Fisher samples via a
rejection sampling scheme from
    "Directional Statistics" (Mardia and Jupp, 1999)

This solution originally appears here:
http://stats.stackexchange.com/questions/156729/sampling-from-von-mises-fisher-distribution-in-python


References:

Sampling from vMF on S^2:
    https://www.mitsuba-renderer.org/~wenzel/files/vmf.pdf
    http://www.stat.pitt.edu/sungkyu/software/randvonMisesFisher3.pdf
'''
import numpy as np


def weight_rejection_sampling(kappa, dim, num_samples):
    dim = dim - 1 # since S^{n-1}
    b = dim / (np.sqrt(4. * kappa**2 + dim**2) + 2 * kappa)
    x = (1 - b) / (1 + b)
    c = kappa * x + dim * np.log(1 - x**2)

    samples = []
    for nn in range(num_samples):
        while True:
            z = np.random.beta(dim / 2., dim / 2.)
            w = (1. - (1. + b) * z) / (1. - (1. - b) * z)
            u = np.random.uniform(low=0, high=1)
            if kappa * w + dim * np.log(1 - x * w) - c >= np.log(u):
                samples.append(w)
                break

    return samples


def sample_tangent_unit(mu):
    mat = np.matrix(mu)

    if mat.shape[1] > mat.shape[0]:
        mat = mat.T

    U, _, _ = np.linalg.svd(mat)
    nu = np.matrix(np.random.randn(mat.shape[0])).T
    x = np.dot(U[:, 1:], nu[1:, :])
    return np.squeeze(np.array(x / np.linalg.norm(x)))


def vMF(mu, kappa, num_samples):
    dim = len(mu)

    # sample some weights
    ws = weight_rejection_sampling(kappa, dim, num_samples)

    result = np.zeros((num_samples, dim))
    for nn in range(num_samples):
        v = sample_tangent_unit(mu)

        # compute new point
        result[nn, :] = v * np.sqrt(1. - ws[nn]**2) + ws[nn] * mu

    return result

@@ -47,98 +47,9 @@
 ples
- via a%0Arejection sampling scheme from%0A    %22Directional Statistics%22 (Mardia and Jupp, 1999)
+.
 %0A%0ATh
@@ -192,19 +192,16 @@
 on%0A%0A
-%0AReferences
+Also see
 :%0A%0AS
@@ -383,32 +383,59 @@
 p%0A%0A%0Adef 
-weight_r
+sample_weight(kappa, dim):%0A    %22%22%22R
 ejection
 _samplin
@@ -422,33 +422,33 @@
 %22%22%22Rejection
-_
+ 
 sampling
 (kappa, dim,
@@ -439,34 +439,87 @@
 ling
-(kappa, dim, num_samples):
+ scheme for sampling distance from center on%0A    surface of the sphere.%0A    %22%22%22
 %0A   
@@ -683,63 +683,8 @@
 2)%0A%0A
-    samples = %5B%5D%0A    for nn in range(num_samples):%0A    
     
@@ -703,20 +703,16 @@
         
-    
 z = np.r
@@ -742,20 +742,16 @@
 m / 2.)%0A
-    
         
@@ -804,20 +804,16 @@
         
-    
 u = np.r
@@ -849,20 +849,16 @@
         
-    
 if kappa
@@ -874,24 +874,25 @@
 m * np.log(1
+.
  - x * w) - 
@@ -923,237 +923,207 @@
     
-    samples.append(w)%0A                break%0A%0A    return samples%0A%0A%0Adef sample_tangent_unit(mu):%0A    mat = np.matrix(mu)
+return w%0A%0A%0Adef _orthonormal_to(mu):%0A    %22%22%22Sample point on sphere orthogonal to mu.
 %0A%0A    
-if mat.shape%5B1%5D %3E mat.shape%5B0%5D:%0A        mat = mat.T%0A%0A    U, _, _ = np.linalg.svd(mat)%0A    nu = np.matrix(
+This implementation is like one iteration of gram-schmidt%0A    between mu and a random vector.%0A    %22%22%22%0A    v = 
 np.r
@@ -1135,18 +1135,17 @@
 .randn(m
-at
+u
 .shape%5B0
@@ -1150,80 +1150,109 @@
 %5B0%5D)
-).T
 %0A    
-x = np.dot(U%5B:, 1:%5D, nu%5B1:, :%5D)%0A    return np.squeeze(np.array(x
+proj_mu_v = mu * np.dot(mu, v) / np.linalg.norm(mu)%0A    orthto = v - proj_mu_v%0A    return orthto
  / n
@@ -1269,11 +1269,14 @@
 orm(
-x))
+orthto
 )%0A%0A%0A
@@ -1316,217 +1316,448 @@
     
-dim = len(mu)%0A%0A    # sample some weights%0A    ws = weight_rejection_sampling(kappa, dim, num_samples)%0A%0A    result = np.zeros((num_samples, dim))%0A    for nn in range(num_samples):%0A        v = sample_tangent_unit
+%22%22%22Generate num_samples N-dimensional samples from von Mises Fisher%0A    distribution around center mu %5Cin R%5EN with concentration kappa.%0A    %22%22%22%0A    dim = len(mu)%0A    result = np.zeros((num_samples, dim))%0A    for nn in range(num_samples):%0A        # sample offset from center (on sphere) with spread kappa%0A        w = sample_weight(kappa, dim)%0A%0A        # sample a point v on the unit sphere that's orthogonal to mu%0A        v = _orthonormal_to
 (mu)
@@ -1832,21 +1832,16 @@
  - w
-s%5Bnn%5D
 **2) + w
 s%5Bnn
@@ -1840,13 +1840,8 @@
  + w
-s%5Bnn%5D
  * m

Update comments, make code a bit cleaner