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)
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
method="cubic", ) sc_data1.values[:] = gaussian_dist(cub_absc, mu=0.0, sigma=2.0) sc_data1.prepare_interpolation() # A.3) optional: demonstrate that interpolation works -> # set plot_data = True! plot_data = False if plot_data: from matplotlib import pyplot as plt z_absc = np.linspace(-8.0, +8.0, num=512, dtype="float64") y_exact1 = gaussian_dist(z_absc, mu=0.0, sigma=1.0) sc_data0 = interpol_buffer.get_object(sc_data0_idx) y_interp_lin = np.array([sc_data0.interpol(zz) for zz in z_absc]) y_exact2 = gaussian_dist(z_absc, mu=0.0, sigma=2.0) sc_data1 = interpol_buffer.get_object(sc_data1_idx) y_interp_cub = np.array([sc_data1.interpol(zz) for zz in z_absc]) plt.figure() plt.subplot(211) plt.plot( lin_absc, sc_data0.values, "bo", z_absc, y_exact1, "k",
st.st_CudaArgument_get_arg_buffer(lattice_arg), line_middle, line_end, ct.c_bool(True), num_blocks, threads_per_block) success=st.st_CudaArgument_receive_buffer(particles_arg, pbuffer) assert success == 0 cmp_pb=CBuffer() cmp_particles=pyst.makeCopy(initial_particles, cbuffer=cmp_pb) cmp_pbuffer=st.st_Buffer_new_mapped_on_cbuffer(cmp_pb) assert cmp_pbuffer != st.st_NullBuffer st.st_Track_all_particles_until_turn( st.st_Particles_buffer_get_particles(cmp_pbuffer, 0), lattice, ct.c_int64(1)) assert pyst.compareParticlesDifference( track_pb.get_object(0, cls=pyst.Particles), cmp_pb.get_object(0, cls=pyst.Particles), abs_treshold=2e-14) == 0 st.st_CudaArgument_delete(particles_arg) st.st_CudaArgument_delete(lattice_arg) st.st_CudaContext_delete(ctx) st.st_Buffer_delete(pbuffer) st.st_Buffer_delete(lattice) st.st_Buffer_delete(cmp_pbuffer) sys.exit(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
st.st_CudaArgument_get_arg_buffer(lattice_arg), line_middle, line_end, ct.c_bool(True), num_blocks, threads_per_block) success = st.st_CudaArgument_receive_buffer(particles_arg, pbuffer) assert success == 0 cmp_pb = CBuffer() cmp_particles = pyst.makeCopy(initial_particles, cbuffer=cmp_pb) cmp_pbuffer = st.st_Buffer_new_mapped_on_cbuffer(cmp_pb) assert cmp_pbuffer != st.st_NullBuffer st.st_Track_all_particles_until_turn( st.st_Particles_buffer_get_particles(cmp_pbuffer, 0), lattice, ct.c_int64(1)) assert pyst.compareParticlesDifference( track_pb.get_object(0, cls=pyst.Particles), cmp_pb.get_object(0, cls=pyst.Particles), abs_treshold=2e-14) == 0 st.st_CudaArgument_delete(particles_arg) st.st_CudaArgument_delete(lattice_arg) st.st_CudaContext_delete(ctx) st.st_Buffer_delete(pbuffer) st.st_Buffer_delete(lattice) st.st_Buffer_delete(cmp_pbuffer) sys.exit(0)