示例#1
0
class OpenPMDRawDataReader(RawDataReaderBase):
    def __init__(self, location, speciesName, dataName, internalName,
                 firstTimeStep):
        # Store an openPMD timeseries object
        # (Its API is used in order to conveniently extract data from the file)
        self.openpmd_ts = OpenPMDTimeSeries(location, check_all_files=False)
        # Initialize the instance
        RawDataReaderBase.__init__(self, location, speciesName, dataName,
                                   internalName, firstTimeStep)

    def _ReadData(self, timeStep):
        data, = self.openpmd_ts.get_particle([self.internalName],
                                             species=self.speciesName,
                                             iteration=timeStep)
        self.currentTime = self._ReadTime(timeStep)
        return data

    def _ReadTime(self, timeStep):
        # The line below sets the attribute `_current_i` of openpmd_ts
        self.openpmd_ts._find_output(None, timeStep)
        # This sets the corresponding time
        self.currentTime = self.openpmd_ts.t[self.openpmd_ts._current_i]

    def _ReadUnits(self):
        # OpenPMD data always provide conversion to SI units
        self.dataUnits = ""  # TODO:
        self.timeUnits = "s"

    def _OpenFile(self, timeStep):
        # The line below sets the attribute `_current_i` of openpmd_ts
        self.openpmd_ts._find_output(None, timeStep)
        # This finds the full path to the corresponding file
        fileName = self.openpmd_ts.h5_files[self.openpmd_ts._current_i]
        file_content = H5File(fileName, 'r')
        return file_content
示例#2
0
 def __init__(self, location, speciesName, dataName, firstTimeStep ):
     # Store an openPMD timeseries object
     # (Its API is used in order to conveniently extract data from the file)
     self.openpmd_ts = OpenPMDTimeSeries( location, check_all_files=False )
     self.openpmd_dataName = dataName
     # Initialize the instance
     FieldReaderBase.__init__(self, location, speciesName, dataName, firstTimeStep)
示例#3
0
def test_beam_focusing( show=False ):
    """
    Runs the simulation of a focusing charged beam, in a boosted-frame,
    with and without the injection through a plane.
    The value of the RMS radius at focus is automatically checked.
    """
    # Simulate beam focusing with injection through plane or not
    simulate_beam_focusing( None, 'direct' )
    simulate_beam_focusing( z_focus, 'through_plane' )

    # Analyze the results and show that the beam reaches
    # the right RMS radius at focus
    ts1 = OpenPMDTimeSeries('./direct/hdf5/')
    r1 = get_rms_radius( ts1 )
    ts2 = OpenPMDTimeSeries('./through_plane/hdf5/')
    r2 = get_rms_radius( ts2 )
    if show:
        import matplotlib.pyplot as plt
        plt.plot( 1.e3*c*ts1.t, 1.e6*r1 )
        plt.plot( 1.e3*c*ts2.t, 1.e6*r2 )
        plt.xlabel( 'z (mm)' )
        plt.ylabel( 'RMS radius (microns)' )
        plt.show()
    # Find the index of the output at z_focus
    i = np.argmin( abs( c*ts2.t - z_focus ) )
    # With injection through plane, we get the right RMS value at focus
    assert abs( r2[i] - sigma_r ) < 0.05e-6
    # Without injection through plane, the RMS value is significantly different
    assert abs( r1[i] - sigma_r ) > 0.5e-6

    # Clean up the data folders
    shutil.rmtree( 'direct' )
    shutil.rmtree( 'through_plane' )
