def add_em_solver(self, deposition_species = []):
grid_EM = picmi.Cartesian3DGrid(number_of_cells=self.number_of_cells,
lower_bound=self.lower_bound,
upper_bound=self.upper_bound,
lower_boundary_conditions=self.pml_bc,
upper_boundary_conditions=self.pml_bc)
n_pass = [[1], [1], [1]]
stride = [[1], [1], [1]]
compensation = [[False], [False], [False]]
alpha = [[0.5], [0.5], [0.5]]
smoother = picmi.BinomialSmoother(n_pass=n_pass,
compensation=compensation,
stride=stride,
alpha=alpha)
self.EM_solver = picmi.ElectromagneticSolver(grid=self.grid_EM,
method=self.EM_method, cfl=self.cfl,
source_smoother=smoother,
warp_l_correct_num_Cherenkov=False,
warp_type_rz_depose=0,
warp_l_setcowancoefs=True,
warp_l_getrho=False,
warp_deposition_species=deposition_species,
warp_conductors = self.chamber.conductors,
warp_conductor_dfill = picmi.warp.largepos)
self.EM_solver.initialize_solver_inputs()
registersolver(self.EM_solver.solver)
def add_ms_solver(self, deposition_species = []):
grid_MS = picmi.Cartesian3DGrid(number_of_cells=self.number_of_cells,
lower_bound=self.lower_bound,
upper_bound=self.upper_bound,
lower_boundary_conditions=self.dir_bc,
upper_boundary_conditions=self.dir_bc,
warp_conductors = self.chamber.conductors,
warp_conductor_dfill = picmi.warp.largepos)
self.MS_solver = picmi.MagnetostaticSolver(grid=grid_MS, warp_deposition_species=deposition_species)
self.MS_solver.initialize_solver_inputs()
registersolver(self.MS_solver.solver)
def measure_SEY(Ekin,
Nmp,
N_elec_p_mp,
sey_params_dict,
xgen=0,
ygen=0,
zgen=0,
thetagen=0,
phigen=0,
tot_t=None,
r_sphere=0.05,
flag_video=False):
import numpy as np
import numpy.random as random
#from pywarpx import picmi
from warp import picmi
from scipy.stats import gaussian_kde
from warp.particles.Secondaries import Secondaries, warp, top, posC, time
import matplotlib.pyplot as plt
import parser
import scipy.io as sio
from scipy.constants import m_e
from scipy.constants import c as c_light
from io import BytesIO as StringIO
from mpi4py import MPI
# Construct PyECLOUD secondary emission object
import PyECLOUD.sec_emission_model_ECLOUD as seec
sey_mod = seec.SEY_model_ECLOUD(
Emax=sey_params_dict['Emax'],
del_max=sey_params_dict['del_max'],
R0=sey_params_dict['R0'],
E_th=sey_params_dict['E_th'],
sigmafit=sey_params_dict['sigmafit'],
mufit=sey_params_dict['mufit'],
secondary_angle_distribution=sey_params_dict[
'secondary_angle_distribution'])
##########################
# physics parameters
##########################
mysolver = 'ES' # solver type ('ES'=electrostatic; 'EM'=electromagnetic)
# --- Beam
##########################
# numerics parameters
##########################
# geometry
# --- grid
dh = r_sphere / 10.
xmin = -r_sphere - 0.001
xmax = -xmin
ymin = -r_sphere - 0.001
ymax = -ymin
zmin = -r_sphere - 0.001
zmax = -zmin
nx = (xmax - xmin) / dh
ny = (xmax - xmin) / dh
nz = (xmax - xmin) / dh
#######################################################
# compute beam size from normalized emittance and beta
# Uncomment if data available
#######################################################
sigmax = 2e-4
sigmay = 2.1e-4
E0 = 0.511 * 1e6
E = E0 + Ekin
beam_gamma = E / E0
beam_beta = np.sqrt(1 - 1 / (beam_gamma * beam_gamma))
v = beam_beta * c_light
elec_beam = picmi.Species(particle_type='electron',
particle_shape='linear',
name='elec_beam')
#initial_distribution = gauss_dist)
secelec = picmi.Species(particle_type='electron',
particle_shape='linear',
name='Secondary electrons')
##########################
# Numeric components
##########################
if mysolver == 'ES':
lower_boundary_conditions = ['dirichlet', 'dirichlet', 'dirichlet']
upper_boundary_conditions = ['dirichlet', 'dirichlet', 'dirichlet']
if mysolver == 'EM':
lower_boundary_conditions = ['open', 'open', 'open']
upper_boundary_conditions = ['open', 'open', 'open']
grid = picmi.Cartesian3DGrid(
number_of_cells=[nx, ny, nz],
lower_bound=[xmin, ymin, zmin],
upper_bound=[xmax, ymax, zmax],
lower_boundary_conditions=lower_boundary_conditions,
upper_boundary_conditions=upper_boundary_conditions) #,
#warpx_max_grid_size=32)
if mysolver == 'ES':
solver = picmi.ElectrostaticSolver(grid=grid)
if mysolver == 'EM':
smoother = picmi.BinomialSmoother(n_pass=[[1], [1], [1]],
compensation=[[False], [False],
[False]],
stride=[[1], [1], [1]],
alpha=[[0.5], [0.5], [0.5]])
solver = picmi.ElectromagneticSolver(
grid=grid,
method='CKC',
cfl=1.,
source_smoother=smoother,
# field_smoother = smoother,
warp_l_correct_num_Cherenkov=False,
warp_type_rz_depose=0,
warp_l_setcowancoefs=True,
warp_l_getrho=False)
##########################
# Simulation setup
##########################
sphere = picmi.warp.Sphere(r_sphere)
box = picmi.warp.Box(xmax - xmin, ymax - ymin, zmax - zmin)
wall = box - sphere
sim = picmi.Simulation(solver=solver,
verbose=1,
cfl=1.0,
warp_initialize_solver_after_generate=1)
sim.conductors = wall
beam_layout = picmi.PseudoRandomLayout(n_macroparticles=10**5, seed=3)
sim.add_species(elec_beam,
layout=beam_layout,
initialize_self_field=solver == 'EM')
sim.add_species(secelec, layout=None, initialize_self_field=False)
#########################
# Add Dipole
########################
vmod = picmi.warp.clight * np.sqrt(1 - 1. / (beam_gamma**2))
vxgen = vmod * np.sin(thetagen) * np.cos(phigen)
vygen = vmod * np.sin(thetagen) * np.sin(phigen)
vzgen = vmod * np.cos(thetagen)
def nonlinearsource():
Nmp_inj = Nmp * (top.it == 3)
x = xgen * np.ones(Nmp_inj)
y = ygen * np.ones(Nmp_inj)
z = zgen * np.ones(Nmp_inj)
vx = vxgen * np.ones(Nmp_inj)
vy = vygen * np.ones(Nmp_inj)
vz = vzgen * np.ones(Nmp_inj)
elec_beam.wspecies.addparticles(x=x,
y=y,
z=z,
vx=vx,
vy=vy,
vz=vz,
gi=1. / beam_gamma,
w=N_elec_p_mp)
picmi.warp.installuserinjection(nonlinearsource)
##########################
# simulation run
##########################
sim.step(1)
solver.solver.installconductor(sim.conductors, dfill=picmi.warp.largepos)
sim.step(1)
pp = warp.ParticleScraper(sim.conductors,
lsavecondid=1,
lsaveintercept=1,
lcollectlpdata=1)
sec = Secondaries(conductors=sim.conductors,
l_usenew=1,
pyecloud_secemi_object=sey_mod,
pyecloud_nel_mp_ref=N_elec_p_mp,
pyecloud_fact_clean=1e-16,
pyecloud_fact_split=1.5)
sec.add(incident_species=elec_beam.wspecies,
emitted_species=secelec.wspecies,
conductor=sim.conductors)
# --- set weights of secondary electrons
#secelec.wspecies.getw()[...] = elecb.wspecies.getw()[0]
# define shortcuts
pw = picmi.warp
em = solver.solver
step = pw.step
top.dt = dh / v
n_step = 0
tot_nsteps = int(np.ceil((tot_t / top.dt)))
n_primelecs = []
n_secelecs = []
n_mp_primelecs = []
n_mp_secelecs = []
secelecs_dist = []
primelecs_dist = []
for n_step in range(tot_nsteps + 1):
step(1)
n_primelecs.append(np.sum(elec_beam.wspecies.getw()))
n_secelecs.append(np.sum(secelec.wspecies.getw()))
n_mp_primelecs.append(np.shape(elec_beam.wspecies.getw())[0])
n_mp_secelecs.append(np.shape(secelec.wspecies.getw())[0])
if flag_video:
secelecs_dist.append(np.sum(secelec.wspecies.get_density(),
axis=1))
primelecs_dist.append(
np.sum(elec_beam.wspecies.get_density(), axis=1))
dict_out = {}
dict_out['SEY'] = np.sum(secelec.wspecies.getw()) / (Nmp * N_elec_p_mp)
dict_out['n_primelecs'] = np.array(n_primelecs)
dict_out['n_secelecs'] = np.array(n_secelecs)
dict_out['n_mp_primelecs'] = np.array(n_mp_primelecs)
dict_out['n_mp_secelecs'] = np.array(n_mp_secelecs)
dict_out['secelecs_dist'] = np.array(secelecs_dist)
dict_out['primelecs_dist'] = np.array(primelecs_dist)
dict_out['xmesh'] = solver.solver.xmesh.copy()
dict_out['ymesh'] = solver.solver.ymesh.copy()
dict_out['zmesh'] = solver.solver.zmesh.copy()
#angle_dist =np.arctan(np.divide(secelec.wspecies.getvy(),secelec.wspecies.getvx()))
#dict_out['angle_dist'] = angle_dist
return dict_out
def measure_SEY(Ekin):
import numpy as np
import numpy.random as random
#from pywarpx import picmi
from warp import picmi
from scipy.stats import gaussian_kde
from warp.particles.Secondaries import Secondaries, warp, top, posC, time
import matplotlib.pyplot as plt
import parser
import scipy.io as sio
from scipy.constants import m_e
from scipy.constants import c as c_light
from io import BytesIO as StringIO
from mpi4py import MPI
##########################
# physics parameters
##########################
mysolver = 'ES' # solver type ('ES'=electrostatic; 'EM'=electromagnetic)
# --- Beam
##########################
# numerics parameters
##########################
# geometry
r = 0.005
l = 0.1
# --- grid
dh = r/10.
xmin = -2*r + 0.001
xmax = -xmin
ymin = -2*r + 0.001
ymax = -ymin
zmin = -l - 0.001
zmax = l
nx = (xmax-xmin)/dh
ny = (xmax-xmin)/dh
nz = (xmax-xmin)/dh
#######################################################
# compute beam size from normalized emittance and beta
# Uncomment if data available
#######################################################
sigmax = 2e-4
sigmay = 2.1e-4
E0 = 0.511*1e6
E = E0 + Ekin
beam_gamma = E/E0
beam_beta = np.sqrt(1-1/(beam_gamma*beam_gamma))
v = beam_beta*c_light
elec_beam = picmi.Species(particle_type = 'electron',
particle_shape = 'linear',
name = 'elec_beam')
#initial_distribution = gauss_dist)
secelec = picmi.Species(particle_type = 'electron',
particle_shape = 'linear',
name = 'Secondary electrons')
##########################
# Numeric components
##########################
if mysolver=='ES':
lower_boundary_conditions = ['dirichlet', 'dirichlet', 'dirichlet']
upper_boundary_conditions = ['dirichlet', 'dirichlet', 'dirichlet']
if mysolver=='EM':
lower_boundary_conditions = ['open', 'open', 'open']
upper_boundary_conditions = ['open', 'open', 'open']
grid = picmi.Cartesian3DGrid(number_of_cells = [nx, ny, nz],
lower_bound = [xmin, ymin, zmin],
upper_bound = [xmax, ymax, zmax],
lower_boundary_conditions = lower_boundary_conditions,
upper_boundary_conditions = upper_boundary_conditions)#,
#warpx_max_grid_size=32)
if mysolver=='ES':
solver = picmi.ElectrostaticSolver(grid = grid)
if mysolver=='EM':
smoother = picmi.BinomialSmoother(n_pass = [[1], [1], [1]],
compensation = [[False], [False], [False]],
stride = [[1], [1], [1]],
alpha = [[0.5], [0.5], [0.5]])
solver = picmi.ElectromagneticSolver(grid = grid,
method = 'CKC',
cfl = 1.,
source_smoother = smoother,
# field_smoother = smoother,
warp_l_correct_num_Cherenkov = False,
warp_type_rz_depose = 0,
warp_l_setcowancoefs = True,
warp_l_getrho = False)
##########################
# Simulation setup
##########################
wall = picmi.warp.YCylinderOut(r,l)
sim = picmi.Simulation(solver = solver,
verbose = 1,
cfl = 1.0,
warp_initialize_solver_after_generate = 1)
sim.conductors = wall
beam_layout = picmi.PseudoRandomLayout(n_macroparticles = 10**5, seed = 3)
sim.add_species(elec_beam, layout=beam_layout,
initialize_self_field = solver=='EM')
sim.add_species(secelec, layout=None, initialize_self_field=False)
#########################
# Add Dipole
#########################
bunch_w = 1
def nonlinearsource():
NP = 1000*(top.it==3)
x = 0*np.ones(NP)
y = 0*np.ones(NP)
z = 0*np.ones(NP)
vx = np.zeros(NP)
vy = np.zeros(NP)
vz = picmi.warp.clight*np.sqrt(1-1./(beam_gamma**2))
elec_beam.wspecies.addparticles(x=x,y=y,z=z,vx=vx,vy=vy,vz=vz,gi = 1./beam_gamma, w=bunch_w)
picmi.warp.installuserinjection(nonlinearsource)
##########################
# simulation run
##########################
sim.step(1)
solver.solver.installconductor(sim.conductors, dfill = picmi.warp.largepos)
sim.step(1)
pp = warp.ParticleScraper(sim.conductors,lsavecondid=1,lsaveintercept=1,lcollectlpdata=1)
####################################
# Set material parameters from file
####################################
def set_params_user(maxsec, matnum):
dict = parser.pos2dic('LHC_inj_72bx5.in')
posC.matsurf = dict['matsurf']
posC.iprob = dict['iprob']
posC.enpar = dict['enpar']
posC.pnpar= dict['pnpar']
posC.dtspk = dict['dtspk']
posC.dtotpk = dict['dtotpk']
posC.pangsec = dict['pangsec']
posC.pr = dict['pr']
posC.sige = dict['sige']
posC.Ecr = dict['Ecr']
posC.E0tspk = dict['E0tspk']
posC.E0epk = dict['E0epk']
posC.E0w = dict['E0w']
posC.rpar1 = dict['rpar'][0]
posC.rpar2 = dict['rpar'][1]
posC.tpar1 = dict['tpar'][0]
posC.tpar2 = dict['tpar'][1]
posC.tpar3 = dict['tpar'][2]
posC.tpar4 = dict['tpar'][3]
posC.tpar5 = dict['tpar'][4]
posC.tpar6 = dict['tpar'][5]
posC.epar1 = dict['epar'][0]
posC.epar2 = dict['epar'][1]
posC.P1rinf = dict['P1rinf']
posC.P1einf = dict['P1einf']
posC.P1epk = dict['P1epk']
posC.powts = dict['powts']
posC.powe = dict['powe']
posC.qr = dict['qr']
sec=Secondaries(conductors=sim.conductors, set_params_user = set_params_user,
l_usenew=1)
sec.add(incident_species = elec_beam.wspecies,
emitted_species = secelec.wspecies,
conductor = sim.conductors)
# --- set weights of secondary electrons
#secelec.wspecies.getw()[...] = elecb.wspecies.getw()[0]
# define shortcuts
pw = picmi.warp
em = solver.solver
step=pw.step
top.dt = dh/v
n_step = 0
tot_t = 2*r/v
tot_nsteps = int(tot_t/top.dt)
for n_step in range(tot_nsteps):
step(1)
secondaries_count = np.sum(secelec.wspecies.getw())
return secondaries_count/1000
def warp_pyecloud_dipole(z_length=None,
nx=None,
ny=None,
nz=None,
n_bunches=None,
b_spac=None,
beam_gamma=None,
sigmax=None,
sigmay=None,
sigmat=None,
bunch_intensity=None,
init_num_elecs=None,
init_num_elecs_mp=None,
By=None,
pyecloud_nel_mp_ref=None,
dt=None,
pyecloud_fact_clean=None,
pyecloud_fact_split=None,
chamber_type=None,
flag_save_video=None,
enable_trap=True,
Emax=None,
del_max=None,
R0=None,
E_th=None,
sigmafit=None,
mufit=None,
secondary_angle_distribution=None,
N_mp_max=None,
N_mp_target=None):
import numpy as np
import numpy.random as random
from warp import picmi
from scipy.stats import gaussian_kde
from warp.particles.Secondaries import Secondaries, top, warp, time
import matplotlib.pyplot as plt
import scipy.io as sio
from io import BytesIO as StringIO
from scipy.constants import c as clight
import sys
import PyECLOUD.myfilemanager as mfm
import os
# Construct PyECLOUD secondary emission object
import PyECLOUD.sec_emission_model_ECLOUD as seec
sey_mod = seec.SEY_model_ECLOUD(Emax=Emax,
del_max=del_max,
R0=R0,
E_th=E_th,
sigmafit=sigmafit,
mufit=mufit,
secondary_angle_distribution='cosine_3D')
##########################
# physics parameters
##########################
mysolver = 'ES' # solver type ('ES'=electrostatic; 'EM'=electromagnetic)
##########################
# numerics parameters
##########################
# --- grid
zs_dipo = -0.5 * z_length
ze_dipo = 0.5 * z_length
r = 23.e-3
h = 18.e-3
l = z_length
unit = 1e-3
ghost = 1e-3
xmin = -r - ghost
xmax = r + ghost
ymin = -r - ghost
ymax = r + ghost
zmin = zs_dipo - 50 * unit
zmax = ze_dipo + 50 * unit
if chamber_type == 'LHC':
chamber_area = 0.0014664200235342726
elif chamber_type == 'rect':
chamber_area = 4 * r * h
##########################
# Beam parameters
##########################
sigmaz = sigmat * picmi.clight
t_offs = b_spac - 6 * sigmat
bunch_w = 1e6
bunch_macro_particles = bunch_intensity / bunch_w
#######################################################
# compute beam size from normalized emittance and beta
# Uncomment if data available
#######################################################
bunch_rms_size = [sigmax, sigmay, sigmaz]
bunch_rms_velocity = [0., 0., 0.]
bunch_centroid_position = [0, 0, zs_dipo - 10 * unit]
bunch_centroid_velocity = [0., 0., beam_gamma * picmi.constants.c]
beam = picmi.Species(particle_type='proton',
particle_shape='linear',
name='beam')
lower_bound = [-r, -h, zs_dipo]
upper_bound = [r, h, ze_dipo]
temp_file_name = 'temp_mps_info.h5'
########################################################
# if checkopoint is found reload it,
# otherwise start from scratch
########################################################
if not os.path.exists(temp_file_name):
x0 = random.uniform(lower_bound[0], upper_bound[0], init_num_elecs_mp)
y0 = random.uniform(lower_bound[1], upper_bound[1], init_num_elecs_mp)
z0 = random.uniform(lower_bound[2], upper_bound[2], init_num_elecs_mp)
vx0 = np.zeros(init_num_elecs_mp)
vy0 = np.zeros(init_num_elecs_mp)
vz0 = np.zeros(init_num_elecs_mp)
w0 = float(init_num_elecs) / float(init_num_elecs_mp)
b_pass_prev = 0
else:
print('############################################################')
print('Temp distribution found. Regenarating and restarting')
print('############################################################')
dict_init_dist = mfm.dict_of_arrays_and_scalar_from_h5(temp_file_name)
# Load particles status
xold = dict_init_dist['x_mp']
yold = dict_init_dist['y_mp']
zold = dict_init_dist['z_mp']
uxold = dict_init_dist['ux_mp']
uyold = dict_init_dist['uy_mp']
uzold = dict_init_dist['uz_mp']
wold = dict_init_dist['w_mp']
# Reload the outputs and other auxiliary stuff
numelecs = dict_init_dist['numelecs']
elecs_density = dict_init_dist['elecs_density']
N_mp = dict_init_dist['N_mp']
b_pass_prev = dict_init_dist['b_pass'] - 1
# compute the velocities
invgamma = np.sqrt(
1 - picmi.clight**2 /
(np.square(uxold) + np.square(uyold) + np.square(uzold)))
vxold = np.multiply(invgamma, uxold)
vyold = np.multiply(invgamma, uxold)
vzold = np.multiply(invgamma, uxold)
# perform the regeneration
x0, y0, z0, vx0, vy0, vz0, w0, N_mpaa = perform_regeneration(
N_mp_target, xold, yold, zold, vxold, vyold, vzold, wold)
print('DEBUG N_mp: %d' % len(x0))
electron_background_dist = picmi.ParticleListDistribution(x=x0,
y=y0,
z=z0,
vx=vx0,
vy=vy0,
vz=vz0,
weight=w0)
elecb = picmi.Species(particle_type='electron',
particle_shape='linear',
name='Electron background',
initial_distribution=electron_background_dist)
secelec = picmi.Species(particle_type='electron',
particle_shape='linear',
name='Secondary electrons')
##########################
# Numeric components
##########################
if mysolver == 'ES':
lower_boundary_conditions = ['dirichlet', 'dirichlet', 'dirichlet']
upper_boundary_conditions = ['dirichlet', 'dirichlet', 'dirichlet']
if mysolver == 'EM':
lower_boundary_conditions = ['open', 'open', 'open']
upper_boundary_conditions = ['open', 'open', 'open']
grid = picmi.Cartesian3DGrid(
number_of_cells=[nx, ny, nz],
lower_bound=[xmin, ymin, zmin],
upper_bound=[xmax, ymax, zmax],
lower_boundary_conditions=lower_boundary_conditions,
upper_boundary_conditions=upper_boundary_conditions)
if mysolver == 'ES':
solver = picmi.ElectrostaticSolver(grid=grid)
if mysolver == 'EM':
smoother = picmi.BinomialSmoother(n_pass=[[1], [1], [1]],
compensation=[[False], [False],
[False]],
stride=[[1], [1], [1]],
alpha=[[0.5], [0.5], [0.5]])
solver = picmi.ElectromagneticSolver(
grid=grid,
method='CKC',
cfl=1.,
source_smoother=smoother,
warp_l_correct_num_Cherenkov=False,
warp_type_rz_depose=0,
warp_l_setcowancoefs=True,
warp_l_getrho=False)
##########################
# Simulation setup
##########################
pipe = picmi.warp.ZCylinderOut(r, l, condid=1)
if chamber_type == 'rect':
upper_box = picmi.warp.YPlane(y0=h, ysign=1, condid=1)
lower_box = picmi.warp.YPlane(y0=-h, ysign=-1, condid=1)
left_box = picmi.warp.XPlane(x0=r, xsign=1, condid=1)
right_box = picmi.warp.XPlane(x0=-r, xsign=-1, condid=1)
sim = picmi.Simulation(solver=solver,
verbose=1,
cfl=1.0,
warp_initialize_solver_after_generate=1)
sim.conductors = upper_box + lower_box + left_box + right_box
elif chamber_type == 'LHC':
pipe = picmi.warp.ZCylinderOut(r, l, condid=1)
upper_box = picmi.warp.YPlane(y0=h, ysign=1, condid=1)
lower_box = picmi.warp.YPlane(y0=-h, ysign=-1, condid=1)
sim = picmi.Simulation(solver=solver,
verbose=1,
cfl=1.0,
warp_initialize_solver_after_generate=1)
sim.conductors = pipe + upper_box + lower_box
beam_layout = picmi.PseudoRandomLayout(n_macroparticles=10**5, seed=3)
sim.add_species(beam,
layout=beam_layout,
initialize_self_field=solver == 'EM')
elecb_layout = picmi.PseudoRandomLayout(n_macroparticles=init_num_elecs_mp,
seed=3)
sim.add_species(elecb,
layout=elecb_layout,
initialize_self_field=solver == 'EM')
sim.add_species(secelec, layout=None, initialize_self_field=False)
#########################
# Add Dipole
#########################
picmi.warp.addnewdipo(zs=zs_dipo, ze=ze_dipo, by=By)
def time_prof(t):
val = 0
sigmat = sigmaz / picmi.clight
for i in range(1, n_bunches + 1):
val += bunch_macro_particles * 1. / np.sqrt(
2 * np.pi * sigmat * sigmat) * np.exp(
-(t - i * b_spac + t_offs) * (t - i * b_spac + t_offs) /
(2 * sigmat * sigmat)) * picmi.warp.top.dt
return val
def nonlinearsource():
NP = int(time_prof(top.time))
x = random.normal(bunch_centroid_position[0], bunch_rms_size[0], NP)
y = random.normal(bunch_centroid_position[1], bunch_rms_size[1], NP)
z = bunch_centroid_position[2]
vx = random.normal(bunch_centroid_velocity[0], bunch_rms_velocity[0],
NP)
vy = random.normal(bunch_centroid_velocity[1], bunch_rms_velocity[1],
NP)
vz = picmi.warp.clight * np.sqrt(1 - 1. / (beam_gamma**2))
beam.wspecies.addparticles(x=x,
y=y,
z=z,
vx=vx,
vy=vy,
vz=vz,
gi=1. / beam_gamma,
w=bunch_w)
picmi.warp.installuserinjection(nonlinearsource)
##########################
# simulation run
##########################
sim.step(1)
solver.solver.installconductor(sim.conductors, dfill=picmi.warp.largepos)
sim.step(1)
pp = warp.ParticleScraper(sim.conductors,
lsavecondid=1,
lsaveintercept=1,
lcollectlpdata=1)
sec = Secondaries(conductors=sim.conductors,
l_usenew=1,
pyecloud_secemi_object=sey_mod,
pyecloud_nel_mp_ref=pyecloud_nel_mp_ref,
pyecloud_fact_clean=pyecloud_fact_clean,
pyecloud_fact_split=pyecloud_fact_split)
sec.add(incident_species=elecb.wspecies,
emitted_species=secelec.wspecies,
conductor=sim.conductors)
sec.add(incident_species=secelec.wspecies,
emitted_species=secelec.wspecies,
conductor=sim.conductors)
# just some shortcuts
pw = picmi.warp
step = pw.step
if mysolver == 'ES':
pw.top.dt = dt
def myplots(l_force=0):
if l_force or pw.top.it % 10 == 0:
plt.close()
(Nx, Ny, Nz) = np.shape(secelec.wspecies.get_density())
fig, axs = plt.subplots(1, 2, figsize=(12, 4.5))
fig.subplots_adjust(left=0.05,
bottom=0.1,
right=0.97,
top=0.94,
wspace=0.15)
d = secelec.wspecies.get_density() + elecb.wspecies.get_density(
) + beam.wspecies.get_density()
d2 = secelec.wspecies.get_density() + elecb.wspecies.get_density()
im1 = axs[0].imshow(d[:, :, Nz / 2].T,
cmap='jet',
origin='lower',
vmin=0.2 * np.min(d2[:, :, Nz / 2]),
vmax=0.8 * np.max(d2[:, :, Nz / 2]),
extent=[xmin, xmax, ymin, ymax])
axs[0].set_xlabel('x [m]')
axs[0].set_ylabel('y [m]')
axs[0].set_title('e- density')
fig.colorbar(
im1,
ax=axs[0],
)
im2 = axs[1].imshow(d[Nx / 2, :, :],
cmap='jet',
origin='lower',
vmin=0.2 * np.min(d2[Nx / 2, :, :]),
vmax=0.8 * np.max(d2[Nx / 2, :, :]),
extent=[zmin, zmax, ymin, ymax],
aspect='auto')
axs[1].set_xlabel('z [m]')
axs[1].set_ylabel('y [m]')
axs[1].set_title('e- density')
fig.colorbar(im2, ax=axs[1])
n_step = top.time / top.dt
figname = 'images/%d.png' % n_step
plt.savefig(figname)
if flag_save_video:
pw.installafterstep(myplots)
myplots(1)
ntsteps_p_bunch = b_spac / top.dt
tot_nsteps = int(np.ceil(b_spac * (n_bunches) / top.dt))
t_start = b_pass_prev * b_spac
tstep_start = int(round(t_start / top.dt))
# pre-allocate outputs
if not os.path.exists('temp_mps_info.h5'):
numelecs = np.zeros(tot_nsteps)
elecs_density = np.zeros((tot_nsteps, nx + 1, ny + 1, 3))
beam_density = np.zeros((tot_nsteps, nx + 1, ny + 1, 3))
N_mp = np.zeros(tot_nsteps)
dict_out = {}
# aux variables
b_pass = 0
perc = 10
original = sys.stdout
text_trap = StringIO()
t0 = time.time()
# trapping warp std output
text_trap = {True: StringIO(), False: sys.stdout}[enable_trap]
original = sys.stdout
for n_step in range(tstep_start, tot_nsteps):
# if a passage is starting...
if n_step / ntsteps_p_bunch >= b_pass + b_pass_prev:
b_pass += 1
perc = 10
if secelec.wspecies.getn() > N_mp_max:
dict_out_temp = {}
print('Number of macroparticles: %d' %
(secelec.wspecies.getn() + elecb.wspecies.getn()))
print('MAXIMUM LIMIT OF MPS HAS BEEN RACHED')
print('Please restart the simulation')
dict_out_temp['x_mp'] = np.concatenate(
(secelec.wspecies.getx(), elecb.wspecies.getx()))
dict_out_temp['y_mp'] = np.concatenate(
(secelec.wspecies.gety(), elecb.wspecies.gety()))
dict_out_temp['z_mp'] = np.concatenate(
(secelec.wspecies.getz(), elecb.wspecies.gety()))
dict_out_temp['ux_mp'] = np.concatenate(
(secelec.wspecies.getux(), elecb.wspecies.getux()))
dict_out_temp['uy_mp'] = np.concatenate(
(secelec.wspecies.getuy(), elecb.wspecies.getuy()))
dict_out_temp['uz_mp'] = np.concatenate(
(secelec.wspecies.getuz(), elecb.wspecies.getuz()))
dict_out_temp['w_mp'] = np.concatenate(
(secelec_w, elecb.wspecies.getw()))
dict_out_temp['numelecs'] = numelecs
dict_out_temp['elecs_density'] = elecs_density
dict_out_temp['N_mp'] = N_mp
dict_out_temp['b_pass'] = b_pass + b_pass_prev
filename = 'temp_mps_info.h5'
mfm.dict_to_h5(dict_out_temp, filename)
return
print('===========================')
print('Bunch passage: %d' % (b_pass + b_pass_prev))
print('Number of electrons in the dipole: %d' %
(np.sum(secelec.wspecies.getw()) +
np.sum(elecb.wspecies.getw())))
print('Number of macroparticles: %d' %
(secelec.wspecies.getn() + elecb.wspecies.getn()))
if n_step % ntsteps_p_bunch / ntsteps_p_bunch * 100 > perc:
print('%d%% of bunch passage' % perc)
perc = perc + 10
sys.stdout = text_trap
step(1)
sys.stdout = original
#print(elecb.wspecies.getvy()[0])
# store outputs
secelec_w = secelec.wspecies.getw()
elecb_w = elecb.wspecies.getw()
numelecs[n_step] = np.sum(secelec_w) + np.sum(elecb_w)
elecs_density[n_step, :, :, :] = secelec.wspecies.get_density(
)[:, :, (nz + 1) / 2 - 1:(nz + 1) / 2 +
2] + elecb.wspecies.get_density()[:, :,
(nz + 1) / 2 - 1:(nz + 1) / 2 + 2]
N_mp[n_step] = len(secelec_w) + len(elecb_w)
t1 = time.time()
totalt = t1 - t0
dict_out['numelecs'] = numelecs
dict_out['elecs_density'] = elecs_density
dict_out['N_mp'] = N_mp
mfm.dict_to_h5(dict_out, 'output.h5')
if os.path.exists('temp_mps_info.h5'):
os.remove('temp_mps_info.h5')
print('Run terminated in %ds' % totalt)
def __init__(self,
z_length=None,
nx=None,
ny=None,
nz=None,
dh_t=None,
n_bunches=None,
b_spac=None,
beam_gamma=None,
sigmax=None,
sigmay=None,
sigmat=None,
bunch_intensity=None,
init_num_elecs=None,
init_num_elecs_mp=None,
By=None,
pyecloud_nel_mp_ref=None,
dt=None,
pyecloud_fact_clean=None,
pyecloud_fact_split=None,
chamber_type=None,
flag_save_video=None,
enable_trap=True,
Emax=None,
del_max=None,
R0=None,
E_th=None,
sigmafit=None,
mufit=None,
secondary_angle_distribution=None,
N_mp_max=None,
N_mp_target=None,
flag_checkpointing=False,
checkpoints=None,
flag_output=False,
bunch_macro_particles=None,
t_offs=None,
width=None,
height=None,
output_name='output.mat',
flag_relativ_tracking=False,
nbins=100,
radius=None,
lattice_elem=None,
ghost=None,
ghost_z=None,
stride_imgs=10,
stride_output=1000):
# Construct PyECLOUD secondary emission object
sey_mod = seec.SEY_model_ECLOUD(
Emax=Emax,
del_max=del_max,
R0=R0,
E_th=E_th,
sigmafit=sigmafit,
mufit=mufit,
secondary_angle_distribution='cosine_3D')
##########################
# physics parameters
##########################
mysolver = 'ES' # solver type ('ES'=electrostatic; 'EM'=electromagnetic)
##########################
# numerics parameters
##########################
self.nbins = nbins
self.N_mp_target = N_mp_target
self.N_mp_max = N_mp_max
self.nx = nx
self.ny = ny
self.nz = nz
self.flag_checkpointing = flag_checkpointing
self.flag_output = flag_output
self.output_name = output_name
self.stride_imgs = stride_imgs
self.stride_output = stride_output
# --- grid
if lattice_elem == 'dipole':
zs_dipo = -0.5 * z_length
ze_dipo = 0.5 * z_length
# Consistency check
if chamber_type == 'rect':
assert (height is not None) and (
width is not None
), 'When using rectangular chamber, height and width must be specified'
if chamber_type == 'circle':
assert radius is not None, 'When using circular chamber r must be specified'
if chamber_type == 'LHC':
radius = 23.e-3
height = 2 * 18.e-3
width = 2 * radius
z_length
unit = 1e-3
if ghost is None:
ghost = 1e-3
if ghost_z is None:
ghost_z = ghost
if chamber_type == 'rect' or chamber_type == 'LHC':
self.xmin = -width / 2. - ghost
self.xmax = -self.xmin
self.ymin = -height / 2 - ghost
self.ymax = -self.ymin
self.zmin = zs_dipo - 50 * unit
self.zmax = ze_dipo + 50 * unit
lower_bound = [-width / 2, -height / 2, zs_dipo]
upper_bound = [width / 2, height / 2, ze_dipo]
elif chamber_type == 'circle':
self.xmin = -radius - ghost
self.xmax = radius + ghost
self.ymin = -radius - ghost
self.ymax = radius + ghost
self.zmin = zs_dipo - 50 * unit
self.zmax = ze_dipo + 50 * unit
lower_bound = [-radius, -radius, zs_dipo]
upper_bound = [radius, radius, ze_dipo]
elif chamber_type == 'crab':
max_z = 500 * unit
l_main_y = 242 * unit
l_main_x = 300 * unit
l_main_z = 350 * unit
l_beam_pipe = 84 * unit
l_int = 62 * unit
l_main_int_y = l_main_y - 0.5 * l_beam_pipe
l_main_int_z = 0.5 * l_main_z - l_int
l_main_int_x = 0.5 * l_main_x - l_int
self.xmin = -0.5 * l_main_x - ghost
self.xmax = 0.5 * l_main_x + ghost
self.ymin = -0.5 * l_main_y - ghost
self.ymax = 0.5 * l_main_y + ghost
self.zmin = -0.5 * max_z - 50 * unit
self.zmax = 0.5 * max_z + 50 * unit
lower_bound = [-0.5 * l_main_x, -0.5 * l_main_y, -0.5 * l_main_z]
upper_bound = [0.5 * l_main_x, 0.5 * l_main_y, 0.5 * l_main_z]
##########################
# Beam parameters
##########################
sigmaz = sigmat * picmi.clight
if bunch_macro_particles > 0:
bunch_w = bunch_intensity / bunch_macro_particles
else:
bunch_w = 0
#######################################################
# compute beam size from normalized emittance and beta
# Uncomment if data available
#######################################################
bunch_rms_size = [sigmax, sigmay, sigmaz]
bunch_rms_velocity = [0., 0., 0.]
if chamber_type == 'crab':
bunch_centroid_position = [0, 0, self.zmin + 10 * unit]
else:
bunch_centroid_position = [0, 0, zs_dipo - 10 * unit]
bunch_centroid_velocity = [0., 0., beam_gamma * picmi.constants.c]
self.beam = picmi.Species(particle_type='proton',
particle_shape='linear',
name='beam')
temp_file_name = 'temp_mps_info.mat'
########################################################
# if checkopoint is found reload it,
# otherwise start from scratch
########################################################
if flag_checkpointing and os.path.exists(temp_file_name):
print(
'############################################################')
print('Temp distribution found. Regenarating and restarting')
print(
'############################################################')
dict_init_dist = sio.loadmat(temp_file_name)
# Load particles status
x0 = dict_init_dist['x_mp'][0]
y0 = dict_init_dist['y_mp'][0]
z0 = dict_init_dist['z_mp'][0]
vx0 = dict_init_dist['vx_mp'][0]
vy0 = dict_init_dist['vy_mp'][0]
vz0 = dict_init_dist['vz_mp'][0]
w0 = dict_init_dist['nel_mp'][0]
# compute the velocities
#invgamma = np.sqrt(1-picmi.clight**2/(np.square(ux0)+np.square(uy0)+np.square(uz0)))
#vx0 = np.multiply(invgamma,ux0)
#vy0 = np.multiply(invgamma,uy0)
#vz0 = np.multiply(invgamma,uy0)
# Reload the outputs and other auxiliary stuff
if flag_output:
self.numelecs = dict_init_dist['numelecs'][0]
self.N_mp = dict_init_dist['N_mp'][0]
self.numelecs_tot = dict_init_dist['numelecs_tot'][0]
self.xhist = dict_init_dist['xhist']
self.bins = dict_init_dist['bins']
self.b_pass_prev = dict_init_dist['b_pass'] - 1
else:
if chamber_type == 'circle':
x0 = random.uniform(lower_bound[0], upper_bound[0],
init_num_elecs_mp)
y0 = random.uniform(lower_bound[1], upper_bound[1],
init_num_elecs_mp)
z0 = random.uniform(lower_bound[2], upper_bound[2],
init_num_elecs_mp)
vx0 = np.zeros(init_num_elecs_mp)
vy0 = np.zeros(init_num_elecs_mp)
vz0 = np.zeros(init_num_elecs_mp)
r0_sq = np.square(x0) + np.square(y0)
flag_out = np.where(r0_sq > radius * radius)[0]
Nout = len(flag_out)
while Nout > 0:
x0[flag_out] = random.uniform(lower_bound[0],
upper_bound[0], Nout)
y0[flag_out] = random.uniform(lower_bound[1],
upper_bound[1], Nout)
r0_sq = np.square(x0) + np.square(y0)
flag_out = np.where(r0_sq > radius * radius)[0]
Nout = len(flag_out)
w0 = float(init_num_elecs) / float(init_num_elecs_mp)
elif chamber_type == 'crab':
x0 = random.uniform(lower_bound[0], upper_bound[0],
init_num_elecs_mp)
y0 = random.uniform(lower_bound[1], upper_bound[1],
init_num_elecs_mp)
z0 = random.uniform(lower_bound[2], upper_bound[2],
init_num_elecs_mp)
vx0 = np.zeros(init_num_elecs_mp)
vy0 = np.zeros(init_num_elecs_mp)
vz0 = np.zeros(init_num_elecs_mp)
flag_out = np.where(
np.logical_and(
np.logical_and(
np.logical_and(x0 < l_main_int_x,
x0 > -l_main_int_x),
np.logical_and(z0 < l_main_int_z,
z0 > -l_main_int_z)),
np.logical_or(y0 > 0.5 * l_beam_pipe,
y0 < -0.5 * l_beam_pipe)))[0]
Nout = len(flag_out)
while Nout > 0:
x0[flag_out] = random.uniform(lower_bound[0],
upper_bound[0], Nout)
y0[flag_out] = random.uniform(lower_bound[1],
upper_bound[1], Nout)
z0[flag_out] = random.uniform(lower_bound[2],
upper_bound[2], Nout)
flag_out = np.where(
np.logical_and(
np.logical_and(
np.logical_and(x0 < l_main_int_x,
x0 > -l_main_int_x),
np.logical_and(z0 < l_main_int_z,
z0 > -l_main_int_z)),
np.logical_or(y0 > 0.5 * l_beam_pipe,
y0 < -0.5 * l_beam_pipe)))[0]
Nout = len(flag_out)
w0 = float(init_num_elecs) / float(init_num_elecs_mp)
else:
x0 = random.uniform(lower_bound[0], upper_bound[0],
init_num_elecs_mp)
y0 = random.uniform(lower_bound[1], upper_bound[1],
init_num_elecs_mp)
z0 = random.uniform(lower_bound[2], upper_bound[2],
init_num_elecs_mp)
vx0 = np.zeros(init_num_elecs_mp)
vy0 = np.zeros(init_num_elecs_mp)
vz0 = np.zeros(init_num_elecs_mp)
w0 = float(init_num_elecs) / float(init_num_elecs_mp)
self.b_pass_prev = 0
electron_background_dist = picmi.ParticleListDistribution(x=x0,
y=y0,
z=z0,
vx=vx0,
vy=vy0,
vz=vz0,
weight=w0)
self.elecb = picmi.Species(
particle_type='electron',
particle_shape='linear',
name='Electron background',
initial_distribution=electron_background_dist)
self.secelec = picmi.Species(particle_type='electron',
particle_shape='linear',
name='Secondary electrons')
##########################
# Numeric components
##########################
if mysolver == 'ES':
lower_boundary_conditions = ['dirichlet', 'dirichlet', 'dirichlet']
upper_boundary_conditions = ['dirichlet', 'dirichlet', 'dirichlet']
if mysolver == 'EM':
lower_boundary_conditions = ['open', 'open', 'open']
upper_boundary_conditions = ['open', 'open', 'open']
grid = picmi.Cartesian3DGrid(
number_of_cells=[self.nx, self.ny, self.nz],
lower_bound=[self.xmin, self.ymin, self.zmin],
upper_bound=[self.xmax, self.ymax, self.zmax],
lower_boundary_conditions=lower_boundary_conditions,
upper_boundary_conditions=upper_boundary_conditions)
if mysolver == 'ES':
solver = picmi.ElectrostaticSolver(grid=grid)
if mysolver == 'EM':
smoother = picmi.BinomialSmoother(n_pass=[[1], [1], [1]],
compensation=[[False], [False],
[False]],
stride=[[1], [1], [1]],
alpha=[[0.5], [0.5], [0.5]])
solver = picmi.ElectromagneticSolver(
grid=grid,
method='CKC',
cfl=1.,
source_smoother=smoother,
warp_l_correct_num_Cherenkov=False,
warp_type_rz_depose=0,
warp_l_setcowancoefs=True,
warp_l_getrho=False)
##########################
# Simulation setup
##########################
if chamber_type == 'rect':
print('Using rectangular chamber with xaper: %1.2e, yaper: %1.2e' %
(width / 2., height / 2.))
upper_box = picmi.warp.YPlane(y0=height / 2., ysign=1, condid=1)
lower_box = picmi.warp.YPlane(y0=-height / 2., ysign=-1, condid=1)
left_box = picmi.warp.XPlane(x0=width / 2., xsign=1, condid=1)
right_box = picmi.warp.XPlane(x0=-width / 2., xsign=-1, condid=1)
sim = picmi.Simulation(solver=solver,
verbose=1,
cfl=1.0,
warp_initialize_solver_after_generate=1)
sim.conductors = upper_box + lower_box + left_box + right_box
elif chamber_type == 'LHC':
print('Using the LHC chamber')
#pipe_annulus = picmi.warp.ZAnnulus(rmin = radius, rmax = radius+ghost/2, length = l, condid =1)
upper_box = picmi.warp.YPlane(y0=height / 2, ysign=1, condid=1)
lower_box = picmi.warp.YPlane(y0=-height / 2, ysign=-1, condid=1)
pipe = picmi.warp.ZCylinderOut(radius=radius,
length=z_length,
condid=1)
sim = picmi.Simulation(solver=solver,
verbose=1,
cfl=1.0,
warp_initialize_solver_after_generate=1)
sim.conductors = pipe + upper_box + lower_box
elif chamber_type == 'circle':
print('Using a circular chamber with radius %1.2e' % radius)
pipe_annulus = picmi.warp.ZAnnulus(rmin=radius,
rmax=radius + ghost,
length=l,
voltage=0.,
xcent=0.,
ycent=0.,
zcent=0.,
condid=1)
pipe = picmi.warp.ZCylinderOut(radius=radius, length=l, condid=1)
sim = picmi.Simulation(solver=solver,
verbose=1,
cfl=1.0,
warp_initialize_solver_after_generate=1)
sim.conductors = pipe_annulus
elif chamber_type == 'crab':
print('Simulating ECLOUD in a crab cavity')
box1 = picmi.warp.Box(zsize=self.zmax - self.zmin,
xsize=self.xmax - self.xmin,
ysize=self.ymax - self.ymin)
box2 = picmi.warp.Box(zsize=self.zmax - self.zmin,
xsize=l_beam_pipe,
ysize=l_beam_pipe)
box3 = picmi.warp.Box(zsize=l_main_z,
xsize=l_main_x,
ysize=l_main_y)
# Shape the cavity
ycen1 = 0.5 * l_beam_pipe + l_main_int_y
ycen2 = -ycen1
box4 = picmi.warp.Box(zsize=2 * l_main_int_z,
xsize=2 * l_main_int_x,
ysize=2 * l_main_int_y,
ycent=ycen1)
box5 = picmi.warp.Box(zsize=2 * l_main_int_z,
xsize=2 * l_main_int_x,
ysize=2 * l_main_int_y,
ycent=ycen2)
sim = picmi.Simulation(solver=solver,
verbose=1,
cfl=1.0,
warp_initialize_solver_after_generate=1)
sim.conductors = box1 - box2 - box3 + box4 + box5
sim.add_species(self.beam,
layout=None,
initialize_self_field=solver == 'EM')
self.elecb_layout = picmi.PseudoRandomLayout(
n_macroparticles=init_num_elecs_mp, seed=3)
sim.add_species(self.elecb,
layout=self.elecb_layout,
initialize_self_field=solver == 'EM')
sim.add_species(self.secelec, layout=None, initialize_self_field=False)
#########################
# Add Dipole
#########################
if lattice_elem == 'dipole':
picmi.warp.addnewdipo(zs=zs_dipo, ze=ze_dipo, by=By)
elif lattice_elem == 'crab':
[self.x, self.y, self.z, self.Ex, self.Ey,
self.Ez] = picmi.getdatafromtextfile("efield.txt",
nskip=1,
dims=[6, None])
[_, _, _, Hx, Hy, Hz] = picmi.getdatafromtextfile("hfield.txt",
nskip=1,
dims=[6, None])
self.Bx = Hx * picmi.mu0
self.By = Hy * picmi.mu0
self.Bz = Hz * picmi.mu0
### Interpolate them at cell centers (as prescribed by Warp doc)
self.d = abs(self.x[1] - self.x[0])
### Number of mesh cells
self.NNx = int(round(2 * np.max(self.x) / self.d))
self.NNy = int(round(2 * np.max(self.y) / self.d))
self.NNz = int(round(2 * np.max(self.z) / self.d))
### Number of mesh vertices
self.nnx = self.NNx + 1
self.nny = self.NNy + 1
self.nnz = self.NNz + 1
self.Ex3d = self.Ex.reshape(self.nnz, self.nny,
self.nnx).transpose(2, 1, 0)
self.Ey3d = self.Ey.reshape(self.nnz, self.nny,
self.nnx).transpose(2, 1, 0)
self.Ez3d = self.Ez.reshape(self.nnz, self.nny,
self.nnx).transpose(2, 1, 0)
self.Bx3d = self.Bx.reshape(self.nnz, self.nny,
self.nnx).transpose(2, 1, 0)
self.By3d = self.By.reshape(self.nnz, self.nny,
self.nnx).transpose(2, 1, 0)
self.Bz3d = self.Bz.reshape(self.nnz, self.nny,
self.nnx).transpose(2, 1, 0)
self.x3d = self.x.reshape(self.nnz, self.nny,
self.nnx).transpose(2, 1, 0)
self.y3d = self.y.reshape(self.nnz, self.nny,
self.nnx).transpose(2, 1, 0)
self.z3d = self.z.reshape(self.nnz, self.nny,
self.nnx).transpose(2, 1, 0)
### Rescale the fields at convenience
maxE = 57 * 1e6
kk = maxE / np.max(abs(self.Ey3d))
self.Ex3d *= kk
self.Ey3d *= kk
self.Ez3d *= kk
self.Bx3d *= kk
self.By3d *= kk
self.Bz3d *= kk
self.Exx = 0.125 * (
self.Ex3d[0:-1, 0:-1, 0:-1] + self.Ex3d[0:-1:, 0:-1, 1:] +
self.Ex3d[0:-1, 1:, 0:-1] + self.Ex3d[0:-1, 1:, 1:] +
self.Ex3d[1:, 1:, 0:-1] + self.Ex3d[1:, 1:, 1:] +
self.Ex3d[1:, 0:-1, 1:] + self.Ex3d[1:, 0:-1, 0:-1])
self.Eyy = 0.125 * (
self.Ey3d[0:-1, 0:-1, 0:-1] + self.Ey3d[0:-1:, 0:-1, 1:] +
self.Ey3d[0:-1, 1:, 0:-1] + self.Ey3d[0:-1, 1:, 1:] +
self.Ey3d[1:, 1:, 0:-1] + self.Ey3d[1:, 1:, 1:] +
self.Ey3d[1:, 0:-1, 1:] + self.Ey3d[1:, 0:-1, 0:-1])
self.Ezz = 0.125 * (
self.Ez3d[0:-1, 0:-1, 0:-1] + self.Ez3d[0:-1:, 0:-1, 1:] +
self.Ez3d[0:-1, 1:, 0:-1] + self.Ez3d[0:-1, 1:, 1:] +
self.Ez3d[1:, 1:, 0:-1] + self.Ez3d[1:, 1:, 1:] +
self.Ez3d[1:, 0:-1, 1:] + self.Ez3d[1:, 0:-1, 0:-1])
self.Bxx = 0.125 * (
self.Bx3d[0:-1, 0:-1, 0:-1] + self.Bx3d[0:-1:, 0:-1, 1:] +
self.Bx3d[0:-1, 1:, 0:-1] + self.Bx3d[0:-1, 1:, 1:] +
self.Bx3d[1:, 1:, 0:-1] + self.Bx3d[1:, 1:, 1:] +
self.Bx3d[1:, 0:-1, 1:] + self.Bx3d[1:, 0:-1, 0:-1])
self.Byy = 0.125 * (
self.By3d[0:-1, 0:-1, 0:-1] + self.By3d[0:-1:, 0:-1, 1:] +
self.By3d[0:-1, 1:, 0:-1] + self.By3d[0:-1, 1:, 1:] +
self.By3d[1:, 1:, 0:-1] + self.By3d[1:, 1:, 1:] +
self.By3d[1:, 0:-1, 1:] + self.By3d[1:, 0:-1, 0:-1])
self.Bzz = 0.125 * (
self.Bz3d[0:-1, 0:-1, 0:-1] + self.Bz3d[0:-1:, 0:-1, 1:] +
self.Bz3d[0:-1, 1:, 0:-1] + self.Bz3d[0:-1, 1:, 1:] +
self.Bz3d[1:, 1:, 0:-1] + self.Bz3d[1:, 1:, 1:] +
self.Bz3d[1:, 0:-1, 1:] + self.Bz3d[1:, 0:-1, 0:-1])
self.xx = 0.125 * (
self.x3d[0:-1, 0:-1, 0:-1] + self.x3d[0:-1:, 0:-1, 1:] +
self.x3d[0:-1, 1:, 0:-1] + self.x3d[0:-1, 1:, 1:] +
self.x3d[1:, 1:, 0:-1] + self.x3d[1:, 1:, 1:] +
self.x3d[1:, 0:-1, 1:] + self.x3d[1:, 0:-1, 0:-1])
self.yy = 0.125 * (
self.y3d[0:-1, 0:-1, 0:-1] + self.y3d[0:-1:, 0:-1, 1:] +
self.y3d[0:-1, 1:, 0:-1] + self.y3d[0:-1, 1:, 1:] +
self.y3d[1:, 1:, 0:-1] + self.y3d[1:, 1:, 1:] +
self.y3d[1:, 0:-1, 1:] + self.y3d[1:, 0:-1, 0:-1])
self.zz = 0.125 * (
self.z3d[0:-1, 0:-1, 0:-1] + self.z3d[0:-1:, 0:-1, 1:] +
self.z3d[0:-1, 1:, 0:-1] + self.z3d[0:-1, 1:, 1:] +
self.z3d[1:, 1:, 0:-1] + self.z3d[1:, 1:, 1:] +
self.z3d[1:, 0:-1, 1:] + self.z3d[1:, 0:-1, 0:-1])
### Lattice spatial parameters
self.zs = np.min(self.z3d) - self.d / 2.
self.ze = np.max(self.z3d) + self.d / 2.
self.xs = np.min(self.x3d) - self.d / 2.
self.ys = np.min(self.y3d) - self.d / 2.
### Lattice temporal parameters
self.Tf = 25e-9
self.freq = 400 * 1e6
self.Nt = 1000
self.phase_disp = 0
self.time_array = np.linspace(0., self.Tf, self.Nt)
self.data_arrayE = np.sin(self.time_array * self.freq * 2 * np.pi +
self.phase_disp)
self.data_arrayB = np.sin(self.time_array * self.freq * 2 * np.pi -
np.pi / 2 + self.phase_disp)
### Create overlapped lattice elements to have E and B in the same region
self.ie, self.egrid = picmi.warp.addnewegrd(self.zs,
self.ze,
dx=self.d,
dy=self.d,
xs=self.xs,
ys=self.ys,
time=self.time_array,
data=self.data_arrayE,
ex=self.Ex3d,
ey=self.Ey3d,
ez=self.Ez3d)
picmi.warp.addnewbgrd(self.zs,
self.ze,
dx=self.d,
dy=self.d,
xs=self.xs,
ys=self.ys,
time=self.time_array,
data=self.data_arrayB,
bx=self.Bx3d,
by=self.By3d,
bz=self.Bz3d)
def time_prof(t):
val = 0
for i in range(0, n_bunches):
val += bunch_macro_particles * 1. / np.sqrt(
2 * np.pi * sigmat * sigmat) * np.exp(
-(t - i * b_spac - t_offs) *
(t - i * b_spac - t_offs) /
(2 * sigmat * sigmat)) * picmi.warp.top.dt
return val
def nonlinearsource():
NP = int(round(time_prof(top.time)))
x = random.normal(bunch_centroid_position[0], bunch_rms_size[0],
NP)
y = random.normal(bunch_centroid_position[1], bunch_rms_size[1],
NP)
z = bunch_centroid_position[2]
vx = random.normal(bunch_centroid_velocity[0],
bunch_rms_velocity[0], NP)
vy = random.normal(bunch_centroid_velocity[1],
bunch_rms_velocity[1], NP)
vz = picmi.warp.clight * np.sqrt(1 - 1. / (beam_gamma**2))
self.beam.wspecies.addparticles(x=x,
y=y,
z=z,
vx=vx,
vy=vy,
vz=vz,
gi=1. / beam_gamma,
w=bunch_w)
picmi.warp.installuserinjection(nonlinearsource)
##########################
# simulation run
##########################
sim.step(1)
solver.solver.installconductor(sim.conductors,
dfill=picmi.warp.largepos)
sim.step(1)
pp = warp.ParticleScraper(sim.conductors,
lsavecondid=1,
lsaveintercept=1,
lcollectlpdata=1)
sec = Secondaries(conductors=sim.conductors,
l_usenew=1,
pyecloud_secemi_object=sey_mod,
pyecloud_nel_mp_ref=pyecloud_nel_mp_ref,
pyecloud_fact_clean=pyecloud_fact_clean,
pyecloud_fact_split=pyecloud_fact_split)
sec.add(incident_species=self.elecb.wspecies,
emitted_species=self.secelec.wspecies,
conductor=sim.conductors)
sec.add(incident_species=self.secelec.wspecies,
emitted_species=self.secelec.wspecies,
conductor=sim.conductors)
#Subcycle(10)
# just some shortcuts
pw = picmi.warp
step = pw.step
if mysolver == 'ES':
pw.top.dt = dt
if flag_relativ_tracking:
pw.top.lrelativ = pw.true
else:
pw.top.lrelativ = pw.false
def myplots(l_force=0):
if l_force or self.n_step % self.stride_imgs == 0:
plt.close()
(Nx, Ny, Nz) = np.shape(self.secelec.wspecies.get_density())
fig, axs = plt.subplots(1, 2, figsize=(12, 4.5))
fig.subplots_adjust(left=0.05,
bottom=0.1,
right=0.97,
top=0.94,
wspace=0.15)
d = self.secelec.wspecies.get_density(
) + self.elecb.wspecies.get_density(
) + self.beam.wspecies.get_density()
d2 = self.secelec.wspecies.get_density(
) + self.elecb.wspecies.get_density()
im1 = axs[0].imshow(
d[:, :, int(Nz / 2)].T,
cmap='jet',
origin='lower',
vmin=0.2 * np.min(d2[:, :, int(Nz / 2)]),
vmax=0.8 * np.max(d2[:, :, int(Nz / 2)]),
extent=[self.xmin, self.xmax, self.ymin, self.ymax])
axs[0].set_xlabel('x [m]')
axs[0].set_ylabel('y [m]')
axs[0].set_title('e- density')
fig.colorbar(
im1,
ax=axs[0],
)
im2 = axs[1].imshow(
d[int(Nx / 2), :, :],
cmap='jet',
origin='lower',
vmin=0.2 * np.min(d2[int(Nx / 2), :, :]),
vmax=0.8 * np.max(d2[int(Nx / 2), :, :]),
extent=[self.zmin, self.zmax, self.ymin, self.ymax],
aspect='auto')
axs[1].set_xlabel('z [m]')
axs[1].set_ylabel('y [m]')
axs[1].set_title('e- density')
fig.colorbar(im2, ax=axs[1])
n_step = top.time / top.dt
figname = 'images/%d.png' % n_step
plt.savefig(figname)
if flag_save_video:
pw.installafterstep(myplots)
myplots(1)
self.ntsteps_p_bunch = b_spac / top.dt
self.tot_nsteps = int(np.ceil(b_spac * (n_bunches) / top.dt))
t_start = self.b_pass_prev * b_spac
tstep_start = int(round(t_start / top.dt))
# pre-allocate outputs
if self.flag_output and not (self.flag_checkpointing
and os.path.exists('temp_mps_info.mat')):
self.numelecs = np.zeros(self.tot_nsteps)
self.numelecs_tot = np.zeros(self.tot_nsteps)
self.N_mp = np.zeros(self.tot_nsteps)
self.xhist = np.zeros((n_bunches, nbins))
dict_out = {}
# aux variables
self.b_pass = 0
self.perc = 10
self.t0 = time.time()
# trapping warp std output
self.text_trap = {True: StringIO(), False: sys.stdout}[enable_trap]
self.original = sys.stdout
self.n_step = 0
name='Secondary electrons')
##########################
# Numeric components
##########################
if mysolver == 'ES':
lower_boundary_conditions = ['dirichlet', 'dirichlet', 'dirichlet']
upper_boundary_conditions = ['dirichlet', 'dirichlet', 'dirichlet']
if mysolver == 'EM':
lower_boundary_conditions = ['open', 'open', 'open']
upper_boundary_conditions = ['open', 'open', 'open']
grid = picmi.Cartesian3DGrid(
number_of_cells=[nx, ny, nz],
lower_bound=[xmin, ymin, zmin],
upper_bound=[xmax, ymax, zmax],
lower_boundary_conditions=lower_boundary_conditions,
upper_boundary_conditions=upper_boundary_conditions) #,
#warpx_max_grid_size=32)
if mysolver == 'ES':
solver = picmi.ElectrostaticSolver(grid=grid)
if mysolver == 'EM':
smoother = picmi.BinomialSmoother(n_pass=[[1], [1], [1]],
compensation=[[False], [False], [False]],
stride=[[1], [1], [1]],
alpha=[[0.5], [0.5], [0.5]])
solver = picmi.ElectromagneticSolver(
grid=grid,
def warp_pyecloud_dipole(z_length=None,
nx=None,
ny=None,
nz=None,
dh_t=None,
n_bunches=None,
b_spac=None,
beam_gamma=None,
sigmax=None,
sigmay=None,
sigmat=None,
bunch_intensity=None,
init_num_elecs=None,
init_num_elecs_mp=None,
By=None,
pyecloud_nel_mp_ref=None,
dt=None,
pyecloud_fact_clean=None,
pyecloud_fact_split=None,
chamber_type=None,
flag_save_video=None,
enable_trap=True,
Emax=None,
del_max=None,
R0=None,
E_th=None,
sigmafit=None,
mufit=None,
secondary_angle_distribution=None,
N_mp_max=None,
N_mp_target=None,
flag_checkpointing=False,
checkpoints=None,
flag_output=False,
bunch_macro_particles=None,
t_offs=None,
width=None,
height=None,
output_name='output.mat',
flag_relativ_tracking=False,
nbins=100,
r=None):
import numpy as np
import numpy.random as random
from warp import picmi
from scipy.stats import gaussian_kde
from warp.particles.Secondaries import Secondaries, top, warp, time
import matplotlib.pyplot as plt
import scipy.io as sio
from io import BytesIO as StringIO
from scipy.constants import c as clight
import sys
import PyECLOUD.myfilemanager as mfm
import os
# Construct PyECLOUD secondary emission object
import PyECLOUD.sec_emission_model_ECLOUD as seec
sey_mod = seec.SEY_model_ECLOUD(Emax=Emax,
del_max=del_max,
R0=R0,
E_th=E_th,
sigmafit=sigmafit,
mufit=mufit,
secondary_angle_distribution='cosine_3D')
##########################
# physics parameters
##########################
mysolver = 'ES' # solver type ('ES'=electrostatic; 'EM'=electromagnetic)
##########################
# numerics parameters
##########################
# --- grid
zs_dipo = -0.5 * z_length
ze_dipo = 0.5 * z_length
# Consistency check
if chamber_type == 'rect':
assert (height is not None) and (
width is not None
), 'When using rectangular chamber, height and width must be specified'
if chamber_type == 'circle':
assert r is not None, 'When using circular chamber r must be specified'
if chamber_type == 'LHC':
h = 18.e-3
l = z_length
unit = 1e-3
ghost = 1e-3
if chamber_type == 'rect':
xmin = -width / 2. - ghost
xmax = -xmin
ymin = -height / 2 - ghost
ymax = -ymin
lower_bound = [-width / 2, -height / 2, zs_dipo]
upper_bound = [width / 2, height / 2, ze_dipo]
else:
xmin = -r - ghost
xmax = r + ghost
ymin = -r - ghost
ymax = r + ghost
lower_bound = [-r, -r, zs_dipo]
upper_bound = [r, r, ze_dipo]
zmin = zs_dipo - 50 * unit
zmax = ze_dipo + 50 * unit
if chamber_type == 'LHC':
chamber_area = 0.0014664200235342726
elif chamber_type == 'rect':
chamber_area = width * height
elif chamber_type == 'circle':
chamber_area = np.pi * r * r
init_dist_area = (upper_bound[0] - lower_bound[0]) * (upper_bound[1] -
lower_bound[1])
##########################
# Beam parameters
##########################
sigmaz = sigmat * picmi.clight
bunch_w = bunch_intensity / bunch_macro_particles
print('DEBUG: bunch_w: %d' % bunch_w)
print('DEBUG: bunch_intensity: %d' % bunch_intensity)
print('DEBUG: bunch_macro_particles: %d' % bunch_macro_particles)
#######################################################
# compute beam size from normalized emittance and beta
# Uncomment if data available
#######################################################
bunch_rms_size = [sigmax, sigmay, sigmaz]
bunch_rms_velocity = [0., 0., 0.]
bunch_centroid_position = [0, 0, zs_dipo - 10 * unit]
bunch_centroid_velocity = [0., 0., beam_gamma * picmi.constants.c]
beam = picmi.Species(particle_type='proton',
particle_shape='linear',
name='beam')
temp_file_name = 'temp_mps_info.mat'
########################################################
# if checkopoint is found reload it,
# otherwise start from scratch
########################################################
if flag_checkpointing and os.path.exists(temp_file_name):
print('############################################################')
print('Temp distribution found. Regenarating and restarting')
print('############################################################')
dict_init_dist = sio.loadmat(temp_file_name)
# Load particles status
x0 = dict_init_dist['x_mp'][0]
y0 = dict_init_dist['y_mp'][0]
z0 = dict_init_dist['z_mp'][0]
ux0 = dict_init_dist['ux_mp'][0]
uy0 = dict_init_dist['uy_mp'][0]
uz0 = dict_init_dist['uz_mp'][0]
w0 = dict_init_dist['w_mp'][0]
# compute the velocities
invgamma = np.sqrt(1 - picmi.clight**2 /
(np.square(ux0) + np.square(uy0) + np.square(uz0)))
vx0 = np.multiply(invgamma, ux0)
vy0 = np.multiply(invgamma, uy0)
vz0 = np.multiply(invgamma, uy0)
# Reload the outputs and other auxiliary stuff
if flag_output:
numelecs = dict_init_dist['numelecs'][0]
N_mp = dict_init_dist['N_mp'][0]
numelecs_tot = dict_init_dist['numelecs_tot'][0]
xhist = dict_init_dist['xhist'][0]
bins = dict_init_dist['bins'][0]
b_pass_prev = dict_init_dist['b_pass'] - 1
else:
x0 = random.uniform(lower_bound[0], upper_bound[0], init_num_elecs_mp)
y0 = random.uniform(lower_bound[1], upper_bound[1], init_num_elecs_mp)
z0 = random.uniform(lower_bound[2], upper_bound[2], init_num_elecs_mp)
vx0 = np.zeros(init_num_elecs_mp)
vy0 = np.zeros(init_num_elecs_mp)
vz0 = np.zeros(init_num_elecs_mp)
w0 = float(init_num_elecs) / float(init_num_elecs_mp)
#Correct the weight to the chambers size
w0 = w0 * init_dist_area / chamber_area
b_pass_prev = 0
electron_background_dist = picmi.ParticleListDistribution(x=x0,
y=y0,
z=z0,
vx=vx0,
vy=vy0,
vz=vz0,
weight=w0)
elecb = picmi.Species(particle_type='electron',
particle_shape='linear',
name='Electron background',
initial_distribution=electron_background_dist)
secelec = picmi.Species(particle_type='electron',
particle_shape='linear',
name='Secondary electrons')
##########################
# Numeric components
##########################
if mysolver == 'ES':
lower_boundary_conditions = ['dirichlet', 'dirichlet', 'dirichlet']
upper_boundary_conditions = ['dirichlet', 'dirichlet', 'dirichlet']
if mysolver == 'EM':
lower_boundary_conditions = ['open', 'open', 'open']
upper_boundary_conditions = ['open', 'open', 'open']
grid = picmi.Cartesian3DGrid(
number_of_cells=[nx, ny, nz],
lower_bound=[xmin, ymin, zmin],
upper_bound=[xmax, ymax, zmax],
lower_boundary_conditions=lower_boundary_conditions,
upper_boundary_conditions=upper_boundary_conditions)
if mysolver == 'ES':
solver = picmi.ElectrostaticSolver(grid=grid)
if mysolver == 'EM':
smoother = picmi.BinomialSmoother(n_pass=[[1], [1], [1]],
compensation=[[False], [False],
[False]],
stride=[[1], [1], [1]],
alpha=[[0.5], [0.5], [0.5]])
solver = picmi.ElectromagneticSolver(
grid=grid,
method='CKC',
cfl=1.,
source_smoother=smoother,
warp_l_correct_num_Cherenkov=False,
warp_type_rz_depose=0,
warp_l_setcowancoefs=True,
warp_l_getrho=False)
##########################
# Simulation setup
##########################
if chamber_type == 'rect':
print 'Using rectangular chamber with xaper: %1.2e, yaper: %1.2e' % (
width / 2., height / 2.)
upper_box = picmi.warp.YPlane(y0=height / 2., ysign=1, condid=1)
lower_box = picmi.warp.YPlane(y0=-height / 2., ysign=-1, condid=1)
left_box = picmi.warp.XPlane(x0=width / 2., xsign=1, condid=1)
right_box = picmi.warp.XPlane(x0=-width / 2., xsign=-1, condid=1)
sim = picmi.Simulation(solver=solver,
verbose=1,
cfl=1.0,
warp_initialize_solver_after_generate=1)
sim.conductors = upper_box + lower_box + left_box + right_box
elif chamber_type == 'LHC':
print 'Using the LHC chamber'
pipe_annulus = picmi.warp.ZAnnulus(rmin=r,
rmax=r + ghost / 2,
length=l,
condid=1)
upper_box = picmi.warp.YPlane(y0=h, ysign=1, condid=1)
lower_box = picmi.warp.YPlane(y0=-h, ysign=-1, condid=1)
pipe = picmi.warp.ZCylinderOut(radius=r, length=l, condid=1)
sim = picmi.Simulation(solver=solver,
verbose=1,
cfl=1.0,
warp_initialize_solver_after_generate=1)
sim.conductors = pipe + upper_box + lower_box
elif chamber_type == 'circle':
print 'Using a circular chamber with radius %d' % r
pipe_annulus = picmi.warp.ZAnnulus(rmin=r,
rmax=r + ghost,
length=l,
voltage=0.,
xcent=0.,
ycent=0.,
zcent=0.,
condid=1)
pipe = picmi.warp.ZCylinderOut(radius=r, length=l, condid=1)
sim = picmi.Simulation(solver=solver,
verbose=1,
cfl=1.0,
warp_initialize_solver_after_generate=1)
sim.conductors = pipe
beam_layout = picmi.PseudoRandomLayout(n_macroparticles=10**5, seed=3)
sim.add_species(beam, layout=None, initialize_self_field=solver == 'EM')
elecb_layout = picmi.PseudoRandomLayout(n_macroparticles=init_num_elecs_mp,
seed=3)
sim.add_species(elecb,
layout=elecb_layout,
initialize_self_field=solver == 'EM')
sim.add_species(secelec, layout=None, initialize_self_field=False)
#########################
# Add Dipole
#########################
picmi.warp.addnewdipo(zs=zs_dipo, ze=ze_dipo, by=By)
def time_prof(t):
val = 0
for i in range(0, n_bunches):
val += bunch_macro_particles * 1. / np.sqrt(
2 * np.pi * sigmat * sigmat) * np.exp(
-(t - i * b_spac - t_offs) * (t - i * b_spac - t_offs) /
(2 * sigmat * sigmat)) * picmi.warp.top.dt
return val
def nonlinearsource():
NP = int(round(time_prof(top.time)))
x = random.normal(bunch_centroid_position[0], bunch_rms_size[0], NP)
y = random.normal(bunch_centroid_position[1], bunch_rms_size[1], NP)
z = bunch_centroid_position[2]
vx = random.normal(bunch_centroid_velocity[0], bunch_rms_velocity[0],
NP)
vy = random.normal(bunch_centroid_velocity[1], bunch_rms_velocity[1],
NP)
vz = picmi.warp.clight * np.sqrt(1 - 1. / (beam_gamma**2))
beam.wspecies.addparticles(x=x,
y=y,
z=z,
vx=vx,
vy=vy,
vz=vz,
gi=1. / beam_gamma,
w=bunch_w)
picmi.warp.installuserinjection(nonlinearsource)
##########################
# simulation run
##########################
sim.step(1)
solver.solver.installconductor(sim.conductors, dfill=picmi.warp.largepos)
sim.step(1)
pp = warp.ParticleScraper(sim.conductors,
lsavecondid=1,
lsaveintercept=1,
lcollectlpdata=1)
sec = Secondaries(conductors=sim.conductors,
l_usenew=1,
pyecloud_secemi_object=sey_mod,
pyecloud_nel_mp_ref=pyecloud_nel_mp_ref,
pyecloud_fact_clean=pyecloud_fact_clean,
pyecloud_fact_split=pyecloud_fact_split)
sec.add(incident_species=elecb.wspecies,
emitted_species=secelec.wspecies,
conductor=sim.conductors)
sec.add(incident_species=secelec.wspecies,
emitted_species=secelec.wspecies,
conductor=sim.conductors)
#Subcycle(10)
# just some shortcuts
pw = picmi.warp
step = pw.step
if mysolver == 'ES':
pw.top.dt = dt
if flag_relativ_tracking:
pw.top.lrelativ = pw.true
else:
pw.top.lrelativ = pw.false
def myplots(l_force=0):
if l_force or pw.top.it % 10 == 0:
plt.close()
(Nx, Ny, Nz) = np.shape(secelec.wspecies.get_density())
fig, axs = plt.subplots(1, 2, figsize=(12, 4.5))
fig.subplots_adjust(left=0.05,
bottom=0.1,
right=0.97,
top=0.94,
wspace=0.15)
d = secelec.wspecies.get_density() + elecb.wspecies.get_density(
) + beam.wspecies.get_density()
d2 = secelec.wspecies.get_density() + elecb.wspecies.get_density()
im1 = axs[0].imshow(d2[:, :, Nz / 2].T,
cmap='jet',
origin='lower',
vmin=0.2 * np.min(d2[:, :, Nz / 2]),
vmax=0.8 * np.max(d2[:, :, Nz / 2]),
extent=[xmin, xmax, ymin, ymax])
axs[0].set_xlabel('x [m]')
axs[0].set_ylabel('y [m]')
axs[0].set_title('e- density')
fig.colorbar(
im1,
ax=axs[0],
)
im2 = axs[1].imshow(d2[Nx / 2, :, :],
cmap='jet',
origin='lower',
vmin=0.2 * np.min(d2[Nx / 2, :, :]),
vmax=0.8 * np.max(d2[Nx / 2, :, :]),
extent=[zmin, zmax, ymin, ymax],
aspect='auto')
axs[1].set_xlabel('z [m]')
axs[1].set_ylabel('y [m]')
axs[1].set_title('e- density')
fig.colorbar(im2, ax=axs[1])
n_step = top.time / top.dt
figname = 'images/%d.png' % n_step
plt.savefig(figname)
if flag_save_video:
pw.installafterstep(myplots)
myplots(1)
ntsteps_p_bunch = b_spac / top.dt
tot_nsteps = int(np.ceil(b_spac * (n_bunches) / top.dt))
t_start = b_pass_prev * b_spac
tstep_start = int(round(t_start / top.dt))
# pre-allocate outputs
if flag_output and not (flag_checkpointing
and os.path.exists('temp_mps_info.mat')):
numelecs = np.zeros(tot_nsteps)
numelecs_tot = np.zeros(tot_nsteps)
N_mp = np.zeros(tot_nsteps)
xhist = np.zeros((n_bunches, nbins))
dict_out = {}
# aux variables
b_pass = 0
perc = 10
original = sys.stdout
text_trap = StringIO()
t0 = time.time()
# trapping warp std output
text_trap = {True: StringIO(), False: sys.stdout}[enable_trap]
original = sys.stdout
for n_step in range(tstep_start, tot_nsteps):
# if a passage is starting...
if n_step / ntsteps_p_bunch >= b_pass + b_pass_prev:
b_pass += 1
perc = 10
# Measure the duration of the previous passage
if b_pass > 1:
t_pass_1 = time.time()
t_pass = t_pass_1 - t_pass_0
t_pass_0 = time.time()
# Dump outputs
if flag_output:
dict_out['numelecs'] = numelecs
dict_out['numelecs_tot'] = numelecs_tot
dict_out['N_mp'] = N_mp
# Compute the x-position histogram
(xhist[b_pass - 1],
bins) = np.histogram(secelec.wspecies.getx(),
range=(xmin, xmax),
bins=nbins,
weights=secelec.wspecies.getw(),
density=False)
dict_out['bins'] = bins
dict_out['xhist'] = xhist
sio.savemat(output_name, dict_out)
# Perform regeneration if needed
if secelec.wspecies.getn() > N_mp_max:
dict_out_temp = {}
print('Number of macroparticles: %d' %
(secelec.wspecies.getn()))
print('MAXIMUM LIMIT OF MPS HAS BEEN RACHED')
perform_regeneration(N_mp_target, secelec.wspecies, sec)
# Save stuff if checkpoint
if flag_checkpointing and np.any(checkpoints == b_pass +
b_pass_prev) and b_pass > 1:
dict_out_temp = {}
print('Saving a checkpoint!')
secelec_w = secelec.wspecies.getw()
dict_out_temp['x_mp'] = np.concatenate(
(secelec.wspecies.getx(), elecb.wspecies.getx()))
dict_out_temp['y_mp'] = np.concatenate(
(secelec.wspecies.gety(), elecb.wspecies.gety()))
dict_out_temp['z_mp'] = np.concatenate(
(secelec.wspecies.getz(), elecb.wspecies.gety()))
dict_out_temp['ux_mp'] = np.concatenate(
(secelec.wspecies.getux(), elecb.wspecies.getux()))
dict_out_temp['uy_mp'] = np.concatenate(
(secelec.wspecies.getuy(), elecb.wspecies.getuy()))
dict_out_temp['uz_mp'] = np.concatenate(
(secelec.wspecies.getuz(), elecb.wspecies.getuz()))
dict_out_temp['w_mp'] = np.concatenate(
(secelec_w, elecb.wspecies.getw()))
if flag_output:
dict_out_temp['numelecs'] = numelecs
dict_out_temp['numelecs_tot'] = numelecs_tot
dict_out_temp['N_mp'] = N_mp
dict_out_temp['xhist'] = xhist
dict_out_temp['bins'] = bins
dict_out_temp['b_pass'] = b_pass + b_pass_prev
filename = 'temp_mps_info.mat'
sio.savemat(filename, dict_out_temp)
print('===========================')
print('Bunch passage: %d' % (b_pass + b_pass_prev))
print('Number of electrons in the dipole: %d' %
(np.sum(secelec.wspecies.getw()) +
np.sum(elecb.wspecies.getw())))
print('Number of macroparticles: %d' %
(secelec.wspecies.getn() + elecb.wspecies.getn()))
if b_pass > 1:
print('Previous passage took %ds' % t_pass)
if n_step % ntsteps_p_bunch / ntsteps_p_bunch * 100 > perc:
print('%d%% of bunch passage' % perc)
perc = perc + 10
# Perform a step
sys.stdout = text_trap
step(1)
sys.stdout = original
#if secelec.wspecies.getn()>0 and elecb.wspecies.getn()>0:
# print(max(max(np.sqrt(np.square(elecb.wspecies.getvx())+np.square(elecb.wspecies.getvy())+np.square(elecb.wspecies.getvz()))), max(np.sqrt(np.square(secelec.wspecies.getvx())+np.square(secelec.wspecies.getvy())+np.square(secelec.wspecies.getvz())))))
# Store stuff to be saved
if flag_output:
secelec_w = secelec.wspecies.getw()
elecb_w = elecb.wspecies.getw()
elecs_density = secelec.wspecies.get_density(
l_dividebyvolume=0)[:, :,
int(nz / 2.)] + elecb.wspecies.get_density(
l_dividebyvolume=0)[:, :,
int(nz / 2.)]
numelecs[n_step] = np.sum(elecs_density)
elecs_density_tot = secelec.wspecies.get_density(
l_dividebyvolume=0)[:, :, :] + elecb.wspecies.get_density(
l_dividebyvolume=0)[:, :, :]
numelecs_tot[n_step] = np.sum(elecs_density_tot)
N_mp[n_step] = len(secelec_w) + len(elecb_w)
# Timer
t1 = time.time()
totalt = t1 - t0
# Dump outputs
if flag_output:
dict_out['numelecs'] = numelecs
dict_out['N_mp'] = N_mp
dict_out['numelecs_tot'] = numelecs_tot
dict_out['xhist'] = xhist
dict_out['bins'] = bins
sio.savemat(output_name, dict_out)
# Delete checkpoint if found
if flag_checkpointing and os.path.exists('temp_mps_info.mat'):
os.remove('temp_mps_info.mat')
print('Run terminated in %ds' % totalt)
def __init__(self,
z_length=None,
nx=None,
ny=None,
nz=None,
solver_type='ES',
n_bunches=None,
b_spac=None,
beam_gamma=None,
sigmax=None,
sigmay=None,
sigmat=None,
bunch_intensity=None,
init_num_elecs=None,
init_num_elecs_mp=None,
By=None,
N_subcycle=None,
pyecloud_nel_mp_ref=None,
dt=None,
pyecloud_fact_clean=None,
pyecloud_fact_split=None,
chamber_type=None,
flag_save_video=None,
enable_trap=True,
Emax=None,
del_max=None,
R0=None,
E_th=None,
sigmafit=None,
mufit=None,
secondary_angle_distribution=None,
N_mp_max=None,
N_mp_target=None,
flag_checkpointing=False,
checkpoints=None,
flag_output=False,
bunch_macro_particles=None,
t_offs=None,
width=None,
height=None,
output_filename='output.mat',
flag_relativ_tracking=False,
nbins=100,
radius=None,
ghost=None,
ghost_z=None,
stride_imgs=10,
stride_output=1000,
chamber=False,
lattice_elem=None,
temps_filename='temp_mps_info.mat',
custom_plot=None):
# Construct PyECLOUD secondary emission object
sey_mod = seec.SEY_model_ECLOUD(
Emax=Emax,
del_max=del_max,
R0=R0,
E_th=E_th,
sigmafit=sigmafit,
mufit=mufit,
secondary_angle_distribution='cosine_3D')
self.nbins = nbins
self.N_mp_target = N_mp_target
self.N_mp_max = N_mp_max
self.nx = nx
self.ny = ny
self.nz = nz
self.flag_checkpointing = flag_checkpointing
self.checkpoints = checkpoints
self.flag_output = flag_output
self.output_filename = output_filename
self.stride_imgs = stride_imgs
self.stride_output = stride_output
self.beam_gamma = beam_gamma
self.chamber = chamber
self.init_num_elecs_mp = init_num_elecs_mp
self.init_num_elecs = init_num_elecs
self.n_bunches = n_bunches
self.bunch_macro_particles = bunch_macro_particles
self.sigmat = sigmat
self.b_spac = b_spac
self.t_offs = t_offs
self.temps_filename = temps_filename
# Just some shortcuts
pw = picmi.warp
step = pw.step
if solver_type == 'ES':
pw.top.dt = dt
if flag_relativ_tracking:
pw.top.lrelativ = pw.true
else:
pw.top.lrelativ = pw.false
self.tot_nsteps = int(np.ceil(b_spac * (n_bunches) / top.dt))
self.saver = Saver(flag_output,
flag_checkpointing,
self.tot_nsteps,
n_bunches,
nbins,
temps_filename=temps_filename,
output_filename=output_filename)
# Beam parameters
sigmaz = sigmat * picmi.clight
if bunch_macro_particles > 0:
self.bunch_w = bunch_intensity / bunch_macro_particles
else:
self.bunch_w = 0
self.bunch_rms_size = [sigmax, sigmay, sigmaz]
self.bunch_rms_velocity = [0., 0., 0.]
self.bunch_centroid_position = [0, 0, chamber.zmin + 10e-3]
self.bunch_centroid_velocity = [0., 0., beam_gamma * picmi.constants.c]
# Instantiate beam
self.beam = picmi.Species(particle_type='proton',
particle_shape='linear',
name='beam')
# If checkopoint is found reload it,
# otherwise start with uniform distribution
self.temp_file_name = 'temp_mps_info.mat'
if self.flag_checkpointing and os.path.exists(self.temp_file_name):
electron_background_dist = self.load_elec_density()
else:
electron_background_dist = self.init_uniform_density()
self.elecb = picmi.Species(
particle_type='electron',
particle_shape='linear',
name='Electron background',
initial_distribution=electron_background_dist)
self.secelec = picmi.Species(particle_type='electron',
particle_shape='linear',
name='Secondary electrons')
# Setup grid and boundary conditions
if solver_type == 'ES':
lower_bc = ['dirichlet', 'dirichlet', 'dirichlet']
upper_bc = ['dirichlet', 'dirichlet', 'dirichlet']
if solver_type == 'EM':
lower_boundary_conditions = ['open', 'open', 'open']
upper_boundary_conditions = ['open', 'open', 'open']
grid = picmi.Cartesian3DGrid(
number_of_cells=[self.nx, self.ny, self.nz],
lower_bound=[chamber.xmin, chamber.ymin, chamber.zmin],
upper_bound=[chamber.xmax, chamber.ymax, chamber.zmax],
lower_boundary_conditions=lower_bc,
upper_boundary_conditions=upper_bc)
if solver_type == 'ES':
solver = picmi.ElectrostaticSolver(grid=grid)
elif solver_type == 'EM':
smoother = picmi.BinomialSmoother(n_pass=[[1], [1], [1]],
compensation=[[False], [False],
[False]],
stride=[[1], [1], [1]],
alpha=[[0.5], [0.5], [0.5]])
solver = picmi.ElectromagneticSolver(
grid=grid,
method='CKC',
cfl=1.,
source_smoother=smoother,
warp_l_correct_num_Cherenkov=False,
warp_type_rz_depose=0,
warp_l_setcowancoefs=True,
warp_l_getrho=False)
# Setup simulation
sim = picmi.Simulation(solver=solver,
verbose=1,
cfl=1.0,
warp_initialize_solver_after_generate=1)
sim.conductors = chamber.conductors
sim.add_species(self.beam,
layout=None,
initialize_self_field=solver == 'EM')
self.elecb_layout = picmi.PseudoRandomLayout(
n_macroparticles=init_num_elecs_mp, seed=3)
sim.add_species(self.elecb,
layout=self.elecb_layout,
initialize_self_field=solver == 'EM')
sim.add_species(self.secelec, layout=None, initialize_self_field=False)
picmi.warp.installuserinjection(self.bunched_beam)
sim.step(1)
solver.solver.installconductor(sim.conductors,
dfill=picmi.warp.largepos)
sim.step(1)
# Setup secondary emission stuff
pp = warp.ParticleScraper(sim.conductors,
lsavecondid=1,
lsaveintercept=1,
lcollectlpdata=1)
sec = Secondaries(conductors=sim.conductors,
l_usenew=1,
pyecloud_secemi_object=sey_mod,
pyecloud_nel_mp_ref=pyecloud_nel_mp_ref,
pyecloud_fact_clean=pyecloud_fact_clean,
pyecloud_fact_split=pyecloud_fact_split)
sec.add(incident_species=self.elecb.wspecies,
emitted_species=self.secelec.wspecies,
conductor=sim.conductors)
sec.add(incident_species=self.secelec.wspecies,
emitted_species=self.secelec.wspecies,
conductor=sim.conductors)
if N_subcycle is not None:
Subcycle(N_subcycle)
if custom_plot is not None:
plot_func = custom_plot
else:
plot_func = self.myplots
pw.installafterstep(plot_func)
plot_func(1)
self.ntsteps_p_bunch = b_spac / top.dt
t_start = self.b_pass_prev * b_spac
tstep_start = int(np.round(t_start / top.dt))
# aux variables
self.b_pass = 0
self.perc = 10
self.t0 = time.time()
# trapping warp std output
self.text_trap = {True: StringIO(), False: sys.stdout}[enable_trap]
self.original = sys.stdout
self.n_step = int(np.round(self.b_pass_prev * b_spac / dt))
upper_bound=[beam_xmax, beam_ymax, beam_zmax],
directed_velocity=[0., 0., beam_uz])
beam = picmi.Species(particle_type='electron',
particle_shape='linear',
name='beam',
initial_distribution=parabolic_beam)
##########################
# numerics components
##########################
grid = picmi.Cartesian3DGrid(
number_of_cells=[nx, ny, nz],
lower_bound=[xmin, ymin, zmin],
upper_bound=[xmax, ymax, zmax],
lower_boundary_conditions=['open', 'open', 'open'],
upper_boundary_conditions=['open', 'open', 'open'],
moving_window_velocity=moving_window_velocity,
warpx_max_grid_size=32)
smoother = picmi.BinomialSmoother(n_pass=[[1], [1], [1]],
compensation=[[False], [False], [False]],
stride=[[1], [1], [1]],
alpha=[[0.5], [0.5], [0.5]])
solver = picmi.ElectromagneticSolver(grid=grid,
method='Yee',
cfl=1.,
source_smoother=smoother,
field_smoother=smoother,
warp_l_correct_num_Cherenkov=True,
def __init__(self, fieldsolver_inputs=None,
beam_inputs=None, ecloud_inputs=None,
antenna_inputs=None,
saving_inputs=None,
simulation_inputs=None):
self.defaultsfromdict(self.__fieldsolver_inputs__, fieldsolver_inputs)
self.defaultsfromdict(self.__beam_inputs__, beam_inputs)
self.defaultsfromdict(self.__ecloud_inputs__, ecloud_inputs)
self.defaultsfromdict(self.__antenna_inputs__, antenna_inputs)
self.defaultsfromdict(self.__saving_inputs__, saving_inputs)
self.defaultsfromdict(self.__simulation_inputs__, simulation_inputs)
# Construct PyECLOUD secondary emission object
self.sey_mod = seec.SEY_model_ECLOUD(Emax=self.Emax,
del_max=self.del_max,
R0=self.R0, E_th=self.E_th,
sigmafit=self.sigmafit,
mufit=self.mufit,
secondary_angle_distribution='cosine_3D')
self.beam_beta = np.sqrt(1 - 1 / (self.beam_gamma ** 2))
if self.n_bunches is not None and self.tot_nsteps is not None:
print("""WARNING: if both n_bunches and tot_nsteps are specified
tot_nsteps is going to be ignored and the number of steps is
going to be determined basing on n_bunches and dt""")
if self.n_bunches is not None and self.t_end is not None:
print("""WARNING: if both n_bunches and t_end are specified
tot_nsteps is going to be ignored and the number of steps is
going to be determined basing on n_bunches and dt""")
if not os.path.exists(self.images_dir) and picmi.warp.me == 0:
os.makedirs(self.images_dir)
# Just some shortcuts
pw = picmi.warp
if self.custom_time_prof is None:
self.time_prof = self.gaussian_time_prof
else:
self.time_prof = self.self_wrapped_custom_time_prof
if self.solver_type == 'EM':
if self.dt is not None:
print('WARNING: dt is going to be ignored for the EM solver')
else:
pw.top.dt = self.dt
if self.flag_relativ_tracking:
pw.top.lrelativ = pw.true
else:
pw.top.lrelativ = pw.false
# Beam parameters
self.sigmaz = self.sigmat * picmi.clight
if self.bunch_macro_particles > 0:
self.bunch_w = self.bunch_intensity / self.bunch_macro_particles
else:
self.bunch_w = 0
self.bunch_rms_size = [self.sigmax, self.sigmay, self.sigmaz]
self.bunch_rms_velocity = [0., 0., 0.]
self.bunch_centroid_position = [0, 0, self.chamber.z_inj_beam]
self.bunch_centroid_velocity = [0., 0., self.beam_beta * picmi.constants.c]
self.species_names = ['beam', 'ecloud']
# Instantiate beam
self.elecs_injected = False
if self.flag_checkpointing and os.path.exists(self.temps_filename):
self.ecloud = picmi.Species(particle_type='electron',
particle_shape='linear',
name=self.species_names[1],
initial_distribution=self.load_elec_density())
self.elecs_injected = True
else:
self.ecloud = picmi.Species(particle_type='electron',
particle_shape='linear',
name=self.species_names[1])
self.beam = picmi.Species(particle_type='proton',
particle_shape='linear',
name=self.species_names[0],
warp_fselfb = self.bunch_centroid_velocity[2])
self.b_pass = 0
# Setup grid and boundary conditions
self.dir_bc = ['dirichlet', 'dirichlet', 'dirichlet']
self.pml_bc = ['open', 'open', 'open']
#self.pml_bc = ['dirichlet', 'dirichlet', 'dirichlet']
self.number_of_cells = [self.nx, self.ny, self.nz]
self.lower_bound = [self.chamber.xmin, self.chamber.ymin, self.chamber.zmin]
self.upper_bound = [self.chamber.xmax, self.chamber.ymax, self.chamber.zmax]
self.grid_EM = picmi.Cartesian3DGrid(number_of_cells=self.number_of_cells,
lower_bound=self.lower_bound,
upper_bound=self.upper_bound,
lower_boundary_conditions=self.pml_bc,
upper_boundary_conditions=self.pml_bc)
grid_ES = picmi.Cartesian3DGrid(number_of_cells=self.number_of_cells,
lower_bound=self.lower_bound,
upper_bound=self.upper_bound,
lower_boundary_conditions=self.dir_bc,
upper_boundary_conditions=self.dir_bc)
if self.solver_type == 'ES':
self.solver = picmi.ElectrostaticSolver(grid=grid_ES,
warp_conductors = self.chamber.conductors,
warp_conductor_dfill = picmi.warp.largepos)
elif self.solver_type == 'EM':
if self.source_smoothing:
n_pass = [[1], [1], [1]]
stride = [[1], [1], [1]]
compensation = [[False], [False], [False]]
alpha = [[0.5], [0.5], [0.5]]
smoother = picmi.BinomialSmoother(n_pass=n_pass,
compensation=compensation,
stride=stride,
alpha=alpha)
else:
smoother = None
if hasattr(self, 'laser_func'):
self.solver = picmi.ElectromagneticSolver(grid=self.grid_EM,
method=self.EM_method, cfl=self.cfl,
source_smoother=smoother,
warp_l_correct_num_Cherenkov=False,
warp_type_rz_depose=0,
warp_l_setcowancoefs=True,
warp_l_getrho=False,
warp_laser_func=self.laser_func,
warp_laser_source_z=self.laser_source_z,
warp_laser_polangle=self.laser_polangle,
warp_laser_emax=self.laser_emax,
warp_laser_xmin=self.laser_xmin,
warp_laser_xmax=self.laser_xmax,
warp_laser_ymin=self.laser_ymin,
warp_laser_ymax=self.laser_ymax,
warp_conductors = self.chamber.conductors,
warp_conductor_dfill = picmi.warp.largepos,
warp_deposition_species =[self.ecloud.wspecies],
warp_iselfb_list = [0])
else:
self.solver = picmi.ElectromagneticSolver(grid=self.grid_EM,
method=self.EM_method, cfl=self.cfl,
source_smoother=smoother,
warp_l_correct_num_Cherenkov=False,
warp_type_rz_depose=0,
warp_l_setcowancoefs=True,
warp_l_getrho=False,
warp_conductors = self.chamber.conductors,
warp_conductor_dfill = picmi.warp.largepos,
warp_deposition_species =[self.ecloud.wspecies],
warp_iselfb_list = [0])
# Setup simulation
self.sim = picmi.Simulation(solver=self.solver, verbose=1,
warp_initialize_solver_after_generate=1)
self.sim.add_species(self.beam, layout=None,
initialize_self_field=False)
self.ecloud_layout = picmi.PseudoRandomLayout(
n_macroparticles=self.init_num_elecs_mp,
seed=3)
#tot_cells = 2 #self.nx*self.ny*self.nz
#self.ecloud_layout = picmi.GriddedLayout(grid=self.grid_EM,
# n_macroparticle_per_cell= np.array([2, 2, 2])) #np.array([int(self.init_num_elecs_mp/tot_cells),
#int(self.init_num_elecs_mp/tot_cells),
#int(self.init_num_elecs_mp/tot_cells)]))
self.sim.add_species(self.ecloud, layout=self.ecloud_layout)
# initialize_self_field= self.init_ecloud_fields)
#init_field = self.solver_type=='EM'
self.sim.step(1)
if self.init_num_elecs_mp > 0 and (not self.elecs_injected):
if self.t_inject_elec == 0:
self.init_uniform_density()
else:
picmi.warp.installuserinjection(self.init_uniform_density)
if self.tot_nsteps is None and self.n_bunches is not None:
self.tot_nsteps = int(np.round(self.b_spac*(self.n_bunches)/top.dt))
if self.tot_nsteps is None and self.t_end is not None:
self.tot_nsteps = int(np.round(self.t_end / top.dt))
# to be fixed
if self.t_end is not None:
self.tot_nsteps = int(np.round(self.t_end / top.dt))
elif self.tot_nsteps is None and self.n_bunches is None:
raise Exception('One between n_bunches, tot_nsteps, t_end has to be specified')
# needed to be consistent with the conductors
self.solver.solver.current_cor = False
self.flag_first_pass = True
if self.bunch_macro_particles > 0:
picmi.warp.installuserinjection(self.bunched_beam)
# Initialize the EM fields
# if self.init_em_fields:
# em = self.solver.solver
# me = pw.me
#
# fields = dict_of_arrays_and_scalar_from_h5_serial(self.folder_em_fields+'/'+str(picmi.warp.me)+'/'+self.file_em_fields)
# em.fields.Ex = fields['ex']*self.em_scale_fac
# em.fields.Ey = fields['ey']*self.em_scale_fac
# em.fields.Ez = fields['ez']*self.em_scale_fac
# em.fields.Bx = fields['bx']*self.em_scale_fac
# em.fields.By = fields['by']*self.em_scale_fac
# em.fields.Bz = fields['bz']*self.em_scale_fac
# em.setebp()
# Setup secondary emission stuff
self.part_scraper = ParticleScraper(self.chamber.conductors, lsavecondid=1, lsaveintercept=1, lcollectlpdata=1)
self.sec = Secondaries(conductors=self.chamber.conductors, l_usenew=1,
pyecloud_secemi_object=self.sey_mod,
pyecloud_nel_mp_ref=self.pyecloud_nel_mp_ref,
pyecloud_fact_clean=self.pyecloud_fact_clean,
pyecloud_fact_split=self.pyecloud_fact_split)
self.sec.add(incident_species=self.ecloud.wspecies,
emitted_species=self.ecloud.wspecies,
conductor=self.chamber.conductors)
self.saver = Saver(self.flag_output, self.flag_checkpointing,
self.nbins,
self.solver, self.sec, temps_filename=self.temps_filename,
output_filename=self.output_filename,
probe_filename = self.probe_filename,
tot_nsteps = self.tot_nsteps, n_bunches = self.n_bunches,
flag_save_ek_impacts = self.flag_save_ek_impacts)
self.ntsteps_p_bunch = int(np.round(self.b_spac / top.dt))
self.n_step = int(np.round(self.b_pass * self.ntsteps_p_bunch))
if self.N_subcycle is not None:
Subcycle(self.N_subcycle)
if self.custom_plot is not None:
pw.installafterstep(self.self_wrapped_custom_plot)
# Install field probes
if len(self.field_probes) > 0:
self.saver.init_field_probes(np.shape(self.field_probes)[0],
self.tot_nsteps,
self.field_probes_dump_stride)
pw.installafterstep(self.self_wrapped_probe_fun)
# Install other user-specified functions
for fun in self.after_step_fun_list:
pw.installafterstep(fun)
# aux variables
self.perc = 10
self.t0 = time.time()
# trapping warp std output
self.text_trap = {True: StringIO(), False: sys.stdout}[self.enable_trap]
self.original = sys.stdout
self.print_solvers_info()
