Skip to content

Commit

Permalink
Merge branch 'feature/PyEC4PyHT_slice_function'
Browse files Browse the repository at this point in the history
  • Loading branch information
@giadarol
giadarol committed Feb 10, 2016
2 parents 15b320f + 9f09e5b commit 892dd36
Showing 1 changed file with 153 additions and 136 deletions.
289 changes: 153 additions & 136 deletions PyEC4PyHT.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -254,8 +254,158 @@ def track(self, beam):
if self.N_tracks>0:
print 'Warning: Track skipped because track_only_first_time is True.'
return

self._reinitialize()

if hasattr(beam.particlenumber_per_mp, '__iter__'):
raise ValueError('ecloud module assumes same size for all beam MPs')

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])

self._track_single_slice(beam, ix, dz)

self._finalize()

self.N_tracks+=1

def replace_with_recorded_field_map(self, delete_ecloud_data=True):

if self.track_only_first_time:
print 'Warning: replace_with_recorded_field_map resets track_only_first_time = False'
self.track_only_first_time=False

if not hasattr(self, 'efieldmap'):
from PyHEADTAIL.field_maps.Transverse_Efield_map import Transverse_Efield_map
self.efieldmap = Transverse_Efield_map(xg = self.spacech_ele.xg, yg = self.spacech_ele.yg,
Ex=self.Ex_ele_last_track, Ey=self.Ey_ele_last_track, n_slices=self.slicer.n_slices,
z_cut=self.slicer.z_cuts, L_interaction=self.L_ecloud, flag_clean_slices = True)

self._ecloud_track = self.track

self.track = self.efieldmap.track

if delete_ecloud_data:
self.spacech_ele=None
self.Mp_e = None
self.init_nel = None
self.init_vx = None
self.init_vy = None
self.init_vz = None
self.init_x = None
self.init_y = None
self.init_z = None


else:
print 'Warning: efieldmap already exists. I do nothing.'

def track_once_and_replace_with_recorded_field_map(self, bunch, delete_ecloud_data=True):
self.save_ele_field = True
self.track_only_first_time = True
self.track(bunch)
self.save_ele_field = False
self.track_only_first_time = False
self.replace_with_recorded_field_map(delete_ecloud_data=delete_ecloud_data)


def _track_single_slice(self, beam, ix, dz):

MP_e = self.MP_e
dynamics = self.dynamics
impact_man = self.impact_man
spacech_ele = self.spacech_ele

dt = dz / (beam.beta * c)

# define substep
if dt>self.Dt_ref:
N_sub_steps = int(np.round(dt/self.Dt_ref))
else:
N_sub_steps=1

Dt_substep = dt/N_sub_steps
#print Dt_substep, N_sub_steps, dt

# beam field
MP_p = MP_light()
MP_p.x_mp = beam.x[ix]+self.x_beam_offset
MP_p.y_mp = beam.y[ix]+self.y_beam_offset
MP_p.nel_mp = beam.x[ix]*0.+beam.particlenumber_per_mp/dz#they have to become cylinders
MP_p.N_mp = len(beam.x[ix])
#compute beam field (it assumes electrons!)
spacech_ele.recompute_spchg_efield(MP_p)
#scatter to electrons
Ex_n_beam, Ey_n_beam = spacech_ele.get_sc_eletric_field(MP_e)
# go to actual beam particles
Ex_n_beam = -Ex_n_beam * beam.charge/e
Ey_n_beam = -Ey_n_beam * beam.charge/e


## compute electron field map
spacech_ele.recompute_spchg_efield(MP_e)

## compute electron field on electrons
Ex_sc_n, Ey_sc_n = spacech_ele.get_sc_eletric_field(MP_e)

## compute electron field on beam particles
Ex_sc_p, Ey_sc_p = spacech_ele.get_sc_eletric_field(MP_p)

## Total electric field on electrons
Ex_n=Ex_sc_n+Ex_n_beam;
Ey_n=Ey_sc_n+Ey_n_beam;

## save position before motion step
old_pos=MP_e.get_positions()

## motion electrons
MP_e = dynamics.stepcustomDt(MP_e, Ex_n,Ey_n, Dt_substep=Dt_substep, N_sub_steps=N_sub_steps)

## impacts: backtracking and secondary emission
MP_e = impact_man.backtrack_and_second_emiss(old_pos, MP_e)

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

if self.save_ele_distributions_last_track:
self.rho_ele_last_track.append(spacech_ele.rho.copy())
#print 'Here'

if self.save_ele_potential:
self.phi_ele_last_track.append(spacech_ele.phi.copy())

if self.save_ele_field:
self.Ex_ele_last_track.append(spacech_ele.efx.copy())
self.Ey_ele_last_track.append(spacech_ele.efy.copy())

if self.save_ele_MP_position:
self.x_MP_last_track.append(MP_e.x_mp.copy())
self.y_MP_last_track.append(MP_e.y_mp.copy())

if self.save_ele_MP_velocity:
self.vx_MP_last_track.append(MP_e.vx_mp.copy())
self.vy_MP_last_track.append(MP_e.vy_mp.copy())

if self.save_ele_MP_size:
self.nel_MP_last_track.append(MP_e.nel_mp.copy())

if self.save_ele_MP_position or self.save_ele_MP_velocity or self.save_ele_MP_size:
self.N_MP_last_track.append(MP_e.N_mp)

def _reinitialize(self):

