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)