Merge pull request #39 from PyCOMPLETE/develop

PyHEADTAIL v1.10.0

field_maps/Transverse_Efield_map.py

@@ -4,76 +4,63 @@
 from PyHEADTAIL.particles.slicing import UniformBinSlicer
 from PyPIC.PyPIC_Scatter_Gather import PyPIC_Scatter_Gather
 
+from . import Element
 
+class Transverse_Efield_map(Element):
+    def __init__(self, xg, yg, Ex, Ey, L_interaction, slicer,
+        	flag_clean_slices=False, wrt_slice_centroid=False,
+        	x_beam_offset=0., y_beam_offset=0., *args, **kwargs):
 
-class Transverse_Efield_map(object):
-	def __init__(self, xg, yg, Ex, Ey, n_slices, z_cut,
-		L_interaction, flag_clean_slices = False, wrt_slice_centroid = False,
-		x_beam_offset = 0., y_beam_offset = 0.):
-
-		if type(z_cut) is float:
-			z_cuts = (-z_cut, z_cut)
-		elif type(z_cut) is tuple:
-			z_cuts = z_cut
-		else:
-			raise ValueError('Type not recognized!')
+        self.slicer = slicer
+        self.L_interaction = L_interaction
+        self.flag_clean_slices = flag_clean_slices
+        self.wrt_slice_centroid = wrt_slice_centroid
 
-		self.slicer = UniformBinSlicer(n_slices = n_slices, z_cuts=z_cuts)
-		self.L_interaction = L_interaction
-		self.flag_clean_slices = flag_clean_slices
-		self.wrt_slice_centroid = wrt_slice_centroid
+        self.Ex = Ex
+        self.Ey = Ey
+        self.pic = PyPIC_Scatter_Gather(xg=xg, yg=yg)
 
-		self.Ex = Ex
-		self.Ey = Ey
-		self.pic = PyPIC_Scatter_Gather(xg=xg, yg=yg)
-
-		self.x_beam_offset = x_beam_offset
-		self.y_beam_offset = y_beam_offset
+        self.x_beam_offset = x_beam_offset
+        self.y_beam_offset = y_beam_offset
 
+    def get_beam_x(self, beam):
+        return beam.x
 
-	def get_beam_x(self, beam):
-		return beam.x
+    def get_beam_y(self, beam):
+        return beam.y
 
-	def get_beam_y(self, beam):
-		return beam.y
+    def track(self, beam):
+        if self.flag_clean_slices:
+            beam.clean_slices()
 
-#	@profile
-	def track(self, beam):
+        slices = beam.get_slices(self.slicer)
 
+        for sid in xrange(slices.n_slices-1, -1, -1):
 
-		if self.flag_clean_slices:
-			beam.clean_slices()
-
-		slices = beam.get_slices(self.slicer)
-
-
-		for i in xrange(slices.n_slices-1, -1, -1):
-
-			# select particles in the slice
-			ix = slices.particle_indices_of_slice(i)
-
-			# slice size and time step
-			dz = (slices.z_bins[i + 1] - slices.z_bins[i])
-			#dt = dz / (beam.beta * c)
-
-			x = self.get_beam_x(beam)[ix]
-			y = self.get_beam_y(beam)[ix]
-
-			self.pic.efx = np.squeeze(self.Ex[i,:,:])
-			self.pic.efy = np.squeeze(self.Ey[i,:,:])
-			if self.wrt_slice_centroid:
-				Ex_sc_p, Ey_sc_p = self.pic.gather(x-np.mean(x)+self.x_beam_offset, y-np.mean(y)+self.y_beam_offset)
-			else:
-				Ex_sc_p, Ey_sc_p = self.pic.gather(x+self.x_beam_offset, y+self.y_beam_offset)
-
-			## kick beam particles
-			fact_kick = beam.charge/(beam.mass*beam.beta*beam.beta*beam.gamma*c*c)*self.L_interaction
-			beam.xp[ix]+=fact_kick*Ex_sc_p
-			beam.yp[ix]+=fact_kick*Ey_sc_p
-
+            # select particles in the slice
+            pid = slices.particle_indices_of_slice(sid)
 
