def run_continuous_injection( gamma_boost, dens_func, p_zmin, p_zmax, show, N_check=3 ): # Chose the time step dt = (zmax-zmin)/Nz/c # Initialize the different structures sim = Simulation( Nz, zmax, Nr, rmax, Nm, dt, p_zmin, p_zmax, 0, p_rmax, p_nz, p_nr, p_nt, 0.5*n, dens_func=dens_func, initialize_ions=False, zmin=zmin, use_cuda=use_cuda, gamma_boost=gamma_boost, boundaries='open' ) # Add another species with a different number of particles per cell sim.add_new_species( -e, m_e, 0.5*n, dens_func, 2*p_nz, 2*p_nr, 2*p_nt, p_zmin, p_zmax, 0, p_rmax ) # Set the moving window, which handles the continuous injection # The moving window has an arbitrary velocity (0.7*c) so as to check # that the injection is correct in this case also sim.set_moving_window( v=c ) # Check that the density is correct after different timesteps N_step = int( Nz/N_check/2 ) for i in range( N_check ): sim.step( N_step, move_momenta=False ) check_density( sim, gamma_boost, dens_func, show )
def simulate_beam_focusing( z_injection_plane, write_dir ): """ Simulate a focusing beam in the boosted frame Parameters ---------- z_injection_plane: float or None when this is not None, the injection through a plane is activated. write_dir: string The directory where the boosted diagnostics are written. """ # Initialize the simulation object sim = Simulation( Nz, zmax, Nr, rmax, Nm, dt, zmin=zmin, gamma_boost=gamma_boost, boundaries='open', use_cuda=use_cuda, v_comoving=v_comoving ) # Note: no macroparticles get created because we do not pass # the density and number of particle per cell # Remove the plasma particles sim.ptcl = [] # Initialize the bunch, along with its space charge add_elec_bunch_gaussian( sim, sigma_r, sigma_z, n_emit, gamma0, sigma_gamma, Q, N, tf=(z_focus-z0)/c, zf=z_focus, boost=boost, z_injection_plane=z_injection_plane ) # Configure the moving window sim.set_moving_window( v=c ) # Add a field diagnostic sim.diags = [ BackTransformedParticleDiagnostic( zmin, zmax, c, dt_snapshot_lab, Ntot_snapshot_lab, gamma_boost, period=100, fldobject=sim.fld, species={'bunch':sim.ptcl[0]}, comm=sim.comm, write_dir=write_dir) ] # Run the simulation sim.step( N_step )
bunch.track(sim.comm) # Add a laser to the fields of the simulation add_laser(sim, a0, w0, ctau, z0, lambda0=lambda0, zf=zfoc, gamma_boost=boost.gamma0) # Convert parameter to boosted frame v_window_boosted, = boost.velocity([v_window]) # Configure the moving window sim.set_moving_window(v=v_window_boosted) # Add a field diagnostic sim.diags = [ # Diagnostics in the boosted frame FieldDiagnostic(dt_period=dt_boosted_diag_period, fldobject=sim.fld, comm=sim.comm), ParticleDiagnostic(dt_period=dt_boosted_diag_period, species={ "electrons": plasma_elec, "bunch": bunch }, comm=sim.comm), # Diagnostics in the lab frame (back-transformed) BackTransformedFieldDiagnostic(zmin,
boost=boost) if track_bunch: sim.ptcl[2].track(sim.comm) # Add a laser to the fields of the simulation add_laser(sim, a0, w0, ctau, z0, lambda0=lambda0, zf=zfoc, gamma_boost=gamma_boost) # Configure the moving window sim.set_moving_window(v=v_window) # Add a field diagnostic sim.diags = [ FieldDiagnostic(diag_period, sim.fld, sim.comm), ParticleDiagnostic(diag_period, { "electrons": sim.ptcl[0], "bunch": sim.ptcl[2] }, sim.comm), BoostedFieldDiagnostic(zmin, zmax, c, dt_snapshot_lab, Ntot_snapshot_lab, gamma_boost, period=diag_period,
def run_simulation(gamma_boost, use_separate_electron_species): """ Run a simulation with a laser pulse going through a gas jet of ionizable N5+ atoms, and check the fraction of atoms that are in the N5+ state. Parameters ---------- gamma_boost: float The Lorentz factor of the frame in which the simulation is carried out. use_separate_electron_species: bool Whether to use separate electron species for each level, or a single electron species for all levels. """ # The simulation box zmax_lab = 20.e-6 # Length of the box along z (meters) zmin_lab = 0.e-6 Nr = 3 # Number of gridpoints along r rmax = 10.e-6 # Length of the box along r (meters) Nm = 2 # Number of modes used # The particles of the plasma p_zmin = 5.e-6 # Position of the beginning of the plasma (meters) p_zmax = 15.e-6 p_rmin = 0. # Minimal radial position of the plasma (meters) p_rmax = 100.e-6 # Maximal radial position of the plasma (meters) n_atoms = 0.2 # The atomic density is chosen very low, # to avoid collective effects p_nz = 2 # Number of particles per cell along z p_nr = 1 # Number of particles per cell along r p_nt = 4 # Number of particles per cell along theta # Boosted frame boost = BoostConverter(gamma_boost) # Boost the different quantities beta_boost = np.sqrt(1. - 1. / gamma_boost**2) zmin, zmax = boost.static_length([zmin_lab, zmax_lab]) p_zmin, p_zmax = boost.static_length([p_zmin, p_zmax]) n_atoms, = boost.static_density([n_atoms]) # Increase the number of particles per cell in order to keep sufficient # statistics for the evaluation of the ionization fraction if gamma_boost > 1: p_nz = int(2 * gamma_boost * (1 + beta_boost) * p_nz) # The laser a0 = 1.8 # Laser amplitude lambda0_lab = 0.8e-6 # Laser wavelength # Boost the laser wavelength before calculating the laser amplitude lambda0, = boost.copropag_length([lambda0_lab], beta_object=1.) # Duration and initial position of the laser ctau = 10. * lambda0 z0 = -2 * ctau # Calculate laser amplitude omega = 2 * np.pi * c / lambda0 E0 = a0 * m_e * c * omega / e B0 = E0 / c def laser_func(F, x, y, z, t, amplitude, length_scale): """ Function that describes a Gaussian laser with infinite waist """ return( F + amplitude * math.cos( 2*np.pi*(z-c*t)/lambda0 ) * \ math.exp( - (z - c*t - z0)**2/ctau**2 ) ) # Resolution and number of timesteps dz = lambda0 / 16. dt = dz / c Nz = int((zmax - zmin) / dz) + 1 N_step = int( (2. * 40. * lambda0 + zmax - zmin) / (dz * (1 + beta_boost))) + 1 # Get the speed of the plasma uz_m, = boost.longitudinal_momentum([0.]) v_plasma, = boost.velocity([0.]) # The diagnostics diag_period = N_step - 1 # Period of the diagnostics in number of timesteps # Initial ionization level of the Nitrogen atoms level_start = 2 # Initialize the simulation object, with the neutralizing electrons # No particles are created because we do not pass the density sim = Simulation(Nz, zmax, Nr, rmax, Nm, dt, zmin=zmin, v_comoving=v_plasma, use_galilean=False, boundaries='open', use_cuda=use_cuda) # Add the charge-neutralizing electrons elec = sim.add_new_species(q=-e, m=m_e, n=level_start * n_atoms, p_nz=p_nz, p_nr=p_nr, p_nt=p_nt, p_zmin=p_zmin, p_zmax=p_zmax, p_rmin=p_rmin, p_rmax=p_rmax, continuous_injection=False, uz_m=uz_m) # Add the N atoms ions = sim.add_new_species(q=0, m=14. * m_p, n=n_atoms, p_nz=p_nz, p_nr=p_nr, p_nt=p_nt, p_zmin=p_zmin, p_zmax=p_zmax, p_rmin=p_rmin, p_rmax=p_rmax, continuous_injection=False, uz_m=uz_m) # Add the target electrons if use_separate_electron_species: # Use a dictionary of electron species: one per ionizable level target_species = {} level_max = 6 # N can go up to N7+, but here we stop at N6+ for i_level in range(level_start, level_max): target_species[i_level] = sim.add_new_species(q=-e, m=m_e) else: # Use the pre-existing, charge-neutralizing electrons target_species = elec level_max = None # Default is going up to N7+ # Define ionization ions.make_ionizable(element='N', level_start=level_start, level_max=level_max, target_species=target_species) # Set the moving window sim.set_moving_window(v=v_plasma) # Add a laser to the fields of the simulation (external fields) sim.external_fields = [ ExternalField(laser_func, 'Ex', E0, 0.), ExternalField(laser_func, 'By', B0, 0.) ] # Add a particle diagnostic sim.diags = [ ParticleDiagnostic( diag_period, {"ions": ions}, particle_data=["position", "gamma", "weighting", "E", "B"], # Test output of fields and gamma for standard # (non-boosted) particle diagnostics write_dir='tests/diags', comm=sim.comm) ] if gamma_boost > 1: T_sim_lab = (2. * 40. * lambda0_lab + zmax_lab - zmin_lab) / c sim.diags.append( BackTransformedParticleDiagnostic(zmin_lab, zmax_lab, v_lab=0., dt_snapshots_lab=T_sim_lab / 2., Ntot_snapshots_lab=3, gamma_boost=gamma_boost, period=diag_period, fldobject=sim.fld, species={"ions": ions}, comm=sim.comm, write_dir='tests/lab_diags')) # Run the simulation sim.step(N_step, use_true_rho=True) # Check the fraction of N5+ ions at the end of the simulation w = ions.w ioniz_level = ions.ionizer.ionization_level # Get the total number of N atoms/ions (all ionization levels together) ntot = w.sum() # Get the total number of N5+ ions n_N5 = w[ioniz_level == 5].sum() # Get the fraction of N5+ ions, and check that it is close to 0.32 N5_fraction = n_N5 / ntot print('N5+ fraction: %.4f' % N5_fraction) assert ((N5_fraction > 0.30) and (N5_fraction < 0.34)) # When different electron species are created, check the fraction of # each electron species if use_separate_electron_species: for i_level in range(level_start, level_max): n_N = w[ioniz_level == i_level].sum() assert np.allclose(target_species[i_level].w.sum(), n_N) # Check consistency in the regular openPMD diagnostics ts = OpenPMDTimeSeries('./tests/diags/hdf5/') last_iteration = ts.iterations[-1] w, q = ts.get_particle(['w', 'charge'], species="ions", iteration=last_iteration) # Check that the openPMD file contains the same number of N5+ ions n_N5_openpmd = np.sum(w[(4.5 * e < q) & (q < 5.5 * e)]) assert np.isclose(n_N5_openpmd, n_N5) # Remove openPMD files shutil.rmtree('./tests/diags/') # Check consistency of the back-transformed openPMD diagnostics if gamma_boost > 1.: ts = OpenPMDTimeSeries('./tests/lab_diags/hdf5/') last_iteration = ts.iterations[-1] w, q = ts.get_particle(['w', 'charge'], species="ions", iteration=last_iteration) # Check that the openPMD file contains the same number of N5+ ions n_N5_openpmd = np.sum(w[(4.5 * e < q) & (q < 5.5 * e)]) assert np.isclose(n_N5_openpmd, n_N5) # Remove openPMD files shutil.rmtree('./tests/lab_diags/')
boost=boost) if track_bunch: sim.ptcl[2].track(sim.comm) # Add a laser to the fields of the simulation add_laser(sim, a0, w0, ctau, z0, lambda0=lambda0, zf=zfoc, gamma_boost=gamma_boost) # Configure the moving window sim.set_moving_window(v=v_window, gamma_boost=gamma_boost) # Add a field diagnostic sim.diags = [ FieldDiagnostic(diag_period, sim.fld, sim.comm), ParticleDiagnostic(diag_period, { "electrons": sim.ptcl[0], "bunch": sim.ptcl[2] }, sim.comm), BoostedFieldDiagnostic(zmin, zmax, c, dt_snapshot_lab, Ntot_snapshot_lab, gamma_boost, period=diag_period,
class Simulation(PICMI_Simulation): # Redefine the `init` method, as required by the picmi `_ClassWithInit` def init(self, kw): self.sim_kw = {} for argname in ['use_ruyten_shapes', 'use_modified_volume']: if f'fbpic_{argname}' in kw: self.sim_kw[argname] = kw.pop(f'fbpic_{argname}') self.step_kw = {} for argname in [ 'correct_currents', 'correct_divE', 'use_true_rho', 'move_positions', 'move_momenta', 'show_progress' ]: if f'fbpic_{argname}' in kw: self.step_kw[argname] = kw.pop(f'fbpic_{argname}') # Get the grid grid = self.solver.grid if not type(grid) == PICMI_CylindricalGrid: raise ValueError('When using fbpic with PICMI, ' 'the grid needs to be a CylindricalGrid object.') # Check rmin and boundary conditions assert grid.lower_bound[0] == 0. assert grid.lower_boundary_conditions[ 1] == grid.upper_boundary_conditions[1] if grid.lower_boundary_conditions[1] == 'reflective': warnings.warn( "FBPIC does not support reflective boundary condition in z.\n" "The z boundary condition was automatically converted to 'open'." ) grid.lower_boundary_conditions[1] = 'open' grid.upper_boundary_conditions[1] = 'open' assert grid.upper_boundary_conditions[1] in ['periodic', 'open'] assert grid.upper_boundary_conditions[0] in ['reflective', 'open'] # Determine timestep if self.solver.cfl is not None: dz = (grid.upper_bound[1] - grid.lower_bound[1]) / grid.number_of_cells[1] dr = (grid.upper_bound[0] - grid.lower_bound[0]) / grid.number_of_cells[0] if self.gamma_boost is not None: beta = np.sqrt(1. - 1. / self.gamma_boost**2) dr = dr / ((1 + beta) * self.gamma_boost) dt = self.solver.cfl * min(dz, dr) / c elif self.time_step_size is not None: dt = self.time_step_size else: raise ValueError('You need to either set the `cfl` of the solver\n' 'or the `timestep_size` of the `Simulation`.') # Convert API for the smoother if self.solver.source_smoother is None: smoother = BinomialSmoother() else: if self.solver.source_smoother.n_pass is None: n_passes = 1 else: n_passes = { 'r': self.solver.source_smoother.n_pass[0], 'z': self.solver.source_smoother.n_pass[1] } if self.solver.source_smoother.compensation is None: compensator = False else: compensator = all(self.solver.source_smoother.compensation) smoother = BinomialSmoother(n_passes=n_passes, compensator=compensator) # Convert verbose level: verbose_level = self.verbose if verbose_level is None: verbose_level = 1 # Order of the stencil for z derivatives in the Maxwell solver if self.solver.stencil_order is None: n_order = -1 else: n_order = self.solver.stencil_order[-1] # Number of guard cells if grid.guard_cells is None: n_guard = None else: n_guard = grid.guard_cells[-1] if self.solver.galilean_velocity is None: v_comoving = None else: v_comoving = self.solver.galilean_velocity[-1] # Initialize and store the FBPIC simulation object self.fbpic_sim = FBPICSimulation(Nz=int(grid.number_of_cells[1]), zmin=grid.lower_bound[1], zmax=grid.upper_bound[1], Nr=int(grid.number_of_cells[0]), rmax=grid.upper_bound[0], Nm=grid.n_azimuthal_modes, dt=dt, use_cuda=True, smoother=smoother, n_order=n_order, boundaries={ 'z': grid.upper_boundary_conditions[1], 'r': grid.upper_boundary_conditions[0] }, n_guard=n_guard, verbose_level=verbose_level, particle_shape=self.particle_shape, v_comoving=v_comoving, gamma_boost=self.gamma_boost, **self.sim_kw) # Set the moving window if grid.moving_window_velocity is not None: self.fbpic_sim.set_moving_window(grid.moving_window_velocity[-1]) # Redefine the method `add_laser` from the PICMI Simulation class def add_laser(self, laser, injection_method): # Call method of parent class PICMI_Simulation.add_laser(self, laser, injection_method) # Handle injection method assert type(injection_method) == PICMI_LaserAntenna # Handle laser profile method if type(laser) == PICMI_GaussianLaser: assert laser.propagation_direction[0] == 0. assert laser.propagation_direction[1] == 0. assert (laser.zeta is None) or (laser.zeta == 0) assert (laser.beta is None) or (laser.beta == 0) phi2_chirp = laser.phi2 if phi2_chirp is None: phi2_chirp = 0 polarization_angle = np.arctan2(laser.polarization_direction[1], laser.polarization_direction[0]) laser_profile = GaussianLaser(a0=laser.a0, waist=laser.waist, z0=laser.centroid_position[-1], zf=laser.focal_position[-1], tau=laser.duration, theta_pol=polarization_angle, phi2_chirp=phi2_chirp) else: raise ValueError('Unknown laser profile: %s' % type(injection_method)) # Inject the laser add_laser_pulse(self.fbpic_sim, laser_profile, method='antenna', z0_antenna=injection_method.position[-1], gamma_boost=self.gamma_boost) # Redefine the method `add_species` from the PICMI Simulation class def add_species(self, species, layout, initialize_self_field=False): # Call method of parent class PICMI_Simulation.add_species(self, species, layout, initialize_self_field) # Call generic method internally self._add_species_generic(species, layout, injection_plane_position=None, injection_plane_normal_vector=None, initialize_self_field=initialize_self_field) def add_species_through_plane(self, species, layout, injection_plane_position, injection_plane_normal_vector, initialize_self_field=False): # Call method of parent class PICMI_Simulation.add_species_through_plane( self, species, layout, injection_plane_position, injection_plane_normal_vector, initialize_self_field=initialize_self_field) # Call generic method internally self._add_species_generic( species, layout, injection_plane_position=injection_plane_position, injection_plane_normal_vector=injection_plane_normal_vector, initialize_self_field=initialize_self_field) def _add_species_generic(self, species, layout, injection_plane_position, injection_plane_normal_vector, initialize_self_field): # Extract list of species if type(species) == PICMI_Species: species_instances_list = [species] elif type(species) == PICMI_MultiSpecies: species_instances_list = species.species_instances_list else: raise ValueError('Unknown type: %s' % type(species)) # Loop over species and create FBPIC species for s in species_instances_list: # Get their charge and mass if s.particle_type is not None: s.charge = particle_charge[s.particle_type] s.mass = particle_mass[s.particle_type] # If `charge_state` is set, redefine the charge and mass if s.charge_state is not None: s.charge = s.charge_state * e s.mass -= s.charge_state * m_e # Add the species to the FBPIC simulation fbpic_species = self._create_new_fbpic_species( s, layout, injection_plane_position, injection_plane_normal_vector, initialize_self_field) # Register a pointer to the FBPIC species in the PICMI species itself # (Useful for particle diagnostics later on) s.fbpic_species = fbpic_species # Loop over species and handle ionization for s in species_instances_list: for interaction in s.interactions: assert interaction[0] == 'ionization' assert interaction[1] == 'ADK' picmi_target = interaction[2] if not hasattr(picmi_target, 'fbpic_species'): raise RuntimeError( 'For ionization with PICMI+FBPIC:\n' 'You need to add the target species to the simulation,' ' before the other species.') fbpic_target = picmi_target.fbpic_species fbpic_source = s.fbpic_species fbpic_source.make_ionizable(element=s.particle_type, level_start=s.charge_state, target_species=fbpic_target) def _create_new_fbpic_species(self, s, layout, injection_plane_position, injection_plane_normal_vector, initialize_self_field): # - For the case of a plasma/beam defined in a gridded layout if type(layout) == PICMI_GriddedLayout: # - Uniform distribution if type(s.initial_distribution) == PICMI_UniformDistribution: n0 = s.initial_distribution.density dens_func = None # - Analytic distribution elif type(s.initial_distribution) == PICMI_AnalyticDistribution: import numexpr density_expression = s.initial_distribution.density_expression if s.density_scale is not None: n0 = s.density_scale else: n0 = 1. def dens_func(x, y, z): d = locals() d.update(s.initial_distribution.user_defined_kw) n = numexpr.evaluate(density_expression, local_dict=d) return n else: raise ValueError( 'Unknown combination of layout and distribution') p_nr = layout.n_macroparticle_per_cell[0] p_nt = layout.n_macroparticle_per_cell[1] p_nz = layout.n_macroparticle_per_cell[2] if initialize_self_field or (injection_plane_position is not None): assert s.initial_distribution.fill_in != True if injection_plane_position is None: z_injection_plane = None else: z_injection_plane = injection_plane_position[-1] gamma0_beta0 = s.initial_distribution.directed_velocity[-1] / c gamma0 = (1 + gamma0_beta0**2)**.5 dist = s.initial_distribution fbpic_species = add_particle_bunch( self.fbpic_sim, q=s.charge, m=s.mass, gamma0=gamma0, n=n0, dens_func=dens_func, p_nz=p_nz, p_nr=p_nr, p_nt=p_nt, p_zmin=dist.lower_bound[-1] if dist.lower_bound[-1] is not None else -np.inf, p_zmax=dist.upper_bound[-1] if dist.upper_bound[-1] is not None else +np.inf, p_rmin=0, p_rmax=dist.upper_bound[0] if dist.upper_bound[0] is not None else +np.inf, boost=self.fbpic_sim.boost, z_injection_plane=z_injection_plane, initialize_self_field=initialize_self_field, boost_positions_in_dens_func=True) else: dist = s.initial_distribution fbpic_species = self.fbpic_sim.add_new_species( q=s.charge, m=s.mass, n=n0, dens_func=dens_func, p_nz=p_nz, p_nr=p_nr, p_nt=p_nt, p_zmin=dist.lower_bound[-1] if dist.lower_bound[-1] is not None else -np.inf, p_zmax=dist.upper_bound[-1] if dist.upper_bound[-1] is not None else +np.inf, p_rmax=dist.upper_bound[0] if dist.upper_bound[0] is not None else +np.inf, continuous_injection=s.initial_distribution.fill_in, boost_positions_in_dens_func=True) # - For the case of a Gaussian beam elif (type(s.initial_distribution)==PICMI_GaussianBunchDistribution) \ and (type(layout) == PICMI_PseudoRandomLayout): dist = s.initial_distribution gamma0_beta0 = dist.centroid_velocity[-1] / c gamma0 = (1 + gamma0_beta0**2)**.5 sig_r = dist.rms_bunch_size[0] sig_z = dist.rms_bunch_size[-1] sig_gamma = dist.rms_velocity[-1] / c sig_vr = dist.rms_velocity[0] / gamma0 if sig_vr != 0: tf = -sig_r**2 / sig_vr**2 * dist.velocity_divergence[0] else: tf = 0. zf = dist.centroid_position[-1] + \ dist.centroid_velocity[-1]/gamma0 * tf # Calculate size at focus and emittance sig_r0 = (sig_r**2 - (sig_vr * tf)**2)**0.5 n_emit = gamma0 * sig_r0 * sig_vr / c # Get the number of physical particles n_physical_particles = dist.n_physical_particles if s.density_scale is not None: n_physical_particles *= s.density_scale fbpic_species = add_particle_bunch_gaussian( self.fbpic_sim, q=s.charge, m=s.mass, gamma0=gamma0, sig_gamma=sig_gamma, sig_r=sig_r0, sig_z=sig_z, n_emit=n_emit, n_physical_particles=n_physical_particles, n_macroparticles=layout.n_macroparticles, zf=zf, tf=tf, boost=self.fbpic_sim.boost, initialize_self_field=initialize_self_field) # - For the case of an empty species elif (s.initial_distribution is None) and (layout is None): fbpic_species = self.fbpic_sim.add_new_species(q=s.charge, m=s.mass) else: raise ValueError('Unknown combination of layout and distribution') return fbpic_species # Redefine the method `add_diagnostic` of the parent class def add_diagnostic(self, diagnostic): # Call method of parent class PICMI_Simulation.add_diagnostic(self, diagnostic) # Handle iteration_min/max in regular diagnostic if type(diagnostic) in [ PICMI_FieldDiagnostic, PICMI_ParticleDiagnostic ]: if diagnostic.step_min is None: iteration_min = 0 else: iteration_min = diagnostic.step_min if diagnostic.step_max is None: iteration_max = np.inf else: iteration_max = diagnostic.step_max # Register field diagnostic if type(diagnostic) in [ PICMI_FieldDiagnostic, PICMI_LabFrameFieldDiagnostic ]: if diagnostic.data_list is None: data_list = ['rho', 'E', 'B', 'J'] else: data_list = set() # Use set to avoid redundancy for data in diagnostic.data_list: if data in ['Ex', 'Ey', 'Ez', 'E']: data_list.add('E') elif data in ['Bx', 'By', 'Bz', 'B']: data_list.add('B') elif data in ['Jx', 'Jy', 'Jz', 'J']: data_list.add('J') elif data == 'rho': data_list.add('rho') data_list = list(data_list) if type(diagnostic) == PICMI_FieldDiagnostic: diag = FieldDiagnostic(period=diagnostic.period, fldobject=self.fbpic_sim.fld, comm=self.fbpic_sim.comm, fieldtypes=data_list, write_dir=diagnostic.write_dir, iteration_min=iteration_min, iteration_max=iteration_max) # Register particle density diagnostic rho_density_list = [] if diagnostic.data_list is not None: for data in diagnostic.data_list: if data.startswith('rho_'): # particle density diagnostics, rho_speciesname rho_density_list.append(data) if rho_density_list: species_dict = {} for data in rho_density_list: sname = data[4:] for s in self.species: if s.name == sname: species_dict[s.name] = s.fbpic_species pdd_diag = ParticleChargeDensityDiagnostic( period=diagnostic.period, sim=self.fbpic_sim, species=species_dict, write_dir=diagnostic.write_dir, iteration_min=iteration_min, iteration_max=iteration_max) self.fbpic_sim.diags.append(pdd_diag) elif type(diagnostic) == PICMI_LabFrameFieldDiagnostic: diag = BackTransformedFieldDiagnostic( zmin_lab=diagnostic.grid.lower_bound[1], zmax_lab=diagnostic.grid.upper_bound[1], v_lab=c, dt_snapshots_lab=diagnostic.dt_snapshots, Ntot_snapshots_lab=diagnostic.num_snapshots, gamma_boost=self.gamma_boost, period=100, fldobject=self.fbpic_sim.fld, comm=self.fbpic_sim.comm, fieldtypes=diagnostic.data_list, write_dir=diagnostic.write_dir) # Register particle diagnostic elif type(diagnostic) in [ PICMI_ParticleDiagnostic, PICMI_LabFrameParticleDiagnostic ]: species_dict = {} for s in diagnostic.species: if s.name is None: raise ValueError('When using a species in a diagnostic, ' 'its name must be set.') species_dict[s.name] = s.fbpic_species if diagnostic.data_list is None: data_list = ['position', 'momentum', 'weighting'] else: data_list = diagnostic.data_list if type(diagnostic) == PICMI_ParticleDiagnostic: diag = ParticleDiagnostic(period=diagnostic.period, species=species_dict, comm=self.fbpic_sim.comm, particle_data=data_list, write_dir=diagnostic.write_dir, iteration_min=iteration_min, iteration_max=iteration_max) else: diag = BackTransformedParticleDiagnostic( zmin_lab=diagnostic.grid.lower_bound[1], zmax_lab=diagnostic.grid.upper_bound[1], v_lab=c, dt_snapshots_lab=diagnostic.dt_snapshots, Ntot_snapshots_lab=diagnostic.num_snapshots, gamma_boost=self.gamma_boost, period=100, fldobject=self.fbpic_sim.fld, species=species_dict, comm=self.fbpic_sim.comm, particle_data=data_list, write_dir=diagnostic.write_dir) # Add it to the FBPIC simulation self.fbpic_sim.diags.append(diag) # Redefine the method `add_diagnostic` of the parent class def add_applied_field(self, applied_field): # Call method of parent class PICMI_Simulation.add_applied_field(self, applied_field) if type(applied_field) == PICMI_Mirror: assert applied_field.z_front_location is not None mirror = Mirror(z_lab=applied_field.z_front_location, n_cells=applied_field.number_of_cells, gamma_boost=self.fbpic_sim.boost.gamma0) self.fbpic_sim.mirrors.append(mirror) elif type(applied_field) == PICMI_ConstantAppliedField: # TODO: Handle bounds for field_name in ['Ex', 'Ey', 'Ez', 'Bx', 'By', 'Bz']: field_value = getattr(applied_field, field_name) if field_value is None: continue def field_func(F, x, y, z, t, amplitude, length_scale): return (F + amplitude * field_value) # Pass it to FBPIC self.fbpic_sim.external_fields.append( ExternalField(field_func, field_name, 1., 0.)) elif type(applied_field) == PICMI_AnalyticAppliedField: # TODO: Handle bounds for field_name in ['Ex', 'Ey', 'Ez', 'Bx', 'By', 'Bz']: # Extract expression and execute it inside a function definition expression = getattr(applied_field, field_name + '_expression') if expression is None: continue fieldfunc = None define_function_code = \ """def fieldfunc( F, x, y, z, t, amplitude, length_scale ):\n return( F + amplitude * %s )""" %expression # Take into account user-defined variables for k in applied_field.user_defined_kw: define_function_code = \ "%s = %s\n" %(k,applied_field.user_defined_kw[k]) \ + define_function_code exec(define_function_code, globals()) # Pass it to FBPIC self.fbpic_sim.external_fields.append( ExternalField(fieldfunc, field_name, 1., 0.)) else: raise ValueError("Unrecognized `applied_field` type.") # Redefine the method `step` of the parent class def step(self, nsteps=None): if nsteps is None: nsteps = self.max_steps self.fbpic_sim.step(nsteps, **self.step_kw)
class Simulation(PICMI_Simulation): # Redefine the `init` method, as required by the picmi `_ClassWithInit` def init(self, kw): # Get the grid grid = self.solver.grid if not type(grid) == PICMI_CylindricalGrid: raise ValueError('When using fbpic with PICMI, ' 'the grid needs to be a CylindricalGrid object.') # Check rmin and boundary conditions assert grid.rmin == 0. assert grid.bc_zmin == grid.bc_zmax assert grid.bc_zmax in ['periodic', 'open'] assert grid.bc_rmax in ['reflective', 'open'] # Determine timestep if self.solver.cfl is not None: dz = (grid.zmax - grid.zmin) / grid.nz dt = self.solver.cfl * dz / c elif self.time_step_size is not None: dt = self.time_step_size else: raise ValueError('You need to either set the `cfl` of the solver\n' 'or the `timestep_size` of the `Simulation`.') # Convert API for the smoother if self.solver.source_smoother is None: smoother = BinomialSmoother() else: smoother = BinomialSmoother( n_passes=self.solver.source_smoother.n_pass, compensator=self.solver.source_smoother.compensation) # Initialize and store the FBPIC simulation object self.fbpic_sim = FBPICSimulation(Nz=int(grid.nz), zmin=grid.zmin, zmax=grid.zmax, Nr=int(grid.nr), rmax=grid.rmax, Nm=grid.n_azimuthal_modes, dt=dt, use_cuda=True, smoother=smoother, n_order=32, boundaries={ 'z': grid.bc_zmax, 'r': grid.bc_rmax }) # Set the moving window if grid.moving_window_zvelocity is not None: self.fbpic_sim.set_moving_window(grid.moving_window_zvelocity) # Redefine the method `add_laser` from the PICMI Simulation class def add_laser(self, laser, injection_method): # Call method of parent class PICMI_Simulation.add_laser(self, laser, injection_method) # Handle injection method assert type(injection_method) == PICMI_LaserAntenna # Handle laser profile method if type(laser) == PICMI_GaussianLaser: assert laser.propagation_direction[0] == 0. assert laser.propagation_direction[1] == 0. assert (laser.zeta is None) or (laser.zeta == 0) assert (laser.beta is None) or (laser.beta == 0) phi2_chirp = laser.phi2 if phi2_chirp is None: phi2_chirp = 0 polarization_angle = np.arctan2(laser.polarization_direction[1], laser.polarization_direction[0]) laser_profile = GaussianLaser(a0=laser.a0, waist=laser.waist, z0=laser.centroid_position[-1], zf=laser.focal_position[-1], tau=laser.duration, theta_pol=polarization_angle, phi2_chirp=phi2_chirp) else: raise ValueError('Unknown laser profile: %s' % type(injection_method)) # Inject the laser add_laser_pulse(self.fbpic_sim, laser_profile, method='antenna', z0_antenna=injection_method.position[-1]) # Redefine the method `add_species` from the PICMI Simulation class def add_species(self, species, layout, initialize_self_field=False): # Call method of parent class PICMI_Simulation.add_species(self, species, layout, initialize_self_field) # Extract list of species if type(species) == PICMI_Species: species_instances_list = [species] elif type(species) == PICMI_MultiSpecies: species_instances_list = species.species_instances_list else: raise ValueError('Unknown type: %s' % type(species)) # Loop over species and create FBPIC species for s in species_instances_list: # Get their charge and mass if s.particle_type is not None: s.charge = particle_charge[s.particle_type] s.mass = particle_mass[s.particle_type] # If `charge_state` is set, redefine the charge and mass if s.charge_state is not None: s.charge = s.charge_state * e s.mass -= s.charge_state * m_e # Add the species to the FBPIC simulation fbpic_species = self._create_new_fbpic_species( s, layout, initialize_self_field) # Register a pointer to the FBPIC species in the PICMI species itself # (Useful for particle diagnostics later on) s.fbpic_species = fbpic_species # Loop over species and handle ionization for s in species_instances_list: for interaction in s.interactions: assert interaction[0] == 'ionization' assert interaction[1] == 'ADK' picmi_target = interaction[2] if not hasattr(picmi_target, 'fbpic_species'): raise RuntimeError( 'For ionization with PICMI+FBPIC:\n' 'You need to add the target species to the simulation,' ' before the other species.') fbpic_target = picmi_target.fbpic_species fbpic_source = s.fbpic_species fbpic_source.make_ionizable(element=s.particle_type, level_start=s.charge_state, target_species=fbpic_target) def _create_new_fbpic_species(self, s, layout, initialize_self_field): # - For the case of a plasma defined in a gridded layout if (type(s.initial_distribution)==PICMI_AnalyticDistribution) and \ (type(layout) == PICMI_GriddedLayout): assert initialize_self_field == False import numexpr density_expression = s.initial_distribution.density_expression if s.density_scale is not None: density_expression = "%f*(%s)" \ %(s.density_scale, density_expression) def dens_func(z, r): n = numexpr.evaluate(density_expression) return n p_nr = layout.n_macroparticle_per_cell[0] p_nt = layout.n_macroparticle_per_cell[1] p_nz = layout.n_macroparticle_per_cell[2] fbpic_species = self.fbpic_sim.add_new_species( q=s.charge, m=s.mass, n=1., dens_func=dens_func, p_nz=p_nz, p_nr=p_nr, p_nt=p_nt, p_zmin=s.initial_distribution.lower_bound[-1], p_zmax=s.initial_distribution.upper_bound[-1], continuous_injection=s.initial_distribution.fill_in) # - For the case of a Gaussian beam elif (type(s.initial_distribution)==PICMI_GaussianBunchDistribution) \ and (type(layout) == PICMI_PseudoRandomLayout): dist = s.initial_distribution gamma0_beta0 = dist.centroid_velocity[-1] / c gamma0 = (1 + gamma0_beta0**2)**.5 sig_r = dist.rms_bunch_size[0] sig_z = dist.rms_bunch_size[-1] sig_gamma = dist.rms_velocity[-1] / c sig_vr = dist.rms_velocity[0] / gamma0 if sig_vr != 0: tf = -sig_r**2 / sig_vr**2 * dist.velocity_divergence[0] else: tf = 0. zf = dist.centroid_position[-1] + \ dist.centroid_velocity[-1]/gamma0 * tf # Calculate size at focus and emittance sig_r0 = (sig_r**2 - (sig_vr * tf)**2)**0.5 n_emit = gamma0 * sig_r0 * sig_vr / c # Get the number of physical particles n_physical_particles = dist.n_physical_particles if s.density_scale is not None: n_physical_particles *= s.density_scale fbpic_species = add_particle_bunch_gaussian( self.fbpic_sim, q=s.charge, m=s.mass, gamma0=gamma0, sig_gamma=sig_gamma, sig_r=sig_r0, sig_z=sig_z, n_emit=n_emit, n_physical_particles=n_physical_particles, n_macroparticles=layout.n_macroparticles, zf=zf, tf=tf, initialize_self_field=initialize_self_field) # - For the case of an empty species elif (s.initial_distribution is None) and (layout is None): fbpic_species = self.fbpic_sim.add_new_species(q=s.charge, m=s.mass) else: raise ValueError('Unknown combination of layout and distribution') return fbpic_species # Redefine the method `add_diagnostic` of the parent class def add_diagnostic(self, diagnostic): # Call method of parent class PICMI_Simulation.add_diagnostic(self, diagnostic) # Handle diagnostic if diagnostic.step_min is None: iteration_min = 0 else: iteration_min = diagnostic.step_min if diagnostic.step_max is None: iteration_max = np.inf else: iteration_max = diagnostic.step_max # Register field diagnostic if type(diagnostic) == PICMI_FieldDiagnostic: diag = FieldDiagnostic(period=diagnostic.period, fldobject=self.fbpic_sim.fld, comm=self.fbpic_sim.comm, fieldtypes=diagnostic.data_list, write_dir=diagnostic.write_dir, iteration_min=iteration_min, iteration_max=iteration_max) # Register particle diagnostic elif type(diagnostic) == PICMI_ParticleDiagnostic: species_dict = {} for s in diagnostic.species: if s.name is None: raise ValueError('When using a species in a diagnostic, ' 'its name must be set.') species_dict[s.name] = s.fbpic_species diag = ParticleDiagnostic(period=diagnostic.period, species=species_dict, comm=self.fbpic_sim.comm, particle_data=diagnostic.data_list, write_dir=diagnostic.write_dir, iteration_min=iteration_min, iteration_max=iteration_max) # Add it to the FBPIC simulation self.fbpic_sim.diags.append(diag) # Redefine the method `step` of the parent class def step(self, nsteps): self.fbpic_sim.step(nsteps)
sim = Simulation(Nz, zmax, Nr, rmax, Nm, dt, 0, 0, 0, 0, 2, 2, 4, 0., zmin=zmin, n_order=n_order, boundaries={ 'z': 'open', 'r': 'reflective' }) # Configure the moving window sim.set_moving_window(v=c) # Suppress the particles that were intialized by default and add the bunch sim.ptcl = [] add_elec_bunch_gaussian(sim, sig_r, sig_z, n_emit, gamma0, sig_gamma, Q, N, tf, zf) # Set the diagnostics sim.diags = [FieldDiagnostic(10, sim.fld, comm=sim.comm)] # Perform one simulation step (essentially in order to write the diags) sim.step(1)
def propagate_pulse(Nz, Nr, Nm, zmin, zmax, Lr, L_prop, zf, dt, N_diag, w0, ctau, k0, E0, m, N_show, n_order, rtol, boundaries, v_window=0, use_galilean=False, v_comoving=0, show=False): """ Propagate the beam over a distance L_prop in Nt steps, and extracts the waist and a0 at each step. Parameters ---------- show : bool Wether to show the fields, so that the user can manually check the agreement with the theory. If True, this will periodically show the map of the fields (with a period N_show), as well as (eventually) the evoluation of a0 and w0. If False, this N_diag : int Number of diagnostic points (i.e. measure of waist and a0) along the propagation Nz, Nr : int The number of points on the grid in z and r respectively Nm : int The number of modes in the azimuthal direction zmin, zmax : float The limits of the box in z Lr : float The size of the box in the r direction (In the case of Lr, this is the distance from the *axis* to the outer boundary) L_prop : float The total propagation distance (in meters) zf : float The position of the focal plane of the laser (only works for m=1) dt : float The timestep of the simulation w0 : float The initial waist of the laser (in meters) ctau : float The initial temporal waist of the laser (in meters) k0 : flat The central wavevector of the laser (in meters^-1) E0 : float The initial E0 of the pulse m : int Index of the mode to be tested For m = 1 : test with a gaussian, linearly polarized beam For m = 0 : test with an annular beam, polarized in E_theta n_order : int Order of the stencil rtol : float Relative precision with which the results are tested boundaries : string Type of boundary condition Either 'open' or 'periodic' v_window : float Speed of the moving window v_comoving : float Velocity at which the currents are assumed to move use_galilean: bool Whether to use a galilean frame that moves at the speed v_comoving Returns ------- A dictionary containing : - 'E' : 1d array containing the values of the amplitude - 'w' : 1d array containing the values of waist - 'fld' : the Fields object at the end of the simulation. """ # Initialize the simulation object sim = Simulation(Nz, zmax, Nr, Lr, Nm, dt, p_zmin=0, p_zmax=0, p_rmin=0, p_rmax=0, p_nz=2, p_nr=2, p_nt=2, n_e=0., n_order=n_order, zmin=zmin, use_cuda=use_cuda, boundaries=boundaries, v_comoving=v_comoving, exchange_period=1, use_galilean=use_galilean) # Remove the particles sim.ptcl = [] # Set the moving window object if v_window != 0: sim.set_moving_window(v=v_window) # Initialize the laser fields z0 = (zmax + zmin) / 2 init_fields(sim, w0, ctau, k0, z0, zf, E0, m) # Create the arrays to get the waist and amplitude w = np.zeros(N_diag) E = np.zeros(N_diag) # Calculate the number of steps to run between each diagnostic Ntot_step = int(round(L_prop / (c * dt))) N_step = int(round(Ntot_step / N_diag)) # Loop over the iterations print('Running the simulation...') for it in range(N_diag): print('Diagnostic point %d/%d' % (it, N_diag)) # Fit the fields to find the waist and a0 w[it], E[it] = fit_fields(sim.fld, m) # Plot the fields during the simulation if show == True and it % N_show == 0: import matplotlib.pyplot as plt plt.clf() sim.fld.interp[m].show('Er') plt.show() # Advance the Maxwell equations sim.step(N_step, show_progress=False) # Get the analytical solution z_prop = c * dt * N_step * np.arange(N_diag) ZR = 0.5 * k0 * w0**2 w_analytic = w0 * np.sqrt(1 + (z_prop - zf)**2 / ZR**2) E_analytic = E0 / (1 + (z_prop - zf)**2 / ZR**2)**(1. / 2) # Either plot the results and check them manually if show is True: import matplotlib.pyplot as plt plt.suptitle('Diffraction of a pulse in the mode %d' % m) plt.subplot(121) plt.plot(1.e6 * z_prop, 1.e6 * w, 'o', label='Simulation') plt.plot(1.e6 * z_prop, 1.e6 * w_analytic, '--', label='Theory') plt.xlabel('z (microns)') plt.ylabel('w (microns)') plt.title('Waist') plt.legend(loc=0) plt.subplot(122) plt.plot(1.e6 * z_prop, E, 'o', label='Simulation') plt.plot(1.e6 * z_prop, E_analytic, '--', label='Theory') plt.xlabel('z (microns)') plt.ylabel('E') plt.legend(loc=0) plt.title('Amplitude') plt.show() # or automatically check that the theoretical and simulated curves # of w and E are close else: assert np.allclose(w, w_analytic, rtol=rtol) assert np.allclose(E, E_analytic, rtol=5.e-3) print('The simulation results agree with the theory to %e.' % rtol) # Return a dictionary of the results return ({'E': E, 'w': w, 'fld': sim.fld})
def test_linear_wakefield(Nm=1, show=False): """ Run a simulation of linear laser-wakefield and compare the fields with the analytical solution. Parameters ---------- Nm: int The number of azimuthal modes used in the simulation (Use 1, 2 or 3) This also determines the profile of the driving laser: - Nm=1: azimuthally-polarized annular laser (laser in mode m=0, wakefield in mode m=0) - Nm=2: linearly-polarized Gaussian laser (laser in mode m=1, wakefield in mode m=0) - Nm=3: linearly-polarized Laguerre-Gauss laser (laser in mode m=0 and m=2, wakefield in mode m=0 and m=2, show: bool Whether to have pop-up windows show the comparison between analytical and simulated results """ # Automatically choose higher number of macroparticles along theta p_nt = 2 * Nm # Initialize the simulation object sim = Simulation(Nz, zmax, Nr, rmax, Nm, dt, p_zmin, p_zmax, p_rmin, p_rmax, p_nz, p_nr, p_nt, n_e, use_cuda=use_cuda, boundaries={ 'z': 'open', 'r': 'reflective' }) # Create the relevant laser profile if Nm == 1: # Build an azimuthally-polarized pulse from 2 Laguerre-Gauss profiles profile = LaguerreGaussLaser( 0, 1, a0=a0, waist=w0, tau=tau, z0=z0, theta_pol=np.pi/2, theta0=0. ) \ + LaguerreGaussLaser( 0, 1, a0=a0, waist=w0, tau=tau, z0=z0, theta_pol=0., theta0=-np.pi/2 ) elif Nm == 2: profile = GaussianLaser(a0=a0, waist=w0, tau=tau, z0=z0, theta_pol=np.pi / 2) elif Nm == 3: profile = LaguerreGaussLaser(0, 1, a0=a0, waist=w0, tau=tau, z0=z0, theta_pol=np.pi / 2) add_laser_pulse(sim, profile) # Configure the moving window sim.set_moving_window(v=c) # Add diagnostics if write_fields: sim.diags.append(FieldDiagnostic(diag_period, sim.fld, sim.comm)) if write_particles: sim.diags.append( ParticleDiagnostic(diag_period, {'electrons': sim.ptcl[0]}, sim.comm)) # Prevent current correction for MPI simulation if sim.comm.size > 1: correct_currents = False else: correct_currents = True # Run the simulation sim.step(N_step, correct_currents=correct_currents) # Compare the fields compare_fields(sim, Nm, show)
def test_boosted_output(gamma_boost=10.): """ # TODO Parameters ---------- gamma_boost: float The Lorentz factor of the frame in which the simulation is carried out. """ # The simulation box Nz = 500 # Number of gridpoints along z zmax_lab = 0.e-6 # Length of the box along z (meters) zmin_lab = -20.e-6 Nr = 10 # Number of gridpoints along r rmax = 10.e-6 # Length of the box along r (meters) Nm = 2 # Number of modes used # Number of timesteps N_steps = 500 diag_period = 20 # Period of the diagnostics in number of timesteps dt_lab = (zmax_lab - zmin_lab) / Nz * 1. / c T_sim_lab = N_steps * dt_lab # Move into directory `tests` os.chdir('./tests') # Initialize the simulation object sim = Simulation( Nz, zmax_lab, Nr, rmax, Nm, dt_lab, 0, 0, # No electrons get created because we pass p_zmin=p_zmax=0 0, rmax, 1, 1, 4, n_e=0, zmin=zmin_lab, initialize_ions=False, gamma_boost=gamma_boost, v_comoving=-0.9999 * c, boundaries='open', use_cuda=use_cuda) sim.set_moving_window(v=c) # Remove the electron species sim.ptcl = [] # Add a Gaussian electron bunch # Note: the total charge is 0 so all fields should remain 0 # throughout the simulation. As a consequence, the motion of the beam # is a mere translation. N_particles = 3000 add_elec_bunch_gaussian(sim, sig_r=1.e-6, sig_z=1.e-6, n_emit=0., gamma0=100, sig_gamma=0., Q=0., N=N_particles, zf=0.5 * (zmax_lab + zmin_lab), boost=BoostConverter(gamma_boost)) sim.ptcl[0].track(sim.comm) # openPMD diagnostics sim.diags = [ BackTransformedParticleDiagnostic(zmin_lab, zmax_lab, v_lab=c, dt_snapshots_lab=T_sim_lab / 3., Ntot_snapshots_lab=3, gamma_boost=gamma_boost, period=diag_period, fldobject=sim.fld, species={"bunch": sim.ptcl[0]}, comm=sim.comm) ] # Run the simulation sim.step(N_steps) # Check consistency of the back-transformed openPMD diagnostics: # Make sure that all the particles were retrived by checking particle IDs ts = OpenPMDTimeSeries('./lab_diags/hdf5/') ref_pid = np.sort(sim.ptcl[0].tracker.id) for iteration in ts.iterations: pid, = ts.get_particle(['id'], iteration=iteration) pid = np.sort(pid) assert len(pid) == N_particles assert np.all(ref_pid == pid) # Remove openPMD files shutil.rmtree('./lab_diags/') os.chdir('../')
def run_fbpic(job: Job) -> None: """ This ``signac-flow`` operation runs a ``fbpic`` simulation. :param job: the job instance is a handle to the data of a unique statepoint """ from fbpic.main import Simulation from fbpic.lpa_utils.laser import add_laser_pulse, GaussianLaser from fbpic.openpmd_diag import FieldDiagnostic, ParticleDiagnostic # The density profile def dens_func(z: np.ndarray, r: np.ndarray) -> np.ndarray: """Returns relative density at position z and r. :param z: longitudinal positions, 1d array :param r: radial positions, 1d array :return: a 1d array ``n`` containing the density (between 0 and 1) at the given positions (z, r) """ # Allocate relative density n = np.ones_like(z) # Make linear ramp n = np.where( z < job.sp.ramp_start + job.sp.ramp_length, (z - job.sp.ramp_start) / job.sp.ramp_length, n, ) # Supress density before the ramp n = np.where(z < job.sp.ramp_start, 0.0, n) return n # plot density profile for checking all_z = np.linspace(job.sp.zmin, job.sp.p_zmax, 1000) dens = dens_func(all_z, 0.0) width_inch = job.sp.p_zmax / 1e-5 major_locator = pyplot.MultipleLocator(10) minor_locator = pyplot.MultipleLocator(5) major_locator.MAXTICKS = 10000 minor_locator.MAXTICKS = 10000 def mark_on_plot(*, ax, parameter: str, y=1.1): ax.annotate(s=parameter, xy=(job.sp[parameter] * 1e6, y), xycoords="data") ax.axvline(x=job.sp[parameter] * 1e6, linestyle="--", color="red") return ax fig, ax = pyplot.subplots(figsize=(width_inch, 4.8)) ax.plot(all_z * 1e6, dens) ax.set_xlabel(r"$%s \;(\mu m)$" % "z") ax.set_ylim(-0.1, 1.2) ax.set_xlim(job.sp.zmin * 1e6 - 20, job.sp.p_zmax * 1e6 + 20) ax.set_ylabel("Density profile $n$") ax.xaxis.set_major_locator(major_locator) ax.xaxis.set_minor_locator(minor_locator) mark_on_plot(ax=ax, parameter="zmin") mark_on_plot(ax=ax, parameter="zmax") mark_on_plot(ax=ax, parameter="p_zmin", y=0.9) mark_on_plot(ax=ax, parameter="z0", y=0.8) mark_on_plot(ax=ax, parameter="zf", y=0.6) mark_on_plot(ax=ax, parameter="ramp_start", y=0.7) mark_on_plot(ax=ax, parameter="L_interact") mark_on_plot(ax=ax, parameter="p_zmax") ax.annotate(s="ramp_start + ramp_length", xy=(job.sp.ramp_start * 1e6 + job.sp.ramp_length * 1e6, 1.1), xycoords="data") ax.axvline(x=job.sp.ramp_start * 1e6 + job.sp.ramp_length * 1e6, linestyle="--", color="red") ax.fill_between(all_z * 1e6, dens, alpha=0.5) fig.savefig(job.fn("check_density.png")) # redirect stdout to "stdout.txt" orig_stdout = sys.stdout f = open(job.fn("stdout.txt"), "w") sys.stdout = f # Initialize the simulation object sim = Simulation( job.sp.Nz, job.sp.zmax, job.sp.Nr, job.sp.rmax, job.sp.Nm, job.sp.dt, n_e=None, # no electrons zmin=job.sp.zmin, boundaries={ "z": "open", "r": "reflective" }, n_order=-1, use_cuda=True, verbose_level=2, ) # Create a Gaussian laser profile laser_profile = GaussianLaser(a0=job.sp.a0, waist=job.sp.w0, tau=job.sp.ctau / c_light, z0=job.sp.z0, zf=job.sp.zf, theta_pol=0., lambda0=job.sp.lambda0, cep_phase=0., phi2_chirp=0., propagation_direction=1) # Add it to the simulation add_laser_pulse(sim, laser_profile, gamma_boost=None, method='direct', z0_antenna=None, v_antenna=0.) # Create the plasma electrons elec = sim.add_new_species(q=-q_e, m=m_e, n=job.sp.n_e, dens_func=dens_func, p_zmin=job.sp.p_zmin, p_zmax=job.sp.p_zmax, p_rmax=job.sp.p_rmax, p_nz=job.sp.p_nz, p_nr=job.sp.p_nr, p_nt=job.sp.p_nt) # Track electrons, useful for betatron radiation # elec.track(sim.comm) # Configure the moving window sim.set_moving_window(v=c_light) # Add diagnostics write_dir = os.path.join(job.ws, "diags") sim.diags = [ FieldDiagnostic(job.sp.diag_period, sim.fld, comm=sim.comm, write_dir=write_dir, fieldtypes=["rho", "E"]), ParticleDiagnostic(job.sp.diag_period, {"electrons": elec}, select={"uz": [1., None]}, comm=sim.comm, write_dir=write_dir, particle_data=["momentum", "weighting"]), ] # Plot the Ex component of the laser # Get the fields in the half-plane theta=0 (Sum mode 0 and mode 1) gathered_grids = [ sim.comm.gather_grid(sim.fld.interp[m]) for m in range(job.sp.Nm) ] rgrid = gathered_grids[0].r zgrid = gathered_grids[0].z # construct the Er field for theta=0 Er = gathered_grids[0].Er.T.real for m in range(1, job.sp.Nm): # There is a factor 2 here so as to comply with the convention in # Lifschitz et al., which is also the convention adopted in Warp Circ Er += 2 * gathered_grids[m].Er.T.real e0 = electric_field_amplitude_norm(lambda0=job.sp.lambda0) fig = pyplot.figure(figsize=(8, 8)) sliceplots.Plot2D( fig=fig, arr2d=Er / e0, h_axis=zgrid * 1e6, v_axis=rgrid * 1e6, zlabel=r"$E_r/E_0$", xlabel=r"$z \;(\mu m)$", ylabel=r"$r \;(\mu m)$", extent=( zgrid[0] * 1e6, # + 40 zgrid[-1] * 1e6, # - 20 rgrid[0] * 1e6, rgrid[-1] * 1e6, # - 15, ), cbar=True, vmin=-3, vmax=3, hslice_val=0.0, # do a 1D slice through the middle of the simulation box ) fig.savefig(job.fn('check_laser.png')) # set deterministic random seed np.random.seed(0) # Run the simulation sim.step(job.sp.N_step, show_progress=False) # redirect stdout back and close "stdout.txt" sys.stdout = orig_stdout f.close()
def run_continuous_injection(gamma_boost, ramp, p_zmin, p_zmax, show, N_check=2): # Chose the time step dt = (zmax - zmin) / Nz / c def dens_func(z, r): dens = np.ones_like(z) # Make the density smooth at rmax dens = np.where(r > rmax - smooth_r, np.cos(0.5 * np.pi * (r - smooth_r) / smooth_r)**2, dens) # Make the density 0 below p_zmin dens = np.where(z < p_zmin, 0., dens) # Make a linear ramp dens = np.where((z >= p_zmin) & (z < p_zmin + ramp), (z - p_zmin) / ramp * dens, dens) return (dens) # Initialize the different structures sim = Simulation(Nz, zmax, Nr, rmax, Nm, dt, p_zmin, p_zmax, 0, p_rmax, p_nz, p_nr, p_nt, 0.5 * n, dens_func=dens_func, initialize_ions=False, zmin=zmin, use_cuda=use_cuda, gamma_boost=gamma_boost, boundaries='open') # Add another species with a different number of particles per cell # and with a finite temperature uth = 0.0001 sim.add_new_species(-e, m_e, 0.5 * n, dens_func, 2 * p_nz, 2 * p_nr, 2 * p_nt, p_zmin, p_zmax, 0, p_rmax, ux_th=uth, uy_th=uth, uz_th=uth) # Set the moving window, which handles the continuous injection # The moving window has an arbitrary velocity (0.7*c) so as to check # that the injection is correct in this case also sim.set_moving_window(v=c) # Check that the density is correct after different timesteps N_step = int(Nz / N_check / 2) for i in range(N_check): sim.step(N_step, move_momenta=False) check_density(sim, gamma_boost, dens_func, show)
def run_simulation(gamma_boost): """ Run a simulation with a laser pulse going through a gas jet of ionizable N5+ atoms, and check the fraction of atoms that are in the N5+ state. Parameters ---------- gamma_boost: float The Lorentz factor of the frame in which the simulation is carried out. """ # The simulation box zmax_lab = 20.e-6 # Length of the box along z (meters) zmin_lab = 0.e-6 Nr = 3 # Number of gridpoints along r rmax = 10.e-6 # Length of the box along r (meters) Nm = 2 # Number of modes used # The particles of the plasma p_zmin = 5.e-6 # Position of the beginning of the plasma (meters) p_zmax = 15.e-6 p_rmin = 0. # Minimal radial position of the plasma (meters) p_rmax = 100.e-6 # Maximal radial position of the plasma (meters) n_e = 1. # The plasma density is chosen very low, # to avoid collective effects p_nz = 2 # Number of particles per cell along z p_nr = 1 # Number of particles per cell along r p_nt = 4 # Number of particles per cell along theta # Boosted frame boost = BoostConverter(gamma_boost) # Boost the different quantities beta_boost = np.sqrt(1. - 1. / gamma_boost**2) zmin, zmax = boost.static_length([zmin_lab, zmax_lab]) p_zmin, p_zmax = boost.static_length([p_zmin, p_zmax]) n_e, = boost.static_density([n_e]) # Increase the number of particles per cell in order to keep sufficient # statistics for the evaluation of the ionization fraction if gamma_boost > 1: p_nz = int(2 * gamma_boost * (1 + beta_boost) * p_nz) # The laser a0 = 1.8 # Laser amplitude lambda0_lab = 0.8e-6 # Laser wavelength # Boost the laser wavelength before calculating the laser amplitude lambda0, = boost.copropag_length([lambda0_lab], beta_object=1.) # Duration and initial position of the laser ctau = 10. * lambda0 z0 = -2 * ctau # Calculate laser amplitude omega = 2 * np.pi * c / lambda0 E0 = a0 * m_e * c * omega / e B0 = E0 / c def laser_func(F, x, y, z, t, amplitude, length_scale): """ Function that describes a Gaussian laser with infinite waist """ return( F + amplitude * math.cos( 2*np.pi*(z-c*t)/lambda0 ) * \ math.exp( - (z - c*t - z0)**2/ctau**2 ) ) # Resolution and number of timesteps dz = lambda0 / 16. dt = dz / c Nz = int((zmax - zmin) / dz) + 1 N_step = int( (2. * 40. * lambda0 + zmax - zmin) / (dz * (1 + beta_boost))) + 1 # Get the speed of the plasma uz_m, = boost.longitudinal_momentum([0.]) v_plasma, = boost.velocity([0.]) # The diagnostics diag_period = N_step - 1 # Period of the diagnostics in number of timesteps # Initialize the simulation object sim = Simulation( Nz, zmax, Nr, rmax, Nm, dt, p_zmax, p_zmax, # No electrons get created because we pass p_zmin=p_zmax p_rmin, p_rmax, p_nz, p_nr, p_nt, n_e, zmin=zmin, initialize_ions=False, v_comoving=v_plasma, use_galilean=False, boundaries='open', use_cuda=use_cuda) sim.set_moving_window(v=v_plasma) # Add the N atoms p_zmin, p_zmax, Npz = adapt_to_grid(sim.fld.interp[0].z, p_zmin, p_zmax, p_nz) p_rmin, p_rmax, Npr = adapt_to_grid(sim.fld.interp[0].r, p_rmin, p_rmax, p_nr) sim.ptcl.append( Particles(q=e, m=14. * m_p, n=0.2 * n_e, Npz=Npz, zmin=p_zmin, zmax=p_zmax, Npr=Npr, rmin=p_rmin, rmax=p_rmax, Nptheta=p_nt, dt=dt, use_cuda=use_cuda, uz_m=uz_m, grid_shape=sim.fld.interp[0].Ez.shape, continuous_injection=False)) sim.ptcl[1].make_ionizable(element='N', level_start=0, target_species=sim.ptcl[0]) # Add a laser to the fields of the simulation (external fields) sim.external_fields = [ ExternalField(laser_func, 'Ex', E0, 0.), ExternalField(laser_func, 'By', B0, 0.) ] # Add a field diagnostic sim.diags = [ ParticleDiagnostic(diag_period, {"ions": sim.ptcl[1]}, write_dir='tests/diags', comm=sim.comm) ] if gamma_boost > 1: T_sim_lab = (2. * 40. * lambda0_lab + zmax_lab - zmin_lab) / c sim.diags.append( BoostedParticleDiagnostic(zmin_lab, zmax_lab, v_lab=0., dt_snapshots_lab=T_sim_lab / 2., Ntot_snapshots_lab=3, gamma_boost=gamma_boost, period=diag_period, fldobject=sim.fld, species={"ions": sim.ptcl[1]}, comm=sim.comm, write_dir='tests/lab_diags')) # Run the simulation sim.step(N_step, use_true_rho=True) # Check the fraction of N5+ ions at the end of the simulation w = sim.ptcl[1].w ioniz_level = sim.ptcl[1].ionizer.ionization_level # Get the total number of N atoms/ions (all ionization levels together) ntot = w.sum() # Get the total number of N5+ ions n_N5 = w[ioniz_level == 5].sum() # Get the fraction of N5+ ions, and check that it is close to 0.32 N5_fraction = n_N5 / ntot print('N5+ fraction: %.4f' % N5_fraction) assert ((N5_fraction > 0.30) and (N5_fraction < 0.34)) # Check consistency in the regular openPMD diagnostics ts = OpenPMDTimeSeries('./tests/diags/hdf5/') last_iteration = ts.iterations[-1] w, q = ts.get_particle(['w', 'charge'], species="ions", iteration=last_iteration) # Check that the openPMD file contains the same number of N5+ ions n_N5_openpmd = np.sum(w[(4.5 * e < q) & (q < 5.5 * e)]) assert np.isclose(n_N5_openpmd, n_N5) # Remove openPMD files shutil.rmtree('./tests/diags/') # Check consistency of the back-transformed openPMD diagnostics if gamma_boost > 1.: ts = OpenPMDTimeSeries('./tests/lab_diags/hdf5/') last_iteration = ts.iterations[-1] w, q = ts.get_particle(['w', 'charge'], species="ions", iteration=last_iteration) # Check that the openPMD file contains the same number of N5+ ions n_N5_openpmd = np.sum(w[(4.5 * e < q) & (q < 5.5 * e)]) assert np.isclose(n_N5_openpmd, n_N5) # Remove openPMD files shutil.rmtree('./tests/lab_diags/')