class warp_pyecloud_sim: 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)) #breakpoint() 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 >= self.b_pass + self.b_pass_prev): self.b_pass += 1 self.perc = 10 # Measure the duration of the previous passage if self.b_pass > 1: 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.secelec.wspecies.getn() > self.N_mp_max: print('Number of macroparticles: %d' % (self.secelec.wspecies.getn())) print('MAXIMUM LIMIT OF MPS HAS BEEN RACHED') perform_regeneration(self.N_mp_target, self.secelec.wspecies, sec) # Save stuff if checkpoint if (self.flag_checkpointing and np.any(self.checkpoints == self.b_pass + self.b_pass_prev)): self.saver.save_checkpoint(self.b_pass + self.b_pass_prev, self.elecb.wspecies, self.secelec.wspecies) print('===========================') print('Bunch passage: %d' % (self.b_pass + self.b_pass_prev)) print('Number of electrons in the dipole: %d' % (np.sum(self.secelec.wspecies.getw()) + np.sum(self.elecb.wspecies.getw()))) print('Number of macroparticles: %d' % (self.secelec.wspecies.getn() + self.elecb.wspecies.getn())) if self.b_pass > 1: print('Previous passage took %ds' % self.t_pass) 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.secelec.wspecies, self.b_pass) # Perform a step sys.stdout = self.text_trap picmi.warp.step(1) sys.stdout = self.original #print(sum(self.secelec.wspecies.getw())+sum(self.elecb.wspecies.getw())) # Store stuff to be saved if self.flag_output: self.saver.update_outputs(self.secelec.wspecies, self.elecb.wspecies, self.nz, self.n_step) self.n_step += 1 if self.n_step > self.tot_nsteps: # Timer t1 = time.time() totalt = t1 - self.t0 # 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 all_steps(self): self.step(self.tot_nsteps) def init_uniform_density(self): chamber = self.chamber lower_bound = chamber.lower_bound upper_bound = chamber.upper_bound init_num_elecs_mp = self.init_num_elecs_mp 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 = 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) flag_out = chamber.isOutside(x0, y0, z0) Nout = np.sum(flag_out) w0 = float(self.init_num_elecs) / float(init_num_elecs_mp) self.b_pass_prev = 0 return picmi.ParticleListDistribution(x=x0, y=y0, z=z0, vx=vx0, vy=vy0, vz=vz0, weight=w0) def load_elec_density(self): print('#############################################################') print('Temp distribution found. Reloading it as initial distribution') print('#############################################################') dict_init_dist = sio.loadmat(self.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] self.b_pass = dict_init_dist['b_pass'][0][0] self.b_pass_prev = dict_init_dist['b_pass'][0][0] - 1 return picmi.ParticleListDistribution(x=x0, y=y0, z=z0, vx=vx0, vy=vy0, vz=vz0, weight=w0) def time_prof(self, t): 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(top.time))) 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 * np.sqrt(1 - 1. / (self.beam_gamma**2)) 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 myplots(self, l_force=0): chamber = self.chamber 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=[ chamber.xmin, chamber.xmax, chamber.ymin, chamber.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=[ chamber.zmin, chamber.zmax, chamber.ymin, chamber.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)
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)