def test_track_LimitRect():
min_x = -0.1
max_x = 0.3
min_y = -0.5
max_y = 0.1
el = pysixtrack.elements.LimitRect(min_x=min_x,
max_x=max_x,
min_y=min_y,
max_y=max_y)
p1 = pysixtrack.Particles()
p1.x = 1
p1.y = 1
ret = el.track(p1)
assert ret == "Particle lost"
arr = np.arange(0, 1, 0.001)
p2 = pysixtrack.Particles(x=arr, y=arr)
survive = np.where((p2.x >= min_x) & (p2.x <= max_x) & (p2.y >= min_y)
& (p2.y <= max_y))
ret = el.track(p2)
assert len(p2.state) == len(survive[0])
p2.x += max_x + 1e-6
ret = el.track(p2)
assert ret == "All particles lost"
def test_track_LimitRectEllipse():
limit_a = 0.1
limit_b = 0.2
max_x = 0.1
max_y = 0.05
el = pysixtrack.elements.LimitRectEllipse(max_x=max_x,
max_y=max_y,
a=limit_a,
b=limit_b)
p1 = pysixtrack.Particles()
p1.x = 1
p1.y = 1
ret = el.track(p1)
assert ret == "Particle lost"
arr = np.arange(0, 1, 0.001)
p2 = pysixtrack.Particles(x=arr, y=arr)
survive = np.where((p2.x**2 / limit_a**2 + p2.y**2 / limit_b**2 <= 1.0)
& (p2.x >= -max_x)
& (p2.x <= max_x)
& (p2.y >= -max_y)
& (p2.y <= max_y))
ret = el.track(p2)
assert len(p2.state) == len(survive[0])
p2.x += limit_a + 1e-6
ret = el.track(p2)
assert ret == "All particles lost"
def remove_dipolar_kicks(self):
temp_lattice = sixtracklib.Elements()
self.elements
sixtracklib.Drift(cbuffer=temp_lattice.cbuffer)
temp_tc_index = temp_lattice.cbuffer.n_objects
temp_tc = sixtracklib.TriCub(cbuffer=temp_lattice.cbuffer)
first_ecloud = list(self.tricubs.keys())[0]
temp_tc.length = self.tricubs[first_ecloud].length
temp_tc.x_shift = 0.
temp_tc.y_shift = 0.
temp_tc.tau_shift = 0.
temp_tc.dipolar_kick_px = 0.
temp_tc.dipolar_kick_py = 0.
temp_tc.dipolar_kick_ptau = 0.
temp_ps = particles_set = sixtracklib.ParticlesSet()
particles = particles_set.Particles(num_particles=1)
temp_part = pysixtrack.Particles(p0c=self.partCO.p0c)
temp_part.x = 0
temp_part.px = 0
temp_part.y = 0
temp_part.py = 0
temp_part.tau = 0
temp_part.ptau = 0
temp_part.state = 1
temp_part.partid = 0
temp_part.elemid = 0
temp_part.turn = 0
particles.from_pysixtrack(temp_part, 0)
temp_job = sixtracklib.TrackJob(temp_lattice, temp_ps, device=None)
temp_tricub_data_buffer_id = temp_job.add_stored_buffer(
cbuffer=self.tricub_data_buffer)
first_tricub_data = list(self.tricub_data.keys())[0]
sixtracklib.TriCub_buffer_create_assign_address_item(
temp_job, temp_tc_index, temp_tricub_data_buffer_id,
self.tricub_data_indices[first_tricub_data])
temp_job.commit_address_assignments()
temp_job.assign_all_addresses()
temp_job.track_until(1)
dipolar_kick_px = particles.px[0]
dipolar_kick_py = particles.py[0]
dipolar_kick_ptau = particles.ptau[0]
print(dipolar_kick_px, dipolar_kick_py, dipolar_kick_ptau)
#dipolar_kick_px = 0.* particles.px[0]
#dipolar_kick_py = 0.*particles.py[0]
#dipolar_kick_ptau = 0.*particles.ptau[0]
for tc in self.tricubs.keys():
tc_index = self.tricub_indices[tc]
tricub = self.job.beam_elements_buffer.get_object(tc_index)
tricub.dipolar_kick_px = dipolar_kick_px
tricub.dipolar_kick_py = dipolar_kick_py
tricub.dipolar_kick_ptau = dipolar_kick_ptau
self.job.push_beam_elements()
return
def mktest(pdata, tol, rtol, elem, **elemdata):
pref = six.Particles(**pdata)
p = sim.Particles(nparticles=1)
for k, v in pdata.items():
setattr(p, k, v)
getattr(six, elem)(**elemdata).track(pref)
elements = sim.Elements()
getattr(elements, elem)(**elemdata)
cljob = sim.TrackJobCL(p, elements, device="0.0", debug=False)
cljob.track(1)
cljob.collect()
tdiff = 0
trdiff = 0
print(elem, elemdata)
for coord in 'x px y py zeta delta rvv rpp beta0'.split():
diff, rdiff = check_diff(coord, p, pref, tol, rtol)
tdiff += diff**2
trdiff += rdiff**2
tdiff = np.sqrt(tdiff)
trdiff = np.sqrt(trdiff)
if tdiff > tol or trdiff > rtol:
print(f"Sum : {tdiff:17.12g} {trdiff:17.12g}")
return False
else:
return True
def trackn(n):
p=pysixtrack.Particles(x=1e-3,px=0.,y=-0.5e-3,py=0.,tau=0.74,pt=0.,
e0=26.01692438e9, m0=0.93827205e9)
for iel,el in enumerate(sps.elems[:n]):
print(iel,el)
el.track(p)
print("%12.9f %12.9f %12.9f %12.9f %12.9f"%(p.s,p.x,p.px,p.tau,p.pt))
if abs(p.x)>1:
break
print((out[n-1][0]))
return p
def test_track_rfmultipole():
p1 = pysixtrack.Particles()
p1.x = 1
p1.y = 1
p2 = p1.copy()
el1 = pysixtrack.elements.RFMultipole(knl=[0.5, 2, 0.2], ksl=[0.5, 3, 0.1])
el2 = pysixtrack.elements.Multipole(knl=el1.knl, ksl=el1.ksl)
el1.track(p1)
el2.track(p2)
assert p1.compare(p2, abs_tol=1e-15)
def track_turn(n):
p=pysixtrack.Particles(x=1e-3,px=np.zeros(5000),
y=-0.5e-3,py=np.zeros(5000),
tau=0.74,pt=np.zeros(5000),
e0=26.01692438e9, m0=0.93827205e9)
out=[]
before=time.time()
for i in range(n):
out.append(p.copy())
sps.track(p)
now=time.time()
print((i,now-before))
before=now
return out
def generate_particle_data_sequ_by_sequ(output_path,
line,
iconv,
sixdump,
conf=dict()):
num_iconv = int(len(iconv))
num_belem = int(len(line))
num_dumps = int(len(sixdump.particles))
assert num_iconv > 0
assert num_belem > iconv[num_iconv - 1]
assert num_dumps >= num_iconv
assert (num_dumps % num_iconv) == 0
num_particles = num_dumps // num_iconv
assert num_particles > 0
assert num_belem > 0
assert num_iconv > 0
NORM_ADDR = conf.get("cbuffer_norm_base_addr", 4096)
pset_buffer = create_particle_set_cbuffer(num_iconv, num_particles, conf)
for ii in range(num_iconv):
assert iconv[ii] < num_belem
assert ii < pset_buffer.num_objects
pset = st.st_Particles.GET(pset_buffer, ii)
assert pset.num_particles == num_particles
for jj in range(0, num_particles):
kk = num_particles * ii + jj
assert kk < num_dumps
in_p = pysix.Particles(**sixdump[kk].get_minimal_beam())
in_p.state = 1
in_p.turn = 0
in_p.partid = jj
in_p.elemid = iconv[ii]
assert pset.num_particles == num_particles
pysix_particle_to_pset(in_p, pset, jj, conf=conf)
path_cobj_pset = os.path.join(output_path, "cobj_particles_sixtrack.bin")
if 0 == pset_buffer.tofile_normalised(path_cobj_pset, NORM_ADDR):
print("**** -> Generated cbuffer of sixtrack particle sequ-by-sequ "
"data:\r\n" + f"**** {path_cobj_pset}")
else:
raise RuntimeError(
"Unable to generate cobjects sixtrack sequency-by-sequence data")
return
def get_init_particles_for_linear_map(part, d):
new_part = pysixtrack.Particles()
new_part.p0c = part.p0c
new_part.x = part.x
new_part.px = part.px
new_part.y = part.y
new_part.py = part.py
new_part.zeta = part.zeta
new_part.delta = part.delta
new_part.x += np.array([0., 1. * d, 0., 0., 0., 0., 0.])
new_part.px += np.array([0., 0., 1. * d, 0., 0., 0., 0.])
new_part.y += np.array([0., 0., 0., 1. * d, 0., 0., 0.])
new_part.py += np.array([0., 0., 0., 0., 1. * d, 0., 0.])
new_part.zeta += np.array([0., 0., 0., 0., 0., 1. * d, 0.])
new_part.delta += np.array([0., 0., 0., 0., 0., 0., 1. * d])
return new_part
def fromSixDump101(cls, input_folder, st_dump_file, **kwargs):
import sixtracktools
import pysixtrack
six = sixtracktools.SixInput(input_folder)
line, rest, iconv = six.expand_struct(convert=pysixtrack.element_types)
sixdump = sixtracktools.SixDump101(st_dump_file)
num_iconv = int(len(iconv))
num_belem = int(len(line))
num_dumps = int(len(sixdump.particles))
assert num_iconv > 0
assert num_belem > iconv[num_iconv - 1]
assert num_dumps >= num_iconv
assert (num_dumps % num_iconv) == 0
num_particles = int(num_dumps / num_iconv)
self = cls(**kwargs)
for ii in range(num_iconv):
elem_id = iconv[ii]
assert elem_id < num_belem
p = self.Particles(num_particles=num_particles)
assert p.num_particles == num_particles
assert len(p.q0) == num_particles
for jj in range(num_particles):
kk = num_particles * ii + jj
assert kk < num_dumps
p.from_pysixtrack(
pysixtrack.Particles(**sixdump[kk].get_minimal_beam()), jj)
p.state[jj] = 1
p.at_element[jj] = elem_id
return self
def linearize_around_closed_orbit(part, line, d):
init_part = get_init_particles_for_linear_map(part, d)
fin_part = init_part.copy()
line.track(fin_part)
X_init = np.empty([6, 7])
X_fin = np.empty([6, 7])
X_init[0, :] = init_part.x
X_init[1, :] = init_part.px
X_init[2, :] = init_part.y
X_init[3, :] = init_part.py
X_init[4, :] = init_part.zeta
X_init[5, :] = init_part.delta
X_fin[0, :] = fin_part.x
X_fin[1, :] = fin_part.px
X_fin[2, :] = fin_part.y
X_fin[3, :] = fin_part.py
X_fin[4, :] = fin_part.zeta
X_fin[5, :] = fin_part.delta
m = X_fin[:, 0] - X_init[:, 0]
M = np.empty([6, 6])
for j in range(6):
M[:, j] = (X_fin[:, j + 1] - X_fin[:, 0]) / d
X_CO = X_init[:, 0] + np.matmul(np.linalg.inv(np.identity(6) - M), m.T)
part_CO = pysixtrack.Particles(p0c=part.p0c)
part_CO.x = X_CO[0]
part_CO.px = X_CO[1]
part_CO.y = X_CO[2]
part_CO.py = X_CO[3]
part_CO.zeta = X_CO[4]
part_CO.delta = X_CO[5]
return M, m, part_CO
def get_sixtracklib_particle_set(init_physical_coordinates: np.ndarray,
p0c_ev: float):
n_particles = init_physical_coordinates.shape[0]
ps = sixtracklib.ParticlesSet()
p = ps.Particles(num_particles=n_particles)
for i_part in range(n_particles):
part = pysixtrack.Particles(p0c=p0c_ev)
part.x = init_physical_coordinates[i_part, 0]
part.px = init_physical_coordinates[i_part, 1]
part.y = init_physical_coordinates[i_part, 2]
part.py = init_physical_coordinates[i_part, 3]
part.tau = init_physical_coordinates[i_part, 4]
part.ptau = init_physical_coordinates[i_part, 5]
part.partid = i_part
part.state = 1
part.elemid = 0
part.turn = 0
p.from_pysixtrack(part, i_part)
return ps
nslices = 4
mad.select(flag="makethin", clear=True)
mad.select(flag="makethin", class_="sbend", slice=nslices)
mad.select(flag="makethin", class_="quadrupole", slice=nslices)
mad.makethin(sequence="ring")
mad.use(sequence="ring")
print(mad.table.summ.q1, mad.table.summ.q2)
twiss = mad.twiss()
print(mad.table.summ.q1, mad.table.summ.q2)
twissout = pysixtrack.Particles.from_madx_twiss(
mad.twiss(betx=1, bety=1, x=0.001))
line = pysixtrack.Line.from_madx_sequence(mad.sequence.ring)
part = pysixtrack.Particles()
part.x = 0.001
pysixout = pysixtrack.Particles.from_list(
line.track_elem_by_elem(part, start=False, end=True))
def mkd(name, t1, t2, ii, jj):
v1 = getattr(t1, name)[ii]
v2 = getattr(t2, name)[jj]
print(f"{name:4} {v2:20} {v2:20} {v2-v1:20}")
return v2 - v1
for ii in range(len(twissout.s)):
sm = twissout.s[ii]
sp = pysixout.s[ii]
import pysixtrack as pyst
for el in pyst.element_list:
p = pyst.Particles()
el().track(p)
def compare(prun,pbench):
out=[]
for att in 'x px y py delta sigma'.split():
vrun=getattr(prun,att)
vbench=getattr(pbench,att)
diff=vrun-vbench
out.append(abs(diff))
print(f"{att:<5} {vrun:22.13e} {vbench:22.13e} {diff:22.13g}")
print(f"max {max(out):21.12e}")
return max(out)
print("")
for ii in range(1,len(iconv)):
jja=iconv[ii-1]
jjb=iconv[ii]
prun=pysixtrack.Particles(**sixdump[ii-1].get_minimal_beam())
print(f"\n-----sixtrack={ii} sixtracklib={jja} --------------")
#print(f"pysixtr {jja}, x={prun.x}, px={prun.px}")
for jj in range(jja+1, jjb+1):
label,elem=line.element_names[jj], line.elements[jj]
elem.track(prun)
print(f"{jj} {label},{str(elem)[:50]}")
pbench=pysixtrack.Particles(**sixdump[ii].get_minimal_beam())
#print(f"sixdump {ii}, x={pbench.x}, px={pbench.px}")
print("-----------------------")
out=compare(prun,pbench)
print("-----------------------\n\n")
if out>1e-13:
print("Too large discrepancy")
break
def test_neutral_errors():
# make sure that some misaligned drifts do not influence particle
cpymad_spec = util.find_spec("cpymad")
if cpymad_spec is None:
print("cpymad is not available - abort test")
sys.exit(0)
from cpymad.madx import Madx
madx = Madx()
madx.input('''
T1: Collimator, L=1.0, apertype=CIRCLE, aperture={0.5};
T2: Collimator, L=1.0, apertype=CIRCLE, aperture={0.5};
T3: Collimator, L=1.0, apertype=CIRCLE, aperture={0.5};
KQ1 = 0.02;
KQ2 = -0.02;
testseq: SEQUENCE, l = 20.0;
T1, at = 5;
T2, at = 12;
T3, at = 18;
ENDSEQUENCE;
!---the usual stuff
BEAM, PARTICLE=PROTON, ENERGY=7000.0, EXN=2.2e-6, EYN=2.2e-6;
USE, SEQUENCE=testseq;
Select, flag=makethin, pattern="T1", slice=2;
makethin, sequence=testseq;
use, sequence=testseq;
!---misalign collimators
select, flag = error, clear;
select, flag = error, pattern = "T1";
ealign, dx = 0.01, dy = 0.01, arex = 0.02, arey = 0.02;
select, flag = error, clear;
select, flag = error, pattern = "T2";
ealign, dx = 0.04, dy = 0.04, dpsi = 0.1;
select, flag = error, clear;
select, flag = error, pattern = "T3";
ealign, dx = 0.02, dy = 0.01, arex = 0.03, arey = 0.02, dpsi = 0.1;
select, flag = error, full;
''')
seq = madx.sequence.testseq
pysixtrack_line = pysixtrack.Line.from_madx_sequence(
seq,
install_apertures=True,
apply_madx_errors=True,
)
madx.input('stop;')
initial_x = 0.025
initial_y = -0.015
particle = pysixtrack.Particles()
particle.x = initial_x
particle.y = initial_y
particle.state = 1
pysixtrack_line.track(particle)
assert abs(particle.x-initial_x) < 1e-14
assert abs(particle.y-initial_y) < 1e-14
def test_track_all():
for el in element_list:
p = pysixtrack.Particles()
el().track(p)
N_theta = int(round((theta_max_rad - theta_min_rad) / theta_step)) + 1
fLine = 'line_from_mad_with_bbCO.pkl'
fParticleCO = 'particle_on_CO_mad_line.pkl'
fOptics = 'optics_mad.pkl'
with open(fLine, 'rb') as fid:
line = pysixtrack.Line.from_dict(pickle.load(fid))
with open(fParticleCO, 'rb') as fid:
partCO = pickle.load(fid)
with open(fOptics, 'rb') as fid:
optics = pickle.load(fid)
part = pysixtrack.Particles(**partCO)
beta_x = optics['betx']
beta_y = optics['bety']
sigmax = np.sqrt(beta_x * epsn_x / part.beta0 / part.gamma0)
sigmay = np.sqrt(beta_y * epsn_y / part.beta0 / part.gamma0)
xy_norm = footprint.initial_xy_polar(r_min=r_min_sigma,
r_max=r_max_sigma,
r_N=N_r,
theta_min=theta_min_rad,
theta_max=theta_max_rad,
theta_N=N_theta)
DpxDpy_wrt_CO = np.zeros_like(xy_norm)
def test_track_spacecharge():
x_co = 0.1
y_co = -0.5
sigma_x = 0.5
sigma_y = 0.1
el1 = pysixtrack.elements.SCQGaussProfile(
number_of_particles=1e11,
bunchlength_rms=0.22,
sigma_x=sigma_x,
sigma_y=sigma_y,
length=2.0,
x_co=x_co,
y_co=y_co,
)
el2 = pysixtrack.elements.SCCoasting(
number_of_particles=el1.number_of_particles,
circumference=el1.bunchlength_rms * np.sqrt(2 * np.pi),
sigma_x=el1.sigma_x,
sigma_y=el1.sigma_y,
length=el1.length,
x_co=el1.x_co,
y_co=el1.y_co,
)
# test absolute kick for sigma_x > sigma_y
x_offset = 0.2
y_offset = -0.5
p1 = pysixtrack.Particles()
p1.x = x_co + x_offset
p1.y = y_co + y_offset
p2 = p1.copy()
el1.track(p1)
el2.track(p2)
assert np.isclose(p1.px, 1.8329795395186613e-07, atol=1e-15)
assert np.isclose(p1.py, -8.540420459001383e-07, atol=1e-15)
assert p1.compare(p2, abs_tol=1e-15)
# test absolute kick for sigma_y > sigma_x
el1.sigma_x = sigma_y
el1.sigma_y = sigma_x
el2.sigma_x = sigma_y
el2.sigma_y = sigma_x
p1 = pysixtrack.Particles()
p1.x = x_co + y_offset
p1.y = y_co + x_offset
p2 = p1.copy()
el1.track(p1)
el2.track(p2)
assert np.isclose(p1.px, -8.540420459001383e-07, atol=1e-15)
assert np.isclose(p1.py, 1.8329795395186613e-07, atol=1e-15)
assert p1.compare(p2, abs_tol=1e-15)
# test no kick for particle on closed orbit
p1 = pysixtrack.Particles()
p1.x = el1.x_co
p1.y = el1.y_co
p2 = p1.copy()
el1.track(p1)
el2.track(p2)
assert np.isclose(p1.px, 0.0, atol=1e-15)
assert np.isclose(p1.py, 0.0, atol=1e-15)
assert p1.compare(p2, abs_tol=1e-15)
# test round beam
p1 = pysixtrack.Particles()
p1.x = el1.x_co + 0.5
p1.y = el1.y_co + 0.1
p2 = p1.copy()
el1.sigma_y = el1.sigma_x
el2.sigma_y = el2.sigma_x
el1.track(p1)
el2.track(p2)
assert np.isclose(p1.px, 1.2895332740238447e-06, atol=1e-15)
assert np.isclose(p1.py, 2.579066548047689e-07, atol=1e-15)
assert p1.compare(p2, abs_tol=1e-15)
#elements=sim.Elements.fromline(line)
elements = sim.Elements()
lhc = sixtracktools.SixInput('.')
line, rest, iconv = lhc.expand_struct(convert=elements.gen_builder())
cljob = sim.TrackJobCL(particles, elements, device="0.0", dump_element=nturns)
cljob.track()
cljob.collect()
out = cljob.dump_element
refline, rest, iconv = lhc.expand_struct(convert=pysixtrack.element_types)
refline = pysixtrack.Line(elements=[l[2] for l in refline])
refline.elements.append(pysixtrack.Monitor())
prun = pysixtrack.Particles(p0c=7000e9, x=x0)
refout = refline.track_elem_by_elem(prun)
refx = np.array([p.x for p in refout])
refzeta = np.array([p.zeta for p in refout])
plot(out.x - refx)
plot(out.zeta - refzeta)
#turn-by-turn
nturn = 128
elements.Monitor(turns=nturn)
particles = sim.Particles(nparticles=npart)
particles.p0c = 7000e9
particles.x = x0
def test_madx_import():
cpymad_spec = util.find_spec("cpymad")
if cpymad_spec is None:
print("cpymad is not available - abort test")
sys.exit(0)
from cpymad.madx import Madx
seq_name = "psb1"
use_aperture = True
n_SCkicks = 120
length_fuzzy = 0.0
p0c = 0.571e6
particle = pysixtrack.Particles(p0c=p0c)
betagamma = particle.beta0 * particle.gamma0
# mass = pysixtrack.Particles.pmass
delta_rms = 1e-3
neps_x = 1.5e-6
neps_y = 1.5e-6
# for space charge bunched
number_of_particles = 1e11
bunchlength_rms = 1.0
# for space charge coasting
line_density = 1e11
for sc_mode in ["Bunched", "Coasting"]:
mad = Madx()
mad.options.echo = False
mad.options.info = False
mad.warn = False
file_path = os.path.realpath(__file__)
path = os.path.dirname(file_path) + "/psb/"
mad.call(path + "psb_fb_lhc.madx", chdir=True)
# Determine space charge locations
temp_line = pysixtrack.Line.from_madx_sequence(mad.sequence[seq_name])
sc_locations, sc_lengths = bt.determine_sc_locations(
temp_line, n_SCkicks, length_fuzzy
)
# Install spacecharge place holders
sc_names = ["sc%d" % number for number in range(len(sc_locations))]
bt.install_sc_placeholders(
mad, seq_name, sc_names, sc_locations, mode=sc_mode
)
# Generate line with spacecharge
line = pysixtrack.Line.from_madx_sequence(
mad.sequence[seq_name], install_apertures=use_aperture
)
# Get sc info from optics
mad_sc_names, sc_twdata = bt.get_spacecharge_names_twdata(
mad, seq_name, mode=sc_mode
)
# Check consistency
if sc_mode == "Bunched":
sc_elements, sc_names = line.get_elements_of_type(
pysixtrack.elements.SpaceChargeBunched
)
elif sc_mode == "Coasting":
sc_elements, sc_names = line.get_elements_of_type(
pysixtrack.elements.SpaceChargeCoasting
)
else:
raise ValueError("mode not understood")
bt.check_spacecharge_consistency(
sc_elements, sc_names, sc_lengths, mad_sc_names
)
# Setup spacecharge in the line
if sc_mode == "Bunched":
bt.setup_spacecharge_bunched_in_line(
sc_elements,
sc_lengths,
sc_twdata,
betagamma,
number_of_particles,
bunchlength_rms,
delta_rms,
neps_x,
neps_y,
)
elif sc_mode == "Coasting":
bt.setup_spacecharge_coasting_in_line(
sc_elements,
sc_lengths,
sc_twdata,
betagamma,
line_density,
delta_rms,
neps_x,
neps_y,
)
else:
raise ValueError("mode not understood")
def test_error_functionality():
# check if errors are actually working as intended
cpymad_spec = util.find_spec("cpymad")
if cpymad_spec is None:
print("cpymad is not available - abort test")
sys.exit(0)
from cpymad.madx import Madx
import numpy as np
madx = Madx()
madx.input('''
T1: Collimator, L=0.0, apertype=CIRCLE, aperture={0.5};
T2: Marker;
T3: Collimator, L=0.0, apertype=CIRCLE, aperture={0.5};
testseq: SEQUENCE, l = 20.0;
T1, at = 5;
T2, at = 10;
T3, at = 15;
ENDSEQUENCE;
!---the usual stuff
BEAM, PARTICLE=PROTON, ENERGY=7000.0, EXN=2.2e-6, EYN=2.2e-6;
USE, SEQUENCE=testseq;
!---assign misalignments and field errors
select, flag = error, clear;
select, flag = error, pattern = "T1";
ealign, dx = 0.01, dy = 0.02, arex = 0.03, arey = 0.04;
select, flag = error, clear;
select, flag = error, pattern = "T3";
ealign, dx = 0.07, dy = 0.08, dpsi = 0.7, arex = 0.08, arey = 0.09;
select, flag = error, full;
''')
seq = madx.sequence.testseq
pysixtrack_line = pysixtrack.Line.from_madx_sequence(
seq,
install_apertures=True,
apply_madx_errors=True,
)
madx.input('stop;')
x_init = 0.1*np.random.rand(10)
y_init = 0.1*np.random.rand(10)
particles = pysixtrack.Particles(
x=x_init.copy(),
y=y_init.copy()
)
T1_checked = False
T1_aper_checked = False
T2_checked = False
T3_checked = False
T3_aper_checked = False
for element, element_name in zip(pysixtrack_line.elements,
pysixtrack_line.element_names):
ret = element.track(particles)
if element_name == 't1':
T1_checked = True
assert np.all(abs(particles.x - (x_init - 0.01)) < 1e-14)
assert np.all(abs(particles.y - (y_init - 0.02)) < 1e-14)
if element_name == 't1_aperture':
T1_aper_checked = True
assert np.all(abs(particles.x - (x_init - 0.01 - 0.03)) < 1e-14)
assert np.all(abs(particles.y - (y_init - 0.02 - 0.04)) < 1e-14)
if element_name == 't2':
T2_checked = True
assert np.all(abs(particles.x - x_init) < 1e-14)
assert np.all(abs(particles.y - y_init) < 1e-14)
cospsi = np.cos(0.7)
sinpsi = np.sin(0.7)
if element_name == 't3':
T3_checked = True
assert np.all(abs(
particles.x
- (x_init - 0.07)*cospsi
- (y_init - 0.08)*sinpsi
) < 1e-14)
assert np.all(abs(
particles.y
+ (x_init - 0.07)*sinpsi
- (y_init - 0.08)*cospsi
) < 1e-14)
if element_name == 't3_aperture':
T3_aper_checked = True
assert np.all(abs(
particles.x
- (x_init - 0.07)*cospsi
- (y_init - 0.08)*sinpsi
- (-0.08)
) < 1e-14)
assert np.all(abs(
particles.y
+ (x_init - 0.07)*sinpsi
- (y_init - 0.08)*cospsi
- (-0.09)
) < 1e-14)
if ret is not None:
break
assert not ret
assert np.all([T1_checked, T1_aper_checked,
T2_checked, T3_checked, T3_aper_checked])
# Load iconv
with open('iconv.pkl', 'rb') as fid:
iconv = pickle.load(fid)
# Load sixtrack tracking data
sixdump_all = sixtracktools.SixDump101('sixtrack/res/dump3.dat')
# Assume first particle to be on the closed orbit
Nele_st = len(iconv)
sixdump_CO = sixdump_all[::2][:Nele_st]
# Compute closed orbit using tracking
closed_orbit = line.track_elem_by_elem(part_on_CO)
# Check that closed orbit is closed
pstart = closed_orbit[0].copy()
pstart_st = pysixtrack.Particles(**sixdump_CO[0].get_minimal_beam())
print('STsigma, Sigma, Stdelta, delta, Stpx, px')
for iturn in range(10):
line.track(pstart)
line.track(pstart_st)
print('%e, %e, %e, %e, %e, %e' %
(pstart_st.sigma, pstart.sigma, pstart_st.delta, pstart.delta,
pstart_st.px, pstart.px))
# Compare closed orbit against sixtrack
for att in 'x px y py delta sigma'.split():
att_CO = np.array([getattr(pp, att) for pp in closed_orbit])
att_CO_at_st_ele = att_CO[iconv]
print('Max C.O. discrepancy in %s %.2e' %
(att, np.max(np.abs(att_CO_at_st_ele - getattr(sixdump_CO, att)))))
def set_optics_CO(self, optics, partCO):
self.optics = optics
self.partCO = pysixtrack.Particles(**partCO)
return
n_turns = 100
with open('line.pkl', 'rb') as fid:
line = pysixtrack.Line.from_dict(pickle.load(fid), keepextra=True)
with open('particle_on_CO.pkl', 'rb') as fid:
partCO = pysixtrack.Particles.from_dict(pickle.load(fid))
with open('DpxDpy_for_footprint.pkl', 'rb') as fid:
temp_data = pickle.load(fid)
xy_norm = temp_data['xy_norm']
DpxDpy_wrt_CO = temp_data['DpxDpy_wrt_CO']
part = partCO.copy() # pysixtrack.Particles(**partCO)
part._m = pysixtrack.Particles()._m # to be sorted out later
part.sigma += 0.05
if track_with == 'PySixtrack':
x_tbt, px_tbt, y_tbt, py_tbt, sigma_tbt, delta_tbt = hp.track_particle_pysixtrack(
line,
part=part,
Dx_wrt_CO_m=0.,
Dpx_wrt_CO_rad=DpxDpy_wrt_CO[:, :, 0].flatten(),
Dy_wrt_CO_m=0,
Dpy_wrt_CO_rad=DpxDpy_wrt_CO[:, :, 1].flatten(),
Dsigma_wrt_CO_m=0.,
Ddelta_wrt_CO=0.,
n_turns=n_turns,
verbose=True)
def track_to_checkpoint(job,
n_particles,
checkpoint=1,
checkpoint_turns=1e6,
monitor1_stores=100,
monitor2_stores=1,
skip_turns=10000):
start_time = time.time()
turn_to_track = checkpoint * checkpoint_turns
job.track_until(turn_to_track)
end_time = time.time()
print(f'Tracking time (1M): {(end_time - start_time)/60.:.4f}mins')
start_time = time.time()
job.collect()
monitor1 = job.output.particles[0]
monitor2 = job.output.particles[1]
p0c = monitor2.p0c[0]
x_skip = monitor1.x.reshape(monitor1_stores, n_particles)
px_skip = monitor1.px.reshape(monitor1_stores, n_particles)
y_skip = monitor1.y.reshape(monitor1_stores, n_particles)
py_skip = monitor1.py.reshape(monitor1_stores, n_particles)
zeta_skip = monitor1.zeta.reshape(monitor1_stores, n_particles)
delta_skip = monitor1.delta.reshape(monitor1_stores, n_particles)
at_turn_skip = monitor1.at_turn.reshape(monitor1_stores, n_particles)
skip_dicts = []
for i in range(monitor1_stores):
zeta = zeta_skip[i]
delta = delta_skip[i]
temp_part = pysixtrack.Particles(p0c=p0c)
temp_part.zeta = zeta
temp_part.delta = delta
tau = temp_part.tau
ptau = temp_part.ptau
x = x_skip[i]
px = px_skip[i]
y = y_skip[i]
py = py_skip[i]
supposed_turn = turn_to_track - checkpoint_turns + (i + 1) * skip_turns
at_turn = at_turn_skip[i] + 1
# print(at_turn, supposed_turn)
mask = at_turn == supposed_turn
not_mask = np.logical_not(mask)
x[not_mask] = 0.
px[not_mask] = 0.
y[not_mask] = 0.
py[not_mask] = 0.
tau[not_mask] = 0.
ptau[not_mask] = 0.
at_turn[not_mask] = 0.
skip_dict = {
'x': x,
'px': px,
'y': y,
'py': py,
'tau': tau,
'ptau': ptau,
'at_turn': at_turn
}
skip_dicts.append((supposed_turn, skip_dict))
x_last = monitor2.x.reshape(monitor2_stores, n_particles)
px_last = monitor2.px.reshape(monitor2_stores, n_particles)
y_last = monitor2.y.reshape(monitor2_stores, n_particles)
py_last = monitor2.py.reshape(monitor2_stores, n_particles)
zeta_last = monitor2.zeta.reshape(monitor2_stores, n_particles)
delta_last = monitor2.delta.reshape(monitor2_stores, n_particles)
at_turn_last = monitor2.at_turn.reshape(monitor2_stores, n_particles)
state_last = monitor2.state.reshape(monitor2_stores, n_particles)
temp_part = pysixtrack.Particles(p0c=p0c)
temp_part.zeta = zeta_last
temp_part.delta = delta_last
tau_last = temp_part.tau
ptau_last = temp_part.ptau
last_dict = {
'x': x_last.flatten(),
'px': px_last.flatten(),
'y': y_last.flatten(),
'py': py_last.flatten(),
'tau': tau_last.flatten(),
'ptau': ptau_last.flatten(),
'at_turn': at_turn_last.flatten(),
'state': state_last.flatten()
}
end_time = time.time()
print(f'Processing time: {(end_time - start_time)/60.:.4f}mins')
return skip_dicts, last_dict
def track_particle_sixtrack(
partCO, Dx_wrt_CO_m, Dpx_wrt_CO_rad,
Dy_wrt_CO_m, Dpy_wrt_CO_rad,
Dsigma_wrt_CO_m, Ddelta_wrt_CO, n_turns
):
Dx_wrt_CO_m, Dpx_wrt_CO_rad,\
Dy_wrt_CO_m, Dpy_wrt_CO_rad,\
Dsigma_wrt_CO_m, Ddelta_wrt_CO = vectorize_all_coords(
Dx_wrt_CO_m, Dpx_wrt_CO_rad,
Dy_wrt_CO_m, Dpy_wrt_CO_rad,
Dsigma_wrt_CO_m, Ddelta_wrt_CO)
n_part = len(Dx_wrt_CO_m)
wfold = 'temp_trackfun'
if not os.path.exists(wfold):
os.mkdir(wfold)
os.system('cp fort.* %s' % wfold)
with open('fort.3', 'r') as fid:
lines_f3 = fid.readlines()
# Set initial coordinates
i_start_ini = None
for ii, ll in enumerate(lines_f3):
if ll.startswith('INITIAL COO'):
i_start_ini = ii
break
lines_f3[i_start_ini + 2] = ' 0.\n'
lines_f3[i_start_ini + 3] = ' 0.\n'
lines_f3[i_start_ini + 4] = ' 0.\n'
lines_f3[i_start_ini + 5] = ' 0.\n'
lines_f3[i_start_ini + 6] = ' 0.\n'
lines_f3[i_start_ini + 7] = ' 0.\n'
lines_f3[i_start_ini + 2 + 6] = ' 0.\n'
lines_f3[i_start_ini + 3 + 6] = ' 0.\n'
lines_f3[i_start_ini + 4 + 6] = ' 0.\n'
lines_f3[i_start_ini + 5 + 6] = ' 0.\n'
lines_f3[i_start_ini + 6 + 6] = ' 0.\n'
lines_f3[i_start_ini + 7 + 6] = ' 0.\n'
lines_f13 = []
for i_part in range(n_part):
temp_part = pysixtrack.Particles(**partCO)
temp_part.x += Dx_wrt_CO_m[i_part]
temp_part.px += Dpx_wrt_CO_rad[i_part]
temp_part.y += Dy_wrt_CO_m[i_part]
temp_part.py += Dpy_wrt_CO_rad[i_part]
temp_part.sigma += Dsigma_wrt_CO_m[i_part]
temp_part.delta += Ddelta_wrt_CO[i_part]
lines_f13.append('%.10e\n' % ((temp_part.x) * 1e3))
lines_f13.append('%.10e\n' % ((temp_part.px) * temp_part.rpp * 1e3))
lines_f13.append('%.10e\n' % ((temp_part.y) * 1e3))
lines_f13.append('%.10e\n' % ((temp_part.py) * temp_part.rpp * 1e3))
lines_f13.append('%.10e\n' % ((temp_part.sigma) * 1e3))
lines_f13.append('%.10e\n' % ((temp_part.delta)))
if i_part % 2 == 1:
lines_f13.append('%.10e\n' % (temp_part.energy0 * 1e-6))
lines_f13.append('%.10e\n' % (prev_part.energy * 1e-6))
lines_f13.append('%.10e\n' % (temp_part.energy * 1e-6))
prev_part = temp_part
with open(wfold + '/fort.13', 'w') as fid:
fid.writelines(lines_f13)
if np.mod(n_part, 2) != 0:
raise ValueError('SixTrack does not like this!')
i_start_tk = None
for ii, ll in enumerate(lines_f3):
if ll.startswith('TRACKING PAR'):
i_start_tk = ii
break
# Set number of turns and number of particles
temp_list = lines_f3[i_start_tk + 1].split(' ')
temp_list[0] = '%d' % n_turns
temp_list[2] = '%d' % (n_part / 2)
lines_f3[i_start_tk + 1] = ' '.join(temp_list)
# Set number of idfor = 2
temp_list = lines_f3[i_start_tk + 2].split(' ')
temp_list[2] = '2'
lines_f3[i_start_tk + 2] = ' '.join(temp_list)
# Setup turn-by-turn dump
i_start_dp = None
for ii, ll in enumerate(lines_f3):
if ll.startswith('DUMP'):
i_start_dp = ii
break
lines_f3[i_start_dp + 1] = 'StartDUMP 1 664 101 dumtemp.dat\n'
with open(wfold + '/fort.3', 'w') as fid:
fid.writelines(lines_f3)
os.system('./runsix_trackfun')
# Load sixtrack tracking data
sixdump_all = sixtracktools.SixDump101('%s/dumtemp.dat' % wfold)
x_tbt = np.zeros((n_turns, n_part))
px_tbt = np.zeros((n_turns, n_part))
y_tbt = np.zeros((n_turns, n_part))
py_tbt = np.zeros((n_turns, n_part))
sigma_tbt = np.zeros((n_turns, n_part))
delta_tbt = np.zeros((n_turns, n_part))
for i_part in range(n_part):
sixdump_part = sixdump_all[i_part::n_part]
x_tbt[:, i_part] = sixdump_part.x
px_tbt[:, i_part] = sixdump_part.px
y_tbt[:, i_part] = sixdump_part.y
py_tbt[:, i_part] = sixdump_part.py
sigma_tbt[:, i_part] = sixdump_part.sigma
delta_tbt[:, i_part] = sixdump_part.delta
return x_tbt, px_tbt, y_tbt, py_tbt, sigma_tbt, delta_tbt
def update_optics(line, eclouds_info, optics, partCO, pinch_path):
new_optics = optics.copy()
d = 1.e-10
init_part = pysixtrack.Particles(**partCO)
init_part.x += np.array([0., 1. * d, 0., 0., 0., 0., 0.])
init_part.px += np.array([0., 0., 1. * d, 0., 0., 0., 0.])
init_part.y += np.array([0., 0., 0., 1. * d, 0., 0., 0.])
init_part.py += np.array([0., 0., 0., 0., 1. * d, 0., 0.])
init_part.tau += np.array([0., 0., 0., 0., 0., 1. * d, 0.])
init_part.ptau += np.array([0., 0., 0., 0., 0., 0., 1. * d])
n_part = 7
ps = sixtracklib.ParticlesSet()
p = ps.Particles(num_particles=n_part)
for i_part in range(n_part):
part = pysixtrack.Particles(p0c=init_part.p0c)
part.x = init_part.x[i_part]
part.px = init_part.px[i_part]
part.y = init_part.y[i_part]
part.py = init_part.py[i_part]
part.tau = init_part.tau[i_part]
part.ptau = init_part.ptau[i_part]
part.partid = i_part
part.state = 1
part.elemid = 0
part.turn = 0
p.from_pysixtrack(part, i_part)
ecloud_lattice = LatticeWithEclouds(line, eclouds_info, ps, device=None)
ecloud_lattice.set_optics_CO(optics, partCO)
ecloud_lattice.add_tricub_data(pinch_path, 'dipole', max_z=0.05)
ecloud_lattice.remove_dipolar_kicks()
tricub_to_tricub_data = {}
for key in eclouds_info['length'].keys():
tricub_to_tricub_data[key] = 'dipole'
ecloud_lattice.finalize_assignments(tricub_to_tricub_data)
job = ecloud_lattice.job
job.track_until(1)
X_init = np.empty([6, 7])
X_fin = np.empty([6, 7])
X_init[0, :] = init_part.x
X_init[1, :] = init_part.px
X_init[2, :] = init_part.y
X_init[3, :] = init_part.py
X_init[4, :] = init_part.tau
X_init[5, :] = init_part.ptau
fin_part = pysixtrack.Particles(p0c=init_part.p0c)
fin_part.x = p.x.flatten()
fin_part.px = p.px.flatten()
fin_part.y = p.y.flatten()
fin_part.py = p.py.flatten()
fin_part.zeta = p.zeta.flatten()
fin_part.delta = p.delta.flatten()
X_fin[0, :] = fin_part.x
X_fin[1, :] = fin_part.px
X_fin[2, :] = fin_part.y
X_fin[3, :] = fin_part.py
X_fin[4, :] = fin_part.tau
X_fin[5, :] = fin_part.ptau
m = X_fin[:, 0] - X_init[:, 0]
M = np.empty([6, 6])
for j in range(6):
M[:, j] = (X_fin[:, j + 1] - X_fin[:, 0]) / d
Ms = normalization.healy_symplectify(M)
W, invW, R = normalization.linear_normal_form(Ms)
q1 = np.arccos(R[0, 0]) / np.pi / 2.
q2 = np.arccos(R[2, 2]) / np.pi / 2.
Qs = np.arccos(R[4, 4]) / np.pi / 2.
new_optics['W'] = W
new_optics['invW'] = invW
new_optics['R'] = R
new_optics['new_Qs'] = Qs
new_optics['new_q1'] = q1
new_optics['new_q2'] = q2
return new_optics
parser.add_argument('--noblock', dest='plot_block', action='store_false')
args = parser.parse_args()
inspect_closed_orbit = True
with open('line_with_ecloud_markers.pkl', 'rb') as fid:
line = pysixtrack.Line.from_dict(pickle.load(fid),keepextra=True)
with open('ecloud_lengths.pkl', 'rb') as fid:
ecloud_lengths = pickle.load(fid)
with open('optics.pkl', 'rb') as fid:
optics = pickle.load(fid)
p0c_eV = optics['p0c_eV']
part = pysixtrack.Particles(p0c = p0c_eV)
part_on_CO, M = normalization.get_CO_and_linear_map(part, line, d=1e-10, tol=1e-10)
print(f'before symplectifying det(M) = {np.linalg.det(M)}')
Ms = normalization.healy_symplectify(M)
print(f'after symplectifying det(M) = {np.linalg.det(Ms)}')
W, invW, R = normalization.linear_normal_form(Ms)
Qs = np.arccos(R[4,4])/np.pi/2.
optics['W'] = W
optics['invW'] = invW
optics['R'] = R
optics['Qs'] = Qs
with open('optics.pkl', 'wb') as fid:
sixinput = sixtracktools.SixInput('sixtrack_input')
# Build a pysixtrack line from pyblep line
import pysixtrack
ps_line, other = pysixtrack.Line.from_sixinput(sixinput)
# Build a pysixtracklib line from pyblep line
import pysixtracklib
pslib_line = pysixtracklib.Elements()
pslib_line.append_line(ps_line)
pslib_line.BeamMonitor(num_stores=1)
# Build a pysixtrack particle
ps_part = pysixtrack.Particles(p0c=7000e9)
ps_part.x = 1e-3
ps_part.px = 2e-4
# Build a pysixtracklib particle
pslib_part_set = pysixtracklib.ParticlesSet()
pslib_part = pslib_part_set.Particles(num_particles=1)
ps_part.partid = 0
ps_part.state = 1
ps_part.elemid = 0
ps_part.turn = 0
pslib_part.from_pysixtrack(ps_part, particle_index=0)
# Track with pysisxtrack
ps_line.track(ps_part)