+            # slice size
+            dz = (slices.z_bins[sid + 1] - slices.z_bins[sid])
 
+            x = self.get_beam_x(beam)[pid]
+            y = self.get_beam_y(beam)[pid]
 
+            self.pic.efx = np.squeeze(self.Ex[sid,:,:])
+            self.pic.efy = np.squeeze(self.Ey[sid,:,:])
 
+            centroid_x = 0
+            centroid_y = 0
+            if self.wrt_slice_centroid:
+                centroid_x = np.mean(x)
+                centroid_y = np.mean(y)
 
+            Ex_sc_p, Ey_sc_p = self.pic.gather(
+                x - centroid_x + self.x_beam_offset,
+                y - centroid_y + self.y_beam_offset
+            )
 
+            # kick beam particles
+            fact_kick = beam.charge / (beam.p0*beam.beta*c) * self.L_interaction
+            beam.xp[pid] += fact_kick * Ex_sc_p
+            beam.yp[pid] += fact_kick * Ey_sc_p

field_maps/__init__.py

@@ -0,0 +1 @@
+from .. import Element

particles/generators.py

@@ -61,6 +61,11 @@ def transverse_linear_matcher(alpha, beta, dispersion=None):
                     phase space
     Returns: Matcher(closure) taking two parameters: coords and direction
     '''
+#    if dispersion and alpha:
+#        raise NotImplementedError('Transverse phase space matching: for '
+#                                  'alpha != 0 we need to match including the '
+#                                  'D\' (dispersion derivative). This is '
+#                                  'currently not implemented.')
     sqrt = np.sqrt
     # build the M matrix: only depends on twiss parameters for the
     # special case of alpha0=0, beta0=1 and phi = 0 (=2pi)
@@ -86,9 +91,9 @@ def _transverse_linear_matcher(beam, direction):
         momentum_coords = getattr(beam, direction[1])
         space_coords =    (M[0, 0]*space_coords + # copy if not using +=, *=..
                            M[0, 1]*momentum_coords)
-        momentum_coords = (M[1 ,0]*space_coords_copy +
+        momentum_coords = (M[1, 0]*space_coords_copy +
                            M[1, 1]*momentum_coords)
-        # add dispersion effects, raise exception of coords['dp'] inexistent
+        # add dispersion effects, raise exception if coords['dp'] inexistent
         if dispersion:
             try:
                 space_coords += dispersion * getattr(beam, 'dp')

trackers/rf_bucket.py

@@ -40,6 +40,10 @@ class RFBucket(Printing):
     Kick objects, i.e. the left and right side of the bucket are not
     accordingly moved w.r.t. the harmonic phase offset.
     """
+
+    """Sampling points to find zero crossings."""
+    sampling_points = 1000
+
     def __init__(self, circumference, gamma, mass,
                  charge, alpha_array, p_increment,
                  harmonic_list, voltage_list, phi_offset_list,
@@ -49,10 +53,11 @@ def __init__(self, circumference, gamma, mass,
         Arguments:
         - mass is the mass of the particle type in the beam
         - charge is the charge of the particle type in the beam
-        - z_offset determines where to start the root finding
-        (of the electric force field to calibrate the separatrix
-        Hamiltonian value to zero). z_offset is per default given by
-        the phase shift of the fundamental RF element.
+        - z_offset determines the centre for the bucket interval
+        over which the root finding (of the electric force field to
+        calibrate the separatrix Hamiltonian value to zero) is done.
+        z_offset is per default determined by the zero crossing located
+        closest to z == 0.
         '''
 
         self.circumference = circumference
