Skip to content

Commit

Permalink
Merge tag 'v1.10.4'
Browse files Browse the repository at this point in the history 
PyHEADTAIL v1.10.4
  • Loading branch information
@aoeftiger
aoeftiger committed Jul 5, 2016
2 parents 6ed1479 + 38ae675 commit cec59ee
Show file tree
Hide file tree
Showing 9 changed files with 171 additions and 153 deletions.
102 changes: 45 additions & 57 deletions field_maps/Transverse_Efield_map.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,76 +4,64 @@
from PyHEADTAIL.particles.slicing import UniformBinSlicer
from PyPIC.PyPIC_Scatter_Gather import PyPIC_Scatter_Gather

from . import Element

class Transverse_Efield_map(Element):
def __init__(self, xg, yg, Ex, Ey, L_interaction, slicer,
flag_clean_slices=False, wrt_slice_centroid=False,
x_beam_offset=0., y_beam_offset=0., verbose=False,
*args, **kwargs):

class Transverse_Efield_map(object):
def __init__(self, xg, yg, Ex, Ey, n_slices, z_cut,
L_interaction, flag_clean_slices = False, wrt_slice_centroid = False,
x_beam_offset = 0., y_beam_offset = 0.):

if type(z_cut) is float:
z_cuts = (-z_cut, z_cut)
elif type(z_cut) is tuple:
z_cuts = z_cut
else:
raise ValueError('Type not recognized!')
self.slicer = slicer
self.L_interaction = L_interaction
self.flag_clean_slices = flag_clean_slices
self.wrt_slice_centroid = wrt_slice_centroid

self.slicer = UniformBinSlicer(n_slices = n_slices, z_cuts=z_cuts)
self.L_interaction = L_interaction
self.flag_clean_slices = flag_clean_slices
self.wrt_slice_centroid = wrt_slice_centroid
self.Ex = Ex
self.Ey = Ey
self.pic = PyPIC_Scatter_Gather(xg=xg, yg=yg, verbose=verbose)

self.Ex = Ex
self.Ey = Ey
self.pic = PyPIC_Scatter_Gather(xg=xg, yg=yg)

self.x_beam_offset = x_beam_offset
self.y_beam_offset = y_beam_offset
self.x_beam_offset = x_beam_offset
self.y_beam_offset = y_beam_offset

def get_beam_x(self, beam):
return beam.x

def get_beam_x(self, beam):
return beam.x
def get_beam_y(self, beam):
return beam.y

def get_beam_y(self, beam):
return beam.y
def track(self, beam):
if self.flag_clean_slices:
beam.clean_slices()

# @profile
def track(self, beam):
slices = beam.get_slices(self.slicer)

for sid in xrange(slices.n_slices-1, -1, -1):

if self.flag_clean_slices:
beam.clean_slices()

slices = beam.get_slices(self.slicer)


for i in xrange(slices.n_slices-1, -1, -1):

# select particles in the slice
ix = slices.particle_indices_of_slice(i)

# slice size and time step
dz = (slices.z_bins[i + 1] - slices.z_bins[i])
#dt = dz / (beam.beta * c)

x = self.get_beam_x(beam)[ix]
y = self.get_beam_y(beam)[ix]

self.pic.efx = np.squeeze(self.Ex[i,:,:])
self.pic.efy = np.squeeze(self.Ey[i,:,:])
if self.wrt_slice_centroid:
Ex_sc_p, Ey_sc_p = self.pic.gather(x-np.mean(x)+self.x_beam_offset, y-np.mean(y)+self.y_beam_offset)
else:
Ex_sc_p, Ey_sc_p = self.pic.gather(x+self.x_beam_offset, y+self.y_beam_offset)

## kick beam particles
fact_kick = beam.charge/(beam.mass*beam.beta*beam.beta*beam.gamma*c*c)*self.L_interaction
beam.xp[ix]+=fact_kick*Ex_sc_p
beam.yp[ix]+=fact_kick*Ey_sc_p

# select particles in the slice
pid = slices.particle_indices_of_slice(sid)

# slice size
dz = (slices.z_bins[sid + 1] - slices.z_bins[sid])

x = self.get_beam_x(beam)[pid]
y = self.get_beam_y(beam)[pid]

self.pic.efx = np.squeeze(self.Ex[sid,:,:])
self.pic.efy = np.squeeze(self.Ey[sid,:,:])

centroid_x = 0
centroid_y = 0
if self.wrt_slice_centroid:
centroid_x = np.mean(x)
centroid_y = np.mean(y)

Ex_sc_p, Ey_sc_p = self.pic.gather(
x - centroid_x + self.x_beam_offset,
y - centroid_y + self.y_beam_offset
)

