diff --git a/field_maps/Transverse_Efield_map.py b/field_maps/Transverse_Efield_map.py index dfdc3d47..24f8d588 100644 --- a/field_maps/Transverse_Efield_map.py +++ b/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 diff --git a/field_maps/__init__.py b/field_maps/__init__.py index e69de29b..3dc80a16 100644 --- a/field_maps/__init__.py +++ b/field_maps/__init__.py @@ -0,0 +1 @@ +from .. import Element diff --git a/particles/generators.py b/particles/generators.py index b73eb73c..c0429a75 100644 --- a/particles/generators.py +++ b/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') diff --git a/trackers/rf_bucket.py b/trackers/rf_bucket.py index d3f62c26..592f2724 100644 --- a/trackers/rf_bucket.py +++ b/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) diff --git a/trackers/simple_long_tracking.py b/trackers/simple_long_tracking.py index 9f57b07d..fa99bd58 100644 --- a/trackers/simple_long_tracking.py +++ b/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 diff --git a/trackers/transverse_tracking.py b/trackers/transverse_tracking.py index 69ee1b36..11e4d716 100644 --- a/trackers/transverse_tracking.py +++ b/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): diff --git a/trackers/transverse_tracking_cython.pyx b/trackers/transverse_tracking_cython.pyx index 8eaa1e3b..83a9db0f 100644 --- a/trackers/transverse_tracking_cython.pyx +++ b/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],