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 __init__(self, line, eclouds_info, particles_set, device=None):
self.tricub_data_buffer = cobjects.CBuffer()
self.tricub_data = {}
self.tricub_data_indices = {}
self.tricubs = {}
self.tricub_indices = {}
self.tricub_data_buffer_ids = {}
self.elements = sixtracklib.Elements()
self.tricub_data_buffer = cobjects.CBuffer()
self.eclouds_info = eclouds_info
self.tune_is_valid_list = []
self.turn_q_list = []
self.q1_list = []
self.q2_list = []
self.qx_list = []
self.qy_list = []
self.n_particles = len(particles_set.particles[0].particle_id)
ecloud_list = eclouds_info['length'].keys()
print(
f'Number of elements in line before cleaning: {len(line.elements)}'
)
self.clean_line(line, ecloud_list)
print(
f'Number of elements in line after cleaning: {len(line.elements)}')
for element, element_name in zip(line.elements, line.element_names):
element_type = element.__class__.__name__
if element_name in ecloud_list:
tc_index = self.elements.cbuffer.n_objects
tc = sixtracklib.TriCub(cbuffer=self.elements.cbuffer)
tc.x_shift = self.eclouds_info['x_CO'][element_name]
tc.y_shift = self.eclouds_info['y_CO'][element_name]
tc.tau_shift = self.eclouds_info['tau_CO'][element_name]
tc.length = self.eclouds_info['length'][element_name]
self.tricubs[element_name] = tc
self.tricub_indices[element_name] = tc_index
else:
getattr(self.elements,
element_type)(**element.to_dict(keepextra=True))
self.job = sixtracklib.TrackJob(self.elements,
particles_set,
device=device)
return
def sixTrackLib(self,
numStores: int = 1,
installBPMs: bool = True,
finalPhaseSpace: bool = False):
"""Export model to SixTrackLib."""
myElem = stl.Elements()
if not installBPMs and not finalPhaseSpace:
raise ValueError("no output specified")
# build map-wise
maps = self.mergeDrifts()
# add maps to SixTrackLib
for m in maps:
if type(m) is Maps.DriftMap:
myElem.Drift(length=m.length)
elif type(m) is Maps.DipoleKick:
k0L = m.weight.item() * m.dipoleLength
myElem.Multipole(length=m.dipoleLength,
hxl=m.angle,
knl=[
k0L,
])
elif type(m) is Maps.EdgeKick:
myElem.DipoleEdge(h=m.curvature, e1=m.edgeAngle)
elif type(m) is Maps.MultipoleKick:
k1nl = m.kickLength.item() * m.k1n.item()
k2nl = m.kickLength.item() * m.k2n.item()
k1sl = m.kickLength.item() * m.k1s.item()
k2sl = m.kickLength.item() * m.k2s.item()
knl = [0.0, k1nl, k2nl]
ksl = [0.0, k1sl, k2sl]
myElem.Multipole(knl=knl, ksl=ksl)
elif type(m) is Model.MonitorDummy:
if installBPMs:
myElem.BeamMonitor(num_stores=numStores)
else:
continue
else:
raise NotImplementedError()
if finalPhaseSpace:
myElem.BeamMonitor(num_stores=numStores)
return myElem
#!/usr/bin/env python
# -*- coding: utf-8 -*-
import sys
import numpy as np
import sixtracklib as pyst
elements = pyst.Elements()
elements.Drift(length=1.2)
elements.Multipole(knl=[0, 0.001])
particles = pyst.Particles.from_ref(num_particles=10, p0c=1e9)
particles.px += np.linspace(0, 1e-2, 10)
job = pyst.TrackJob(elements, particles)
status = job.track_until(1)
print(particles.x)
print(particles.px)
if pyst.supports('opencl'):
jobcl = pyst.TrackJob(elements, particles, device="opencl:0.0")
status = jobcl.track_until(2)
jobcl.collect()
print(particles.x)
print(particles.px)
"k",
z_absc,
y_interp_cub,
"r-",
)
plt.show()
# -------------------------------------------------------------------------
# B) Init the particle set
beam = st.ParticlesSet()
particles = beam.Particles(num_particles=100, p0c=6.5e12)
# -------------------------------------------------------------------------
# C) Build the lattice. We add three interpolated space charge elements
# and keep track of the indices at which they are available
lattice = st.Elements()
sc0_index = lattice.cbuffer.n_objects # index of sc0 element
sc0 = lattice.SCInterpolatedProfile(
number_of_particles=particles.num_particles)
dr0 = lattice.Drift(length=1.0)
q0 = lattice.Multipole(knl=[0.0, 0.1])
sc1_index = lattice.cbuffer.n_objects # index of sc1 element
sc1 = lattice.SCInterpolatedProfile(
number_of_particles=particles.num_particles)
dr1 = lattice.Drift(length=1.0)
q1 = lattice.Multipole(knl=[0.0, -0.1])
sc2_index = lattice.cbuffer.n_objects # index of sc2 element
sc2 = lattice.SCInterpolatedProfile(
mad.call(file="fodo.madx")
mad.command.beam(particle='proton', energy=str(Etot))
mad.use(sequence="FODO")
mad.twiss()
mad.command.select(flag="makethin", class_="quadrupole", slice='8')
mad.command.select(flag="makethin", class_="sbend", slice='8')
mad.command.makethin(makedipedge=False, style="teapot", sequence="fodo")
mad.twiss()
sis18 = mad.sequence.FODO
nturns = 1
ps_line, _ = pysixtrack.Line.from_madx_sequence(sis18)
elements = pystlib.Elements()
elements.append_line(ps_line)
def prepare(npart=int(1e6), p0c=p0c, elements=elements, device='cpu'):
particles = pystlib.Particles.from_ref(npart, p0c=p0c)
particles.x += np.linspace(0, 1e-6, npart)
job = pystlib.TrackJob(elements, particles, device=device)
return job
class Timer(object):
def __init__(self):
self.interval = 0
def track_particle_sixtracklib(
line, 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,
device=None):
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)
import sixtracklib
elements = sixtracklib.Elements()
elements.BeamMonitor(num_stores=n_turns)
elements.append_line(line)
n_part = len(Dx_wrt_CO_m)
# Build PyST particle
ps = sixtracklib.ParticlesSet()
p = ps.Particles(num_particles=n_part)
for i_part in range(n_part):
part = partCO.copy()
part.x += Dx_wrt_CO_m[i_part]
part.px += Dpx_wrt_CO_rad[i_part]
part.y += Dy_wrt_CO_m[i_part]
part.py += Dpy_wrt_CO_rad[i_part]
part.sigma += Dsigma_wrt_CO_m[i_part]
part.delta += Ddelta_wrt_CO[i_part]
part.partid = i_part
part.state = 1
part.elemid = 0
part.turn = 0
p.from_pysixtrack(part, i_part)
if device is None:
job = sixtracklib.TrackJob(elements, ps)
else:
job = sixtracklib.TrackJob(elements, ps, device=device)
job.track_until(n_turns)
job.collect()
res = job.output
x_tbt = res.particles[0].x.reshape(n_turns, n_part)
px_tbt = res.particles[0].px.reshape(n_turns, n_part)
y_tbt = res.particles[0].y.reshape(n_turns, n_part)
py_tbt = res.particles[0].py.reshape(n_turns, n_part)
sigma_tbt = res.particles[0].sigma.reshape(n_turns, n_part)
delta_tbt = res.particles[0].delta.reshape(n_turns, n_part)
# For now data are saved at the end of the turn by STlib and at the beginning by the others
#x_tbt[1:, :] = x_tbt[:-1, :]
#px_tbt[1:, :] = px_tbt[:-1, :]
#y_tbt[1:, :] = y_tbt[:-1, :]
#py_tbt[1:, :] = py_tbt[:-1, :]
#sigma_tbt[1:, :] = sigma_tbt[:-1, :]
#delta_tbt[1:, :] = delta_tbt[:-1, :]
#x_tbt[0, :] = p.x
#px_tbt[0, :] = p.px
#y_tbt[0, :] = p.y
#py_tbt[0, :] = p.py
#sigma_tbt[0, :] = p.sigma
#delta_tbt[0, :] = p.delta
print('Done loading!')
return x_tbt, px_tbt, y_tbt, py_tbt, sigma_tbt, delta_tbt
f_rf = 400e6 # frequency of the CC RF in Hz (from Mad-x) hardcoded
f_rev = 43.45e3 # revolution frequency of SPS in Hz (doubts about this value)
mod_signal = modulated_rf_phase(A_sec, mod_period_sec, f_rf, pp.n_turns_max+1, f_rev) # amplitude in rad
#####################
parser = ArgumentParser()
parser.add_argument("--device", help="set opencl device")
args = parser.parse_args()
os.makedirs(pp.output_dir, exist_ok=True)
with open(pp.input_dir + 'line.pkl', 'rb') as fid:
line = pysixtrack.Line.from_dict(pickle.load(fid), keepextra=True)
elements = sixtracklib.Elements()
elements.append_line(line)
n_part = pp.n_macroparticles
circumference = line.get_length()
# directory to save the final distribution
parts_distribution_dict = {'x': [], 'px':[], 'y': [], 'py': [], 'sigma': [], 'delta': []}
# directory to save the tbt emittances
tbt_dict = {'turn':[], 'time':[], 'intensity':[], 'neps_x':[], 'neps_y':[], 'std_sigma':[]}
time_cum = 0
if pp.track_with == 'sixtracklib':
ps = sixtracklib.ParticlesSet().fromfile('input/sixtracklib.particles')
def track_particle_sixtracklib_firstlast(
line, 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,
device=None):
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)
#if type(partCO) is pysixtrack.Particles:
# part = partCO.copy()
#else:
# part = pysixtrack.Particles(**partCO)
n_turns_to_store=1000
n_turns_tbt=1000
#skip_turns=1000
import sixtracklib
elements=sixtracklib.Elements()
#sixtracklib.append_beam_monitors_to_lattice(beam_elements_buffer=elements.cbuffer,
# until_turn_elem_by_elem=0,
# until_turn_turn_by_turn=n_turns_tbt,
# until_turn=n_turns,
# skip_turns=skip_turns
# )
elements.BeamMonitor(num_stores=n_turns_tbt,start=0,skip=1,is_rolling=False)
elements.BeamMonitor(num_stores=n_turns_to_store,start=0,skip=1,is_rolling=True)
print(elements.get_elements())
#elements.BeamMonitor(num_stores=n_turns)
#elements.BeamMonitor(num_stores=n_turns_to_store)
elements.append_line(line)
n_stores=elements.get_elements()[1].num_stores
n_part = len(Dx_wrt_CO_m)
# Build PyST particle
ps = sixtracklib.ParticlesSet()
p = ps.Particles(num_particles=n_part)
for i_part in range(n_part):
if type(partCO) is pysixtrack.Particles:
part = partCO.copy()
else:
part = pysixtrack.Particles(**partCO)
part.x += Dx_wrt_CO_m[i_part]
part.px += Dpx_wrt_CO_rad[i_part]
part.y += Dy_wrt_CO_m[i_part]
part.py += Dpy_wrt_CO_rad[i_part]
part.sigma += Dsigma_wrt_CO_m[i_part]
part.delta += Ddelta_wrt_CO[i_part]
part.partid = i_part
part.state = 1
part.elemid = 0
part.turn = 0
p.from_pysixtrack(part, i_part)
if device is None:
job = sixtracklib.TrackJob(elements, ps)
else:
job = sixtracklib.TrackJob(elements, ps, device=device)
start_tracking_time = time.time()
job.track(n_turns)
end_tracking_time = time.time()
job.collect()
end_collecting_time = time.time()
res = job.output
print(res.particles[0])
print(res.particles[1])
x_tbt_first = res.particles[0].x.reshape(n_turns_tbt,n_part)
px_tbt_first = res.particles[0].px.reshape(n_turns_tbt,n_part)
y_tbt_first = res.particles[0].y.reshape(n_turns_tbt,n_part)
py_tbt_first = res.particles[0].py.reshape(n_turns_tbt,n_part)
zeta_tbt_first = res.particles[0].zeta.reshape(n_turns_tbt,n_part)
delta_tbt_first = res.particles[0].delta.reshape(n_turns_tbt,n_part)
at_turn_tbt_first = res.particles[0].at_turn.reshape(n_turns_tbt,n_part)
state_tbt_first = res.particles[0].state.reshape(n_turns_tbt,n_part)
x_tbt_last = res.particles[1].x.reshape(n_stores,n_part)
px_tbt_last = res.particles[1].px.reshape(n_stores,n_part)
y_tbt_last = res.particles[1].y.reshape(n_stores,n_part)
py_tbt_last = res.particles[1].py.reshape(n_stores,n_part)
zeta_tbt_last = res.particles[1].zeta.reshape(n_stores,n_part)
delta_tbt_last = res.particles[1].delta.reshape(n_stores,n_part)
at_turn_tbt_last = res.particles[1].at_turn.reshape(n_stores,n_part)
state_tbt_last = res.particles[1].state.reshape(n_stores,n_part)
output_dict = {'x_tbt_first' : x_tbt_first,
'px_tbt_first' : px_tbt_first,
'y_tbt_first' : y_tbt_first,
'py_tbt_first' : py_tbt_first,
'zeta_tbt_first' : zeta_tbt_first,
'delta_tbt_first' : delta_tbt_first,
'at_turn_tbt_first' : at_turn_tbt_first,
# 'state_tbt_first' : state_tbt_first,
'x_tbt_last' : x_tbt_last,
'px_tbt_last' : px_tbt_last,
'y_tbt_last' : y_tbt_last,
'py_tbt_last' : py_tbt_last,
'zeta_tbt_last' : zeta_tbt_last,
'delta_tbt_last' : delta_tbt_last,
'at_turn_tbt_last' : at_turn_tbt_last,
# 'state_tbt_last' : state_tbt_last,
'tracking_time_mins' : (end_tracking_time - start_tracking_time)/60.,
'collecting_time_mins' : (end_collecting_time - end_tracking_time)/60.,
}
print('Done loading!')
return output_dict
in_knl = np.random.uniform(0.0, 100.0, order + 1)
in_ksl = np.random.uniform(0.0, 100.0, order + 1)
bal_length = 2 * order + 2
in_bal = np.zeros(bal_length)
for ii in range(0, len(in_ksl)):
in_bal[2 * ii] = in_knl[ii] / factorial(ii, exact=True)
in_bal[2 * ii + 1] = in_ksl[ii] / factorial(ii, exact=True)
pyst_line.append(pyst.elements.Multipole(knl=in_knl, ksl=in_ksl))
in_knl_data.append(in_knl)
in_ksl_data.append(in_ksl)
in_bal_data.append(in_bal)
temp = pyst.Line(pyst_line)
st_line = st.Elements()
st_line.append_line(temp)
assert st_line.cbuffer.n_objects == num_tests
assert len(in_knl_data) == num_tests
assert len(in_ksl_data) == num_tests
assert len(in_bal_data) == num_tests
for ii in range(0, num_tests):
assert pyst_line[ii].order == st_line.get(ii).order
assert np.allclose(in_knl_data[ii], st_line.get(ii).knl, rtol=1e-15)
assert np.allclose(in_ksl_data[ii], st_line.get(ii).ksl, rtol=1e-15)
assert np.allclose(in_bal_data[ii], st_line.get(ii).bal, rtol=1e-15)
for ii in range(0, num_tests):
order = pyst_line[ii].order
def test_kicks(
cmp_file_name="precomputed_kicks.pickle",
device_str=None,
abs_tol=1e-15,
rel_tol=0.0,
):
####### load file with correct kicks #######
path_to_testdir = sttest.config.PATH_TO_TESTDATA_DIR
assert path_to_testdir is not None
assert os.path.exists(path_to_testdir)
assert os.path.isdir(path_to_testdir)
path_to_cmp_file = os.path.join(path_to_testdir, "tricub", cmp_file_name)
assert os.path.exists(path_to_cmp_file)
with open(path_to_cmp_file, "rb") as fp:
n_part, prng_seed, kicks = pickle.load(fp)
assert n_part > 0
assert prng_seed is not None
assert kicks is not None
np.random.seed(int(prng_seed))
lattice = st.Elements()
tc_index = lattice.cbuffer.n_objects
tc = st.TriCub(cbuffer=lattice.cbuffer)
tc.length = 1.0
particles_set = st.ParticlesSet()
particles = particles_set.Particles(num_particles=n_part)
nx = 5
ny = 7
nz = 9
A = np.random.rand(nx, ny, nz, 8) * 1.0e-3
dx = 0.001
dy = 0.002
dz = 0.003
x0 = -(nx // 2) * dx
y0 = -(ny // 2) * dy
z0 = -(nz // 2) * dz
test_x = x0 + (nx - 2) * dx * np.random.rand(n_part)
test_y = y0 + (ny - 2) * dy * np.random.rand(n_part)
test_z = z0 + (nz - 2) * dz * np.random.rand(n_part)
for i_part in range(n_part):
part = pysixtrack.Particles()
part.x = test_x[i_part]
part.y = test_y[i_part]
part.tau = test_z[i_part]
part.partid = i_part
part.state = 1
part.elemid = 0
part.turn = 0
particles.from_pysixtrack(part, i_part)
job = st.TrackJob(lattice, particles_set, device=device_str)
tricub_data_buffer = cobjects.CBuffer()
tc_data_index = tricub_data_buffer.n_objects
tc_data = st.TriCubData(cbuffer=tricub_data_buffer, nx=nx, ny=ny, nz=nz)
tc_data.x0 = x0
tc_data.y0 = y0
tc_data.z0 = z0
tc_data.dx = dx
tc_data.dy = dy
tc_data.dz = dz
tc_data.mirror_x = 0
tc_data.mirror_y = 0
tc_data.mirror_z = 0
scale = [1.0, dx, dy, dz, dx * dy, dx * dz, dy * dz, (dx * dy) * dz]
for ii in range(nx):
for jj in range(ny):
for kk in range(nz):
for ll in range(8):
tc_data.table_addr[ll + 8 *
(ii + nx *
(jj + ny * kk))] = (A[ii, jj, kk, ll] *
scale[ll])
tricub_data_buffer_id = job.add_stored_buffer(cbuffer=tricub_data_buffer)
st.TriCub_buffer_create_assign_address_item(job, tc_index,
tricub_data_buffer_id,
tc_data_index)
job.commit_address_assignments()
job.assign_all_addresses()
job.track_until(1)
job.collect()
assert np.allclose(kicks[:, 0], particles.px, rel_tol, abs_tol)
assert np.allclose(kicks[:, 1], particles.py, rel_tol, abs_tol)
assert np.allclose(kicks[:, 2], particles.ptau, rel_tol, abs_tol)
st_NullParticles, \
st_Buffer_new_from_copy, st_Buffer_new_mapped_on_cbuffer, \
st_Buffer_get_num_of_objects, st_Particles_copy
from sixtracklib_test.stcommon import st_Particles_print_out, \
st_Particles_compare_values_with_treshold, st_Particles_compare_values, \
st_Particles_buffers_compare_values_with_treshold, \
st_Particles_random_init
import ctypes as ct
from cobjects import CBuffer
if __name__ == '__main__':
EPS = np.finfo(float).eps
line = pyst.Elements()
line.Drift(length=1.0)
line.Cavity(voltage=100e3, frequency=400e6, lag=0.0)
line.Drift(length=3.0)
line.LimitRect(min_x=-1.0, max_x=1.0, min_y=-2.0, max_y=2.0)
line.Drift(length=5.0)
line.BeamMonitor(num_stores=10,
start=0,
skip=0,
out_address=0,
max_particle_id=0,
min_particle_id=0,
is_rolling=False,
is_turn_ordered=False)
line.Drift(length=7.0)
line.LimitEllipse(a=0.5, b=0.35)
from __future__ import absolute_import, print_function import numpy as np import pyopencl as cl import sixtracklib as sl particles = sl.Particles(nparticles=2560, ndim=10) elements = sl.Elements() elements.add_drift(3.3) elements.add_multipole([0, -0.01]) sim = sl.SixTrackCL(particles, elements)