@@ -79,25 +84,40 @@ def __init__(self, circumference, gamma, mass,
         """
         self._add_potentials = []
 
+        zmax = self.circumference / (2*np.amin(self.h))
+
         if z_offset is None:
-            i_fund = np.argmin(self.h) # index of fundamental RF element
-            phi_offset = self.dphi[i_fund]
-            # account for bucket size between -pi and pi.
-            # below transition there should be no relocation of the
-            # bucket interval by an offset of pi! we only need relative
-            # offset w.r.t. normal phi setting at given gamma (0 resp. pi)
-            if self.eta0 < 0:
-                phi_offset -= np.pi
-            z_offset = -phi_offset * self.R / self.h[i_fund]
+            # i_fund = np.argmin(self.h) # index of fundamental RF element
+            # phi_offset = self.dphi[i_fund]
+            # # account for bucket size between -pi and pi.
+            # # below transition there should be no relocation of the
+            # # bucket interval by an offset of pi! we only need relative
+            # # offset w.r.t. normal phi setting at given gamma (0 resp. pi)
+            # if self.eta0 < 0:
+            #     phi_offset -= np.pi
+            # z_offset = -phi_offset * self.R / self.h[i_fund]
+            ### the above approach does not take into account higher harmonics!
+            ### Within a 2 bucket length interval we find all zero crossings
+            ### of the non-accelerated total_force and identify the outermost
+            ### separatrix UFPs via their minimal (convexified) potential value
+            domain_to_find_bucket_centre = np.linspace(-1.999*zmax, 1.999*zmax,
+                                                       self.sampling_points)
+            z0 = self.zero_crossings(self.make_total_force(),
+                                     domain_to_find_bucket_centre)
+            convex_pot0 = (np.array(self.make_total_potential()(z0)) *
+                           np.sign(self.eta0) / self.charge) # charge for numerical reasons
+            outer_separatrix_pot0 = np.min(convex_pot0)
+            outer_separatrix_z0 = z0[np.isclose(convex_pot0,
+                                                outer_separatrix_pot0)]
+            # outer_separatrix_z0 should contain exactly 2 entries
+            z_offset = np.mean(outer_separatrix_z0)
         self.z_offset = z_offset
 
-        zmax = self.circumference / (2*np.amin(harmonic_list))
-
         """Minimum and maximum z values on either side of the
         stationary bucket to cover the maximally possible bucket area,
         defined by the fundamental harmonic.
-        (Does not necessarily coincide with the unstable fix points
-        of the real bucket including self.p_increment .)
+        (This range is always larger than the outmost unstable fix
+        points of the real bucket including self.p_increment .)
         """
         self.interval = (z_offset - 1.01*zmax, z_offset + 1.01*zmax)
 
@@ -141,8 +161,8 @@ def z_sfp(self):
 
     @property
     def z_ufp_separatrix(self):
-        '''Return the (left-most) unstable fix point at the separatrix
-        of the bucket.
+        '''Return the (left-most) unstable fix point at the outermost
+        separatrix of the bucket.
         (i.e. a bucket boundary defining unstable fix point)
         '''
         if self.eta0 * self.p_increment > 0:
@@ -340,11 +360,13 @@ def acc_potential(self, z, ignore_add_potentials=False,
 
     # ROOT AND BOUNDARY FINDING ROUTINES
     # ==================================
-    def zero_crossings(self, f, x=None, subintervals=1000):
+    def zero_crossings(self, f, x=None, subintervals=None):
         '''Determine roots of f along x.
         If x is not explicitely given, take stationary bucket interval.
         '''
         if x is None:
+            if subintervals is None:
+                subintervals = self.sampling_points
             x = np.linspace(*self.interval, num=subintervals)
 
         y = f(x)

trackers/simple_long_tracking.py

@@ -18,8 +18,8 @@
 from types import MethodType
 import weakref
 
-sin = np.sin
-cos = np.cos
+#from ..general import pmath as pm
+import numpy as pm # as soon as pmath is in develop, replace this line with above line.
 
 # @TODO
 # think about flexible design to separate numerical methods
@@ -202,7 +202,7 @@ def track_without_dispersion(self, beam):
                + self.phi_offset + self._phi_lock)
 
         delta_p = beam.dp * beam.p0
-        delta_p += amplitude * sin(phi) - self.p_increment
+        delta_p += amplitude * pm.sin(phi) - self.p_increment
         beam.p0 += self.p_increment
         beam.dp = delta_p / beam.p0
 
@@ -349,10 +349,10 @@ def __init__(self, circumference, harmonic_list, voltage_list,
           map. See the docstring of the Kick class for a more detailed
           description.
         """
-        
+
         self.charge = charge
         self.mass = mass
-        
+
         super(RFSystems, self).__init__(
 			alpha_array, circumference, *args, **kwargs)
 
@@ -494,13 +494,13 @@ def get_bucket(self, bunch=None, gamma=None, mass=None, charge=None,
         (gamma, mass, charge) explicitely to return a bucket
         defined by these.
         '''
-        
+
         if charge is None:
             charge = self.charge
-            
+
         if mass is None:
             mass = self.mass
-        
+
         try:
             bunch_signature = (bunch.gamma, bunch.mass, bunch.charge)
         except AttributeError:
@@ -701,6 +701,7 @@ def __init__(self, alpha_array, circumference, Qs, D_x=0, D_y=0,*args, **kwargs)
         '''
         super(LinearMap, self).__init__(alpha_array, circumference,
                                         *args, **kwargs)
+        assert (np.isscalar(Qs)), "Qs has to be a scalar"
         self.Qs = Qs
         self.D_x = D_x
         self.D_y = D_y
@@ -717,8 +718,8 @@ def track(self, beam):
         omega_s = self.Qs * omega_0
 
         dQs = 2 * np.pi * self.Qs
-        cosdQs = cos(dQs)
-        sindQs = sin(dQs)
+        cosdQs = np.cos(dQs) # use np because dQs is always a scalar
+        sindQs = np.sin(dQs) # use np because dQs is always a scalar
 
         z0 = beam.z
         dp0 = beam.dp

trackers/transverse_tracking.py

@@ -315,15 +315,16 @@ def _generate_segment_maps(self):
             self.segment_maps.append(transverse_segment_map)
 
     def get_injection_optics(self):
-        """ Return a tuple with the transverse TWISS parameters
-        (alpha_x, beta_x, alpha_y, beta_y) from the beginning of the
-        first segment (injection point). """
+        """Return a dict with the transverse TWISS parameters
+        alpha_x, beta_x, D_x, alpha_y, beta_y, D_y from the 
+        beginning of the first segment (injection point).
+        """
         return {
-            'alpha_x': self.alpha_x[0], 
-            'beta_x': self.beta_x[0], 
+            'alpha_x': self.alpha_x[0],
+            'beta_x': self.beta_x[0],
             'D_x': self.D_x[0],
-            'alpha_y': self.alpha_y[0], 
-            'beta_y': self.beta_y[0], 
+            'alpha_y': self.alpha_y[0],
+            'beta_y': self.beta_y[0],
             'D_y': self.D_y[0]}
 
     def __len__(self):

trackers/transverse_tracking_cython.pyx

@@ -447,9 +447,10 @@ class TransverseMap(Printing):
             self.segment_maps.append(transverse_segment_map)
 
     def get_injection_optics(self):
-        """ Return a tuple with the transverse TWISS parameters
-        (alpha_x, beta_x, alpha_y, beta_y) from the beginning of the
-        first segment (injection point). """
+         """Return a dict with the transverse TWISS parameters
+        alpha_x, beta_x, D_x, alpha_y, beta_y, D_y from the 
+        beginning of the first segment (injection point).
+        """
         return {
             'alpha_x': self.alpha_x[0], 
             'beta_x': self.beta_x[0],