Skip to content

Commit 844dfef

Browse files
aoeftigeraoeftiger
committed
Merge pull request #7 from PyCOMPLETE/develop
Develop with fixes of fencepost error in length of optics functions (need n_segments+1!)
2 parents 2da3983 + fc821c6 commit 844dfef

File tree

2 files changed

+134
-98
lines changed

2 files changed

+134
-98
lines changed

‎CERNmachines.py‎

Lines changed: 54 additions & 54 deletions
Original file line numberDiff line numberDiff line change
@@ -39,12 +39,12 @@ def __init__(self, *args, **kwargs):
3939
self.app_y = 0.0000e-9
4040
self.app_xy = 0
4141

42-
self.alpha_x = 0 * np.ones(self.n_segments)
43-
self.beta_x = self.circumference/(2*np.pi*self.Q_x) * np.ones(self.n_segments)
44-
self.D_x = 0 * np.ones(self.n_segments)
45-
self.alpha_y = 0 * np.ones(self.n_segments)
46-
self.beta_y = self.circumference/(2*np.pi*self.Q_y) * np.ones(self.n_segments)
47-
self.D_y = 0 * np.ones(self.n_segments)
42+
self.alpha_x = 0 * np.ones(self.n_segments + 1)
43+
self.beta_x = self.circumference/(2*np.pi*self.Q_x) * np.ones(self.n_segments + 1)
44+
self.D_x = 0 * np.ones(self.n_segments + 1)
45+
self.alpha_y = 0 * np.ones(self.n_segments + 1)
46+
self.beta_y = self.circumference/(2*np.pi*self.Q_y) * np.ones(self.n_segments + 1)
47+
self.D_y = 0 * np.ones(self.n_segments + 1)
4848

4949
self.alpha = 0.06
5050
self.h1 = 1
@@ -73,12 +73,12 @@ def __init__(self, *args, **kwargs):
7373
self.app_y = 0.0000e-9
7474
self.app_xy = 0
7575

76-
self.alpha_x = 0 * np.ones(self.n_segments)
77-
self.beta_x = self.circumference/(2*np.pi*self.Q_x) * np.ones(self.n_segments)
78-
self.D_x = 0 * np.ones(self.n_segments)
79-
self.alpha_y = 0 * np.ones(self.n_segments)
80-
self.beta_y = self.circumference/(2*np.pi*self.Q_y) * np.ones(self.n_segments)
81-
self.D_y = 0 * np.ones(self.n_segments)
76+
self.alpha_x = 0 * np.ones(self.n_segments + 1)
77+
self.beta_x = self.circumference/(2*np.pi*self.Q_x) * np.ones(self.n_segments + 1)
78+
self.D_x = 0 * np.ones(self.n_segments + 1)
79+
self.alpha_y = 0 * np.ones(self.n_segments + 1)
80+
self.beta_y = self.circumference/(2*np.pi*self.Q_y) * np.ones(self.n_segments + 1)
81+
self.D_y = 0 * np.ones(self.n_segments + 1)
8282

8383
self.alpha = 0.06
8484
self.h1 = 1
@@ -107,12 +107,12 @@ def __init__(self, *args, **kwargs):
107107
self.app_y = 0.0000e-9
108108
self.app_xy = 0
109109

110-
self.alpha_x = 0 * np.ones(self.n_segments)
111-
self.beta_x = self.circumference/(2*np.pi*self.Q_x) * np.ones(self.n_segments)
112-
self.D_x = 0 * np.ones(self.n_segments)
113-
self.alpha_y = 0 * np.ones(self.n_segments)
114-
self.beta_y = self.circumference/(2*np.pi*self.Q_y) * np.ones(self.n_segments)
115-
self.D_y = 0 * np.ones(self.n_segments)
110+
self.alpha_x = 0 * np.ones(self.n_segments + 1)
111+
self.beta_x = self.circumference/(2*np.pi*self.Q_x) * np.ones(self.n_segments + 1)
112+
self.D_x = 0 * np.ones(self.n_segments + 1)
113+
self.alpha_y = 0 * np.ones(self.n_segments + 1)
114+
self.beta_y = self.circumference/(2*np.pi*self.Q_y) * np.ones(self.n_segments + 1)
115+
self.D_y = 0 * np.ones(self.n_segments + 1)
116116

117117
self.alpha = 0.06
118118
self.h1 = 1
@@ -153,12 +153,12 @@ def __init__(self, *args, **kwargs):
153153

