Skip to content

Commit

Permalink
Merge pull request #16 from rdemaria/master
Browse files Browse the repository at this point in the history 
Merge from Riccardo
  • Loading branch information
@giadarol
giadarol authored Nov 16, 2018
2 parents 4137836 + e69959d commit 3b9084f
Show file tree
Hide file tree
Showing 3 changed files with 176 additions and 70 deletions.
221 changes: 153 additions & 68 deletions python/pysixtracklib/beam_elements.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)
4 changes: 2 additions & 2 deletions python/pysixtracklib/export_to_cobjects.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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':
Expand Down
21 changes: 21 additions & 0 deletions python/pysixtracklib/particles.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)



0 comments on commit 3b9084f

Please sign in to comment.