def __init__(self, cbuffer=None):
if cbuffer is None:
self.cbuffer = CBuffer()
else:
self.cbuffer = cbuffer
for name, cls in self.element_types.items():
setattr(self, name, self._mk_fun(self.cbuffer, cls))
self.cbuffer.typeids[cls._typeid] = cls
class Elements(object):
element_types = {'Cavity': Cavity,
'Drift': Drift,
'DriftExact': DriftExact,
'Multipole': Multipole,
# 'RFMultipole': RFMultipole,
'SRotation': SRotation,
'XYShift': XYShift,
'BeamBeam6D': BeamBeam6D,
'BeamBeam4D': BeamBeam4D,
# 'Line': Line,
'BeamMonitor': BeamMonitor,
}
def _mk_fun(self, buff, cls):
def fun(*args, **nargs):
# print(cls.__name__,nargs)
return cls(cbuffer=buff, **nargs)
return fun
@classmethod
def fromfile(cls, filename):
cbuffer = CBuffer.fromfile(filename)
return cls(cbuffer=cbuffer)
@classmethod
def fromline(cls, line):
self = cls()
for label, element_name, element in line:
getattr(self, element_name)(**element._asdict())
return self
def tofile(self, filename):
self.cbuffer.tofile(filename)
def __init__(self, cbuffer=None):
if cbuffer is None:
self.cbuffer = CBuffer()
else:
self.cbuffer = cbuffer
for name, cls in self.element_types.items():
setattr(self, name, self._mk_fun(self.cbuffer, cls))
self.cbuffer.typeids[cls._typeid] = cls
def gen_builder(self):
out = {}
for name, cls in self.element_types.items():
out[name] = getattr(self, name)
return out
def get_elements(self):
n = self.cbuffer.n_objects
return [self.cbuffer.get_object(i) for i in range(n)]
def get(self, objid):
return self.cbuffer.get_object(objid)
def __init__(self, cbuffer=None):
if cbuffer is None:
self.cbuffer = CBuffer()
else:
self.cbuffer = cbuffer
for name, cls in self.element_types.items():
self.cbuffer.typeids[cls._typeid] = cls
def __init__(self, cbuffer=None):
if cbuffer is None:
self.cbuffer = CBuffer()
else:
self.cbuffer = cbuffer
for name, cls in self.element_types.items():
setattr(self, name, self._mk_fun(self.cbuffer, cls))
self.cbuffer.typeids[cls._typeid] = cls
self._builder = self.gen_builder()
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.tofile(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.tofile(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.tofile(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.toPySixTrack( 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.fromPySixTrack( 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.fromPySixTrack( 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.tofile( particles_dump )
raise SystemExit("support for cuda required for this test")
path_to_testdir = testlib.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_particle_data = os.path.join(
path_to_testdir, "beambeam", "particles_dump.bin")
assert os.path.exists(path_to_particle_data)
path_to_beam_elements_data = os.path.join(
path_to_testdir, "beambeam", "beam_elements.bin")
assert os.path.exists(path_to_beam_elements_data)
pb = CBuffer.fromfile(path_to_particle_data)
initial_particles = pb.get_object(0, cls=pyst.Particles)
track_pb = CBuffer()
particles = pyst.makeCopy(initial_particles, cbuffer=track_pb)
eb = CBuffer.fromfile(path_to_beam_elements_data)
num_beam_elements = eb.n_objects
ctx = st.st_CudaContext_create()
assert ctx != st.st_NullCudaContext
lattice = st.st_Buffer_new_mapped_on_cbuffer(eb)
assert lattice != st.st_NullBuffer
lattice_arg = st.st_CudaArgument_new(ctx)
tc = st.TriCub(cbuffer=lattice.cbuffer) # b) the particle set particle_sets = st.ParticlesSet() particles = particle_sets.Particles(num_particles=100) # ------------------------------------------------------------------------------ # 2) Create the track_job; currently only CPU is supported job = st.TrackJob(lattice, particle_sets) # ------------------------------------------------------------------------------ # 3) Create the data buffer for the TriCubData instances and hand it over to # the track_job for management: tricub_data_buffer = CBuffer() tc_data_0_index = tricub_data_buffer.n_objects tc_data_0 = st.TriCubData(cbuffer=tricub_data_buffer, nx=100, ny=100, nz=100) tc_data_1_index = tricub_data_buffer.n_objects tc_data_1 = st.TriCubData(cbuffer=tricub_data_buffer, nx=10, ny=16, nz=8) tricub_data_buffer_id = job.add_stored_buffer(cbuffer=tricub_data_buffer) # ------------------------------------------------------------------------------ # 4) Create the mappings connecting the two TriCubData instances to the three # TriCub beam elements # # tc_data_0 -> tc2 # tc_data_1 -> tc5
if __name__ == '__main__':
path_to_testdir = testlib.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_particle_data = os.path.join(path_to_testdir, "beambeam",
"particles_dump.bin")
assert os.path.exists(path_to_particle_data)
path_to_beam_elements_data = os.path.join(path_to_testdir, "beambeam",
"beam_elements.bin")
assert os.path.exists(path_to_beam_elements_data)
pb = CBuffer.fromfile(path_to_particle_data)
num_elem_by_elem_turns = 1
eb = CBuffer.fromfile(path_to_beam_elements_data)
until_turn_elem_by_elem = 1
until_turn_turn_by_turn = 5
until_turn = 100
skip_turns = 10
# ------------------------------------------------------------------------
initial_num_beam_elements = eb.n_objects
num_beam_monitors = st.beam_elements.append_beam_monitors_to_lattice(
eb, until_turn_elem_by_elem, until_turn_turn_by_turn, until_turn,
skip_turns)
def fromfile(cls, filename):
cbuffer = CBuffer.fromfile(filename)
return cls(cbuffer=cbuffer)
path_to_testdir, "beambeam", "particles_dump.bin")
assert(os.path.exists(path_to_particle_data))
pb = st.st_Buffer_new_from_file(path_to_particle_data.encode('utf-8'))
assert(pb != st.st_NullBuffer)
num_particle_sets = st.st_Particles_buffer_get_num_of_particle_blocks(pb)
assert(num_particle_sets > 0)
total_num_particles = st.st_Particles_buffer_get_total_num_of_particles(pb)
assert(total_num_particles > num_particle_sets)
particles = st.st_Particles_buffer_get_particles(pb, 0)
assert(particles != st.st_NullParticles)
num_particles = st.st_Particles_get_num_of_particles(particles)
assert(num_particles > 0)
# Load the same data file into a CBuffer instance
cobj_pb = CBuffer.fromfile(path_to_particle_data)
assert(cobj_pb.n_objects > 0)
assert(cobj_pb.n_objects == num_particle_sets)
cmp_particles = cobj_pb.get_object(0, cls=pyst.Particles)
cmp_num_particles = cmp_particles.num_particles
assert(cmp_particles.num_particles == num_particles)
# Provide a buffer for calculating the difference
diff_buffer = st.st_Buffer_new(0)
assert(diff_buffer != st.st_NullBuffer)
diff = st.st_Particles_new(diff_buffer, num_particles)
assert(diff != st.st_NullParticles)
assert(num_particles == st.st_Particles_get_num_of_particles(diff))
# Calculate the difference between the particles stored on the NS(Buffer)
def fromfile(cls, filename):
cbuffer = CBuffer.fromfile(filename)
return cls(cbuffer=cbuffer)
raise SystemExit("cuda support required for this test")
try:
num_nodes = st.CudaController.NUM_AVAILABLE_NODES()
except RuntimeError as e:
num_nodes = 0
if num_nodes <= 0:
print("No CUDA nodes available -> skip test")
sys.exit(0)
num_d_values = 10
num_e_values = 10
num_obj = 10
obj_buffer = CBuffer()
for ii in range(0, num_obj):
obj = GenericObj(cbuffer=obj_buffer,
type_id=ii,
a=ii,
b=float(ii),
c=[1.0, 2.0, 3.0, 4.0],
num_d=num_d_values,
num_e=num_e_values)
c_obj_buffer = st.Buffer(cbuffer=obj_buffer)
assert c_obj_buffer.pointer != st_NullBuffer
assert c_obj_buffer.slot_size > 0
assert c_obj_buffer.capacity > 0
assert c_obj_buffer.size > 0 and c_obj_buffer.size <= c_obj_buffer.capacity
assert c_obj_buffer.num_objects == obj_buffer.n_objects
if __name__ == '__main__':
path_to_testdir = testlib.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_particle_data = os.path.join(path_to_testdir, "beambeam",
"particles_dump.bin")
assert (os.path.exists(path_to_particle_data))
path_to_beam_elements_data = os.path.join(path_to_testdir, "beambeam",
"beam_elements.bin")
assert (os.path.exists(path_to_beam_elements_data))
# -------------------------------------------------------------------------
eb = CBuffer.fromfile(path_to_beam_elements_data)
eb_data_begin = eb.base
eb_data_size = eb.size
eb_num_objects = eb.n_objects
pb = CBuffer.fromfile(path_to_particle_data)
assert (pb.n_objects > 0)
particle_type_id = pb.get_object_typeid(0)
pb_data_begin = pb.base
pb_data_size = pb.size
pb_num_objects = pb.n_objects
# Testcase 1: only default parameters, no elem by elem output, no
# beam monitors
def __init__(self,
beam_elements_buffer,
particles_buffer,
until_turn_elem_by_elem=0,
arch='cpu',
device_id=None,
device=None,
output_buffer=None,
config_str=None):
self.ptr_st_track_job = st.st_NullTrackJob
self._particles_buffer = None
self._ptr_c_particles_buffer = st.st_NullBuffer
self._beam_elements_buffer = None
self._ptr_c_beam_elements_buffer = st.st_NullBuffer
self._output_buffer = None
self._ptr_c_output_buffer = st.st_NullBuffer
base_addr_t = ct.POINTER(ct.c_ubyte)
success = False
if particles_buffer is not None:
particles_buffer = _get_buffer(particles_buffer)
self._particles_buffer = particles_buffer
self._ptr_c_particles_buffer = \
st.st_Buffer_new_mapped_on_cbuffer(particles_buffer)
if self._ptr_c_particles_buffer == st.st_NullBuffer:
raise ValueError("Issues with input particles buffer")
if beam_elements_buffer is not None:
beam_elements_buffer = _get_buffer(beam_elements_buffer)
self._beam_elements_buffer = beam_elements_buffer
self._ptr_c_beam_elements_buffer = \
st.st_Buffer_new_mapped_on_cbuffer(beam_elements_buffer)
if self._ptr_c_beam_elements_buffer == st.st_NullBuffer:
raise ValueError("Issues with input beam elements buffer")
particles = st.st_Particles_buffer_get_particles(
self._ptr_c_particles_buffer, 0)
if particles == st.st_NullParticles:
raise ValueError("Required particle sets not available")
until_turn_elem_by_elem = ct.c_uint64(until_turn_elem_by_elem)
out_buffer_flags = st.st_OutputBuffer_required_for_tracking(
particles, self._ptr_c_beam_elements_buffer, until_turn_elem_by_elem)
needs_output_buffer = st.st_OutputBuffer_requires_output_buffer(
ct.c_int32(out_buffer_flags))
if needs_output_buffer:
num_objects = ct.c_uint64(0)
num_slots = ct.c_uint64(0)
num_dataptrs = ct.c_uint64(0)
num_garbage = ct.c_uint64(0)
slot_size = st.st_Buffer_get_slot_size(
self._ptr_c_particles_buffer)
ret = st.st_OutputBuffer_calculate_output_buffer_params(
self._ptr_c_beam_elements_buffer, particles,
until_turn_elem_by_elem, ct.byref(num_objects),
ct.byref(num_slots), ct.byref(num_dataptrs),
ct.byref(num_garbage), slot_size)
if ret == 0:
if num_objects.value > 0 and num_slots.value > 0 and \
num_dataptrs.value > 0 and num_garbage.value >= 0:
if output_buffer is None:
output_buffer = CBuffer(
max_slots=num_slots.value,
max_objects=num_objects.value,
max_pointers=num_dataptrs.value,
max_garbage=num_garbage.value)
else:
output_buffer.reallocate(
max_slots=num_slots.value,
max_objects=num_objects.value,
max_pointers=num_dataptrs.value,
max_garbage=num_garbage.value)
if output_buffer is None:
raise ValueError("Could not provide output buffer")
self._output_buffer = output_buffer
self._ptr_c_output_buffer = \
st.st_Buffer_new_mapped_on_cbuffer(output_buffer)
if self._ptr_c_output_buffer == st.st_NullBuffer:
raise ValueError("Unable to map (optional) output buffer")
else:
raise ValueError("Error pre-calculating out buffer params")
elif output_buffer is not None:
self._output_buffer = output_buffer
self._ptr_c_output_buffer = \
st.st_Buffer_new_mapped_on_cbuffer(self._output_buffer)
if self._ptr_c_output_buffer == st.st_NullBuffer:
raise ValueError("Unable to map (optional) output buffer")
assert((needs_output_buffer and
self._ptr_c_output_buffer != st.st_NullBuffer) or
(not needs_output_buffer))
if device is not None:
arch, device_id = device.split(':')
arch = arch.strip().lower()
if not(stconf.SIXTRACKLIB_MODULES.get(arch, False) is not False
or arch == 'cpu'):
raise ValueError("Unknown architecture {0}".format(arch, ))
if device_id is not None:
if config_str is None:
config_str = device_id
else:
config_str = device_id + ";" + config_str
else:
config_str = ""
arch = arch.encode('utf-8')
config_str = config_str.encode('utf-8')
self.ptr_st_track_job = st.st_TrackJob_new_with_output(
ct.c_char_p(arch), self._ptr_c_particles_buffer,
self._ptr_c_beam_elements_buffer, self._ptr_c_output_buffer,
until_turn_elem_by_elem, ct.c_char_p(config_str))
if self.ptr_st_track_job == st.st_NullTrackJob:
raise ValueError('unable to construct TrackJob from arguments')
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 track_particles_pysixtrack(
input_particles,
line,
until_turn_elem_by_elem,
until_turn_turn_by_turn,
until_turn=0,
skip_turns=1,
output_buffer=None,
):
num_particles = input_particles.num_particles
num_beam_elements = len(line)
initial_at_turn = sys.maxsize
min_particle_id = sys.maxsize
max_particle_id = -(sys.maxsize) - 1
for ii in range(num_particles):
if input_particles.at_turn[ii] < initial_at_turn:
initial_at_turn = input_particles.at_turn[ii]
if min_particle_id > input_particles.particle_id[ii]:
min_particle_id = input_particles.particle_id[ii]
if max_particle_id < input_particles.particle_id[ii]:
max_particle_id = input_particles.particle_id[ii]
nn = max_particle_id - min_particle_id + 1
num_particles_per_turn = nn * num_beam_elements
if output_buffer is None:
output_buffer = CBuffer()
if initial_at_turn < until_turn_elem_by_elem:
num_particles_to_store = num_particles_per_turn * (
until_turn_elem_by_elem - initial_at_turn)
out = st_Particles(num_particles=num_particles_to_store,
cbuffer=output_buffer)
for ii in range(num_particles):
particle = pysix_Particle()
input_particles.to_pysixtrack(particle, ii)
assert particle.partid >= min_particle_id
assert particle.partid <= max_particle_id
assert particle.turn >= initial_at_turn
assert particle.turn < until_turn_elem_by_elem
delta_part_id = particle.partid - min_particle_id
while particle.turn < until_turn_elem_by_elem:
assert particle.elemid == 0
offset = (num_beam_elements * nn *
(particle.turn - initial_at_turn))
for elem in line:
out.from_pysixtrack(
particle,
offset + particle.elemid * nn + delta_part_id)
elem.track(particle)
particle.elemid += 1
particle.turn += 1
particle.elemid = 0
input_particles.from_pysixtrack(particle, ii)
start_turn_by_turn = max(until_turn_elem_by_elem, initial_at_turn)
if start_turn_by_turn < until_turn_turn_by_turn:
max_num_turns_to_store = until_turn_turn_by_turn - start_turn_by_turn
num_particles_to_store = nn * max_num_turns_to_store
out = st_Particles(num_particles=num_particles_to_store,
cbuffer=output_buffer)
for ii in range(num_particles):
particle = pysix_Particle()
input_particles.to_pysixtrack(particle, ii)
assert particle.partid >= min_particle_id
assert particle.partid <= max_particle_id
assert particle.turn >= start_turn_by_turn
delta_part_id = particle.partid - min_particle_id
turns_tracked = particle.turn - start_turn_by_turn
while particle.turn < until_turn_turn_by_turn:
assert particle.elemid == 0
for elem in line:
elem.track(particle)
particle.elemid += 1
out.from_pysixtrack(particle,
nn * turns_tracked + delta_part_id)
turns_tracked += 1
particle.turn += 1
particle.elemid = 0
input_particles.from_pysixtrack(particle, ii)
start_out_turns_turn = max(start_turn_by_turn, until_turn_turn_by_turn)
if start_out_turns_turn < until_turn:
max_num_turns_to_store = until_turn - start_out_turns_turn
if skip_turns <= 0:
skip_turns = 1
remainder = max_num_turns_to_store % skip_turns
max_num_turns_to_store = max_num_turns_to_store // skip_turns
if remainder != 0:
max_num_turns_to_store += 1
num_particles_to_store = int(nn * max_num_turns_to_store)
out = st_Particles(num_particles=num_particles_to_store,
cbuffer=output_buffer)
for ii in range(num_particles):
particle = pysix_Particle()
input_particles.to_pysixtrack(particle, ii)
assert particle.partid >= min_particle_id
assert particle.partid <= max_particle_id
assert particle.turn >= start_out_turns_turn
delta_part_id = particle.partid - min_particle_id
turns_tracked = particle.turn - start_turn_by_turn
while particle.turn < until_turn:
assert particle.elemid == 0
for elem in line:
elem.track(particle)
particle.elemid += 1
if turns_tracked % skip_turns == 0:
jj = nn * (turns_tracked // skip_turns) + delta_part_id
if jj > out.num_particles:
jj = jj % out.num_particles
out.from_pysixtrack(particle, jj)
turns_tracked += 1
particle.turn += 1
particle.elemid = 0
input_particles.from_pysixtrack(particle, ii)
return output_buffer
raise SystemExit("support for cuda required for this test")
path_to_testdir = testlib.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_particle_data = os.path.join(
path_to_testdir, "beambeam", "particles_dump.bin")
assert os.path.exists(path_to_particle_data)
path_to_beam_elements_data = os.path.join(
path_to_testdir, "beambeam", "beam_elements.bin")
assert os.path.exists(path_to_beam_elements_data)
pb = CBuffer.fromfile(path_to_particle_data)
initial_particles = pb.get_object(0, cls=pyst.Particles)
track_pb = CBuffer()
particles = pyst.makeCopy(initial_particles, cbuffer=track_pb)
eb = CBuffer.fromfile(path_to_beam_elements_data)
num_beam_elements = eb.n_objects
ctx = st.st_CudaContext_create()
assert ctx != st.st_NullCudaContext
lattice = st.st_Buffer_new_mapped_on_cbuffer(eb)
assert lattice != st.st_NullBuffer
lattice_arg = st.st_CudaArgument_new(ctx)
class ParticlesSet(object):
element_types = {'Particles': Particles}
@property
def particles(self):
return self.cbuffer.get_objects()
def __init__(self, cbuffer=None):
if cbuffer is None:
self.cbuffer = CBuffer()
else:
self.cbuffer = cbuffer
for name, cls in self.element_types.items():
self.cbuffer.typeids[cls._typeid] = cls
def Particles(self, **nargs):
particles = Particles(cbuffer=self.cbuffer, **nargs)
return particles
@classmethod
def fromfile(cls, filename):
cbuffer = CBuffer.fromfile(filename)
return cls(cbuffer=cbuffer)
@classmethod
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.fromPySixTrack(
pysixtrack.Particles(**sixdump[kk].get_minimal_beam()), jj)
p.state[jj] = 1
p.at_element[jj] = elem_id
return self
def tofile(self, filename):
self.cbuffer.tofile(filename)
import sixtracklib as st
import cobjects
import numpy as np
from cobjects import CBuffer
def gaussian_dist(z, mu=0.0, sigma=1.0):
assert np.all(np.abs(sigma) > 0.0)
norm_z = (z - mu) / (np.sqrt(2.0) * sigma)
return 1.0 / np.sqrt(2 * np.pi * sigma) * np.exp(-norm_z * norm_z)
if __name__ == "__main__":
# -------------------------------------------------------------------------
# A) Create buffer for storing the line density datasets:
interpol_buffer = CBuffer()
# Prepare two line density datasets from gaussian distributions:
# NOTE: We keep track on the indices of the datasets -> this will be used
# later during assignment
# A.1) gaussian, sigma=1.0, 24 points, linear interpolation
lin_absc = np.linspace(-8.0, 8.0, num=24, dtype="float64")
sc_data0_idx = interpol_buffer.n_objects
sc_data0 = st.LineDensityProfileData(
cbuffer=interpol_buffer,
capacity=len(lin_absc),
z0=lin_absc[0],
dz=lin_absc[1] - lin_absc[0],
method="linear",
)
class ParticlesSet(object):
element_types = {"Particles": Particles}
@property
def particles(self):
return self.cbuffer.get_objects()
def __init__(self, cbuffer=None):
if cbuffer is None:
self.cbuffer = CBuffer()
else:
self.cbuffer = cbuffer
for name, cls in self.element_types.items():
self.cbuffer.typeids[cls._typeid] = cls
def Particles(self, **nargs):
particles = Particles(cbuffer=self.cbuffer, **nargs)
return particles
@classmethod
def fromfile(cls, filename):
cbuffer = CBuffer.fromfile(filename)
return cls(cbuffer=cbuffer)
@classmethod
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 to_file(self, filename):
self.cbuffer.tofile(filename)
c=CField(3,'real',length=4,default=1.1)
d=CField(4,'uint8',length=4,default=1,pointer=True)
e=CField(5,'real',length=2,default=1.3,pointer=True)
obj=MyObj()
print(obj)
obj._buffer.info()
obj.a=2
obj.b=3.5
obj.c[2]=1.5
obj.d=2
obj.e=[3,4]
print(obj)
b=CBuffer()
obj1=MyObj(cbuffer=b)
b.check_pointers()
obj2=MyObj(cbuffer=b)
b.check_pointers()
obj2.e=[3,4]
print(obj1)
print(obj2)
b.tofile('test.np')
import numpy as np
data=np.fromfile('test.np',dtype='uint64')
print(data)
c=CBuffer.fromfile('test.np')
obj0=c.get_object(MyObj,0)
class Elements(object):
element_types = {
"Cavity": Cavity,
"Drift": Drift,
"DriftExact": DriftExact,
"Multipole": Multipole,
"RFMultipole": RFMultipole,
"SRotation": SRotation,
"XYShift": XYShift,
"BeamMonitor": BeamMonitor,
"LimitRect": LimitRect,
"LimitEllipse": LimitEllipse,
"LimitRectEllipse": LimitRectEllipse,
"BeamBeam4D": BeamBeam4D,
"BeamBeam6D": BeamBeam6D,
"SCCoasting": SCCoasting,
"SCQGaussProfile": SCQGaussProfile,
"SCInterpolatedProfile": SCInterpolatedProfile,
"LimitRect": LimitRect,
"LimitEllipse": LimitEllipse,
"LimitRectEllipse": LimitRectEllipse,
"DipoleEdge": DipoleEdge,
# 'Line': Line,
# 'Solenoid': Solenoid,
# 'Wire': Wire,
}
@staticmethod
def add_element_type(cls, cls_name):
Elements.element_types.update({cls_name: cls})
def _mk_fun(self, buff, cls):
def fun(*args, **nargs):
# print(cls.__name__,nargs)
return cls(cbuffer=buff, **nargs)
return fun
@classmethod
def fromfile(cls, filename):
cbuffer = CBuffer.fromfile(filename)
return cls(cbuffer=cbuffer)
@classmethod
def from_line(cls, line):
self = cls()
return self.append_line(line)
def append_line(self, line):
for element in line.elements:
element_name = element.__class__.__name__
getattr(self, element_name)(**element.to_dict(keepextra=True))
return self
def to_file(self, filename):
self.cbuffer.tofile(filename)
return self
def __init__(self, cbuffer=None):
if cbuffer is None:
self.cbuffer = CBuffer()
else:
self.cbuffer = cbuffer
for name, cls in self.element_types.items():
setattr(self, name, self._mk_fun(self.cbuffer, cls))
self.cbuffer.typeids[cls._typeid] = cls
self._builder = self.gen_builder()
def gen_builder(self):
out = {}
for name, cls in self.element_types.items():
out[name] = getattr(self, name)
return out
def get_elements(self):
n = self.cbuffer.n_objects
return [self.cbuffer.get_object(i) for i in range(n)]
def get(self, objid):
return self.cbuffer.get_object(objid)
@classmethod
def from_mad(cls, seq, exact_drift=False):
# temporary
self = cls()
for label, element_name, element in madseq_to_generator(seq):
if exact_drift and element_name == "Drift":
element_name = "DriftExact"
getattr(self, element_name)(**element._asdict())
return self
slot_size = st.st_Buffer_get_slot_size(pb)
num_elem_by_elem_turns = 10
ret = st.st_OutputBuffer_calculate_output_buffer_params(
eb, particles, num_elem_by_elem_turns, ct.byref(num_objects),
ct.byref(num_slots), ct.byref(num_dataptrs), ct.byref(num_garbage),
slot_size)
assert (ret == 0)
assert (num_objects.value > 0)
assert (num_slots.value > 0)
assert (num_dataptrs.value > 0)
output_buffer = CBuffer(max_objects=num_objects.value,
max_slots=num_slots.value,
max_pointers=num_dataptrs.value,
max_garbage=num_garbage.value)
output_buffer_base_addr = output_buffer.base
output_buffer_size = output_buffer.size
saved_max_num_objects = output_buffer.max_objects
saved_max_num_slots = output_buffer.max_slots
saved_max_num_dataptrs = output_buffer.max_pointers
saved_max_num_garbage = output_buffer.max_garbage
assert (saved_max_num_objects >= num_objects.value)
assert (saved_max_num_slots >= num_slots.value)
assert (saved_max_num_dataptrs >= num_dataptrs.value)
assert (saved_max_num_garbage >= num_garbage.value)
mapped_output_buffer = st.st_Buffer_new_mapped_on_cbuffer(output_buffer)
class Elements(object):
element_types = {
'Cavity': Cavity,
'Drift': Drift,
'DriftExact': DriftExact,
'Multipole': Multipole,
'RFMultipole': RFMultipole,
'SRotation': SRotation,
'XYShift': XYShift,
'BeamBeam6D': BeamBeam6D,
'BeamBeam4D': BeamBeam4D,
'Line': Line,
'Monitor': Monitor,
}
def _mk_fun(self, buff, cls):
def fun(*args, **nargs):
# print(cls.__name__,nargs)
return cls(cbuffer=buff, **nargs)
return fun
@classmethod
def fromfile(cls, filename):
cbuffer = CBuffer.fromfile(filename)
return cls(cbuffer=cbuffer)
@classmethod
def fromline(cls, line):
self = cls()
for label, element_name, element in line:
getattr(self, element_name)(**element._asdict())
return self
def tofile(self, filename):
self.cbuffer.tofile(filename)
def __init__(self, cbuffer=None):
if cbuffer is None:
self.cbuffer = CBuffer()
else:
self.cbuffer = cbuffer
for name, cls in self.element_types.items():
setattr(self, name, self._mk_fun(self.cbuffer, cls))
self.cbuffer.typeids[cls._typeid] = cls
def gen_builder(self):
out = {}
for name, cls in self.element_types.items():
out[name] = getattr(self, name)
return out
def gen_builder_class(self):
class out:
pass
for name, cls in self.element_types.items():
setattr(out, name, getattr(self, name))
return out
def get_elements(self):
n = self.cbuffer.n_objects
return [self.cbuffer.get_object(i) for i in range(n)]
def get(self, objid):
return self.cbuffer.get_object(objid)
def set_monitors(self, offset=0):
monitorid = self.element_types['Monitor']._typeid
monitors = []
nmonitor = 0
for i in range(self.cbuffer.n_objects):
if self.cbuffer.get_object_typeid(i) == monitorid:
monitor = self.cbuffer.get_object(i)
monitor.ref = nmonitor + offset
monitors.append(monitor)
nmonitor += 1
return monitors
if __name__ == '__main__':
path_to_testdir = testlib.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_particle_data = os.path.join(
path_to_testdir, "beambeam", "particles_dump.bin")
assert(os.path.exists(path_to_particle_data))
path_to_beam_elements_data = os.path.join(
path_to_testdir, "beambeam", "beam_elements.bin")
assert(os.path.exists(path_to_beam_elements_data))
pb = CBuffer.fromfile(path_to_particle_data)
num_elem_by_elem_turns = 1
eb = CBuffer.fromfile(path_to_beam_elements_data)
until_turn_elem_by_elem = 1
until_turn_turn_by_turn = 5
until_turn = 100
skip_turns = 10
# ------------------------------------------------------------------------
initial_num_beam_elements = eb.n_objects
num_beam_monitors = pyst.beam_elements.append_beam_monitors_to_lattice(
eb, until_turn_elem_by_elem, until_turn_turn_by_turn,
until_turn, skip_turns)
def __init__(self):
self.cbuffer = CBuffer()
self.particles = []
from cobjects import CBuffer
b = CBuffer()
b.info()
b.new_object(24, 2, [])
b.info()
obj = b.get_object_buffer(0).view('uint64')
obj[0] = 31
b.reallocate(100, 100, 100, 100)
b.info()
b._test_cffilib()