def generate_real_bunch(self):
#beam parameters
intensity = 1.234e9
circumference = 111.
gamma = 20.1
#simulation parameters
macroparticlenumber = 2048
particlenumber_per_mp = intensity / macroparticlenumber
x = np.random.uniform(-1, 1, macroparticlenumber)
y = np.random.uniform(-1, 1, macroparticlenumber)
z = np.random.uniform(-1, 1, macroparticlenumber)
xp = np.random.uniform(-0.5, 0.5, macroparticlenumber)
yp = np.random.uniform(-0.5, 0.5, macroparticlenumber)
dp = np.random.uniform(-0.5, 0.5, macroparticlenumber)
coords_n_momenta_dict = {
'x': x,
'y': y,
'z': z,
'xp': xp,
'yp': yp,
'dp': dp
}
return Particles(macroparticlenumber, particlenumber_per_mp, e, m_p,
circumference, gamma, coords_n_momenta_dict)
def create_unif_bunch(self,
xmin=-1.,
xmax=1.,
ymin=-1.,
ymax=1.,
zmin=-1.,
zmax=1,
xpmin=-1.,
xpmax=1.,
ypmin=-1.,
ypmax=1.,
dpmin=-1.,
dpmax=1.):
x = np.random.uniform(xmin, xmax, self.macroparticlenumber)
y = np.random.uniform(ymin, ymax, self.macroparticlenumber)
z = np.random.uniform(zmin, zmax, self.macroparticlenumber)
xp = np.random.uniform(xpmin, xpmax, self.macroparticlenumber)
yp = np.random.uniform(ypmin, ypmax, self.macroparticlenumber)
dp = np.random.uniform(dpmin, dpmax, self.macroparticlenumber)
coords_n_momenta_dict = {
'x': x,
'y': y,
'z': z,
'xp': xp,
'yp': yp,
'dp': dp
}
return Particles(self.macroparticlenumber, self.particlenumber_per_mp,
e, m_p, self.circumference, self.gamma,
coords_n_momenta_dict)
def create_transverse_only_bunch(self):
x = np.random.uniform(-1, 1, self.macroparticlenumber)
y = np.random.uniform(-1, 1, self.macroparticlenumber)
xp = np.random.uniform(-0.5, 0.5, self.macroparticlenumber)
yp = np.random.uniform(-0.5, 0.5, self.macroparticlenumber)
coords_n_momenta_dict = {'x': x, 'y': y, 'xp': xp, 'yp': yp}
return Particles(self.macroparticlenumber, self.particlenumber_per_mp,
e, m_p, self.circumference, self.gamma,
coords_n_momenta_dict)
def generate(self):
''' Returns a particle object with the parameters specified
in the constructor of the Generator object
'''
coords = self._create_phase_space()
particles = Particles(self.macroparticlenumber,
self.intensity / self.macroparticlenumber,
self.charge,
self.mass,
self.circumference,
self.gamma,
coords_n_momenta_dict=coords)
self._linear_match_phase_space(particles)
return particles
def create_bunch(self, zmin=-1., zmax=1.):
z = np.linspace(zmin, zmax, num=self.macroparticlenumber)
y = np.copy(z)
x = np.copy(z)
xp = np.linspace(-0.5, 0.5, num=self.macroparticlenumber)
yp = np.copy(xp)
dp = np.copy(xp)
coords_n_momenta_dict = {
'x': x, 'y': y, 'z': z,
'xp': xp, 'yp': yp, 'dp': dp
}
return Particles(
self.macroparticlenumber, self.particlenumber_per_mp, e, m_p,
self.circumference, self.gamma, coords_n_momenta_dict
)
def create_all1_bunch(self):
x = np.ones(self.macroparticlenumber)
y = x.copy()
z = x.copy()
xp = x.copy()
yp = x.copy()
dp = x.copy()
coords_n_momenta_dict = {
'x': x, 'y': y, 'z': z,
'xp': xp, 'yp': yp, 'dp': dp
}
return Particles(
self.macroparticlenumber, 1, e, 1, #never mind the other params
1, 18., coords_n_momenta_dict
)
def create_bunch(self):
np.random.seed(0)
x = np.random.uniform(-1, 1, self.macroparticlenumber)
y = np.copy(x)
z = np.copy(x)
xp = np.random.uniform(-0.5, 0.5, self.macroparticlenumber)
yp = np.copy(xp)
dp = np.copy(xp)
coords_n_momenta_dict = {
'x': x,
'y': y,
'z': z,
'xp': xp,
'yp': yp,
'dp': dp
}
return Particles(self.macroparticlenumber, self.particlenumber_per_mp,
e, m_p, self.circumference, self.gamma,
coords_n_momenta_dict)
def create_all1_bunch(self):
np.random.seed(1)
x = np.ones(self.n_macroparticles)
y = x.copy()
z = x.copy()
xp = x.copy()
yp = x.copy()
dp = x.copy()
coords_n_momenta_dict = {
'x': x,
'y': y,
'z': z,
'xp': xp,
'yp': yp,
'dp': dp
}
return Particles(macroparticlenumber=len(x),
particlenumber_per_mp=1000,
charge=e,
mass=m_p,
circumference=self.circumference,
gamma=self.gamma,
coords_n_momenta_dict=coords_n_momenta_dict)
def create_bunch(self):
np.random.seed(0) #set seed to make results reproducible
x = np.random.uniform(-0.1, 0.1, self.macroparticlenumber)
y = np.random.uniform(-0.1, 0.1, self.macroparticlenumber)
z = x.copy()
xp = np.random.uniform(0., 1., self.macroparticlenumber)
yp = np.random.uniform(0., 1., self.macroparticlenumber)
dp = x.copy()
coords_n_momenta_dict = {
'x': x,
'y': y,
'z': z,
'xp': xp,
'yp': yp,
'dp': dp
}
return Particles(
self.macroparticlenumber,
1,
e,
1, #never mind the other params
1,
18.,
coords_n_momenta_dict)
def slice_a_bunch(this_bunch, z_cut, n_slices):
# Slice bunch if populated
if this_bunch.slice_info['slice_4_EC']:
bunch_center = this_bunch.slice_info['z_bin_center']
this_slicer = UniformBinSlicer(z_cuts=(bunch_center - z_cut,
bunch_center + z_cut),
n_slices=n_slices)
this_slices = this_bunch.extract_slices(this_slicer,
include_non_sliced='always')
sliced = this_slices[:-1]
unsliced = this_slices[-1]
if unsliced.slice_info != 'unsliced':
raise ValueError("Something went wrong")
for ss in sliced:
ss.slice_info['interact_with_EC'] = True
# Build head and tail slices
mask_head = unsliced.z >= sliced[-1].slice_info['z_bin_right']
mask_tail = unsliced.z <= sliced[0].slice_info['z_bin_left']
slice_tail = Particles(macroparticlenumber=np.sum(mask_tail),
particlenumber_per_mp=unsliced.particlenumber_per_mp,
charge=unsliced.charge,
mass=unsliced.mass, circumference=unsliced.circumference,
gamma=unsliced.gamma,
coords_n_momenta_dict={\
'x': np.atleast_1d(unsliced.x[mask_tail]),
'xp':np.atleast_1d(unsliced.xp[mask_tail]),
'y':np.atleast_1d(unsliced.y[mask_tail]),
'yp':np.atleast_1d(unsliced.yp[mask_tail]),
'z':np.atleast_1d(unsliced.z[mask_tail]),
'dp':np.atleast_1d(unsliced.dp[mask_tail])})
slice_tail.slice_info = {
'z_bin_center':
0.5 * (this_bunch.slice_info['z_bin_left'] +
sliced[0].slice_info['z_bin_left']),
'z_bin_left':
this_bunch.slice_info['z_bin_left'],
'z_bin_right':
sliced[0].slice_info['z_bin_left'],
'interact_with_EC':
False
}
slice_head = Particles(macroparticlenumber=np.sum(mask_head),
particlenumber_per_mp=unsliced.particlenumber_per_mp,
charge=unsliced.charge,
mass=unsliced.mass, circumference=unsliced.circumference,
gamma=unsliced.gamma,
coords_n_momenta_dict={\
'x': np.atleast_1d(unsliced.x[mask_head]),
'xp':np.atleast_1d(unsliced.xp[mask_head]),
'y':np.atleast_1d(unsliced.y[mask_head]),
'yp':np.atleast_1d(unsliced.yp[mask_head]),
'z':np.atleast_1d(unsliced.z[mask_head]),
'dp':np.atleast_1d(unsliced.dp[mask_head])})
slice_head.slice_info = {
'z_bin_center':
0.5 * (this_bunch.slice_info['z_bin_right'] +
sliced[-1].slice_info['z_bin_right']),
'z_bin_left':
sliced[-1].slice_info['z_bin_right'],
'z_bin_right':
this_bunch.slice_info['z_bin_right'],
'interact_with_EC':
False
}
list_slices_this_bunch = [slice_tail] + sliced + [slice_head]
else:
# Build a copy of the bunch
copy_this_bunch = Particles(
macroparticlenumber=this_bunch.macroparticlenumber,
particlenumber_per_mp=this_bunch.particlenumber_per_mp,
charge=this_bunch.charge,
mass=this_bunch.mass,
circumference=this_bunch.circumference,
gamma=this_bunch.gamma,
coords_n_momenta_dict=this_bunch.get_coords_n_momenta_dict())
copy_this_bunch.slice_info = {
kk: this_bunch.slice_info[kk]
for kk in list(this_bunch.slice_info.keys())
}
list_slices_this_bunch = [copy_this_bunch]
for i_sl, ss in enumerate(
list_slices_this_bunch[::-1]): # I want slice 0 to be at the head
ss.slice_info['info_parent_bunch'] = {
kk: this_bunch.slice_info[kk]
for kk in list(this_bunch.slice_info.keys())
}
ss.slice_info['i_slice'] = i_sl
ss.slice_info['N_slices_tot_bunch'] = len(list_slices_this_bunch)
return list_slices_this_bunch
def slice_a_bunch(this_bunch, z_cut, n_slices):
# Slice bunch if populated
if this_bunch.slice_info['slice_4_EC']:
bunch_center = this_bunch.slice_info['z_bin_center']
this_slicer = UniformBinSlicer(z_cuts=(bunch_center-z_cut, bunch_center+z_cut), n_slices=n_slices)
this_slices = this_bunch.extract_slices(this_slicer, include_non_sliced='always')
sliced = this_slices[:-1]
unsliced = this_slices[-1]
if unsliced.slice_info!='unsliced':
raise ValueError("Something went wrong")
for ss in sliced:
ss.slice_info['interact_with_EC'] = True
# Build head and tail slices
mask_head = unsliced.z>=sliced[-1].slice_info['z_bin_right']
mask_tail = unsliced.z<=sliced[0].slice_info['z_bin_left']
slice_tail = Particles(macroparticlenumber=np.sum(mask_tail),
particlenumber_per_mp=unsliced.particlenumber_per_mp,
charge=unsliced.charge,
mass=unsliced.mass, circumference=unsliced.circumference,
gamma=unsliced.gamma,
coords_n_momenta_dict={\
'x': np.atleast_1d(unsliced.x[mask_tail]),
'xp':np.atleast_1d(unsliced.xp[mask_tail]),
'y':np.atleast_1d(unsliced.y[mask_tail]),
'yp':np.atleast_1d(unsliced.yp[mask_tail]),
'z':np.atleast_1d(unsliced.z[mask_tail]),
'dp':np.atleast_1d(unsliced.dp[mask_tail])})
slice_tail.slice_info={
'z_bin_center': 0.5*(this_bunch.slice_info['z_bin_left']+sliced[0].slice_info['z_bin_left']),
'z_bin_left': this_bunch.slice_info['z_bin_left'],
'z_bin_right': sliced[0].slice_info['z_bin_left'],
'interact_with_EC': False}
slice_head = Particles(macroparticlenumber=np.sum(mask_head),
particlenumber_per_mp=unsliced.particlenumber_per_mp,
charge=unsliced.charge,
mass=unsliced.mass, circumference=unsliced.circumference,
gamma=unsliced.gamma,
coords_n_momenta_dict={\
'x': np.atleast_1d(unsliced.x[mask_head]),
'xp':np.atleast_1d(unsliced.xp[mask_head]),
'y':np.atleast_1d(unsliced.y[mask_head]),
'yp':np.atleast_1d(unsliced.yp[mask_head]),
'z':np.atleast_1d(unsliced.z[mask_head]),
'dp':np.atleast_1d(unsliced.dp[mask_head])})
slice_head.slice_info={
'z_bin_center': 0.5*(this_bunch.slice_info['z_bin_right']+sliced[-1].slice_info['z_bin_right']),
'z_bin_left': sliced[-1].slice_info['z_bin_right'],
'z_bin_right': this_bunch.slice_info['z_bin_right'],
'interact_with_EC': False}
list_slices_this_bunch = [slice_tail] + sliced + [slice_head]
else:
# Build a copy of the bunch
copy_this_bunch = Particles(macroparticlenumber=this_bunch.macroparticlenumber,
particlenumber_per_mp=this_bunch.particlenumber_per_mp,
charge=this_bunch.charge,
mass=this_bunch.mass, circumference=this_bunch.circumference,
gamma=this_bunch.gamma,
coords_n_momenta_dict=this_bunch.get_coords_n_momenta_dict())
copy_this_bunch.slice_info = {kk:this_bunch.slice_info[kk] for kk in this_bunch.slice_info.keys()}
list_slices_this_bunch = [copy_this_bunch]
for i_sl, ss in enumerate(list_slices_this_bunch[::-1]): # I want slice 0 to be at the head
ss.slice_info['info_parent_bunch'] = {kk: this_bunch.slice_info[kk] for kk in this_bunch.slice_info.keys()}
ss.slice_info['i_slice'] = i_sl
ss.slice_info['N_slices_tot_bunch'] = len(list_slices_this_bunch)
return list_slices_this_bunch
def buffer_2_beam(buf, mode='pickle'):
if buf[0]<0:
beam=None
else:
i_buf = 0
macroparticlenumber = int(buf[i_buf])
i_buf += 1
particlenumber_per_mp = buf[i_buf]
i_buf += 1
charge = buf[i_buf]
i_buf += 1
mass = buf[i_buf]
i_buf += 1
circumference = buf[i_buf]
i_buf += 1
gamma = buf[i_buf]
i_buf += 1
id_ = buf[i_buf:i_buf+macroparticlenumber].astype(np.int32)
i_buf += macroparticlenumber
x = buf[i_buf:i_buf+macroparticlenumber]
i_buf += macroparticlenumber
xp = buf[i_buf:i_buf+macroparticlenumber]
i_buf += macroparticlenumber
y = buf[i_buf:i_buf+macroparticlenumber]
i_buf += macroparticlenumber
yp = buf[i_buf:i_buf+macroparticlenumber]
i_buf += macroparticlenumber
z = buf[i_buf:i_buf+macroparticlenumber]
i_buf += macroparticlenumber
dp = buf[i_buf:i_buf+macroparticlenumber]
i_buf += macroparticlenumber
slice_info_buf_size = int(buf[i_buf])
i_buf += 1
slice_info_buf = buf[i_buf:i_buf+slice_info_buf_size]
i_buf += slice_info_buf_size
beam = Particles(macroparticlenumber=macroparticlenumber,
particlenumber_per_mp=particlenumber_per_mp, charge=charge,
mass=mass, circumference=circumference, gamma=gamma,
coords_n_momenta_dict={\
'x': np.atleast_1d(x),
'xp':np.atleast_1d(xp),
'y':np.atleast_1d(y),
'yp':np.atleast_1d(yp),
'z':np.atleast_1d(z),
'dp':np.atleast_1d(dp)})
beam.id = np.atleast_1d(id_)
if mode=='json':
si_int = slice_info_buf.astype(np.int)
si_str = ''.join(map(unichr, list(si_int)))
beam.slice_info = json.loads(si_str)
elif mode=='pickle':
# Get length in bytes
llrec = int(slice_info_buf[0])
s1back_padded = np.frombuffer(slice_info_buf[1:].tobytes(), dtype='S1')
s1back = s1back_padded[:llrec]
pss_rec = s1back.tobytes()
beam.slice_info = pickle.loads(pss_rec)
else:
raise ValueError('Unknown mode!')
# if slice_info_buf[0] < 0.:
# beam.slice_info = None
# elif slice_info_buf[0] == 0.:
# beam.slice_info = 'unsliced'
# else:
# beam.slice_info = {\
# 'z_bin_center':slice_info_buf[1] ,
# 'z_bin_right':slice_info_buf[2],
# 'z_bin_left':slice_info_buf[3]}
return beam
def buffer_2_beam(buf, mode='pickle'):
if buf[0]<0:
beam=None
else:
i_buf = 0
macroparticlenumber = int(buf[i_buf])
i_buf += 1
particlenumber_per_mp = buf[i_buf]
i_buf += 1
charge = buf[i_buf]
i_buf += 1
mass = buf[i_buf]
i_buf += 1
circumference = buf[i_buf]
i_buf += 1
gamma = buf[i_buf]
i_buf += 1
id_ = buf[i_buf:i_buf+macroparticlenumber].astype(np.int32)
i_buf += macroparticlenumber
x = buf[i_buf:i_buf+macroparticlenumber]
i_buf += macroparticlenumber
xp = buf[i_buf:i_buf+macroparticlenumber]
i_buf += macroparticlenumber
y = buf[i_buf:i_buf+macroparticlenumber]
i_buf += macroparticlenumber
yp = buf[i_buf:i_buf+macroparticlenumber]
i_buf += macroparticlenumber
z = buf[i_buf:i_buf+macroparticlenumber]
i_buf += macroparticlenumber
dp = buf[i_buf:i_buf+macroparticlenumber]
i_buf += macroparticlenumber
slice_info_buf_size = int(buf[i_buf])
i_buf += 1
slice_info_buf = buf[i_buf:i_buf+slice_info_buf_size]
i_buf += slice_info_buf_size
beam = Particles(macroparticlenumber=macroparticlenumber,
particlenumber_per_mp=particlenumber_per_mp, charge=charge,
mass=mass, circumference=circumference, gamma=gamma,
coords_n_momenta_dict={\
'x': np.atleast_1d(x),
'xp':np.atleast_1d(xp),
'y':np.atleast_1d(y),
'yp':np.atleast_1d(yp),
'z':np.atleast_1d(z),
'dp':np.atleast_1d(dp)})
beam.id = np.atleast_1d(id_)
if mode=='json':
si_int = slice_info_buf.astype(np.int)
si_str = ''.join(map(chr, list(si_int)))
beam.slice_info = json.loads(si_str)
elif mode=='pickle':
# Get length in bytes
llrec = int(slice_info_buf[0])
s1back_padded = np.frombuffer(slice_info_buf[1:].tobytes(), dtype='S1')
s1back = s1back_padded[:llrec]
pss_rec = s1back.tobytes()
beam.slice_info = pickle.loads(pss_rec)
else:
raise ValueError('Unknown mode!')
# if slice_info_buf[0] < 0.:
# beam.slice_info = None
# elif slice_info_buf[0] == 0.:
# beam.slice_info = 'unsliced'
# else:
# beam.slice_info = {\
# 'z_bin_center':slice_info_buf[1] ,
# 'z_bin_right':slice_info_buf[2],
# 'z_bin_left':slice_info_buf[3]}
return beam
def buffer_2_beam(buf):
if buf[0] < 0:
beam = None
else:
i_buf = 0
macroparticlenumber = int(buf[i_buf])
i_buf += 1
particlenumber_per_mp = buf[i_buf]
i_buf += 1
charge = buf[i_buf]
i_buf += 1
mass = buf[i_buf]
i_buf += 1
circumference = buf[i_buf]
i_buf += 1
gamma = buf[i_buf]
i_buf += 1
id_ = np.int_(buf[i_buf:i_buf + macroparticlenumber])
i_buf += macroparticlenumber
x = buf[i_buf:i_buf + macroparticlenumber]
i_buf += macroparticlenumber
xp = buf[i_buf:i_buf + macroparticlenumber]
i_buf += macroparticlenumber
y = buf[i_buf:i_buf + macroparticlenumber]
i_buf += macroparticlenumber
yp = buf[i_buf:i_buf + macroparticlenumber]
i_buf += macroparticlenumber
z = buf[i_buf:i_buf + macroparticlenumber]
i_buf += macroparticlenumber
dp = buf[i_buf:i_buf + macroparticlenumber]
i_buf += macroparticlenumber
slice_info_buf = buf[i_buf:i_buf + 4]
i_buf += 4
beam = Particles(macroparticlenumber=macroparticlenumber,
particlenumber_per_mp=particlenumber_per_mp, charge=charge,
mass=mass, circumference=circumference, gamma=gamma,
coords_n_momenta_dict={\
'x': np.atleast_1d(x),
'xp':np.atleast_1d(xp),
'y':np.atleast_1d(y),
'yp':np.atleast_1d(yp),
'z':np.atleast_1d(z),
'dp':np.atleast_1d(dp)})
beam.id = np.atleast_1d(id_)
if slice_info_buf[0] < 0.:
beam.slice_info = None
elif slice_info_buf[0] == 0.:
beam.slice_info = 'unsliced'
else:
beam.slice_info = {\
'z_bin_center':slice_info_buf[1] ,
'z_bin_right':slice_info_buf[2],
'z_bin_left':slice_info_buf[3]}
return beam
def buffer_2_beam(buf):
if buf[0]<0:
beam=None
else:
i_buf = 0
macroparticlenumber = int(buf[i_buf])
i_buf += 1
particlenumber_per_mp = buf[i_buf]
i_buf += 1
charge = buf[i_buf]
i_buf += 1
mass = buf[i_buf]
i_buf += 1
circumference = buf[i_buf]
i_buf += 1
gamma = buf[i_buf]
i_buf += 1
id_ = np.int_(buf[i_buf:i_buf+macroparticlenumber])
i_buf += macroparticlenumber
x = buf[i_buf:i_buf+macroparticlenumber]
i_buf += macroparticlenumber
xp = buf[i_buf:i_buf+macroparticlenumber]
i_buf += macroparticlenumber
y = buf[i_buf:i_buf+macroparticlenumber]
i_buf += macroparticlenumber
yp = buf[i_buf:i_buf+macroparticlenumber]
i_buf += macroparticlenumber
z = buf[i_buf:i_buf+macroparticlenumber]
i_buf += macroparticlenumber
dp = buf[i_buf:i_buf+macroparticlenumber]
i_buf += macroparticlenumber
slice_info_buf = buf[i_buf:i_buf+4]
i_buf += 4
beam = Particles(macroparticlenumber=macroparticlenumber,
particlenumber_per_mp=particlenumber_per_mp, charge=charge,
mass=mass, circumference=circumference, gamma=gamma,
coords_n_momenta_dict={\
'x': np.atleast_1d(x),
'xp':np.atleast_1d(xp),
'y':np.atleast_1d(y),
'yp':np.atleast_1d(yp),
'z':np.atleast_1d(z),
'dp':np.atleast_1d(dp)})
beam.id = np.atleast_1d(id_)
if slice_info_buf[0] < 0.:
beam.slice_info = None
elif slice_info_buf[0] == 0.:
beam.slice_info = 'unsliced'
else:
beam.slice_info = {\
'z_bin_center':slice_info_buf[1] ,
'z_bin_right':slice_info_buf[2],
'z_bin_left':slice_info_buf[3]}
return beam