diff --git a/python/pysixtracklib/beam_elements.py b/python/pysixtracklib/beam_elements.py index eb8088f9..2dfb54e5 100644 --- a/python/pysixtracklib/beam_elements.py +++ b/python/pysixtracklib/beam_elements.py @@ -5,107 +5,192 @@ import numpy as np -class Drift( CObject ): - _typeid = 2 - length = CField( 0, 'real', default=0.0, alignment=8 ) +class Drift(CObject): + _typeid = 2 + length = CField(0, 'real', default=0.0, alignment=8) -class DriftExact( CObject ): + +class DriftExact(CObject): _typeid = 3 - length = CField( 0, 'real', default=0.0, alignment=8 ) + length = CField(0, 'real', default=0.0, alignment=8) + -class MultiPole( CObject ): +class Multipole(CObject): _typeid = 4 - order = CField( 0, 'int64', default=0, alignment=8 ) - length = CField( 1, 'real', default=0.0, alignment=8 ) - hxl = CField( 2, 'real', default=0.0, alignment=8 ) - hyl = CField( 3, 'real', default=0.0, alignment=8 ) - bal = CField( 4, 'real', default=0.0, - length='2 * order + 2', pointer=True, alignment=8 ) - - def _factorial( self, x ): - if not isinstance( x, int ): + order = CField(0, 'int64', default=0, alignment=8) + length = CField(1, 'real', default=0.0, alignment=8) + hxl = CField(2, 'real', default=0.0, alignment=8) + hyl = CField(3, 'real', default=0.0, alignment=8) + bal = CField(4, 'real', default=0.0, + length='2 * order + 2', pointer=True, alignment=8) + + def _factorial(self, x): + if not isinstance(x, int): return 0 - return ( x > 0 ) and ( x * self._factorial( x - 1 ) ) or 1 + return (x > 0) and (x * self._factorial(x - 1)) or 1 - def __init__( self, order=None, knl=None, ksl=None, bal=None, **kwargs ): + def __init__(self, order=None, knl=None, ksl=None, bal=None, **kwargs): - if bal is None and ( not( knl is None ) or not( ksl is None ) ): - if knl is None: knl = [] - if ksl is None: ksl = [] - if order is None: order = 0 + if bal is None and (not(knl is None) or not(ksl is None)): + if knl is None: + knl = [] + if ksl is None: + ksl = [] + if order is None: + order = 0 - n = max( ( order + 1 ), max( len( knl ), len( ksl ) ) ) - _knl = np.array( knl ) - nknl = np.zeros( n, dtype=_knl.dtype ) + n = max((order + 1), max(len(knl), len(ksl))) + _knl = np.array(knl) + nknl = np.zeros(n, dtype=_knl.dtype) nknl[:len(knl)] = knl knl = nknl - del( _knl ) + del(_knl) - _ksl = np.array( ksl ) - nksl = np.zeros( n, dtype=_ksl.dtype ) + _ksl = np.array(ksl) + nksl = np.zeros(n, dtype=_ksl.dtype) nksl[:len(ksl)] = ksl ksl = nksl - del( _ksl ) + del(_ksl) - assert( n > 0 ) + assert(n > 0) order = n - 1 - bal = np.zeros( 2 * order + 2 ) - assert( len( knl ) == len( ksl ) ) + bal = np.zeros(2 * order + 2) + assert(len(knl) == len(ksl)) - for ii in range( 0, len( knl ) ): - inv_factorial = 1.0 / float( self._factorial( ii ) ) + for ii in range(0, len(knl)): + inv_factorial = 1.0 / float(self._factorial(ii)) jj = 2 * ii - bal[ jj ] = knl[ ii ] * inv_factorial - bal[ jj + 1 ] = ksl[ ii ] * inv_factorial + bal[jj] = knl[ii] * inv_factorial + bal[jj + 1] = ksl[ii] * inv_factorial - elif not( bal is None ) and bal and \ - len( bal ) > 2 and ( ( len( bal ) % 2 ) == 0 ): + elif not(bal is None) and bal and \ + len(bal) > 2 and ((len(bal) % 2) == 0): - order = ( len( bal ) - 2 ) / 2 - assert( order > 0 ) + order = (len(bal) - 2) / 2 + assert(order > 0) elif bal is None and knl is None and ksl is None and \ - not( order is None ) and order > 0: - bal = np.zeros( 2 * order + 2 ) - + not(order is None) and order > 0: + bal = np.zeros(2 * order + 2) - if not( bal is None or order is None ): - CObject.__init__( self, bal=bal, order=order, **kwargs ) + if not(bal is None or order is None): + CObject.__init__(self, bal=bal, order=order, **kwargs) else: - CObject.__init__( self, bal=[], order=0, **kwargs ) + CObject.__init__(self, bal=[], order=0, **kwargs) -class Cavity( CObject ): +class Cavity(CObject): _typeid = 5 - voltage = CField( 0, 'real', default=0.0, alignment=8 ) - frequency = CField( 1, 'real', default=0.0, alignment=8 ) - lag = CField( 2, 'real', default=0.0, alignment=8 ) + voltage = CField(0, 'real', default=0.0, alignment=8) + frequency = CField(1, 'real', default=0.0, alignment=8) + lag = CField(2, 'real', default=0.0, alignment=8) -class XYShift( CObject ): + +class XYShift(CObject): _typeid = 6 - dx = CField( 0, 'real', default=0.0, alignment=8 ) - dy = CField( 1, 'real', default=0.0, alignment=8 ) + dx = CField(0, 'real', default=0.0, alignment=8) + dy = CField(1, 'real', default=0.0, alignment=8) + -class SRotation( CObject ): +class SRotation(CObject): _typeid = 7 - cos_z = CField( 0, 'real', default=1.0, alignment=8 ) - sin_z = CField( 1, 'real', default=0.0, alignment=8 ) + cos_z = CField(0, 'real', default=1.0, alignment=8) + sin_z = CField(1, 'real', default=0.0, alignment=8) + + def __init__(self, angle=0, **nargs): + anglerad = angle/180*np.pi + cos_z = np.cos(anglerad) + sin_z = np.sin(anglerad) + CObject.__init__(self, + cos_z=cos_z, sin_z=sin_z, **nargs) -class BeamBeam4D( CObject ): + +class BeamBeam4D(CObject): _typeid = 8 - size = CField( 0, 'uint64', default=0, alignment=8 ) - data = CField( 1, 'real', default=0.0, - length='size', pointer=True, alignment=8 ) + size = CField(0, 'uint64', const=True, default=0) + data = CField(1, 'float64', default=0.0, + length='size', pointer=True) + + def __init__(self, data=None, **kwargs): + if data is None: + slots = ('q_part', 'N_part', 'sigma_x', 'sigma_y', 'beta_s', + 'min_sigma_diff', 'Delta_x', 'Delta_y', 'Dpx_sub', 'Dpy_sub', 'enabled') + data = [kwargs[ss] for ss in slots] + CObject.__init__(self, size=len(data), data=data, **kwargs) + else: + CObject.__init__(self, **kwargs) - def __init__(self, data, **kwargs): - CObject.__init__( self, size=len(data), data=data, **kwargs) -class BeamBeam6D( CObject ): +class BeamBeam6D(CObject): _typeid = 9 - size = CField( 0, 'uint64', default=0, alignment=8 ) - data = CField( 1, 'real', default=0.0, - length='size', pointer=True, alignment=8 ) - - def __init__(self, data, **kwargs): - CObject.__init__( self, size=len(data), data=data, **kwargs) + size = CField(0, 'uint64', const=True, default=0) + data = CField(1, 'float64', default=0.0, + length='size', pointer=True) + + def __init__(self, data=None, **kwargs): + if data is None: + import pysixtrack + data = pysixtrack.BB6Ddata.BB6D_init( + **{kk: kwargs[kk] for kk in kwargs.keys() if kk != 'cbuffer'}).tobuffer() + CObject.__init__(self, size=len(data), data=data, **kwargs) + else: + CObject.__init__(self, **kwargs) + + +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 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) diff --git a/python/pysixtracklib/export_to_cobjects.py b/python/pysixtracklib/export_to_cobjects.py index 3db4f696..06a05969 100644 --- a/python/pysixtracklib/export_to_cobjects.py +++ b/python/pysixtracklib/export_to_cobjects.py @@ -10,7 +10,7 @@ from cobjects import CBuffer, CObject, CField -from .beam_elements import Drift, DriftExact, MultiPole, Cavity, XYShift, SRotation +from .beam_elements import Drift, DriftExact, Multipole, Cavity, XYShift, SRotation from .beam_elements import BeamBeam4D, BeamBeam6D from .particles import Particles as IOParticles @@ -39,7 +39,7 @@ def line2cobject( line, cbuffer=None ): e = DriftExact( cbuffer=cbuffer, length=elem.length ) elif elem_type == 'Multipole': - e = MultiPole( cbuffer=cbuffer, knl=elem.knl, ksl=elem.ksl, + e = Multipole( cbuffer=cbuffer, knl=elem.knl, ksl=elem.ksl, length=elem.length, hxl=elem.hxl, hyl=elem.hyl ) elif elem_type == 'XYShift': diff --git a/python/pysixtracklib/particles.py b/python/pysixtracklib/particles.py index 5061484a..d0298339 100644 --- a/python/pysixtracklib/particles.py +++ b/python/pysixtracklib/particles.py @@ -167,3 +167,24 @@ def compareParticlesDifference( lhs, rhs, abs_treshold=None ): break return cmp_result + + +class ParticlesSet(object): + def __init__(self): + self.cbuffer=CBuffer() + self.particles=[] + def Particles(self,**nargs): + particles=Particles(cbuffer=self.cbuffer,**nargs) + self.particles.append(particles) + return particles + + @classmethod + def fromfile(cls, filename): + cbuffer = CBuffer.fromfile(filename) + return cls(cbuffer=cbuffer) + + def tofile(self, filename): + self.cbuffer.tofile(filename) + + +