#reinitialize
self.MP_e.x_mp[:self.init_N_mp] = self.init_x #it is a mutation and not a binding (and we have tested it :-))
self.MP_e.y_mp[:self.init_N_mp] = self.init_y
self.MP_e.z_mp[:self.init_N_mp] = self.init_z
Expand Down Expand Up @@ -297,101 +447,10 @@ def track(self, beam):

if self.save_ele_MP_position or self.save_ele_MP_velocity or self.save_ele_MP_size:
self.N_MP_last_track = []


if hasattr(beam.particlenumber_per_mp, '__iter__'):
raise ValueError('ecloud module assumes same size for all beam MPs')

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)
def _finalize(self):

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

# define substep
if dt>self.Dt_ref:
N_sub_steps = int(np.round(dt/self.Dt_ref))
else:
N_sub_steps=1

Dt_substep = dt/N_sub_steps
#print Dt_substep, N_sub_steps, dt

# beam field
MP_p = MP_light()
MP_p.x_mp = beam.x[ix]+self.x_beam_offset
MP_p.y_mp = beam.y[ix]+self.y_beam_offset
MP_p.nel_mp = beam.x[ix]*0.+beam.particlenumber_per_mp/dz#they have to become cylinders
MP_p.N_mp = slices.n_macroparticles_per_slice[i]
#compute beam field (it assumes electrons!)
spacech_ele.recompute_spchg_efield(MP_p)
#scatter to electrons
Ex_n_beam, Ey_n_beam = spacech_ele.get_sc_eletric_field(MP_e)
# go to actual beam particles
Ex_n_beam = -Ex_n_beam * beam.charge/e
Ey_n_beam = -Ey_n_beam * beam.charge/e


## compute electron field map
spacech_ele.recompute_spchg_efield(MP_e)

## compute electron field on electrons
Ex_sc_n, Ey_sc_n = spacech_ele.get_sc_eletric_field(MP_e)

## compute electron field on beam particles
Ex_sc_p, Ey_sc_p = spacech_ele.get_sc_eletric_field(MP_p)

## Total electric field on electrons
Ex_n=Ex_sc_n+Ex_n_beam;
Ey_n=Ey_sc_n+Ey_n_beam;

## save position before motion step
old_pos=MP_e.get_positions()

## motion electrons
MP_e = dynamics.stepcustomDt(MP_e, Ex_n,Ey_n, Dt_substep=Dt_substep, N_sub_steps=N_sub_steps)

## impacts: backtracking and secondary emission
MP_e = impact_man.backtrack_and_second_emiss(old_pos, MP_e)

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

if self.save_ele_distributions_last_track:
self.rho_ele_last_track.append(spacech_ele.rho.copy())
#print 'Here'

if self.save_ele_potential:
self.phi_ele_last_track.append(spacech_ele.phi.copy())

if self.save_ele_field:
self.Ex_ele_last_track.append(spacech_ele.efx.copy())
self.Ey_ele_last_track.append(spacech_ele.efy.copy())

if self.save_ele_MP_position:
self.x_MP_last_track.append(MP_e.x_mp.copy())
self.y_MP_last_track.append(MP_e.y_mp.copy())

if self.save_ele_MP_velocity:
self.vx_MP_last_track.append(MP_e.vx_mp.copy())
self.vy_MP_last_track.append(MP_e.vy_mp.copy())

if self.save_ele_MP_size:
self.nel_MP_last_track.append(MP_e.nel_mp.copy())

if self.save_ele_MP_position or self.save_ele_MP_velocity or self.save_ele_MP_size:
self.N_MP_last_track.append(MP_e.N_mp)

if self.save_ele_distributions_last_track:
self.rho_ele_last_track = np.array(self.rho_ele_last_track[::-1])

Expand All @@ -414,49 +473,7 @@ def track(self, beam):
self.nel_MP_last_track = np.array(self.nel_MP_last_track[::-1])

if self.save_ele_MP_position or self.save_ele_MP_velocity or self.save_ele_MP_size:
self.N_MP_last_track = np.array(self.N_MP_last_track[::-1])

self.N_tracks+=1

def replace_with_recorded_field_map(self, delete_ecloud_data=True):

if self.track_only_first_time:
print 'Warning: replace_with_recorded_field_map resets track_only_first_time = False'
self.track_only_first_time=False

if not hasattr(self, 'efieldmap'):
from PyHEADTAIL.field_maps.Transverse_Efield_map import Transverse_Efield_map
self.efieldmap = Transverse_Efield_map(xg = self.spacech_ele.xg, yg = self.spacech_ele.yg,
Ex=self.Ex_ele_last_track, Ey=self.Ey_ele_last_track, n_slices=self.slicer.n_slices,
z_cut=self.slicer.z_cuts, L_interaction=self.L_ecloud, flag_clean_slices = True)

self._ecloud_track = self.track

self.track = self.efieldmap.track

if delete_ecloud_data:
self.spacech_ele=None
self.Mp_e = None
self.init_nel = None
self.init_vx = None
self.init_vy = None
self.init_vz = None
self.init_x = None
self.init_y = None
self.init_z = None


else:
print 'Warning: efieldmap already exists. I do nothing.'

def track_once_and_replace_with_recorded_field_map(self, bunch, delete_ecloud_data=True):
self.save_ele_field = True
self.track_only_first_time = True
self.track(bunch)
self.save_ele_field = False
self.track_only_first_time = False
self.replace_with_recorded_field_map(delete_ecloud_data=delete_ecloud_data)

self.N_MP_last_track = np.array(self.N_MP_last_track[::-1])

def read_parameter_files_pyhdtl(pyecl_input_folder):
switch_model=0
Expand Down

0 comments on commit 892dd36

Please sign in to comment.