154154
self.gamma = 27.7
155155

156-
self.alpha_x = 0 * np.ones(self.n_segments)
157-
self.beta_x = 54.6 * np.ones(self.n_segments)
158-
self.D_x = 0 * np.ones(self.n_segments)
159-
self.alpha_y = 0 * np.ones(self.n_segments)
160-
self.beta_y = 54.6 * np.ones(self.n_segments)
161-
self.D_y = 0 * np.ones(self.n_segments)
156+
self.alpha_x = 0 * np.ones(self.n_segments + 1)
157+
self.beta_x = 54.6 * np.ones(self.n_segments + 1)
158+
self.D_x = 0 * np.ones(self.n_segments + 1)
159+
self.alpha_y = 0 * np.ones(self.n_segments + 1)
160+
self.beta_y = 54.6 * np.ones(self.n_segments + 1)
161+
self.D_y = 0 * np.ones(self.n_segments + 1)
162162

163163
self.Q_x = 20.13
164164
self.Q_y = 20.18
@@ -184,12 +184,12 @@ def __init__(self, *args, **kwargs):
184184

185185
self.gamma = 27.7
186186

187-
self.alpha_x = 0 * np.ones(self.n_segments)
188-
self.beta_x = 42. * np.ones(self.n_segments)
189-
self.D_x = 0 * np.ones(self.n_segments)
190-
self.alpha_y = 0 * np.ones(self.n_segments)
191-
self.beta_y = 42. * np.ones(self.n_segments)
192-
self.D_y = 0 * np.ones(self.n_segments)
187+
self.alpha_x = 0 * np.ones(self.n_segments + 1)
188+
self.beta_x = 42. * np.ones(self.n_segments + 1)
189+
self.D_x = 0 * np.ones(self.n_segments + 1)
190+
self.alpha_y = 0 * np.ones(self.n_segments + 1)
191+
self.beta_y = 42. * np.ones(self.n_segments + 1)
192+
self.D_y = 0 * np.ones(self.n_segments + 1)
193193

194194
self.Q_x = 26.13
195195
self.Q_y = 26.18
@@ -215,12 +215,12 @@ def __init__(self, *args, **kwargs):
215215

216216
self.gamma = np.sqrt((450e9*e/(m_p*c**2))**2+1)
217217

218-
self.alpha_x = 0 * np.ones(self.n_segments)
219-
self.beta_x = 54.6 * np.ones(self.n_segments)
220-
self.D_x = 0 * np.ones(self.n_segments)
221-
self.alpha_y = 0 * np.ones(self.n_segments)
222-
self.beta_y = 54.6 * np.ones(self.n_segments)
223-
self.D_y = 0 * np.ones(self.n_segments)
218+
self.alpha_x = 0 * np.ones(self.n_segments + 1)
219+
self.beta_x = 54.6 * np.ones(self.n_segments + 1)
220+
self.D_x = 0 * np.ones(self.n_segments + 1)
221+
self.alpha_y = 0 * np.ones(self.n_segments + 1)
222+
self.beta_y = 54.6 * np.ones(self.n_segments + 1)
223+
self.D_y = 0 * np.ones(self.n_segments + 1)
224224

225225
self.Q_x = 20.13
226226
self.Q_y = 20.18
@@ -272,12 +272,12 @@ def __init__(self, *args, **kwargs):
272272
self.Q_x = 64.28
273273
self.Q_y = 59.31
274274

275-
self.alpha_x = 0 * np.ones(self.n_segments)
276-
self.beta_x = self.R/self.Q_x * np.ones(self.n_segments)
277-
self.D_x = 0 * np.ones(self.n_segments)
278-
self.alpha_y = 0 * np.ones(self.n_segments)
279-
self.beta_y = self.R/self.Q_y * np.ones(self.n_segments)
280-
self.D_y = 0 * np.ones(self.n_segments)
275+
self.alpha_x = 0 * np.ones(self.n_segments + 1)
276+
self.beta_x = self.R/self.Q_x * np.ones(self.n_segments + 1)
277+
self.D_x = 0 * np.ones(self.n_segments + 1)
278+
self.alpha_y = 0 * np.ones(self.n_segments + 1)
279+
self.beta_y = self.R/self.Q_y * np.ones(self.n_segments + 1)
280+
self.D_y = 0 * np.ones(self.n_segments + 1)
281281

