Merge branch 'develop'

 - SliceSet:Introduced particles_outside_cut
 - Particles: indeces handled correcly by: __add__, __radd__, and extract_slices
 - Particles: Add include_non_sliced functionality to extract_slices

From tests with leptons:
 - Better estimate of Q_s in generate_6D_Gaussian_bunch
 - Longitudinal tracking and matching work with negative charge
 - Added CLIC-DR example with electrons

Author: giadarol

machines/synchrotron.py

@@ -101,15 +101,19 @@ def generate_6D_Gaussian_bunch(self, n_macroparticles, intensity,
         '''
         if self.longitudinal_mode == 'linear':
             check_inside_bucket = lambda z,dp : np.array(len(z)*[True])
+            Qs = self.longitudinal_map.Qs
         elif self.longitudinal_mode == 'non-linear':
-            check_inside_bucket = self.longitudinal_map.get_bucket(
-                gamma=self.gamma).make_is_accepted(margin=0.05)
+            bucket = self.longitudinal_map.get_bucket(
+                gamma=self.gamma, mass=self.mass, charge=self.charge)
+            check_inside_bucket = bucket.make_is_accepted(margin=0.05)
+            Qs = bucket.Qs
+
         else:
             raise NotImplementedError(
                 'Something wrong with self.longitudinal_mode')
 
         eta = self.longitudinal_map.alpha_array[0] - self.gamma**-2
-        beta_z    = np.abs(eta)*self.circumference/2./np.pi/self.longitudinal_map.Qs
+        beta_z    = np.abs(eta)*self.circumference/2./np.pi/Qs
         sigma_dp  = sigma_z/beta_z
         epsx_geo = epsn_x/self.betagamma
         epsy_geo = epsn_y/self.betagamma

particles/particles.py

@@ -151,34 +151,64 @@ def get_slices(self, slicer, *args, **kwargs):
         return self._slice_sets[slicer]
 
 
-    def extract_slices(self, slicer, *args, **kwargs):
-        '''Return a list Particles object with the different slices.
+    def extract_slices(self, slicer, include_non_sliced='if_any', *args, **kwargs):
+        '''Return a list Particles object with the different slices. 
+        The last element of the list contains particles not assigned to any slice.
+        
+        include_non_sliced : {'always', 'never', 'if_any'}, optional
+        'always':
+          extra element in the list with particles not belonging to any slice 
+          is always inserted (it can be empty).
+        'never':
+          extra element in the list with particles not belonging to any slice 
+          is never inserted.
+        'if_any':
+          extra element in the list with particles not belonging to any slice 
+          is inserted only if such particles exist.       
         '''
-
+        
+        if include_non_sliced not in ['if_any', 'always', 'never']:
+        	raise ValueError("include_non_sliced=%s is not valid!\n"%include_non_sliced+
+        					"Possible values are {'always', 'never', 'if_any'}" )
+        
         slices = self.get_slices(slicer, *args, **kwargs)
         self_coords_n_momenta_dict = self.get_coords_n_momenta_dict()
         slice_object_list = []
-
+        
         for i_sl in xrange(slices.n_slices):
-
+        
             ix = slices.particle_indices_of_slice(i_sl)
             macroparticlenumber = len(ix)
-
+            
             slice_object = Particles(macroparticlenumber=macroparticlenumber, 
                 particlenumber_per_mp=self.particlenumber_per_mp, charge=self.charge,
                 mass=self.mass, circumference=self.circumference, gamma=self.gamma, coords_n_momenta_dict={})
 
             for coord in self_coords_n_momenta_dict.keys():
                 slice_object.update({coord: self_coords_n_momenta_dict[coord][ix]})
-
+            
+            slice_object.id[:] = self.id[ix]
+            
             slice_object.slice_info = {\
                     'z_bin_center': slices.z_centers[i_sl],\
                     'z_bin_right':slices.z_bins[i_sl+1],\
                     'z_bin_left':slices.z_bins[i_sl]}
 
             slice_object_list.append(slice_object)
-            
-
+        
+        # handle unsliced
+        if include_non_sliced is not 'never':
+            ix = slices.particles_outside_cuts
+            if len(ix)>0 or include_non_sliced is 'always':
+                slice_object = Particles(macroparticlenumber=len(ix), 
+                    particlenumber_per_mp=self.particlenumber_per_mp, charge=self.charge,
+                    mass=self.mass, circumference=self.circumference, gamma=self.gamma, coords_n_momenta_dict={})
+                for coord in self_coords_n_momenta_dict.keys():
+                	slice_object.update({coord: self_coords_n_momenta_dict[coord][ix]})
+                slice_object.id[:] = self.id[ix]
+                slice_object.slice_info = 'unsliced' 
+            	slice_object_list.append(slice_object)
+        
         return slice_object_list
 
     def clean_slices(self):
@@ -251,6 +281,8 @@ def __add__(self, other):
             #setattr(result, coord, np.concatenate((self_coords_n_momenta_dict[coord].copy(), other_coords_n_momenta_dict[coord].copy())))
             result.update({coord: np.concatenate((self_coords_n_momenta_dict[coord].copy(), other_coords_n_momenta_dict[coord].copy()))})
 
+        result.id = np.concatenate((self.id.copy(), other.id.copy()))
+        
         return result
 
     def __radd__(self, other):
@@ -263,6 +295,7 @@ def __radd__(self, other):
             for coord in self_coords_n_momenta_dict.keys():
                 #setattr(result, coord, np.concatenate((self_coords_n_momenta_dict[coord].copy(), other_coords_n_momenta_dict[coord].copy())))
                 result.update({coord: self_coords_n_momenta_dict[coord].copy()})
+            result.id = self.id.copy()
         else:
             result = self.__add__(other)
 

particles/slicing.py

@@ -174,6 +174,17 @@ def particles_within_cuts(self):
                 (self.slice_index_of_particle < self.n_slices)
             )[0])
         return particles_within_cuts_
