Merge pull request #22 from PyCOMPLETE/develop

PyHEADTAIL v1.5.0

Merry Christmas! :-)

New features:
- UniformBinSlicer allows for specified sampling points (for wakes)
- RFSystems can remove kicks
- KV 1D and 2D distribution included into generators
- Particles can sort their arrays and apply given permutations
- a first version of the GPU module enters PyHEADTAIL (Particles, Slicing etc on GPU)
- CUDA thrust library interface for specific functions
- ParticlesGPU implements sorting via thrust
- slicing with PyPIC (e.g. for GPU usage)
- RFBucketMatcher allows for margin along separatrix
- optimised RFBucketMatcher generation (much quicker!)
- longitudinal wrapping to keep z within circumference
- pre-commit hooks (make sure to include these in your GIT config!)
- ParticleGenerator interface change (linear_matcher removed)
- pure python RFQ implementation (besides existing cython version)
- adaptation of Bunch-/SliceMonitor for GPU usage

Author: aoeftiger

README.md

@@ -34,3 +34,11 @@ In [1]: import PyHEADTAIL
 
 PyHEADTAIL v1.4.0-0-g8422af8081
 
+
+-------------------------------------------------------------------------------
+
+Please use the pre-push script 'pre-push.py' if you want to contribute
+to the repository. It only lets you push to the develop and master branch if
+no unit tests fail.
+
+To install (creates a symlink): ln -s ../../pre-push.py .git/hooks/pre-push

field_maps/Transverse_Efield_map.py

@@ -1,5 +1,5 @@
 import numpy as np
-from scipy.constants import c, e
+from scipy.constants import c
 
 from PyHEADTAIL.particles.slicing import UniformBinSlicer
 from PyPIC.PyPIC_Scatter_Gather import PyPIC_Scatter_Gather

general/utils.py

@@ -68,3 +68,8 @@ def __repr__(self):
 
     def __len__(self):
         return len(self._list_of_objects)
+
+    def pop(self, index):
+        '''Remove the object from the internal list and return the
+        corresponding attribute, analogous to list.pop .'''
+        return getattr(self._list_of_objects.pop(index), self._attr_name)

gpu/__init__.py

@@ -0,0 +1,43 @@
+from .. import Element
+from .. import __version__
+
+# monkey patching default classes with GPU variants
+# is done on importing this module!
+
+from skcuda.misc import mean, std, diff
+from pycuda import gpuarray
+from pycuda import cumath
+
+# I. subjects to monkey patching:
+from ..particles import particles as def_particles
+from ..particles import slicing as def_slicing
+from ..trackers import simple_long_tracking as def_simple_long_tracking
+from ..trackers import wrapper as def_wrapper
+
+# II. actual monkey patching
+
+# a) Particles rebindings for GPU
+from .particles import ParticlesGPU
+
+def_particles.Particles = ParticlesGPU
+def_particles.mean = lambda *args, **kwargs: mean(*args, **kwargs).get()
+def_particles.std = lambda *args, **kwargs: std(*args, **kwargs).get()
+
+# b) Slicing rebindings for GPU
+# def_slicing.min_ = lambda *args, **kwargs: gpuarray.min(*args, **kwargs).get()
+# def_slicing.max_ = lambda *args, **kwargs: gpuarray.max(*args, **kwargs).get()
+# def_slicing.diff = diff
+
+from .slicing import SlicerGPU
+
+# # to be replaced: find a better solution than monkey patching base classes!
+# # (to solve the corresponding need to replace all inheriting classes' parent!)
+# def_slicing.Slicer = SlicerGPU
+# def_slicing.UniformBinSlicer.__bases__ = (SlicerGPU,)
+
+# c) Longitudinal tracker rebindings for GPU
+def_simple_long_tracking.sin = cumath.sin
+
+# d) Wrapper rebindings for GPU
+from .wrapper import LongWrapperGPU
+def_wrapper.LongWrapper = LongWrapperGPU

gpu/particles.py

@@ -0,0 +1,54 @@
+# default classes imports from modules as assigned in gpu/__init__.py
+from . import def_particles
+
+import numpy as np
+
+from pycuda import gpuarray
+
+import thrust_interface
+thrust = thrust_interface.compiled_module
+
+class ParticlesGPU(def_particles.Particles):
+    '''Implementation of the Particles with data stored on the GPU.'''
+    def __init__(self, *args, **kwargs):
+        super(ParticlesGPU, self).__init__(*args, **kwargs)
+
+        for coord in self.coords_n_momenta:
+            device_array = gpuarray.to_gpu(getattr(self, coord))
+            setattr(self, coord, device_array)
+        self.id = gpuarray.to_gpu(self.id)
+
+    def transfer_to_host(self):
+        '''Transfers all GPU device data back to the host RAM.'''
+        for coord in self.coords_n_momenta:
+            device_array = getattr(self, coord).get()
+            setattr(self, coord, device_array)
+        self.id = self.id.get()
+
+    def sort_for(self, attr):
+        '''Sort the named particle attribute (coordinate / momentum)
+        array and reorder all particles accordingly.
+        '''
+        permutation = gpuarray.empty(self.macroparticlenumber, dtype=np.int32)
+        thrust.get_sort_perm_double(getattr(self, attr), permutation)
+        self.reorder(permutation, except_for_attrs=[attr])
+
+    def reorder(self, permutation, except_for_attrs=[]):
+        '''Reorder all particle coordinate and momentum arrays
+        (in self.coords_n_momenta) and ids except for except_for_attrs
+        according to the given index array permutation.
+        '''
+        tmp = gpuarray.empty(self.macroparticlenumber, dtype=np.float64)
+        for attr in self.coords_n_momenta:
+            if attr in except_for_attrs:
+                continue
+            unordered = getattr(self, attr)
+            thrust.apply_sort_perm_double(unordered, tmp, permutation)
+            setattr(self, attr, tmp)
+            tmp = unordered
+        del tmp
+        tmp = gpuarray.empty(self.macroparticlenumber, dtype=self.id.dtype)
+        thrust.apply_sort_perm_int(self.id, tmp, permutation)
+        self.id = tmp
+
+