282282
self.Qp_x = 0
283283
self.Qp_y = 0
@@ -306,12 +306,12 @@ def __init__(self, *args, **kwargs):
306306
self.Q_x = 64.31
307307
self.Q_y = 59.32
308308

309-
self.alpha_x = 0 * np.ones(self.n_segments)
310-
self.beta_x = self.R/self.Q_x * np.ones(self.n_segments)
311-
self.D_x = 0 * np.ones(self.n_segments)
312-
self.alpha_y = 0 * np.ones(self.n_segments)
313-
self.beta_y = self.R/self.Q_y * np.ones(self.n_segments)
314-
self.D_y = 0 * np.ones(self.n_segments)
309+
self.alpha_x = 0 * np.ones(self.n_segments + 1)
310+
self.beta_x = self.R/self.Q_x * np.ones(self.n_segments + 1)
311+
self.D_x = 0 * np.ones(self.n_segments + 1)
312+
self.alpha_y = 0 * np.ones(self.n_segments + 1)
313+
self.beta_y = self.R/self.Q_y * np.ones(self.n_segments + 1)
314+
self.D_y = 0 * np.ones(self.n_segments + 1)
315315

316316
self.Qp_x = 0
317317
self.Qp_y = 0
@@ -360,12 +360,12 @@ def __init__(self, *args, **kwargs):
360360
self.Q_x = 62.31
361361
self.Q_y = 60.32
362362

363-
self.alpha_x = 0 * np.ones(self.n_segments)
364-
self.beta_x = self.R/self.Q_x * np.ones(self.n_segments)
365-
self.D_x = 0 * np.ones(self.n_segments)
366-
self.alpha_y = 0 * np.ones(self.n_segments)
367-
self.beta_y = self.R/self.Q_y * np.ones(self.n_segments)
368-
self.D_y = 0 * np.ones(self.n_segments)
363+
self.alpha_x = 0 * np.ones(self.n_segments + 1)
364+
self.beta_x = self.R/self.Q_x * np.ones(self.n_segments + 1)
365+
self.D_x = 0 * np.ones(self.n_segments + 1)
366+
self.alpha_y = 0 * np.ones(self.n_segments + 1)
367+
self.beta_y = self.R/self.Q_y * np.ones(self.n_segments + 1)
368+
self.D_y = 0 * np.ones(self.n_segments + 1)
369369

370370
self.Qp_x = 0
371371
self.Qp_y = 0

‎machines.py‎

Lines changed: 80 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -5,26 +5,45 @@
55

66
from PyHEADTAIL.general.element import Element
77
import PyHEADTAIL.particles.generators as gen
8-
from PyHEADTAIL.trackers.transverse_tracking import TransverseMap
9-
from PyHEADTAIL.trackers.detuners import Chromaticity, AmplitudeDetuning
8+
try:
9+
from PyHEADTAIL.trackers.transverse_tracking_cython import TransverseMap
10+
from PyHEADTAIL.trackers.detuners_cython import (Chromaticity,
11+
AmplitudeDetuning)
12+
except ImportError as e:
13+
print ("*** Warning: could not import cython variants of trackers, "
14+
"did you cythonize (use the following command)?\n"
15+
"$ ./install \n"
16+
"Falling back to (slower) python version.")
17+
from PyHEADTAIL.trackers.transverse_tracking import TransverseMap
18+
from PyHEADTAIL.trackers.detuners import Chromaticity, AmplitudeDetuning
1019
from PyHEADTAIL.trackers.simple_long_tracking import LinearMap, RFSystems
1120

1221

1322
class Synchrotron(Element):
1423

1524
def __init__(self, *args, **kwargs):
1625
'''
17-
Currently (because the RFSystems tracking uses a verlet integrator)
18-
the RFSystems element will be installed at s=circumference/2,
19-
which is correct for the smooth approximation
26+
Currently (because the RFSystems tracking uses a Verlet
27+
velocity integrator) the RFSystems element will be installed at
28+
s == circumference/2, which is correct for the smooth
29+
approximation.
2030
'''
31+
self.chromaticity_on = kwargs.pop('chromaticity_on', True)
32+
self.amplitude_detuning_on = kwargs.pop('amplitude_detuning_on', True)
33+
34+
if not hasattr(self, 'longitudinal_focusing'):
35+
self.longitudinal_focusing = kwargs.pop('longitudinal_focusing')
36+
if self.longitudinal_focusing not in ['linear', 'non-linear']:
37+
raise ValueError('longitudinal_focusing not recognized!!!')
38+
2139
for attr in kwargs.keys():
2240
if kwargs[attr] is not None:
2341
self.prints('Synchrotron init. From kwargs: %s = %s'
2442
% (attr, repr(kwargs[attr])))
2543
setattr(self, attr, kwargs[attr])
2644

