Add Waterbag and KV generator based on Knuth's algorithm.

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):