class warp_pyecloud_sim(object): __fieldsolver_inputs__ = {'nx': None, 'ny': None, 'nz': None, 'solver_type': 'ES', 'N_subcycle': None, 'EM_method': 'Yee', 'cfl': 1.0, 'dt': None, 'source_smoothing': True} __beam_inputs__ = {'n_bunches': None, 'b_spac': None, 'beam_gamma': None, 'sigmax': None, 'sigmay': None, 'sigmat': None, 'bunch_intensity': None, 't_offs': None, 'bunch_macro_particles': 0, 'custom_time_prof': None} __ecloud_inputs__ = {'init_num_elecs': None, 'init_num_elecs_mp': 0, 'pyecloud_nel_mp_ref': None, 'pyecloud_fact_clean': None, 'pyecloud_fact_split': None, '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, 't_inject_elec': 0, 'init_ecloud_fields': False, } __antenna_inputs__ = {'laser_source_z': None, 'laser_polangle': None, 'laser_emax': None, 'laser_xmin': None, 'laser_xmax': None, 'laser_ymin': None, 'laser_ymax': None, 'laser_func': None } __saving_inputs__ = {'flag_checkpointing': False, 'checkpoints': None, 'flag_output': True, 'output_filename': 'output.h5', 'nbins': 100, 'radius': None, 'stride_imgs': 10, 'stride_output': 1000, 'temps_filename': 'temp_mps_info.h5', 'custom_plot': None, 'images_dir': None, 'field_probes': [], 'field_probes_dump_stride': 1000, 'probe_filename': 'probe.h5', 'flag_save_ek_impacts': False } __simulation_inputs__ = {'enable_trap': True, 'flag_relativ_tracking': True, 'chamber': False, 'lattice_elem': None, 'after_step_fun_list': [], 'tot_nsteps': None, 't_end': None } 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() def print_solvers_info(self): # printing info about the sim print('dx = %e' %self.solver.solver.dx) print('dy = %e' %self.solver.solver.dy) print('dz = %e' %self.solver.solver.dz) dtCFL = self.cfl/(c_light*np.sqrt(1/self.solver.solver.dx**2 + 1/self.solver.solver.dy**2 + 1/self.solver.solver.dz**2)) print('EM solver: dt = %e' %dtCFL) print('ES solver: dt = %e' %picmi.warp.top.dt) def add_es_solver(self, deposition_species = []): 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) self.ES_solver = picmi.ElectrostaticSolver(grid=grid_ES, warp_deposition_species=deposition_species, warp_conductors = self.chamber.conductors, warp_conductor_dfill = picmi.warp.largepos) self.ES_solver.initialize_solver_inputs() registersolver(self.ES_solver.solver) 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 distribute_species(self, primary_species=None, es_species=None, em_species=None, ms_species=None): if primary_species is not None: self.solver.solver.deposition_species = primary_species if es_species is not None: self.ES_solver.solver.deposition_species = es_species if em_species is not None: self.EM_solver.solver.deposition_species = em_species if ms_species is not None: self.MS_solver.solver.deposition_species = ms_species def self_wrapped_probe_fun(self): self.saver.update_field_probes(self.field_probes) def step(self, u_steps=1): for u_step in range(u_steps): # if a passage is starting... if self.n_step % self.ntsteps_p_bunch == 0: self.b_pass += 1 self.perc = 10 # Measure the duration of the previous passage if not self.flag_first_pass: self.t_pass_1 = time.time() self.t_pass = self.t_pass_1 - self.t_pass_0 self.t_pass_0 = time.time() # Perform regeneration if needed if self.ecloud.wspecies.getn() > self.N_mp_max: print('Number of macroparticles: %d' % (self.ecloud.wspecies.getn())) print('MAXIMUM LIMIT OF MPS HAS BEEN RACHED') perform_regeneration(self.N_mp_target, self.ecloud.wspecies, self.sec) # Save stuff if checkpoint if (self.flag_checkpointing and np.any(self.checkpoints == self.b_pass)): self.saver.save_checkpoint(self.b_pass, self.ecloud.wspecies) print('===========================') print('Bunch passage: %d' %self.b_pass) print('Number of electrons: %d' % (np.sum(self.ecloud.wspecies.getw()))) print('Number of macroparticles: %d' % (self.ecloud.wspecies.getn())) if not self.flag_first_pass: print('Previous passage took %ds' % self.t_pass) self.flag_first_pass = False if ((self.n_step % self.ntsteps_p_bunch) / self.ntsteps_p_bunch * 100 > self.perc): print('%d%% of bunch passage' % self.perc) self.perc = self.perc + 10 # Dump outputs if self.flag_output and self.n_step % self.stride_output == 0: self.saver.dump_outputs(self.chamber.xmin, self.chamber.xmax, self.ecloud.wspecies, self.b_pass) # Perform a step sys.stdout = self.text_trap picmi.warp.step() sys.stdout = self.original # Store stuff to be saved if self.flag_output: self.saver.update_outputs(self.ecloud.wspecies, self.beam.wspecies, self.nz, self.n_step) if self.n_step > self.tot_nsteps: # Timer t1 = time.time() totalt = t1 - self.t0 print('Run terminated in %ds' % totalt) def all_steps(self): for i in range(self.n_step, self.tot_nsteps): self.step() self.n_step += 1 #print(self.beam.wspecies.getn()) def all_steps_no_ecloud(self): if picmi.warp.me == 0: for i in tqdm(range(self.n_step, self.tot_nsteps)): sys.stdout = self.text_trap picmi.warp.step(1) sys.stdout = self.original if self.flag_output: self.saver.update_outputs(self.ecloud.wspecies, self.beam.wspecies, self.nz, self.n_step) if self.flag_output and self.n_step % self.stride_output == 0: self.saver.dump_outputs(self.chamber.xmin, self.chamber.xmax, self.ecloud.wspecies, self.b_pass) # Perform regeneration if needed if self.ecloud.wspecies.getn() > self.N_mp_max: print('Number of macroparticles: %d' %self.ecloud.wspecies.getn()) print('MAXIMUM LIMIT OF MPS HAS BEEN RACHED') perform_regeneration(self.N_mp_target, self.ecloud.wspecies, self.sec) self.n_step += 1 else: for i in range(self.n_step, self.tot_nsteps): sys.stdout = self.text_trap picmi.warp.step() sys.stdout = self.original if self.flag_output: self.saver.update_outputs(self.ecloud.wspecies, self.beam.wspecies, self.nz, self.n_step) if self.flag_output and self.n_step % self.stride_output == 0: self.saver.dump_outputs(self.chamber.xmin, self.chamber.xmax, self.ecloud.wspecies, self.b_pass) # Perform regeneration if needed if self.ecloud.wspecies.getn() > self.N_mp_max: print('Number of macroparticles: %d' %self.ecloud.wspecies.getn()) print('MAXIMUM LIMIT OF MPS HAS BEEN RACHED') perform_regeneration(self.N_mp_target, self.ecloud.wspecies, self.sec) self.n_step += 1 def uniform_density(self): init_num_elecs_mp = self.init_num_elecs_mp x0 = np.empty(init_num_elecs_mp, dtype=float) y0 = np.empty(init_num_elecs_mp, dtype=float) z0 = np.empty(init_num_elecs_mp, dtype=float) vx0 = np.empty(init_num_elecs_mp, dtype=float) vy0 = np.empty(init_num_elecs_mp, dtype=float) vz0 = np.empty(init_num_elecs_mp, dtype=float) gi0 = np.empty(init_num_elecs_mp, dtype=float) if picmi.warp.me == 0: chamber = self.chamber lower_bound = chamber.lower_bound upper_bound = chamber.upper_bound 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) gi0 = c_light/np.sqrt(c_light**2 - np.square(vx0) - np.square(vy0) - np.square(vz0)) flag_out = chamber.is_outside(x0, y0, z0) Nout = np.sum(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 = chamber.is_outside(x0, y0, z0) Nout = np.sum(flag_out) comm = MPI.COMM_WORLD comm.Bcast(x0, root=0) comm.Bcast(y0, root=0) comm.Bcast(z0, root=0) comm.Bcast(vx0, root=0) comm.Bcast(vy0, root=0) comm.Bcast(vz0, root=0) comm.Bcast(gi0, root=0) w0 = float(self.init_num_elecs) / float(init_num_elecs_mp) return x0, y0, z0, vx0, vy0, vz0, gi0, w0 def init_uniform_density(self): pwt = picmi.warp.top if np.isclose(pwt.time, self.t_inject_elec, rtol=0, atol=pwt.dt) and (pwt.time - self.t_inject_elec) > 0: x0, y0, z0, vx0, vy0, vz0, gi0, w0 = self.uniform_density() self.ecloud.wspecies.addparticles(x=x0, y=y0, z=z0, vx=vx0, vy=vy0, vz=vz0, gi=gi0, w=w0) if self.init_ecloud_fields: initialize_beam_fields(self.solver.solver, '3d', self.ecloud.wspecies, w3d, top) print('injected %d electrons' % np.sum(self.ecloud.wspecies.getw())) print('injected %d MPs' % self.ecloud.wspecies.getn()) def load_elec_density(self): print('#############################################################') print('Temp distribution found. Reloading it as initial distribution') print('#############################################################') dict_init_dist = dict_of_arrays_and_scalar_from_h5(self.temps_filename) # Load particles status x0 = dict_init_dist['x_mp'] y0 = dict_init_dist['y_mp'] z0 = dict_init_dist['z_mp'] vx0 = dict_init_dist['vx_mp'] vy0 = dict_init_dist['vy_mp'] vz0 = dict_init_dist['vz_mp'] w0 = dict_init_dist['nel_mp'] beta2 = (vx0**2 + vy0**2 + vz0**2)/(clight**2) gamma = 1/np.sqrt(1-beta2) self.b_pass = dict_init_dist['b_pass'] - 1 return picmi.ParticleListDistribution(x=x0, y=y0, z=z0, ux=vx0*gamma, uy=vy0*gamma, uz=vz0*gamma, weight=w0) def gaussian_time_prof(self): t = picmi.warp.top.time val = 0 for i in range(0, self.n_bunches): val += (self.bunch_macro_particles * 1. / np.sqrt(2 * np.pi * self.sigmat ** 2) * np.exp(-(t - i * self.b_spac - self.t_offs) ** 2 / (2 * self.sigmat ** 2)) * picmi.warp.top.dt) return val def bunched_beam(self): NP = int(np.round(self.time_prof())) if NP > 0: x = random.normal(self.bunch_centroid_position[0], self.bunch_rms_size[0], NP) y = random.normal(self.bunch_centroid_position[1], self.bunch_rms_size[1], NP) z = self.bunch_centroid_position[2] vx = random.normal(self.bunch_centroid_velocity[0], self.bunch_rms_velocity[0], NP) vy = random.normal(self.bunch_centroid_velocity[1], self.bunch_rms_velocity[1], NP) vz = picmi.warp.clight * self.beam_beta self.beam.wspecies.addparticles(x=x, y=y, z=z, vx=vx, vy=vy, vz=vz, gi=1. / self.beam_gamma, w=self.bunch_w) def self_wrapped_custom_plot(self, l_force=0): self.custom_plot(self, l_force=l_force) def self_wrapped_custom_time_prof(self): return self.custom_time_prof(self) def defaultsfromdict(self, dic, kw): if kw is not None: for name, defvalue in dic.items(): if name not in self.__dict__: self.__dict__[name] = kw.get(name, defvalue) if name in kw: del kw[name] if kw: raise TypeError("""Keyword argument '%s' is out of place""" % list(kw)[0]) def dump(self, filename): #self.solver.solver.laser_func = None if hasattr(self, 'laser_func'): del self.solver.em3dfft_args['laser_func'] self.laser_func = None self.solver.solver.laser_func = None self.text_trap = None self.original = None self.custom_plot = None self.custom_time_prof = None warpdump(filename)