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)
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)
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 )
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)
ffi,lib,cobj=obj1._cdebug()
#c=CBuffer.from_file('test_buffer_common.bin')