示例#4
0
def test_lpa_sim_twoproc_restart():
    "Test the checkpoint/restart mechanism with two proc"
    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/lwfa_script.py', temporary_dir )

    # Enter the temporary directory
    os.chdir( temporary_dir )

    # Read the script and check that the targeted lines are present
    with open('lwfa_script.py') as f:
        script = f.read()
        # Check that the targeted lines are present
        if script.find('save_checkpoints = False') == -1 \
            or script.find('use_restart = False') == -1 \
            or script.find('track_electrons = False') == -1 \
            or script.find('N_step = 200') == -1:
            raise RuntimeError('Did not find expected lines in lwfa_script.py')

    # Modify the script so as to enable checkpoints and particle tracking
    script = script.replace('save_checkpoints = False',
                                'save_checkpoints = True')
    script = script.replace('track_electrons = False',
                                'track_electrons = True')
    script = script.replace('N_step = 200', 'N_step = 101')
    with open('lwfa_script.py', 'w') as f:
        f.write(script)
    # Launch the modified script from the OS, with 2 proc
    response = os.system( 'mpirun -np 2 python lwfa_script.py' )
    assert response==0

    # Modify the script so as to enable restarts
    script = script.replace('use_restart = False',
                                'use_restart = True')
    script = script.replace('save_checkpoints = True',
                                'save_checkpoints = False')
    with open('lwfa_script.py', 'w') as f:
        f.write(script)
    # Launch the modified script from the OS, with 2 proc
    response = os.system( 'mpirun -np 2 python lwfa_script.py' )
    assert response==0

    # Check that the particle ids are unique at each iterations
    ts = OpenPMDTimeSeries('./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))

   # Exit the temporary directory and suppress it
    os.chdir('../../')
    shutil.rmtree( temporary_dir )
示例#5
0
 def __init__(self, location, speciesName, dataName, internalName, firstTimeStep):
     # First check whether openPMD is installed
     if not openpmd_installed:
         raise RunTimeError("You need to install openPMD-viewer, e.g. with:\n"
             "pip install openPMD-viewer")
     # Store an openPMD timeseries object
     # (Its API is used in order to conveniently extract data from the file)
     self.openpmd_ts = OpenPMDTimeSeries( location, check_all_files=False )
     # Initialize the instance
     RawDataReaderBase.__init__(self, location, speciesName, dataName, internalName, firstTimeStep)
示例#6
0
    def __init__( self, path_to_dir ):
        """
        Initialize an OpenPMD time series with various methods to diagnose the
        data

        Parameter
        ---------
        path_to_dir : 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)
        """
        OpenPMDTimeSeries.__init__( self, path_to_dir )
示例#7
0
class OpenPMDRawDataReader(RawDataReaderBase):
    def __init__(self, location, speciesName, dataName, internalName, firstTimeStep):
        # First check whether openPMD is installed
        if not openpmd_installed:
            raise RunTimeError("You need to install openPMD-viewer, e.g. with:\n"
                "pip install openPMD-viewer")
        # Store an openPMD timeseries object
        # (Its API is used in order to conveniently extract data from the file)
        self.openpmd_ts = OpenPMDTimeSeries( location, check_all_files=False )
        # Initialize the instance
        RawDataReaderBase.__init__(self, location, speciesName, dataName, internalName, firstTimeStep)

    def _ReadData(self, timeStep):
        data, = self.openpmd_ts.get_particle( [self.internalName],
                    species=self.speciesName, iteration=timeStep )
        self.currentTime = self._ReadTime(timeStep)
        return data

    def _ReadTime(self, timeStep):
        # The line below sets the attribute `_current_i` of openpmd_ts
        self.openpmd_ts._find_output( None, timeStep )
        # This sets the corresponding time
        self.currentTime = self.openpmd_ts.t[ self.openpmd_ts._current_i ]

    def _ReadUnits(self):
        # OpenPMD data always provide conversion to SI units
        # TODO: Get the units from file
        self.dataUnits = "arb.u." 
        self.timeUnits = "s"

    def _OpenFile(self, timeStep):
        # The line below sets the attribute `_current_i` of openpmd_ts
        self.openpmd_ts._find_output( None, timeStep )
        # This finds the full path to the corresponding file
        fileName = self.openpmd_ts.h5_files[ self.openpmd_ts._current_i ]
        file_content = H5File(fileName, 'r')
        return file_content

    def _ReadSimulationProperties(self, file_content):
        # TODO: Add the proper resolution
        self.grid_resolution = None
        self.grid_size = None
        self.grid_units = 'm'

    def _ReadBasicData(self):
        file_content = self._OpenFile(self.firstTimeStep)
        self._ReadSimulationProperties(file_content)
        file_content.close()
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 = int(T_interact/sim.dt)',
                            '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)
示例#9
0
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 and check that the targeted lines are present
    with open(script_filename) as f:
        script = f.read()
        # Check that the targeted lines are present
        if script.find('n_order = -1') == -1 \
            or script.find('track_bunch = False') == -1:
            raise RuntimeError('Did not find expected lines in \
                boosted_frame_script.py')

    # Modify the script so as to enable finite order
    script = script.replace('n_order = -1', 'n_order = 16')
    script = script.replace('track_bunch = False', 'track_bunch = True')
    with open(script_filename, 'w') as f:
        f.write(script)

    # Launch the script from the OS
    response = os.system(
        'cd %s; mpirun -np 2 python boosted_frame_script.py' %temporary_dir )
    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 )