# kick beam particles
fact_kick = beam.charge / (beam.p0*beam.beta*c) * self.L_interaction
beam.xp[pid] += fact_kick * Ex_sc_p
beam.yp[pid] += fact_kick * Ey_sc_p
1 change: 1 addition & 0 deletions field_maps/__init__.py
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
from .. import Element
9 changes: 7 additions & 2 deletions particles/generators.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -61,6 +61,11 @@ def transverse_linear_matcher(alpha, beta, dispersion=None):
phase space
Returns: Matcher(closure) taking two parameters: coords and direction
'''
# if dispersion and alpha:
# raise NotImplementedError('Transverse phase space matching: for '
# 'alpha != 0 we need to match including the '
# 'D\' (dispersion derivative). This is '
# 'currently not implemented.')
sqrt = np.sqrt
# build the M matrix: only depends on twiss parameters for the
# special case of alpha0=0, beta0=1 and phi = 0 (=2pi)
Expand All @@ -86,9 +91,9 @@ def _transverse_linear_matcher(beam, direction):
momentum_coords = getattr(beam, direction[1])
space_coords = (M[0, 0]*space_coords + # copy if not using +=, *=..
M[0, 1]*momentum_coords)
momentum_coords = (M[1 ,0]*space_coords_copy +
momentum_coords = (M[1, 0]*space_coords_copy +
M[1, 1]*momentum_coords)
# add dispersion effects, raise exception of coords['dp'] inexistent
# add dispersion effects, raise exception if coords['dp'] inexistent
if dispersion:
try:
space_coords += dispersion * getattr(beam, 'dp')
Expand Down
44 changes: 22 additions & 22 deletions particles/particles.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -104,7 +104,7 @@ def p0(self):
return self._p0
@p0.setter
def p0(self, value):
self.gamma = value / (self.mass * self.beta * c)
self.gamma = np.sqrt(1 + (value / (self.mass * c))**2)

@property
def z_beamframe(self):
Expand Down Expand Up @@ -152,63 +152,63 @@ def get_slices(self, slicer, *args, **kwargs):


def extract_slices(self, slicer, include_non_sliced='if_any', *args, **kwargs):
'''Return a list Particles object with the different slices.
'''Return a list Particles object with the different slices.
The last element of the list contains particles not assigned to any slice.
include_non_sliced : {'always', 'never', 'if_any'}, optional
'always':
extra element in the list with particles not belonging to any slice
extra element in the list with particles not belonging to any slice
is always inserted (it can be empty).
'never':
extra element in the list with particles not belonging to any slice
extra element in the list with particles not belonging to any slice
is never inserted.
'if_any':
extra element in the list with particles not belonging to any slice
is inserted only if such particles exist.
extra element in the list with particles not belonging to any slice
is inserted only if such particles exist.
'''

if include_non_sliced not in ['if_any', 'always', 'never']:
raise ValueError("include_non_sliced=%s is not valid!\n"%include_non_sliced+
"Possible values are {'always', 'never', 'if_any'}" )

slices = self.get_slices(slicer, *args, **kwargs)
self_coords_n_momenta_dict = self.get_coords_n_momenta_dict()
slice_object_list = []

for i_sl in xrange(slices.n_slices):