27-
self.create_transverse_map()
45+
self.create_transverse_map(self.chromaticity_on,
46+
self.amplitude_detuning_on)
2847

2948
# create the one_turn map: install the longitudinal map at
3049
# s = circumference/2
@@ -34,12 +53,11 @@ def __init__(self, *args, **kwargs):
3453

3554
# compute the index of the element before which to insert
3655
# the longitudinal map
37-
if (len(self.one_turn_map) % 2 == 0):
38-
insert_before = len(self.one_turn_map) // 2
39-
else:
40-
insert_before = len(self.one_turn_map) // 2 + 1
56+
for insert_before, si in enumerate(self.s):
57+
if si > 0.5 * self.circumference:
58+
break
4159
n_segments = len(self.transverse_map)
42-
self.create_longitudinal_map(insert_before % n_segments)
60+
self.create_longitudinal_map(insert_before)
4361
self.one_turn_map.insert(insert_before, self.longitudinal_map)
4462

4563
def install_after_each_transverse_segment(self, element_to_add):
@@ -101,26 +119,29 @@ def beta_z(self):
101119
def R(self):
102120
return self.circumference/(2*np.pi)
103121

104-
def track(self, bunch, verbose = False):
122+
def track(self, bunch, verbose=False):
105123
for m in self.one_turn_map:
106124
if verbose:
107125
self.prints('Tracking through:\n' + str(m))
108126
m.track(bunch)
109127

110-
def create_transverse_map(self):
111-
112-
chromaticity = Chromaticity(self.Qp_x, self.Qp_y)
113-
amplitude_detuning = AmplitudeDetuning(self.app_x, self.app_y, self.app_xy)
128+
def create_transverse_map(self, chromaticity_on=True,
129+
amplitude_detuning_on=True):
130+
detuners = []
131+
if chromaticity_on:
132+
detuners.append(Chromaticity(self.Qp_x, self.Qp_y))
133+
if amplitude_detuning_on:
134+
detuners.append(
135+
AmplitudeDetuning(self.app_x, self.app_y, self.app_xy)
136+
)
114137

115138
self.transverse_map = TransverseMap(
116139
self.circumference, self.s,
117140
self.alpha_x, self.beta_x, self.D_x,
118141
self.alpha_y, self.beta_y, self.D_y,
119-
self.Q_x, self.Q_y,
120-
chromaticity, amplitude_detuning)
142+
self.Q_x, self.Q_y, *detuners)
121143