示例#10
0
def check_theory_gaussian():
    """
    Check that the transverse E and B field are close to the high-gamma
    theory for a gaussian bunch
    """
    ts = OpenPMDTimeSeries( os.path.join(temporary_dir, 'diags_serial/hdf5/') )
    Ex, info = ts.get_field( 'E', 'x', iteration=0 )
    By, info = ts.get_field( 'B', 'y', iteration=0 )
    r, z = np.meshgrid( info.r, info.z, indexing='ij' )
    # High-gamma theory for Gaussian bunch
    Eth = -Q/(2*np.pi)**1.5/sig_z/epsilon_0/r * \
        (1 - np.exp(-0.5*r**2/sig_r**2)) * \
        np.exp( -0.5*(z-zf)**2/sig_z**2)
    Bth = Eth/c
    # Check that the fields agree
    assert np.allclose( Ex, Eth, atol=0.1*Eth.max() )
    assert np.allclose( By, Bth, atol=0.1*Bth.max() )
示例#11
0
def _opmd_time_series(data_file):
    prev = None
    try:
        prev = main.list_h5_files
        main.list_h5_files = lambda x: ([data_file.filename], [data_file.iteration])
        return OpenPMDTimeSeries(py.path.local(data_file.filename).dirname)
    finally:
        if prev:
            main.list_h5_files = prev
示例#12
0
    def __init__( self, path_to_dir, check_all_files=True ):
        """
        Initialize an OpenPMD time series with various methods to diagnose the
        data

        Parameter
        ---------
        path_to_dir : 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)

        check_all_files: bool, optional
            Check that all the files in the timeseries are consistent
            (i.e. that they contain the same fields and particles,
            with the same metadata)
            For fast access to the files, this can be changed to False.
        """
        OpenPMDTimeSeries.__init__( self, path_to_dir,
                                    check_all_files=check_all_files )
示例#13
0
    def __init__( self, path_to_dir, check_all_files=True ):
        """
        Initialize an OpenPMD time series with various methods to diagnose the
        data

        Parameter
        ---------
        path_to_dir : 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)

        check_all_files: bool, optional
            Check that all the files in the timeseries are consistent
            (i.e. that they contain the same fields and particles,
            with the same metadata)
            For fast access to the files, this can be changed to False.
        """
        OpenPMDTimeSeries.__init__( self, path_to_dir,
                                    check_all_files=check_all_files )
