def generate_lattice_data(input_path, output_path, conf=dict()):
print("**** Generating Lattice Data From SixTrack Input:")
print("**** -> Reading sixtrack input data from:\r\n" +
f"**** {input_path}")
six = sixtracktools.SixInput(input_path)
slot_size = st.CBufferView.DEFAULT_SLOT_SIZE
line = pysix.Line.from_sixinput(six)
n_slots, n_objs, n_ptrs = calc_cbuffer_params_for_pysix_line(
line, slot_size=slot_size, conf=conf)
cbuffer = st.CBuffer(n_slots, n_objs, n_ptrs, 0, slot_size)
pysix_line_to_cbuffer(line, cbuffer)
path_to_lattice = os.path.join(output_path, "cobj_lattice.bin")
if 0 == cbuffer.tofile_normalised(path_to_lattice,
conf.get("cbuffer_norm_base_addr",
4096)):
print("**** -> Generated cobjects lattice data at:\r\n" +
f"**** {path_to_lattice}")
else:
raise RuntimeError("Problem during creation of lattice data")
path_to_pysix_lattice = os.path.join(output_path,
"pysixtrack_lattice.pickle")
try:
pickle.dump(line.elements, open(path_to_pysix_lattice, "wb"))
print("**** -> Generated pysixtrack lattice as python pickle:\r\n" +
f"**** {path_to_pysix_lattice}")
except:
raise RuntimeError("Unable to generate pysixtrack lattice data")
if conf.get( 'make_demotrack_data', False ) and \
st.Demotrack_enabled() and st.Demotrack_belems_can_convert( cbuffer ):
dt_lattice = st.Demotrack_belems_convert(cbuffer)
if isinstance( dt_lattice, np.ndarray ) and \
st.Demotrack_belems_num_stored_objects( dt_lattice ) == \
cbuffer.num_objects:
path_dt_lattice = os.path.join(output_path,
"demotrack_lattice.bin")
with open(path_dt_lattice, "wb") as fp_out:
fp_out.write(float_to_bytes(len(dt_lattice)))
fp_out.write(dt_lattice.tobytes())
print(
"**** -> Generated demotrack lattice as flat array:\r\n" +
f"**** {path_dt_lattice}")
return
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
sys.path.append('..')
import normalization
import os
os.system('(cd sixtrack; ./runsix)')
import numpy as np
##############
# Build line #
##############
# Read sixtrack input
sixinput = sixtracktools.SixInput('./sixtrack')
p0c_eV = sixinput.initialconditions[-3] * 1e6
# Build pysixtrack line from sixtrack input
line = pysixtrack.Line.from_sixinput(sixinput)
other_data = line.other_info
# Info on sixtrack->pyblep conversion
iconv = other_data['iconv']
########################################################
# Search closed orbit #
# (for comparison purposes we the orbit from sixtrack) #
########################################################
# Load sixtrack tracking data
import pysixtrack
npart = 1
nturns = 1
x0 = 0.0001
particles = sim.Particles(nparticles=npart)
particles.p0c = 7000e9
particles.x = x0
#lhc = sixtracktools.SixInput('.')
#line, rest, iconv = lhc.expand_struct()
#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])
import sixtracktools
import pysixtrack
import numpy as np
##############
# Build line #
##############
# Read sixtrack input
sixinput = sixtracktools.SixInput(".")
p0c_eV = sixinput.initialconditions[-3] * 1e6
# Build pysixtrack line from sixtrack input
line = pysixtrack.Line.from_sixinput(sixinput)
# Info on sixtrack->pyblep conversion
iconv = line.other_info["iconv"]
########################################################
# Search closed orbit #
# (for comparison purposes we the orbit from sixtrack) #
########################################################
# Load sixtrack tracking data
sixdump_all = sixtracktools.SixDump101("res/dump3.dat")
# Assume first particle to be on the closed orbit
Nele_st = len(iconv)
sixdump_CO = sixdump_all[::2][:Nele_st]
# Disable BB elements
# Read sixtrack input using sixtracktools
import sixtracktools
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)
import sixtracktools
import pysixtrack
six = sixtracktools.SixInput('sixtrack')
line, rest, iconv = six.expand_struct(convert=pysixtrack.element_types)
sixdump = sixtracktools.SixDump101('sixtrack/dump3.dat')[::2]
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_type, elem = line[jj]
elem.track(prun)
print(f"{jj} {label},{str(elem)[:50]}")
import sixtracktools
import pysixtrack
import simpletrack as sim
elements = sim.Elements()
six = sixtracktools.SixInput(".")
line = pysixtrack.Line.from_sixinput(six, classes=elements.gen_builder_class())
elements.tofile('line.bin')
def generate_testdata(pyst_example, pysixtrack_line_from_pickle=True):
# -------------------------------------------------------------------------
# Step 1: convert the input data into CObject files
input_folder = os.path.join( testdata.PATH_TO_TESTDATA_DIR, pyst_example )
output_folder = os.path.join( testdata.PATH_TO_TESTDATA_DIR, pyst_example )
st_beam_elem_dump = os.path.join( output_folder, 'beam_elements_sixtrack.bin' )
beam_elem_dump = os.path.join( output_folder, 'beam_elements.bin' )
# Dump the unmodified SixTrack machine description to CBuffer data file
six = sixtracktools.SixInput(input_folder)
st_line, rest, iconv = six.expand_struct(convert=pysixtrack.element_types)
st_elements = pystlib.Elements.fromline(st_line)
st_elements.to_file(st_beam_elem_dump)
# Dump the pysixtrack machine description to CBuffer data file
if pysixtrack_line_from_pickle:
with open(os.path.join(input_folder, 'line.pkl'), 'rb') as fid:
line = pickle.load(fid)
else:
line = st_line
elements = pystlib.Elements.fromline(line)
elements.to_file(beam_elem_dump)
# -------------------------------------------------------------------------
# Step 2: Dump particle state into an element by element I/O buffer
# before tracking happens at each beam element:
# Dump the unmodified SixTrack element-by-element data to CBuffer data file
st_particles_dump = os.path.join(output_folder, 'particles_dump_sixtrack.bin')
st_particles = pystlib.ParticlesSet.fromSixDump101(input_folder,
os.path.join(input_folder, 'dump3.dat'))
st_particles.to_file(st_particles_dump)
# Reload from file
input_particles_buffer = pystlib.CBuffer.fromfile(st_particles_dump)
assert(input_particles_buffer.n_objects > 0)
ebe_particles_buffer = CBuffer()
input_particles = input_particles_buffer.get_object(0, cls=Particles)
npart = input_particles.num_particles
pystlib_particles = pystlib.particles.makeCopy( input_particles )
# one particle is used for the fix-point calculation, thus we would need at
# least two particles for any kind of non-trivial testdata
track_particles = []
for jj in range( npart ):
track_particles.append( pysixtrack.Particles() )
input_particles.to_pysixtrack( track_particles[ jj ], jj )
track_particles[ jj ].turn = 0 #Override turn in case it's not 0
for ii, elem in enumerate( line ):
label, be_type, beam_element = elem
before = Particles( num_particles=npart, cbuffer=ebe_particles_buffer )
for jj in range( 0, npart ):
before.from_pysixtrack( track_particles[ jj ], jj )
beam_element.track( track_particles[ jj ] )
before.at_element[:] = ii
last = Particles( num_particles=npart, cbuffer=ebe_particles_buffer )
for jj in range( npart ):
last.from_pysixtrack( track_particles[ jj ], jj )
last.at_turn[:] = 1
last.at_element[:] = 0
# -------------------------------------------------------------------------
# Step 3: Write the element by element I/O buffer to the output file
# in the output_folder location:
assert( ( len( line ) + 1 ) == ebe_particles_buffer.n_objects )
particles_dump = os.path.join( output_folder, 'particles_dump.bin' )
ebe_particles_buffer.to_file( particles_dump )
def generate_particle_data(input_path, output_path, conf=dict()):
print("**** Generating Particles Data From SixTrack Input:")
path_to_dump_file = os.path.join(input_path, "dump3.dat")
print("**** -> Reading sixtrack input data from:\r\n" +
f"**** {path_to_dump_file}")
six = sixtracktools.SixInput(input_path)
#line, rest, iconv = six.expand_struct( convert=pysixtrack.elements )
line = pysix.Line.from_sixinput(six)
iconv = line.other_info["iconv"]
sixdump = sixtracktools.SixDump101(path_to_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 = num_dumps // num_iconv
print(f"**** Info :: num sixtrack sequences : {num_iconv}")
print(f"**** Info :: num beam elements : {num_belem}")
print(f"**** Info :: num particles : {num_particles}")
assert num_particles > 0
assert num_belem > 0
assert num_iconv > 0
# =========================================================================
# Get initial particle distribution:
# Generate the initial particle disitribution buffers
print("**** -> Generating initial particle distributions ...")
generate_particle_data_initial(output_path, iconv, sixdump, conf=conf)
# =========================================================================
# Make sixtrack sequency-by-sequence data:
if conf.get("make_sixtrack_sequ_by_sequ", False):
print("**** -> Generating SixTrack sequ-by-sequ particle data ...")
generate_particle_data_sequ_by_sequ(output_path,
line,
iconv,
sixdump,
conf=conf)
# =========================================================================
# Make elem-by-elem data using pysixtrack:
if conf.get("make_elem_by_elem_data", False):
print(
"**** -> Generating elem-by-elem particle data using pysixtrack ..."
)
generate_particle_data_elem_by_elem(output_path,
line,
iconv,
sixdump,
conf=conf)
# =========================================================================
# Make until turn data using pysixtrack:
if conf.get( "make_until_num_turn_data", False ) and \
conf.get( "until_num_turns", 1 ) > 0:
print(
"**** -> Generating until_turn tracked data using pysixtrack ...")
until_turn = conf.get("until_num_turns", 1)
generate_particle_data_until_turn(output_path,
line,
iconv,
sixdump,
until_turn,
conf=conf)
import sixtracktools
import pysixtrack
six = sixtracktools.SixInput("sixtrack")
line, rest, iconv = six.expand_struct(convert=pysixtrack.element_types)
sixdump = sixtracktools.SixDump101("sixtrack/dump3.dat")[::2]
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_type, elem = line[jj]
elem.track(prun)
print(f"{jj} {label},{str(elem)[:50]}")
