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)
