Skip to content

Commit aabd655

Browse files
giadarolgiadarol
committed
Merge pull request #28 from PyCOMPLETE/feature/synchrotron
Feature/synchrotron
2 parents 89a9f35 + 17da4cc commit aabd655

File tree

3 files changed

+79
-54
lines changed

3 files changed

+79
-54
lines changed

‎machines/synchrotron.py‎

Lines changed: 55 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -6,18 +6,6 @@
66
from PyHEADTAIL.general.element import Element
77
import PyHEADTAIL.particles.generators as gen
88

9-
try:
10-
from PyHEADTAIL.trackers.transverse_tracking_cython import TransverseMap
11-
from PyHEADTAIL.trackers.detuners_cython import (Chromaticity,
12-
AmplitudeDetuning)
13-
except ImportError as e:
14-
print ("*** Warning: could not import cython variants of trackers, "
15-
"did you cythonize (use the following command)?\n"
16-
"$ ./install \n"
17-
"Falling back to (slower) python version.")
18-
from PyHEADTAIL.trackers.transverse_tracking import TransverseMap
19-
from PyHEADTAIL.trackers.detuners import Chromaticity, AmplitudeDetuning
20-
219
from PyHEADTAIL.trackers.simple_long_tracking import LinearMap, RFSystems
2210

2311
class BasicSynchrotron(Element):
@@ -26,7 +14,8 @@ def __init__(self, optics_mode, circumference=None, n_segments=None, s=None, nam
2614
accQ_x=None, accQ_y=None, Qp_x=0, Qp_y=0, app_x=0, app_y=0, app_xy=0,
2715
alpha_mom_compaction=None, longitudinal_mode=None, Q_s=None,
2816
h_RF=None, V_RF=None, dphi_RF=None, p0=None, p_increment=None,
29-
charge=None, mass=None, **kwargs):
17+
charge=None, mass=None, RF_at='middle', other_detuners=[],
18+
use_cython = True, **kwargs):
3019

3120

3221
self.optics_mode = optics_mode
@@ -36,23 +25,16 @@ def __init__(self, optics_mode, circumference=None, n_segments=None, s=None, nam
3625
self.p0 = p0
3726

3827
self.one_turn_map = []
39-
40-
detuners = []
41-
if Qp_x != 0 or Qp_y != 0:
42-
detuners.append(Chromaticity(Qp_x, Qp_y))
43-
if app_x != 0 or app_y != 0 or app_xy != 0:
44-
detuners.append(AmplitudeDetuning(app_x, app_y, app_xy))
4528

4629
# construct transverse map
4730
self._contruct_transverse_map(optics_mode=optics_mode, circumference=circumference, n_segments=n_segments, s=s, name=name,
4831
alpha_x=alpha_x, beta_x=beta_x, D_x=D_x, alpha_y=alpha_y, beta_y=beta_y, D_y=D_y,
49-
accQ_x=accQ_x, accQ_y=accQ_y, detuners=detuners)
32+
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, other_detuners=other_detuners,
33+
use_cython=use_cython)
5034

5135
# construct longitudinal map
5236
self._contruct_longitudinal_map(alpha_mom_compaction=alpha_mom_compaction, longitudinal_mode=longitudinal_mode, Q_s=Q_s,
53-
h_RF=h_RF, V_RF=V_RF, dphi_RF=dphi_RF, p_increment=p_increment)
54-
55-
37+
h_RF=h_RF, V_RF=V_RF, dphi_RF=dphi_RF, p_increment=p_increment, RF_at=RF_at)
5638