122144
def create_longitudinal_map(self, one_turn_map_insert_idx=0):
123-
124145
if self.longitudinal_focusing == 'linear':
125146
self.longitudinal_map = LinearMap(
126147
[self.alpha],
@@ -137,23 +158,26 @@ def create_longitudinal_map(self, one_turn_map_insert_idx=0):
137158
D_y=self.D_y[one_turn_map_insert_idx]
138159
)
139160
else:
140-
raise ValueError('ERROR: unknown focusing', self.longitudinal_focusing)
141-
142-
def generate_6D_Gaussian_bunch(self, n_macroparticles, intensity, epsn_x, epsn_y, sigma_z):
143-
'''
144-
Generates a 6D Gaussian distribution of particles which is transversely
145-
matched to the Synchrotron. Longitudinally, the distribution is matched
146-
only in terms of linear focusing. For a non-linear bucket, the Gaussian
147-
distribution is cut along the separatrix (with some margin) and will
148-
gradually filament into the bucket. This will change the specified bunch
149-
length.
161+
raise NotImplementedError(
162+
'Something wrong with self.longitudinal_focusing')
163+
164+
def generate_6D_Gaussian_bunch(self, n_macroparticles, intensity,
165+
epsn_x, epsn_y, sigma_z):
166+
'''Generate a 6D Gaussian distribution of particles which is
167+
transversely matched to the Synchrotron. Longitudinally, the
168+
distribution is matched only in terms of linear focusing.
169+
For a non-linear bucket, the Gaussian distribution is cut along
170+
the separatrix (with some margin). It will gradually filament
171+
into the bucket. This will change the specified bunch length.
150172
'''
151173
if self.longitudinal_focusing == 'linear':
152174
check_inside_bucket = lambda z,dp : np.array(len(z)*[True])
153175
elif self.longitudinal_focusing == 'non-linear':
154-
check_inside_bucket = self.longitudinal_map.get_bucket(gamma=self.gamma).make_is_accepted(margin = 0.05)
176+
check_inside_bucket = self.longitudinal_map.get_bucket(
177+
gamma=self.gamma).make_is_accepted(margin=0.05)
155178
else:
156-
raise ValueError('Longitudinal_focusing not recognized!!!')
179+
raise NotImplementedError(
180+
'Something wrong with self.longitudinal_focusing')
157181

158182
beta_z = np.abs(self.eta)*self.circumference/2./np.pi/self.Q_s
159183
sigma_dp = sigma_z/beta_z
@@ -166,25 +190,34 @@ def generate_6D_Gaussian_bunch(self, n_macroparticles, intensity, epsn_x, epsn_y
166190
distribution_x=gen.gaussian2D(epsx_geo),
167191
distribution_y=gen.gaussian2D(epsy_geo),
168192
distribution_z=gen.cut_distribution(
169-
gen.gaussian2D_asymmetrical(sigma_u=sigma_z, sigma_up=sigma_dp),
193+
gen.gaussian2D_asymmetrical(
194+
sigma_u=sigma_z, sigma_up=sigma_dp),
170195
is_accepted=check_inside_bucket),
171196
linear_matcher_x=gen.transverse_linear_matcher(
172-
alpha=self.alpha_x[0], beta=self.beta_x[0], dispersion=self.D_x[0]),
197+
alpha=self.alpha_x[0], beta=self.beta_x[0],
198+
dispersion=self.D_x[0]),
173199
linear_matcher_y=gen.transverse_linear_matcher(
174-
alpha=self.alpha_y[0], beta=self.beta_y[0], dispersion=self.D_y[0])
200+
alpha=self.alpha_y[0], beta=self.beta_y[0],
201+
dispersion=self.D_y[0])
175202
).generate()
176203

177204
return bunch
178205

179-
def generate_6D_Gaussian_bunch_matched(self, n_macroparticles, intensity, epsn_x, epsn_y, sigma_z=None, epsn_z=None):
180-
'''
181-
Generates a 6D Gaussian distribution of particles which is transversely
182-
as well as longitudinally matched. The distribution is found iteratively
183-
to exactly yield the given bunch length while at the same time being
184-
stationary in the non-linear bucket. Thus, the bunch length should amount
185-
to the one specificed and should not change significantly during the
186-
Synchrotron motion.
206+
def generate_6D_Gaussian_bunch_matched(
207+
self, n_macroparticles, intensity, epsn_x, epsn_y,
208+
sigma_z=None, epsn_z=None):
209+
'''Generate a 6D Gaussian distribution of particles which is
210+
transversely as well as longitudinally matched.
211+
The distribution is found iteratively to exactly yield the
212+
given bunch length while at the same time being stationary in
213+
the non-linear bucket. Thus, the bunch length should amount
214+
to the one specificed and should not change significantly
215+
during the synchrotron motion.
216+
217+
Requires self.longitudinal_focusing == 'non-linear'
218+
for the bucket.
187219
'''
220+
assert self.longitudinal_focusing == 'non-linear'
188221
epsx_geo = epsn_x/self.betagamma
189222
epsy_geo = epsn_y/self.betagamma
190223

@@ -197,9 +230,12 @@ def generate_6D_Gaussian_bunch_matched(self, n_macroparticles, intensity, epsn_x
197230
rfbucket=self.longitudinal_map.get_bucket(gamma=self.gamma),
198231
sigma_z=sigma_z, epsn_z=epsn_z),
199232
linear_matcher_x=gen.transverse_linear_matcher(
200-
alpha=self.alpha_x[0], beta=self.beta_x[0], dispersion=self.D_x[0]),
233+
alpha=self.alpha_x[0], beta=self.beta_x[0],
234+
dispersion=self.D_x[0]),
201235
linear_matcher_y=gen.transverse_linear_matcher(
202-
alpha=self.alpha_y[0], beta=self.beta_y[0], dispersion=self.D_y[0])
236+
alpha=self.alpha_y[0], beta=self.beta_y[0],
237+
dispersion=self.D_y[0])
203238
).generate()
204239

205240
return bunch
241+

0 commit comments

Comments
 (0)