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')