5739
@property
5840
def gamma(self):
@@ -176,8 +158,27 @@ def generate_6D_Gaussian_bunch_matched(
176158

177159
def _contruct_transverse_map(self, optics_mode=None, circumference=None, n_segments=None, s=None, name=None,
178160
alpha_x=None, beta_x=None, D_x=None, alpha_y=None, beta_y=None, D_y=None,
179-
accQ_x=None, accQ_y=None, detuners=[]):
180-
161+
accQ_x=None, accQ_y=None, Qp_x=None, Qp_y=None, app_x=None, app_y=None, app_xy=None, other_detuners=None,
162+
use_cython=None):
163+
164+
if use_cython:
165+
try:
166+
from PyHEADTAIL.trackers.transverse_tracking_cython import TransverseMap
167+
from PyHEADTAIL.trackers.detuners_cython import (Chromaticity,
168+
AmplitudeDetuning)
169+
except ImportError as e:
170+
print ("*** Warning: could not import cython variants of trackers, "
171+
"did you cythonize (use the following command)?\n"
172+
"$ make \n"
173+
"Falling back to (slower) python version.")
174+
from PyHEADTAIL.trackers.transverse_tracking import TransverseMap
175+
from PyHEADTAIL.trackers.detuners import Chromaticity, AmplitudeDetuning
176+
else:
177+
"Transverse tracking: forcing python implementation."
178+
from PyHEADTAIL.trackers.transverse_tracking import TransverseMap
179+
from PyHEADTAIL.trackers.detuners import Chromaticity, AmplitudeDetuning
180+
181+
181182
if optics_mode == 'smooth':
182183
if circumference is None:
183184
raise ValueError('circumference has to be specified if optics_mode = "smooth"')
@@ -211,6 +212,13 @@ def _contruct_transverse_map(self, optics_mode=None, circumference=None, n_segme
211212

212213
else:
213214
raise ValueError('optics_mode not recognized')
215+
216+
detuners = []
217+
if Qp_x != 0 or Qp_y != 0:
218+
detuners.append(Chromaticity(Qp_x, Qp_y))
219+
if app_x != 0 or app_y != 0 or app_xy != 0:
220+
detuners.append(AmplitudeDetuning(app_x, app_y, app_xy))
221+
detuners += other_detuners
214222

215223
self.transverse_map = TransverseMap(s=s,
216224
alpha_x=alpha_x,
@@ -247,15 +255,24 @@ def _contruct_transverse_map(self, optics_mode=None, circumference=None, n_segme
247255
self.one_turn_map.append(m)
248256

249257
def _contruct_longitudinal_map(self, alpha_mom_compaction=None, longitudinal_mode=None, Q_s=None,
250-
h_RF=None, V_RF=None, dphi_RF=None, p_increment=None):
258+
h_RF=None, V_RF=None, dphi_RF=None, p_increment=None, RF_at=None):
251259

252-
# compute the index of the element before which to insert
253-
# the longitudinal map
254-
if longitudinal_mode is not None:
255-
for insert_before, si in enumerate(self.transverse_map.s):
256-
if si > 0.5 * self.circumference:
257-
break
260+
if longitudinal_mode is None:
261+
return
258262

263+
if RF_at == 'middle':
264+
# compute the index of the element before which to insert
265+
# the longitudinal map
266+
if longitudinal_mode is not None:
267+
for insert_before, si in enumerate(self.transverse_map.s):
268+
if si > 0.5 * self.circumference:
269+
break
270+
elif RF_at == 'end_of_transverse':
271+
insert_before = -1
272+
else:
273+
raise ValueError('RF_at=%s not recognized!')
274+
275+
259276
if longitudinal_mode == 'linear':
260277

261278
eta = alpha_mom_compaction - self.gamma**-2
@@ -272,8 +289,7 @@ def _contruct_longitudinal_map(self, alpha_mom_compaction=None, longitudinal_mod
272289
self.circumference, Q_s,
273290
D_x=self.transverse_map.D_x[insert_before],
274291
D_y=self.transverse_map.D_y[insert_before])
275-
276-
292+
277293
elif longitudinal_mode == 'non-linear':
278294
self.longitudinal_map = RFSystems(
279295
self.circumference, np.atleast_1d(h_RF),
@@ -286,5 +302,8 @@ def _contruct_longitudinal_map(self, alpha_mom_compaction=None, longitudinal_mod
286302
else:
287303
raise NotImplementedError(
288304
'Something wrong with longitudinal_mode')
289-
290-
self.one_turn_map.insert(insert_before, self.longitudinal_map)
305+
306+
if insert_before==-1:
307+
self.one_turn_map.append(self.longitudinal_map)
308+
else:
309+
self.one_turn_map.insert(insert_before, self.longitudinal_map)

‎testing/script-tests/LHC.py‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
from PyMachines.synchrotron import BasicSynchrotron
1+
from PyHEADTAIL.machines.synchrotron import BasicSynchrotron
22
import numpy as np
33
from scipy.constants import c, e, m_p
44

@@ -12,6 +12,7 @@ def __init__(self, machine_configuration=None, optics_mode='smooth', **kwargs):
1212
h_RF = 35640
1313
mass = m_p
1414
charge = e
15+
RF_at = 'middle'
1516

1617
if machine_configuration=='Injection':
1718
p0 = 450e9 * e /c
@@ -124,7 +125,7 @@ def __init__(self, machine_configuration=None, optics_mode='smooth', **kwargs):
124125
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,
125126
alpha_mom_compaction=alpha, longitudinal_mode=longitudinal_mode,
126127
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,
127-
charge=charge, mass=mass)
128+
charge=charge, mass=mass, RF_at=RF_at)
128129

129130
def _anharmonicities_from_octupole_current_settings(self, i_octupole_focusing, i_octupole_defocusing):
130131
"""Calculate the constants of proportionality app_x, app_y and

‎testing/script-tests/synchrotron_test_LHC.py‎

Lines changed: 21 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -15,8 +15,8 @@
1515

1616
intensity = 1e11
1717

18-
#mode = 'smooth'
19-
mode = 'non-smooth'
18+
mode = 'smooth'
19+
#mode = 'non-smooth'
2020

2121
import pickle
2222
from LHC import LHC
@@ -62,11 +62,13 @@
6262
fig, axes = plt.subplots(2, sharex=True)
6363

6464
axes[0].plot(np.array(beam_beta_x), 'bo')
65-
axes[0].plot(machine.transverse_map.beta_x, 'b-')
65+
axes[0].plot(machine.transverse_map.beta_x, 'b-', label='x')
6666
axes[0].plot(np.array(beam_beta_y), 'ro')
67-
axes[0].plot(machine.transverse_map.beta_y, 'r-')
67+
axes[0].plot(machine.transverse_map.beta_y, 'r-', label='y')
6868
axes[0].grid('on')
6969
axes[0].set_ylabel('beta_x, beta_y')
70+
axes[0].legend(bbox_to_anchor=(1, 1),loc='upper left',prop={'size':12})
71+
plt.subplots_adjust(right=.86)
7072

7173
axes[1].plot(np.array(beam_alpha_x), 'bo')
7274
axes[1].plot(machine.transverse_map.alpha_x, 'b-')
@@ -134,28 +136,31 @@
134136

135137
fig, axes = plt.subplots(3, figsize=(16, 8), tight_layout=True)
136138
twax = [plt.twinx(ax) for ax in axes]
137-
axes[0].plot(sx)
138-
twax[0].plot(epsx, '-g')
139+
axes[0].plot(sx, label=r'$\sigma_x$' )
140+
twax[0].plot(epsx, '-g', label=r'$\varepsilon_x$')
139141
axes[0].set_xlabel('Turns')
140142
axes[0].set_ylabel(r'$\sigma_x$')
141-
twax[0].set_ylabel(r'$\varepsilon_y$')
142-
axes[1].plot(sy)
143-
twax[1].plot(epsy, '-g')
143+
twax[0].set_ylabel(r'$\varepsilon_x$')
144+
axes[1].plot(sy, label=r'$\sigma_y$' )
145+
twax[1].plot(epsy, '-g', label=r'$\varepsilon_y$')
144146
axes[1].set_xlabel('Turns')
145-
axes[1].set_ylabel(r'$\sigma_x$')
147+
axes[1].set_ylabel(r'$\sigma_y$')
146148
twax[1].set_ylabel(r'$\varepsilon_y$')
147-
axes[2].plot(sz)
148-
twax[2].plot(epsz, '-g')
149+
axes[2].plot(sz, label=r'$\sigma_z$' )
150+
twax[2].plot(epsz, '-g', label=r'$\varepsilon_z$')
149151
axes[2].set_xlabel('Turns')
150-
axes[2].set_ylabel(r'$\sigma_x$')
151-
twax[2].set_ylabel(r'$\varepsilon_y$')
152+
axes[2].set_ylabel(r'$\sigma_z$')
153+
twax[2].set_ylabel(r'$\varepsilon_z$')
152154
axes[0].grid()
153155
axes[1].grid()
154156
axes[2].grid()
155157
for ax in list(axes)+list(twax):
156158
ax.ticklabel_format(useOffset=False, style='sci', scilimits=(0,0),axis='y')
157-
158-
159+
for ax in list(axes):
160+
ax.legend(loc='upper right',prop={'size':12})
161+
for ax in list(twax):
162+
ax.legend(loc='lower right',prop={'size':12})
163+
159164
#~ plt.figure(100)
160165
#~ plt.plot(optics['s'][:],optics['beta_x'][:], '-o')
161166

0 commit comments

Comments
 (0)