示例#14
0
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 )
    # Enter the temporary directory
    os.chdir( temporary_dir )

    # Read the script and check that the targeted lines are present
    with open('boosted_frame_script.py') as f:
        script = f.read()
    # Modify the script so as to enable particle tracking
    script = script.replace('track_bunch = False', 'track_bunch = True')
    with open('boosted_frame_script.py', 'w') as f:
        f.write(script)

    # Launch the script from the OS
    response = os.system( 'mpirun -np 2 python boosted_frame_script.py' )
    assert response==0

    # Check that the particle ids are unique at each iterations
    ts = OpenPMDTimeSeries('./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))

    # Exit the temporary directory and suppress it
    os.chdir('../../')
    shutil.rmtree( temporary_dir )
示例#15
0
    def LoadOpenPMDData(self):
        """OpenPMD Loader"""
        # Scan the folder using openPMD-viewer
        ts = OpenPMDTimeSeries(self._dataLocation, check_all_files=False)

        # Register the available fields
        if ts.avail_fields is not None:
            for field in ts.avail_fields:
                # Vector field
                if ts.fields_metadata[field]['type'] == 'vector':
                    available_coord = ['x', 'y', 'z']
                    # Register each coordinate of the vector
                    for coord in available_coord:
                        fieldName = field + '/' + coord
                        standardName = self.GiveStandardNameForOpenPMDQuantity(
                            fieldName)
                        self.AddDomainField(
                            FolderField("openPMD", fieldName, standardName,
                                        self._dataLocation, ts.iterations))
                # Scalar field
                if ts.fields_metadata[field]['type'] == 'scalar':
                    fieldName = field
                    standardName = self.GiveStandardNameForOpenPMDQuantity(
                        field)
                    self.AddDomainField(
                        FolderField("openPMD", fieldName, standardName,
                                    self._dataLocation, ts.iterations))

        # Register the available species
        if ts.avail_species is not None:
            for species in ts.avail_species:
                self.AddSpecies(Species(species))
                for species_quantity in ts.avail_record_components[species]:
                    if species_quantity == "id":
                        self.AddRawDataTagsToSpecies(
                            species,
                            RawDataTags("openPMD", species_quantity,
                                        self._dataLocation, ts.iterations,
                                        species, species_quantity))
                    else:
                        self.AddRawDataToSpecies(
                            species,
                            FolderRawDataSet(
                                "openPMD", species_quantity,
                                self.GiveStandardNameForOpenPMDQuantity(
                                    species_quantity), self._dataLocation,
                                ts.iterations, species, species_quantity))
示例#16
0
def check_identical_fields( folder1, folder2 ):
    ts1 = OpenPMDTimeSeries( folder1 )
    ts2 = OpenPMDTimeSeries( folder2 )
    # Check the vector fields
    for field, coord in [("J", "z"), ("E","r"), ("E","z"), ("B","t")]:
        print("Checking %s%s" %(field, coord))
        field1, info = ts1.get_field(field, coord, iteration=0)
        field2, info = ts2.get_field(field, coord, iteration=0)
        # For 0 fields, do not use allclose
        if abs(field1).max() == 0:
            assert abs(field2).max() == 0
        else:
            assert np.allclose(
                field1/abs(field1).max(), field2/abs(field2).max() )
    # Check the rho field
    print("Checking rho")
    field1, info = ts1.get_field("rho", iteration=0)
    field2, info = ts2.get_field("rho", iteration=0)
    assert np.allclose( field1/abs(field1).max(), field2/abs(field2).max() )
示例#17
0
def check_theory_pml(show, boundaries):
    """
    Check that the transverse E and B field are close to the high-gamma
    theory for a gaussian bunch
    """
    ts = OpenPMDTimeSeries(os.path.join(temporary_dir, 'diags/hdf5/'))
    for iteration in ts.iterations:
        compare_E(ts,
                  'x',
                  m=0,
                  iteration=iteration,
                  rtol=rtol0,
                  show=show,
                  boundaries=boundaries)
        compare_E(ts,
                  'x',
                  m=1,
                  iteration=iteration,
                  rtol=rtol1,
                  show=show,
                  boundaries=boundaries)
示例#18
0
 def GetTimeStepsInOpenPMDLocation(self, location):
     ts = OpenPMDTimeSeries(location, check_all_files=False)
     return ts.iterations
示例#19
0
文件: bunch.py 项目: zcl-maker/fbpic
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)
    """
    # Import openPMD viewer
    try:
        from opmd_viewer import OpenPMDTimeSeries
    except ImportError:
        raise ImportError(
            'The package `opmd_viewer` is required to restart from checkpoints.'
            '\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)
    # 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
示例#20
0
class OpenPMDFieldReader(FieldReaderBase):
    def __init__(self, location, speciesName, dataName, firstTimeStep):
        # First check whether openPMD is installed
        if not openpmd_installed:
            raise RunTimeError(
                "You need to install openPMD-viewer, e.g. with:\n"
                "pip install openPMD-viewer")
        # Store an openPMD timeseries object
        # (Its API is used in order to conveniently extract data from the file)
        self.openpmd_ts = OpenPMDTimeSeries(location, check_all_files=False)
        self.openpmd_dataName = dataName
        # Initialize the instance
        FieldReaderBase.__init__(self, location, speciesName, dataName,
                                 firstTimeStep)

    def _ReadBasicData(self):
        file_content = self._OpenFile(self.firstTimeStep)
        self._ReadInternalName(file_content)
        self._DetermineFieldDimension(file_content)
        self._GetMatrixShape(file_content)
        self._ReadSimulationProperties(file_content)
        file_content.close()

    def _GetMatrixShape(self, file_content):
        _, dataset = openpmd_find_dataset(file_content, self.internalName)
        self.matrixShape = openpmd_get_shape(dataset)

    def _ReadInternalName(self, file_content):
        self.internalName = self.openpmd_dataName

    def _DetermineFieldDimension(self, file_content):
        # Find the name of the field ; vector fields like E are encoded as "E/x"
        fieldname = self.internalName.split("/")[0]
        geometry = self.openpmd_ts.fields_metadata[fieldname]['geometry']
        if geometry == '3dcartesian':
            self.fieldDimension = "3D"
        elif geometry == '2dcartesian':
            self.fieldDimension = "2D"
        else:
            raise ValueError("Unsupported geometry: %s" % geometry)

    def _Read1DSlice(self, timeStep, slicePositionX, slicePositionY=None):
        # Find the name of the field ; vector fields like E are encoded as "E/x"
        field_and_coord = self.internalName.split("/")
        if self.fieldDimension == '2D':
            fieldData, _ = self.openpmd_ts.get_field(*field_and_coord,
                                                     iteration=timeStep)
            elementsX = self.matrixShape[-2]
            selectedRow = round(elementsX * (float(slicePositionX) / 100))
            sliceData = np.array(fieldData[selectedRow])
        elif self.fieldDimension == '3D':
            # Slice first along X
            fieldData = self._Read2DSlice(None, slicePositionX, timeStep)
            # Then slice along Y
            elementsY = self.matrixShape[-2]
            selectedY = round(elementsY * (float(slicePositionY) / 100))
            sliceData = np.array(fieldData[selectedY])
        return sliceData

    def _Read2DSlice(self, sliceAxis, slicePosition, timeStep):
        # Find the name of the field ; vector fields like E are encoded as "E/x"
        field_and_coord = self.internalName.split("/")
        # Convert the `slicePosition` from a 0-to-100 number to -1 to 1
        slicing = -1. + 2 * slicePosition / 100.
        # Extract the slice
        sliceData, _ = self.openpmd_ts.get_field(*field_and_coord,
                                                 iteration=timeStep,
                                                 slicing_dir="x",
                                                 slicing=slicing)
        return sliceData

    def _ReadAllFieldData(self, timeStep):
        # Find the name of the field ; vector fields like E are encoded as "E/x"
        field_and_coord = self.internalName.split("/")
        fieldData, _ = self.openpmd_ts.get_field(*field_and_coord,
                                                 iteration=timeStep,
                                                 slicing=None)
        return fieldData

    def _ReadAxisData(self, timeStep):
        # Find the name of the field ; vector fields like E are encoded as "E/x"
        field_and_coord = self.internalName.split("/")
        # Note: the code below has very bad performance because
        # it automatically reads the field data, just to extract the metadata
        # TODO: improve in the future
        _, field_meta_data = self.openpmd_ts.get_field(*field_and_coord,
                                                       iteration=timeStep,
                                                       slicing=None)
        # Construct the `axisData` from the object `field_meta_data`
        axisData = {}
        axisData["x"] = getattr(field_meta_data, "z")
        axisData["y"] = getattr(field_meta_data, "x")
        if self.fieldDimension == "3D":
            axisData["z"] = getattr(field_meta_data, "y")
        return axisData

    def _ReadTime(self, timeStep):
        # The line below sets the attribute `_current_i` of openpmd_ts
        self.openpmd_ts._find_output(None, timeStep)
        # This sets the corresponding time
        self.currentTime = self.openpmd_ts.t[self.openpmd_ts._current_i]

    def _ReadUnits(self):
        file_content = self._OpenFile(self.firstTimeStep)
        # OpenPMD data always provide conversion to SI units
        self.axisUnits["x"] = "m"
        self.axisUnits["y"] = "m"
        self.axisUnits["z"] = "m"
        self.timeUnits = "t"
        self.dataUnits = "arb.u."  # TODO find the exact unit; needs navigation in file
        file_content.close()

    def _ReadSimulationProperties(self, file_content):
        # TODO: add the proper resolution
        self.grid_resolution = None
        self.grid_size = None
        self.grid_units = "m"

    def _OpenFile(self, timeStep):
        # The line below sets the attribute `_current_i` of openpmd_ts
        self.openpmd_ts._find_output(None, timeStep)
        # This finds the full path to the corresponding file
        fileName = self.openpmd_ts.h5_files[self.openpmd_ts._current_i]
        file_content = H5File(fileName, 'r')
        return file_content
示例#21
0
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/')
import numpy as np
from os.path import basename
import matplotlib.pyplot as plt
from opmd_viewer import OpenPMDTimeSeries
from opmd_viewer.addons import LpaDiagnostics
from scipy.ndimage.filters import gaussian_filter
from scipy import fftpack
from scipy.signal import savgol_filter

from mpl_toolkits import mplot3d
from matplotlib import animation
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm

ts_laser = LpaDiagnostics('./diags/hdf5/', check_all_files=False)
ts_particles = OpenPMDTimeSeries('./diags/hdf5/', check_all_files=False)

fields = {'cmap': 'coolwarm'}
particles = {'cmap': 'jet'}

SimulationName = 'TrailingPulse_PulseTrain'

#Make directories for plots
Path = 'Plots_' + SimulationName
if not os.path.exists(Path):
    os.makedirs(Path)
if not os.path.exists(Path + '/EzLineout'):
    os.makedirs(Path + '/EzLineout')
if not os.path.exists(Path + '/ExLineout'):
    os.makedirs(Path + '/ExLineout')
if not os.path.exists(Path + '/ExLineoutTrailingPulse'):
示例#23
0
import os
import math
import os.path
import numpy as np
from os.path import basename
import matplotlib.pyplot as plt
from opmd_viewer import OpenPMDTimeSeries
from opmd_viewer.addons import LpaDiagnostics
from scipy.ndimage.filters import gaussian_filter

ts_laser = LpaDiagnostics('./diags/hdf5/', check_all_files=False)
ts_particles = OpenPMDTimeSeries('./diags/hdf5/', check_all_files=False)

fields = {'cmap': 'coolwarm'}
particles = {'cmap': 'jet'}

SimulationName = 'BeamLoading_PulseTrain'

#Make directories for plots
Path = 'Plots_' + SimulationName
if not os.path.exists(Path):
    os.makedirs(Path)
if not os.path.exists(Path + '/EzLineout'):
    os.makedirs(Path + '/EzLineout')
if not os.path.exists(Path + '/LaserEnvelope'):
    os.makedirs(Path + '/LaserEnvelope')
if not os.path.exists(Path + '/LaserEnvelope/Lineout'):
    os.makedirs(Path + '/LaserEnvelope/Lineout')
if not os.path.exists(Path + '/rho'):
    os.makedirs(Path + '/rho')
if not os.path.exists('TextFiles'):
import os
import math
import os.path
import numpy as np
from os.path import basename
import matplotlib.pyplot as plt
from opmd_viewer import OpenPMDTimeSeries
from opmd_viewer.addons import LpaDiagnostics
from scipy.ndimage.filters import gaussian_filter

ts_laser = LpaDiagnostics('./diags/hdf5/', check_all_files=False)
ts_particles = OpenPMDTimeSeries('./diags/hdf5/', check_all_files=False)

fields = {'cmap': 'coolwarm'}
particles = {'cmap': 'jet'}

SimulationName = 'PulseTrainSimulation'

#Make directories for plots
Path = 'Plots' + SimulationName
if not os.path.exists(Path):
    os.makedirs(Path)
if not os.path.exists(Path + '/EzLineout'):
    os.makedirs(Path + '/EzLineout')

if not os.path.exists(Path + '/LaserEnvelope'):
    os.makedirs(Path + '/LaserEnvelope')
if not os.path.exists(Path + '/LaserEnvelope/Lineout'):
    os.makedirs(Path + '/LaserEnvelope/Lineout')
if not os.path.exists(Path + '/rho'):
    os.makedirs(Path + '/rho')
示例#25
0
from __future__ import absolute_import, division, print_function, unicode_literals

# OpenPMD imports
from opmd_viewer import OpenPMDTimeSeries

# read in the longitdunal electric field component
path_to_data = '../../rsfbpic/package_data'
field_name = 'E'
coord = 'z'
iteration = 280
mode = 0
theta = 0.
plot = False

# open the file; instantiate "time series" object
time_series = OpenPMDTimeSeries(path_to_data)

# read the specified field
ez, info_ez = time_series.get_field(iteration=iteration, \
                                    field=field_name, \
                                    coord=coord, \
                                    m=mode, \
                                    theta=theta, \
                                    plot=plot)

print()
print("axes = ", info_ez.axes)
print()
print("rmin, rmax = ", info_ez.rmin, "; ", info_ez.rmax)
print("zmin, zmax = ", info_ez.zmin, "; ", info_ez.zmax)
print()
示例#26
0
def add_elec_bunch_openPMD(sim,
                           ts_path,
                           z_off=0.,
                           species=None,
                           select=None,
                           iteration=None,
                           boost=None,
                           filter_currents=True):
    """
    Introduce a relativistic electron 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.

    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.

    filter_currents : bool, optional
        Whether to filter the currents in k space (True by default)
    """
    # Import openPMD viewer
    try:
        from opmd_viewer import OpenPMDTimeSeries
    except ImportError:
        raise ImportError(
            'The package `opmd_viewer` is required to restart from checkpoints.'
            '\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)
    # 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.amax(z) + np.amin(z)) / 2 + z_off

    # Add the electrons to the simulation, and calculate the space charge
    add_elec_bunch_from_arrays(sim,
                               x,
                               y,
                               z,
                               ux,
                               uy,
                               uz,
                               w,
                               boost=boost,
                               filter_currents=filter_currents)
示例#27
0
import os
import math
import os.path
import numpy as np
from os.path import basename
import matplotlib.pyplot as plt
from opmd_viewer import OpenPMDTimeSeries
from opmd_viewer.addons import LpaDiagnostics
from scipy.ndimage.filters import gaussian_filter

ts_laser = LpaDiagnostics('./diags/hdf5/', check_all_files=False)
ts_particles = OpenPMDTimeSeries('./diags/hdf5/', check_all_files=False)

fields = {'cmap': 'coolwarm'}
particles = {'cmap': 'jet'}

SimulationName = 'BeamLoading_PulseTrain'

#Make directories for plots
Path = 'Plots_' + SimulationName
if not os.path.exists(Path):
    os.makedirs(Path)
if not os.path.exists(Path + '/EzLineout'):
    os.makedirs(Path + '/EzLineout')
if not os.path.exists(Path + '/LaserEnvelope'):
    os.makedirs(Path + '/LaserEnvelope')
if not os.path.exists(Path + '/LaserEnvelope/Lineout'):
    os.makedirs(Path + '/LaserEnvelope/Lineout')
if not os.path.exists(Path + '/rho'):
    os.makedirs(Path + '/rho')
if not os.path.exists('TextFiles'):
示例#28
0
def restart_from_checkpoint( sim, iteration=None ):
    """
    Fills the Simulation object `sim` with data saved in a checkpoint.

    More precisely, the following data from `sim` is overwritten:

    - Current time and iteration number of the simulation
    - Position of the boundaries of the simulation box
    - Values of the field arrays
    - Size and values of the particle arrays

    Any other information (e.g. diagnostics of the simulation, presence of a
    moving window, presence of a laser antenna, etc.) need to be set by hand.

    For this reason, a successful restart will often require to modify the
    original input script that produced the checkpoint, rather than to start
    a new input script from scratch.

    NB: This function should always be called *before* the initialization
    of the moving window, since the moving window infers the position of
    particle injection from the existing particle data.

    Parameters
    ----------
    sim: a Simulation object
       The Simulation object into which the checkpoint should be loaded

    iteration: integer (optional)
       The iteration number of the checkpoint from which to restart
       If None, the latest checkpoint available will be used.
    """
    # Import openPMD-viewer
    try:
        from opmd_viewer import OpenPMDTimeSeries
    except ImportError:
        raise ImportError(
        'The package `opmd_viewer` is required to restart from checkpoints.'
        '\nPlease install it from https://github.com/openPMD/openPMD-viewer')

    # Verify that the restart is valid (only for the first processor)
    # (Use the global MPI communicator instead of the `BoundaryCommunicator`,
    # so that this also works for `use_all_ranks=False`)
    if comm.rank == 0:
        check_restart( sim, iteration )
    comm.barrier()

    # Choose the name of the directory from which to restart:
    # one directory per processor
    checkpoint_dir = 'checkpoints/proc%d/hdf5' %comm.rank
    ts = OpenPMDTimeSeries( checkpoint_dir )
    # Select the iteration, and its index
    if iteration is None:
        iteration = ts.iterations[-1]
    # Find the index of the closest iteration
    i_iteration = np.argmin( abs(np.array(ts.iterations) - iteration) )

    # Modify parameters of the simulation
    sim.iteration = iteration
    sim.time = ts.t[ i_iteration ]

    # Export available species as a list
    avail_species = ts.avail_species
    if avail_species is None:
        avail_species = []

    # Load the particles
    # Loop through the different species
    if len(avail_species) == len(sim.ptcl):
        for i in range(len(sim.ptcl)):
            name = 'species %d' %i
            load_species( sim.ptcl[i], name, ts, iteration, sim.comm )
    else:
        raise RuntimeError( \
"""Species numbers in checkpoint and simulation should be same, but
got {:d} and {:d}. Use add_new_species method to add species to
simulation or sim.ptcl = [] to remove them""".format(len(avail_species),
                                                     len(sim.ptcl)) )
    # Record position of grid before restart
    zmin_old = sim.fld.interp[0].zmin

    # Load the fields
    # Loop through the different modes
    for m in range( sim.fld.Nm ):
        # Load the fields E and B
        for fieldtype in ['E', 'B']:
            for coord in ['r', 't', 'z']:
                load_fields( sim.fld.interp[m], fieldtype,
                             coord, ts, iteration )
    # Record position after restart (`zmin` is modified by `load_fields`)
    # and shift the global domain position in the BoundaryCommunicator
    zmin_new = sim.fld.interp[0].zmin
    sim.comm.shift_global_domain_positions( zmin_new - zmin_old )
示例#29
0
def test_cpu_gpu_deposition(show=False):
    "Function that is run by py.test, when doing `python setup.py test`"

    # Skip this test if cuda is not installed
    if not cuda_installed:
        return

    # Perform deposition for a few timesteps, with both the CPU and GPU
    for hardware in ['cpu', 'gpu']:
        if hardware == 'cpu':
            use_cuda = False
        elif hardware == 'gpu':
            use_cuda = True

        # 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,
                         zmin=zmin,
                         use_cuda=use_cuda)

        # Tweak the velocity of the electron bunch
        gamma0 = 10.
        sim.ptcl[0].ux = sim.ptcl[0].x
        sim.ptcl[0].uy = sim.ptcl[0].y
        sim.ptcl[0].uz = np.sqrt(gamma0**2 - sim.ptcl[0].ux**2 -
                                 sim.ptcl[0].uy**2 - 1)
        sim.ptcl[0].inv_gamma = 1. / gamma0 * np.ones_like(sim.ptcl[0].x)
        sim.ptcl[0].m = 1e10

        # Add a field diagnostic
        sim.diags = [
            FieldDiagnostic(diag_period,
                            sim.fld,
                            fieldtypes=['rho', 'J'],
                            comm=sim.comm,
                            write_dir=os.path.join('tests', hardware))
        ]

        ### Run the simulation
        sim.step(N_step)

    # Check that the results are identical
    ts_cpu = OpenPMDTimeSeries('tests/cpu/hdf5')
    ts_gpu = OpenPMDTimeSeries('tests/gpu/hdf5')
    for iteration in ts_cpu.iterations:
        for field, coord in [('rho', ''), ('J', 'x'), ('J', 'z')]:
            # Jy is not tested because it is zero
            print('Testing %s at iteration %d' % (field + coord, iteration))
            F_cpu, info = ts_cpu.get_field(field, coord, iteration=iteration)
            F_gpu, info = ts_gpu.get_field(field, coord, iteration=iteration)
            tolerance = 1.e-13 * (abs(F_cpu).max() + abs(F_gpu).max())
            if not show:
                assert np.allclose(F_cpu, F_gpu, atol=tolerance)
            else:
                if not np.allclose(F_cpu, F_gpu, atol=tolerance):
                    plot_difference(field, coord, iteration, F_cpu, F_gpu,
                                    info)

    # Remove the files used
    shutil.rmtree('tests/cpu')
    shutil.rmtree('tests/gpu')
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 = int(T_interact/sim.dt)',
                            '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)
示例#31
0
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 = [
        BoostedParticleDiagnostic(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('../')