From 9854773d3387cd73af17fec950cea9372b1fa97f Mon Sep 17 00:00:00 2001
From: Kevin Shing Bruce Li <Kevin.Shing.Bruce.Li@cern.ch>
Date: Tue, 21 Feb 2017 13:03:12 +0100
Subject: [PATCH] Add Waterbag and KV generator based on Knuth's algorithm,
 with literature sources.

---
 PyHEADTAIL/particles/generators.py | 87 ++++++++++++++++--------------
 1 file changed, 47 insertions(+), 40 deletions(-)

diff --git a/PyHEADTAIL/particles/generators.py b/PyHEADTAIL/particles/generators.py
index a32baad1..83bc2f96 100644
--- a/PyHEADTAIL/particles/generators.py
+++ b/PyHEADTAIL/particles/generators.py
@@ -375,60 +375,67 @@ def _uniform2D(n_particles):
     return _uniform2D
 
 
-'''
-Why we have this ll algo in here? We had a better one (Knuth):
-http://citeseerx.ist.psu.edu/viewdoc/download?doi=10.1.1.138.8671&rep=rep1&type=pdf
+def _knuth_uniform(order=0, d=2, *args, **kwargs):
+    '''Algorithm based on Don Knuth:
 
+    Jan Poland: "Three Different Algorithms for Generating Uniformly
+    Distributed Random Points on the N-Sphere",
+    http://citeseerx.ist.psu.edu/viewdoc/download?doi=10.1.1.138.8671&rep=rep1&type=pdf
 
-import numpy as np
+    '''
+    def _uniform(n_macroparticles):
+        u = np.random.normal(size=(d, n_macroparticles))
+        r = np.sqrt(np.sum(u**2, axis=0)) # sum across dimensions for radius
+
+        if order == 0: # K-V or airbag
+            s = 1./r
+        elif order==1: # Waterbag
+            s = 1./r * np.random.rand(n_macroparticles)**(1./d)
+        u *= s
 
+        a = args
+        print order, d, args
+        # print(u.shape, r.shape, s.shape)
 
-def transform_into_kv(bunch, a_x, a_xp, a_y, a_yp):
+        u *= np.array([args]).T
+        return u
+    return _uniform
 
-    # KV distribution
-    # ==================================================
-    d = 4
-    for i in range(bunch.macroparticlenumber):
-       u = np.random.normal(size=d)
-       r = np.sqrt(np.sum(u**2))
-       u *= 1./r
 
-       bunch.x[i]  = u[0]
-       bunch.xp[i] = u[1]
-       bunch.y[i]  = u[2]
-       bunch.yp[i] = u[3]
-    # ==================================================
+def KV2D(r_u, r_up):
+    '''Using Knuth's algorithm
 
-    bunch.x  *= a_x
-    bunch.xp *= a_xp
-    bunch.y  *= a_y
-    bunch.yp *= a_yp
+    ..Args: r_u, r_up are +/- extensions in u and u_p
+
+    '''
+    return _knuth_uniform(0, 2, r_u, r_up)
 
 
-def transform_into_waterbag(bunch, a_x, a_xp, a_y, a_yp):
+def KV4D(r_x, r_xp, r_y, r_yp):
+    '''Using Knuth's algorithm
 
-    # waterbag distribution
-    # ==================================================
-    d = 4
-    for i in range(bunch.macroparticlenumber):
-        u = np.random.normal(size=d)
-        r = np.sqrt(np.sum(u**2))
-        u *= (np.random.rand(1))**(1./d)/r
+    ..Args: r_u, r_up are +/- extensions in u and u_p
 
-        bunch.x[i]  = u[0]
-        bunch.xp[i] = u[1]
-        bunch.y[i]  = u[2]
-        bunch.yp[i] = u[3]
-    # ==================================================
+    '''
+    return _knuth_uniform(0, 4, r_x, r_xp, r_y, r_yp)
 
-    bunch.x  *= a_x
-    bunch.xp *= a_xp
-    bunch.y  *= a_y
-    bunch.yp *= a_yp
 
+def waterbag2D(r_u, r_up):
+    '''Using Knuth's algorithm
 
-Want to get this back in the near future...
-'''
+    ..Args: r_u, r_up are +/- extensions in u and u_p
+
+    '''
+    return _knuth_uniform(1, 2, r_u, r_up)
+
+
+def waterbag4D(r_x, r_xp, r_y, r_yp):
+    '''Using Knuth's algorithm
+
+    ..Args: r_u, r_up are +/- extensions in u and u_p
+
+    '''
+    return _knuth_uniform(1, 4, r_x, r_xp, r_y, r_yp)
 
 
 def kv2D(r_u, r_up):