+    
+    @property
+    # @memoize
+    def particles_outside_cuts(self):
+        '''All particle indices which are situated outside the slicing
+        region defined by [z_cut_tail, z_cut_head).'''
+        particles_ouside_cuts_ = make_int32(np.where(np.logical_not(
+                (self.slice_index_of_particle > -1) &
+                (self.slice_index_of_particle < self.n_slices))
+            )[0])
+        return particles_ouside_cuts_
 
     @property
     # @memoize

testing/interactive-tests/SlicingTest.ipynb

@@ -0,0 +1,115 @@
+from PyHEADTAIL.machines.synchrotron import BasicSynchrotron
+import numpy as np
+from scipy.constants import c, e, m_e
+
+class CLIC_DR(BasicSynchrotron):
+
+	def __init__(self, machine_configuration=None, optics_mode='smooth', **kwargs):
+		
+		 
+		alpha		= 0.0001276102729
+		h_RF       	= 1465
+		mass 		= m_e
+		charge		= e
+		RF_at		= 'middle'
+		
+		if machine_configuration=='3TeV':
+			longitudinal_mode = 'non-linear'				
+			p0 			= 2.86e9 * e /c
+			p_increment = 0.
+			accQ_x		= 48.35
+			accQ_y		= 10.4
+			V_RF		= 5.1e6
+			dphi_RF		= 0.
+                        Q_s             = None
+		elif machine_configuration=='3TeV_linear':
+			Q_s		= 0.0059	
+			longitudinal_mode = 'linear'				
+			p0 			= 2.86e9 * e /c
+			p_increment = 0.
+			accQ_x		= 48.35
+			accQ_y		= 10.4
+			V_RF		= 5.1e6
+			dphi_RF		= 0.
+		else:
+			raise ValueError('machine_configuration not recognized!')
+			
+		
+		
+		if optics_mode=='smooth':
+			if 's' in kwargs.keys(): raise ValueError('s vector cannot be provided if optics_mode = "smooth"')
+			
+			n_segments = kwargs['n_segments']
+			circumference = 427.5
+			
+			name = None
+			
+			beta_x 		= 3.475 
+			D_x 		= 0
+			beta_y 		= 9.233
+			D_y 		= 0 
+			
+			alpha_x 	= None
+			alpha_y 	= None
+			
+			s = None
+			
+		elif optics_mode=='non-smooth':
+			if 'n_segments' in kwargs.keys(): raise ValueError('n_segments cannot be provided if optics_mode = "non-smooth"')
+			n_segments = None
+			circumference = None
+			
+			name		= kwargs['name']
+			
+			beta_x 		= kwargs['beta_x']
+			beta_y 		= kwargs['beta_y'] 
+
+			try:
+				D_x 		= kwargs['D_x']	
+			except KeyError:
+				D_x 		= 0*np.array(kwargs['s'])
+			try:
+				D_y 		= kwargs['D_y']	
+			except KeyError:
+				D_y 		= 0*np.array(kwargs['s'])			
+					
+			alpha_x 	= kwargs['alpha_x']
+			alpha_y 	= kwargs['alpha_y']
+			
+			s = kwargs['s']
+		
+		else:
+			raise ValueError('optics_mode not recognized!')		
+
+		
+		# detunings
+		Qp_x		= 0
+		Qp_y		= 0
+		
+		app_x		= 0
+		app_y		= 0
+		app_xy		= 0
+		
+
+		
+		for attr in kwargs.keys():
+			if kwargs[attr] is not None:
+				if type(kwargs[attr]) is list or type(kwargs[attr]) is np.ndarray:
+					str2print = '[%s ...]'%repr(kwargs[attr][0])
+				else:
+					str2print = repr(kwargs[attr])
+				self.prints('Synchrotron init. From kwargs: %s = %s'
+							% (attr, str2print))
+				temp =  kwargs[attr]
+				exec('%s = temp'%attr)
+				
+		
+		
+		super(CLIC_DR, self).__init__(optics_mode=optics_mode, circumference=circumference, n_segments=n_segments, s=s, name=name,
+             alpha_x=alpha_x, beta_x=beta_x, D_x=D_x, alpha_y=alpha_y, beta_y=beta_y, D_y=D_y,
+             accQ_x=accQ_x, accQ_y=accQ_y, Qp_x=Qp_x, Qp_y=Qp_y, app_x=app_x, app_y=app_y, app_xy=app_xy,
+             alpha_mom_compaction=alpha, longitudinal_mode=longitudinal_mode,
+             h_RF=np.atleast_1d(h_RF), V_RF=np.atleast_1d(V_RF), dphi_RF=np.atleast_1d(dphi_RF), p0=p0, p_increment=p_increment,
+             charge=charge, mass=mass, RF_at=RF_at, Q_s=Q_s)
+		
+