class Backend: ''' Use openPMD-viewer as the backend reader to read openPMD files ''' def __init__(self, filename): ''' Constructor: store the dataset object ''' self.dataset = OpenPMDTimeSeries(filename) def fields_list(self): ''' Return the list of fields defined on the grid ''' return self.dataset.avail_fields def species_list(self): ''' Return the list of species in the dataset ''' return self.dataset.avail_species def n_levels(self): ''' Return the number of MR levels in the dataset ''' return 1 def get_field_checksum(self, lev, field, test_name): ''' Calculate the checksum for a given field at a given level in the dataset ''' Q = self.dataset.get_field(field=field, iteration=self.dataset.iterations[-1])[0] return np.sum(np.abs(Q)) def get_species_attributes(self, species): ''' Return the list of attributes for a given species in the dataset ''' return self.dataset.avail_record_components[species] def get_species_checksum(self, species, attribute): ''' Calculate the checksum for a given attribute of a given species in the dataset ''' Q = self.dataset.get_particle(var_list=[attribute], species=species, iteration=self.dataset.iterations[-1]) # JSON complains with numpy integers, so if the quantity is a np.int64, convert to int checksum = np.sum(np.abs(Q)) if type(checksum) in [np.int64, np.uint64]: return int(checksum) return checksum
def test_boosted_frame_sim_twoproc(): "Test the example input script with two procs in `docs/source/example_input`" temporary_dir = './tests/tmp_test_dir' # Create a temporary directory for the simulation # and copy the example script into this directory if os.path.exists(temporary_dir): shutil.rmtree(temporary_dir) os.mkdir(temporary_dir) shutil.copy('./docs/source/example_input/boosted_frame_script.py', temporary_dir) # Shortcut for the script file, which is repeatedly changed script_filename = os.path.join(temporary_dir, 'boosted_frame_script.py') # Read the script with open(script_filename) as f: script = f.read() # Change default N_step script = replace_string(script, 'N_step = .*', 'N_step = 101') # Modify the script so as to enable finite order script = replace_string(script, 'n_order = -1', 'n_order = 16') script = replace_string(script, 'track_bunch = False', 'track_bunch = True') with open(script_filename, 'w') as f: f.write(script) # Launch the script from the OS command_line = 'cd %s; NUMBA_NUM_THREADS=1 MKL_NUM_THREADS=1 ' % temporary_dir command_line += 'mpirun -np 2 python boosted_frame_script.py' response = os.system(command_line) assert response == 0 # Check that the particle ids are unique at each iterations ts = OpenPMDTimeSeries(os.path.join(temporary_dir, 'lab_diags/hdf5')) print('Checking particle ids...') start_time = time.time() for iteration in ts.iterations: pid, = ts.get_particle(["id"], iteration=iteration) assert len(np.unique(pid)) == len(pid) end_time = time.time() print("%.2f seconds" % (end_time - start_time)) # Suppress the temporary directory shutil.rmtree(temporary_dir)
def particle_energy_histogram( tseries: OpenPMDTimeSeries, it: int, energy_min=1, energy_max=800, delta_energy=1, cutoff=1, # CHANGEME ): """ Compute the weighted particle energy histogram from ``tseries`` at step ``iteration``. :param tseries: whole simulation time series :param it: time step in the simulation :param energy_min: lower energy threshold (MeV) :param energy_max: upper energy threshold (MeV) :param delta_energy: size of each energy bin (MeV) :param cutoff: upper threshold for the histogram, in pC / MeV :return: histogram values and bin edges """ nbins = (energy_max - energy_min) // delta_energy energy_bins = np.linspace(start=energy_min, stop=energy_max, num=nbins + 1) ux, uy, uz, w = tseries.get_particle(["ux", "uy", "uz", "w"], iteration=it) energy = mc2 * np.sqrt(1 + ux**2 + uy**2 + uz**2) # Explanation of weights: # 1. convert electron charge from C to pC (factor 1e12) # 2. multiply by weight w to get real number of electrons # 3. divide by energy bin size delta_energy to get charge / MeV hist, _ = np.histogram(energy, bins=energy_bins, weights=q_e * 1e12 / delta_energy * w) # cut off histogram np.clip(hist, a_min=None, a_max=cutoff, out=hist) return hist, energy_bins, nbins
def run_sim(script_name, n_MPI, checked_fields, test_checkpoint_dir=False): """ Runs the script `script_name` from the folder docs/source/example_input, with `n_MPI` MPI processes. The simulation is then restarted with the same number of processes ; the code checks that the restarted results are identical. More precisely: - The first simulation is run for N_step, then the random seed is reset (for reproducibility) and the code runs for N_step more steps. - Then a second simulation is launched, which reruns the last N_step. """ temporary_dir = './tests/tmp_test_dir' # Create a temporary directory for the simulation # and copy the example script into this directory if os.path.exists(temporary_dir): shutil.rmtree(temporary_dir) os.mkdir(temporary_dir) shutil.copy('./docs/source/example_input/%s' % script_name, temporary_dir) # Shortcut for the script file, which is repeatedly changed script_filename = os.path.join(temporary_dir, script_name) # Read the script and check with open(script_filename) as f: script = f.read() # Change default N_step, diag_period and checkpoint_period script = replace_string(script, 'N_step = .*', 'N_step = 200') script = replace_string(script, 'diag_period = 50', 'diag_period = 10') script = replace_string(script, 'checkpoint_period = 100', 'checkpoint_period = 50') # For MPI simulations: modify the script to use finite-order if n_MPI > 1: script = replace_string(script, 'n_order = -1', 'n_order = 16') # Modify the script so as to enable checkpoints script = replace_string(script, 'save_checkpoints = False', 'save_checkpoints = True') if test_checkpoint_dir: # Try to change the name of the checkpoint directory checkpoint_dir = './test_chkpt' script = replace_string( script, 'set_periodic_checkpoint\( sim, checkpoint_period \)', 'set_periodic_checkpoint( sim, checkpoint_period, checkpoint_dir="%s" )' % checkpoint_dir) script = replace_string( script, 'restart_from_checkpoint\( sim \)', 'restart_from_checkpoint( sim, checkpoint_dir="%s" )' % checkpoint_dir) else: checkpoint_dir = './checkpoints' script = replace_string(script, 'track_electrons = False', 'track_electrons = True') # Modify the script to perform N_step, enforce the random seed # (should be the same when restarting, for exact comparison), # and perform again N_step. script = replace_string( script, 'sim.step\( N_step \)', 'sim.step( N_step ); np.random.seed(0); sim.step( N_step )') with open(script_filename, 'w') as f: f.write(script) # Launch the script from the OS command_line = 'cd %s' % temporary_dir if n_MPI == 1: command_line += '; python %s' % script_name else: # Use only one thread for multiple MPI command_line += '; NUMBA_NUM_THREADS=1 MKL_NUM_THREADS=1 ' command_line += 'mpirun -np %d python %s' % (n_MPI, script_name) response = os.system(command_line) assert response == 0 # Move diagnostics (for later comparison with the restarted simulation) shutil.move(os.path.join(temporary_dir, 'diags'), os.path.join(temporary_dir, 'original_diags')) # Keep only the checkpoints from the first N_step N_step = int(get_string('N_step = (\d+)', script)) period = int(get_string('checkpoint_period = (\d+)', script)) for i_MPI in range(n_MPI): for step in range(N_step + period, 2 * N_step + period, period): os.remove( os.path.join( temporary_dir, '%s/proc%d/hdf5/data%08d.h5' % (checkpoint_dir, i_MPI, step))) # Modify the script so as to enable restarts script = replace_string(script, 'use_restart = False', 'use_restart = True') # Redo only the last N_step script = replace_string( script, 'sim.step\( N_step \); np.random.seed\(0\); sim.step\( N_step \)', 'np.random.seed(0); sim.step( N_step )', ) with open(script_filename, 'w') as f: f.write(script) # Launch the modified script from the OS, with 2 proc response = os.system(command_line) assert response == 0 # Check that restarted simulation gives the same results # as the original simulation print('Checking restarted simulation...') start_time = time.time() ts1 = OpenPMDTimeSeries(os.path.join(temporary_dir, 'diags/hdf5')) ts2 = OpenPMDTimeSeries(os.path.join(temporary_dir, 'original_diags/hdf5')) compare_simulations(ts1, ts2, checked_fields) end_time = time.time() print("%.2f seconds" % (end_time - start_time)) # Check that the particle IDs are unique print('Checking particle ids...') start_time = time.time() for iteration in ts1.iterations: pid, = ts1.get_particle(["id"], iteration=iteration, species="electrons") assert len(np.unique(pid)) == len(pid) end_time = time.time() print("%.2f seconds" % (end_time - start_time)) # Suppress the temporary directory shutil.rmtree(temporary_dir)
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 add_particle_bunch_openPMD(sim, q, m, ts_path, z_off=0., species=None, select=None, iteration=None, boost=None, z_injection_plane=None, initialize_self_field=True): """ Introduce a relativistic particle bunch in the simulation, along with its space charge field, loading particles from an openPMD timeseries. Parameters ---------- sim : a Simulation object The structure that contains the simulation. q : float (in Coulomb) Charge of the particle species m : float (in kg) Mass of the particle species ts_path : string The path to the directory where the openPMD files are. For the moment, only HDF5 files are supported. There should be one file per iteration, and the name of the files should end with the iteration number, followed by '.h5' (e.g. data0005000.h5) z_off: float (in meters) Shift the particle positions in z by z_off. By default the initialized phasespace is centered at z=0. species: string A string indicating the name of the species This is optional if there is only one species select: dict, optional Either None or a dictionary of rules to select the particles, of the form 'x' : [-4., 10.] (Particles having x between -4 and 10 microns) 'ux' : [-0.1, 0.1] (Particles having ux between -0.1 and 0.1 mc) 'uz' : [5., None] (Particles with uz above 5 mc) iteration: integer (optional) The iteration number of the openPMD file from which to extract the particles. boost : a BoostConverter object, optional A BoostConverter object defining the Lorentz boost of the simulation. z_injection_plane: float (in meters) or None When `z_injection_plane` is not None, then particles have a ballistic motion for z<z_injection_plane. This is sometimes useful in boosted-frame simulations. `z_injection_plane` is always given in the lab frame. initialize_self_field: bool, optional Whether to calculate the initial space charge fields of the bunch and add these fields to the fields on the grid (Default: True) """ # Try to import openPMD-viewer, version 1 try: from openpmd_viewer import OpenPMDTimeSeries openpmd_viewer_version = 1 except ImportError: # If not available, try to import openPMD-viewer, version 0 try: from opmd_viewer import OpenPMDTimeSeries openpmd_viewer_version = 0 except ImportError: openpmd_viewer_version = None # Otherwise, raise an error if openpmd_viewer_version is None: raise ImportError( 'The package openPMD-viewer is required to load a particle bunch from on openPMD file.' '\nPlease install it from https://github.com/openPMD/openPMD-viewer' ) ts = OpenPMDTimeSeries(ts_path) # Extract phasespace and particle weights x, y, z, ux, uy, uz, w = ts.get_particle( ['x', 'y', 'z', 'ux', 'uy', 'uz', 'w'], iteration=iteration, species=species, select=select) if openpmd_viewer_version == 0: # Convert the positions from microns to meters x *= 1.e-6 y *= 1.e-6 z *= 1.e-6 # Shift the center of the phasespace to z_off z = z - np.average(z, weights=w) + z_off # Add the electrons to the simulation, and calculate the space charge ptcl_bunch = add_particle_bunch_from_arrays( sim, q, m, x, y, z, ux, uy, uz, w, boost=boost, z_injection_plane=z_injection_plane, initialize_self_field=initialize_self_field) return ptcl_bunch
parser = argparse.ArgumentParser(description='Script to analyze the correctness of the beam in vacuum') parser.add_argument('--output-dir', dest='output_dir', default='diags/hdf5', help='Path to the directory containing output files') args = parser.parse_args() # Numerical parameters of the simulation field_strength = 0.5 gamma = 1000. x_std_initial = 1./2. omega_beta = np.sqrt(field_strength/gamma) # Load beam particle data ts = OpenPMDTimeSeries(args.output_dir) xp, yp, uzp, wp = ts.get_particle(species='beam', iteration=ts.iterations[-1], var_list=['x', 'y', 'uz', 'w']) std_theory = x_std_initial * np.abs(np.cos(omega_beta * ts.current_t)) std_sim_x = np.sqrt(np.sum(xp**2*wp)/np.sum(wp)) std_sim_y = np.sqrt(np.sum(yp**2*wp)/np.sum(wp)) if do_plot: plt.figure() plt.plot(xp, yp, '.') plt.xlim(-2, 2) plt.ylim(-2, 2) plt.xlabel('x') plt.ylabel('y') plt.savefig('image.pdf', bbox_inches='tight') print("beam width theory : " + str(std_theory))
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/')