ix = slices.particle_indices_of_slice(i_sl)
macroparticlenumber = len(ix)
slice_object = Particles(macroparticlenumber=macroparticlenumber,

slice_object = Particles(macroparticlenumber=macroparticlenumber,
particlenumber_per_mp=self.particlenumber_per_mp, charge=self.charge,
mass=self.mass, circumference=self.circumference, gamma=self.gamma, coords_n_momenta_dict={})

for coord in self_coords_n_momenta_dict.keys():
slice_object.update({coord: self_coords_n_momenta_dict[coord][ix]})

slice_object.id[:] = self.id[ix]

slice_object.slice_info = {\
'z_bin_center': slices.z_centers[i_sl],\
'z_bin_right':slices.z_bins[i_sl+1],\
'z_bin_left':slices.z_bins[i_sl]}

slice_object_list.append(slice_object)

# handle unsliced
if include_non_sliced is not 'never':
ix = slices.particles_outside_cuts
if len(ix)>0 or include_non_sliced is 'always':
slice_object = Particles(macroparticlenumber=len(ix),
slice_object = Particles(macroparticlenumber=len(ix),
particlenumber_per_mp=self.particlenumber_per_mp, charge=self.charge,
mass=self.mass, circumference=self.circumference, gamma=self.gamma, coords_n_momenta_dict={})
for coord in self_coords_n_momenta_dict.keys():
slice_object.update({coord: self_coords_n_momenta_dict[coord][ix]})
slice_object.id[:] = self.id[ix]
slice_object.slice_info = 'unsliced'
slice_object.slice_info = 'unsliced'
slice_object_list.append(slice_object)

return slice_object_list

def clean_slices(self):
Expand Down Expand Up @@ -282,7 +282,7 @@ def __add__(self, other):
result.update({coord: np.concatenate((self_coords_n_momenta_dict[coord].copy(), other_coords_n_momenta_dict[coord].copy()))})

result.id = np.concatenate((self.id.copy(), other.id.copy()))

return result

def __radd__(self, other):
Expand Down
63 changes: 31 additions & 32 deletions spacecharge/transverse_spacecharge.py
View file
Original file line number Diff line number Diff line change
@@ -1,10 +1,13 @@
import numpy as np
from scipy.constants import c
from __future__ import division

class TransverseSpaceCharge(object):
def __init__(self, L_interaction, slicer, pyPICsolver, flag_clean_slices = False):
from . import Element

import numpy as np
from scipy.constants import c

class TransverseSpaceCharge(Element):
def __init__(self, L_interaction, slicer, pyPICsolver,
flag_clean_slices=False, *args, **kwargs):
self.slicer = slicer
self.L_interaction = L_interaction
self.pyPICsolver = pyPICsolver
Expand All @@ -20,10 +23,7 @@ def get_beam_x(self, beam):
def get_beam_y(self, beam):
return beam.y

# @profile
def track(self, beam):


if self.save_distributions_last_track:
self.rho_last_track = []

Expand All @@ -33,54 +33,53 @@ def track(self, beam):
self.Ey_last_track = []

if hasattr(beam.particlenumber_per_mp, '__iter__'):
raise ValueError('spacecharge module assumes same size for all beam MPs')
raise ValueError('spacecharge module assumes same number of charges'
' for all beam macroparticles!')

if self.flag_clean_slices:
beam.clean_slices()

slices = beam.get_slices(self.slicer)

for i in xrange(slices.n_slices-1, -1, -1):
for sid in xrange(slices.n_slices-1, -1, -1):

# select particles in the slice
ix = slices.particle_indices_of_slice(i)
pid = slices.particle_indices_of_slice(sid)

# slice size and time step
dz = (slices.z_bins[i + 1] - slices.z_bins[i])
dt = dz / (beam.beta * c)
# slice size
dz = slices.slice_widths[sid]

# beam field
x_mp = self.get_beam_x(beam)[ix]
y_mp = self.get_beam_y(beam)[ix]
n_mp = x_mp*0.+beam.particlenumber_per_mp/dz#they have to become cylinders
N_mp = slices.n_macroparticles_per_slice[i]
x = self.get_beam_x(beam)[pid]
y = self.get_beam_y(beam)[pid]

#compute beam field (it assumes electrons!)
self.pyPICsolver.scatter_and_solve(x_mp, y_mp, n_mp, beam.charge)
# the particles have to become cylinders:
n_mp = np.zeros_like(x) + beam.particlenumber_per_mp / dz

#gather to beam particles
Ex_n_beam, Ey_n_beam = self.pyPICsolver.gather(x_mp, y_mp)
# compute beam field (it assumes electrons!)
self.pyPICsolver.scatter_and_solve(x, y, n_mp, beam.charge)

# interpolate beam field to particles
Ex_n_beam, Ey_n_beam = self.pyPICsolver.gather(x, y)

# go to actual beam particles and add relativistic (B field) factor
Ex_n_beam = Ex_n_beam/beam.gamma/beam.gamma
Ey_n_beam = Ey_n_beam/beam.gamma/beam.gamma
Ex_n_beam = Ex_n_beam / beam.gamma / beam.gamma
Ey_n_beam = Ey_n_beam / beam.gamma / beam.gamma


## kick beam particles
fact_kick = beam.charge/(beam.mass*beam.beta*beam.beta*beam.gamma*c*c)*self.L_interaction
beam.xp[ix]+=fact_kick*Ex_n_beam
beam.yp[ix]+=fact_kick*Ex_n_beam
# kick beam particles
fact_kick = beam.charge / (beam.p0*beam.beta*c) * self.L_interaction
beam.xp[pid] += fact_kick * Ex_n_beam
beam.yp[pid] += fact_kick * Ey_n_beam

if self.save_distributions_last_track:
self.rho_last_track.append(self.pyPICsolver.rho.copy())
#print 'Here'

if self.save_potential_and_field:
self.phi_last_track.append(self.pyPICsolver.phi.copy())
self.Ex_last_track.append(self.pyPICsolver.efx.copy()/beam.gamma/beam.gamma)
self.Ey_last_track.append(self.pyPICsolver.efy.copy()/beam.gamma/beam.gamma)

self.Ex_last_track.append(
self.pyPICsolver.efx.copy()/beam.gamma/beam.gamma)
self.Ey_last_track.append(
self.pyPICsolver.efy.copy()/beam.gamma/beam.gamma)

if self.save_distributions_last_track:
self.rho_last_track = np.array(self.rho_last_track[::-1])
Expand Down
Loading

0 comments on commit cec59ee

Please sign in to comment.