def run_gpt_with_particlegroup(settings=None,
                               gpt_input_file=None,
                               input_particle_group=None,
                               workdir=None,
                               use_tempdir=True,
                               gpt_bin='$GPT_BIN',
                               timeout=2500,
                               auto_phase=False,
                               verbose=False,
                               gpt_verbose=False,
                               asci2gdf_bin='$ASCI2GDF_BIN'):
    """
    Run gpt with particles from ParticleGroup. 
    
        settings: dict with keys that are in gpt input file.    
        
    """

    # Call simpler evaluation if there is no input_particle_group:
    if (input_particle_group is None):
        return run_gpt(settings=settings,
                       gpt_input_file=gpt_input_file,
                       workdir=workdir,
                       use_tempdir=use_tempdir,
                       gpt_bin=gpt_bin,
                       timeout=timeout,
                       verbose=verbose)

    if (verbose):
        print('Run GPT with ParticleGroup:')

    unit_registry = UnitRegistry()

    # Make gpt and generator objects
    G = GPT(gpt_bin=gpt_bin,
            input_file=gpt_input_file,
            initial_particles=input_particle_group,
            workdir=workdir,
            use_tempdir=use_tempdir)
    G.timeout = timeout
    G.verbose = verbose

    # Set inputs
    if settings:
        for k, v in settings.items():
            G.set_variable(k, v)

    if ('final_charge' in settings):
        raise ValueError(
            'final_charge is deprecated, please specify value and units instead.'
        )

    # Run
    if (auto_phase):

        if (verbose):
            print('\nAuto Phasing >------\n')
        t1 = time.time()

        # Create the distribution used for phasing
        if (verbose):
            print('****> Creating initial distribution for phasing...')

        phasing_beam = get_distgen_beam_for_phasing_from_particlegroup(
            input_particle_group, n_particle=10, verbose=verbose)
        phasing_particle_file = os.path.join(G.path,
                                             'gpt_particles.phasing.gdf')
        write_gpt(phasing_beam,
                  phasing_particle_file,
                  verbose=verbose,
                  asci2gdf_bin=asci2gdf_bin)

        if (verbose):
            print('<**** Created initial distribution for phasing.\n')

        G.write_input_file()  # Write the unphased input file

        phased_file_name, phased_settings = gpt_phasing(
            G.input_file,
            path_to_gpt_bin=G.gpt_bin[:-3],
            path_to_phasing_dist=phasing_particle_file,
            verbose=verbose)
        G.set_variables(phased_settings)
        t2 = time.time()

        if (verbose):
            print(f'Time Ellapsed: {t2-t1} sec.')
            print('------< Auto Phasing\n')

    # If here, either phasing successful, or no phasing requested
    G.run(gpt_verbose=gpt_verbose)

    if ('final_charge:value' in settings and 'final_charge:units' in settings
            and len(G.screen) > 0):
        final_charge = settings[
            'final_charge:value'] * unit_registry.parse_expression(
                settings['final_charge:units'])
        final_charge = final_charge.to('coulomb').magnitude
        clip_to_charge(G.screen[-1], final_charge, make_copy=False)

    if ('final_radius:value' in settings and 'final_radius:units' in settings
            and len(G.screen) > 0):
        final_radius = settings[
            'final_radius:value'] * unit_registry.parse_expression(
                settings['final_radius:units'])
        final_radius = final_radius.to('meter').magnitude
        take_range(G.screen[-1], 'r', 0, final_radius)

    if (input_particle_group['sigma_t'] == 0.0):
        # Initial distribution is a tout
        if (G.output['n_tout'] > 0):
            G.output['particles'].insert(0, input_particle_group)
            G.output['n_tout'] = G.output['n_tout'] + 1
    else:
        # Initial distribution is a screen
        if (G.output['n_screen'] > 0):
            G.output['particles'].insert(G.output['n_tout'],
                                         input_particle_group)
            G.output['n_screen'] = G.output['n_screen'] + 1

    return G
Beispiel #2
0
def run_gpt_with_distgen(settings=None,
                         gpt_input_file=None,
                         distgen_input_file=None,
                         workdir=None,
                         use_tempdir=True,
                         gpt_bin='$GPT_BIN',
                         timeout=2500,
                         auto_phase=False,
                         verbose=False,
                         gpt_verbose=False,
                         asci2gdf_bin='$ASCI2GDF_BIN'):
    """
    Run gpt with particles generated by distgen. 
    
        settings: dict with keys that can appear in an gpt or distgen Generator input file. 
        
    Example usage:
        G = run_gpt_with_distgen({'lspch':False},
                       gpt_input_file='$LCLS_LATTICE/gpt/models/gunb_eic/gpt.in',
                       distgen_input_file='$LCLS_LATTICE/distgen/models/gunb_gaussian/gunb_gaussian.json',
                       verbose=True,
                       timeout=None
                      )        
        
    """

    # Call simpler evaluation if there is no generator:
    if not distgen_input_file:
        return run_gpt(settings=settings,
                       gpt_input_file=gpt_input_file,
                       workdir=workdir,
                       use_tempdir=use_tempdir,
                       gpt_bin=gpt_bin,
                       timeout=timeout,
                       verbose=verbose)

    if (verbose):
        print('Run GPT with Distgen:')

    # Make gpt and generator objects
    G = GPT(gpt_bin=gpt_bin,
            input_file=gpt_input_file,
            workdir=workdir,
            use_tempdir=use_tempdir)
    G.timeout = timeout
    G.verbose = verbose

    # Distgen generator
    gen = Generator(verbose=verbose)
    f = full_path(distgen_input_file)
    distgen_params = yaml.safe_load(open(f))

    # Set inputs
    if settings:
        G, distgen_params = set_gpt_and_distgen(G,
                                                distgen_params,
                                                settings,
                                                verbose=verbose)

    # Link particle files
    particle_file = os.path.join(G.path, G.get_dist_file())

    if (verbose):
        print('Linking particle files, distgen output will point to -> "' +
              os.path.basename(particle_file) + '" in working directory.')

    G.set_dist_file(particle_file)

    if ('output' in distgen_params and verbose):
        print('Replacing Distgen output params')

    distgen_params['output'] = {'type': 'gpt', 'file': particle_file}

    if (verbose):
        print('\nDistgen >------\n')
    # Configure distgen
    gen.parse_input(distgen_params)

    # Run
    beam = gen.beam()
    write_gpt(beam, particle_file, verbose=verbose, asci2gdf_bin=asci2gdf_bin)

    if (verbose):
        print('------< Distgen\n')

    if (auto_phase):

        if (verbose):
            print('\nAuto Phasing >------\n')
        t1 = time.time()

        # Create the distribution used for phasing
        if (verbose):
            print('****> Creating intiial distribution for phasing...')

        phasing_beam = get_distgen_beam_for_phasing(beam,
                                                    n_particle=10,
                                                    verbose=verbose)
        phasing_particle_file = os.path.join(G.path,
                                             'gpt_particles.phasing.gdf')
        write_gpt(phasing_beam,
                  phasing_particle_file,
                  verbose=verbose,
                  asci2gdf_bin=asci2gdf_bin)

        if (verbose):
            print('<**** Created intiial distribution for phasing.\n')

        G.write_input_file()  # Write the unphased input file
        phased_file_name, phased_settings = gpt_phasing(
            G.input_file,
            path_to_gpt_bin=G.gpt_bin[:-3],
            path_to_phasing_dist=phasing_particle_file,
            verbose=verbose)
        G.set_variables(phased_settings)
        t2 = time.time()
        if (verbose):
            print(f'Time Ellapsed: {t2-t1} sec.')
            print('------< Auto Phasing\n')

    G.run(gpt_verbose=gpt_verbose)

    return G
def multirun_gpt_with_particlegroup(settings=None,
                                    gpt_input_file=None,
                                    input_particle_group=None,
                                    workdir=None,
                                    use_tempdir=True,
                                    gpt_bin='$GPT_BIN',
                                    timeout=2500,
                                    auto_phase=False,
                                    verbose=False,
                                    gpt_verbose=False,
                                    asci2gdf_bin='$ASCI2GDF_BIN'):
    """
    Run gpt with particles from ParticleGroup. 
    
        settings: dict with keys that are in gpt input file.    
        
    """

    unit_registry = UnitRegistry()

    # Call simpler evaluation if there is no input_particle_group:
    if (input_particle_group is None):
        raise ValueError('Must supply input_particle_group')

    if (verbose):
        print('Run GPT with ParticleGroup:')

    if ('clipping_charge' in settings):
        raise ValueError(
            'clipping_charge is deprecated, please specify value and units instead.'
        )
    if ('final_charge' in settings):
        raise ValueError(
            'final_charge is deprecated, please specify value and units instead.'
        )
    if ('t_restart' not in settings):
        raise ValueError('t_restart must be supplied')

    t_restart = settings['t_restart']

    if ('restart_file' not in settings):
        # Make gpt and generator objects
        G = GPT(gpt_bin=gpt_bin,
                input_file=gpt_input_file,
                initial_particles=input_particle_group,
                workdir=workdir,
                use_tempdir=use_tempdir)
        G.timeout = timeout
        G.verbose = verbose

        # Set inputs
        if settings:
            for k, v in settings.items():
                G.set_variable(k, v)
        else:
            raise ValueError('Must supply settings')

        G.set_variable('multi_run', 0)
        if (auto_phase):

            if (verbose):
                print('\nAuto Phasing >------\n')
            t1 = time.time()

            # Create the distribution used for phasing
            if (verbose):
                print('****> Creating initial distribution for phasing...')

            phasing_beam = get_distgen_beam_for_phasing_from_particlegroup(
                input_particle_group, n_particle=10, verbose=verbose)
            phasing_particle_file = os.path.join(G.path,
                                                 'gpt_particles.phasing.gdf')
            write_gpt(phasing_beam,
                      phasing_particle_file,
                      verbose=verbose,
                      asci2gdf_bin=asci2gdf_bin)

            if (verbose):
                print('<**** Created initial distribution for phasing.\n')

            G.write_input_file()  # Write the unphased input file

            phased_file_name, phased_settings = gpt_phasing(
                G.input_file,
                path_to_gpt_bin=G.gpt_bin[:-3],
                path_to_phasing_dist=phasing_particle_file,
                verbose=verbose)
            G.set_variables(phased_settings)
            t2 = time.time()

            if (verbose):
                print(f'Time Ellapsed: {t2-t1} sec.')
                print('------< Auto Phasing\n')

        G.set_variable('multi_run', 1)
        G.set_variable('last_run', 2)
        G.set_variable('t_start', 0.0)
        G.set_variable('t_restart', t_restart)

        # If here, either phasing successful, or no phasing requested
        G.run(gpt_verbose=gpt_verbose)
    else:
        G = GPT()
        G.load_archive(settings['restart_file'])
        if settings:
            for k, v in settings.items():
                G.set_variable(k, v)

    # Remove touts and screens that are after t_restart
    t_restart_with_fudge = t_restart + 1.0e-18  # slightly larger that t_restart to avoid floating point comparison problem
    G.output['n_tout'] = np.count_nonzero(
        G.stat('mean_t', 'tout') <= t_restart_with_fudge)
    G.output['n_screen'] = np.count_nonzero(
        G.stat('mean_t', 'screen') <= t_restart_with_fudge)
    for p in reversed(G.particles):
        if (p['mean_t'] > t_restart_with_fudge):
            G.particles.remove(p)

    G_all = G  # rename it, and then overwrite G

    if (verbose):
        print(f'Looking for tout at t = {t_restart}')
    restart_particles = get_screen_data(G,
                                        tout_t=t_restart,
                                        use_extension=False,
                                        verbose=verbose)[0]

    if ('clipping_charge:value' in settings
            and 'clipping_charge:units' in settings):
        clipping_charge = settings[
            'clipping_charge:value'] * unit_registry.parse_expression(
                settings['clipping_charge:units'])
        clipping_charge = clipping_charge.to('coulomb').magnitude
        clip_to_charge(restart_particles, clipping_charge, make_copy=False)

    G = GPT(gpt_bin=gpt_bin,
            input_file=gpt_input_file,
            initial_particles=restart_particles,
            workdir=workdir,
            use_tempdir=use_tempdir)
    G.timeout = timeout
    G.verbose = verbose

    for k, v in G_all.input["variables"].items():
        G.set_variable(k, v)

    G.set_variable('multi_run', 2)
    G.set_variable('last_run', 2)
    G.set_variable('t_start', t_restart)
    G.run(gpt_verbose=gpt_verbose)

    G_all.output['particles'][G_all.output['n_tout']:G_all.
                              output['n_tout']] = G.tout
    G_all.output['particles'] = G_all.output['particles'] + G.screen
    G_all.output['n_tout'] = G_all.output['n_tout'] + G.output['n_tout']
    G_all.output['n_screen'] = G_all.output['n_screen'] + G.output['n_screen']

    if ('final_charge:value' in settings and 'final_charge:units' in settings
            and len(G_all.screen) > 0):
        final_charge = settings[
            'final_charge:value'] * unit_registry.parse_expression(
                settings['final_charge:units'])
        final_charge = final_charge.to('coulomb').magnitude
        clip_to_charge(G_all.screen[-1], final_charge, make_copy=False)

    if (input_particle_group['sigma_t'] == 0.0):
        # Initial distribution is a tout
        if (G_all.output['n_tout'] > 0):
            # Don't include the cathode if there are no other screens. Screws up optimizations of "final" screen when there is an error
            G_all.output['particles'].insert(0, input_particle_group)
            G_all.output['n_tout'] = G_all.output['n_tout'] + 1
    else:
        # Initial distribution is a screen
        if (G_all.output['n_screen'] > 0):
            # Don't include the cathode if there are no other screens. Screws up optimizations of "final" screen when there is an error
            G_all.output['particles'].insert(G_all.output['n_tout'],
                                             input_particle_group)
            G_all.output['n_screen'] = G_all.output['n_screen'] + 1

    return G_all
    def run(self, inputs, verbose=False):
       
        tag = f'vb24@{self.id}:'

        #----------------------------------------------------------------------------
        # Get laser distribution, cathode quantities, and gun current
        #----------------------------------------------------------------------------
        r_params = {'sigma_xy': dunits(str(inputs[f'{tag}laser:sigma_xy'])),
                    'alpha':    dunits(str(inputs[f'{tag}laser:alpha_xy']))}

        count = self.pvdefs[f'{tag}laser:r']['count']
        laser_wavelength = inputs[f'{tag}laser:wavelength']
        laser_power = inputs[f'{tag}laser:power']
        laser_sigma_xy = inputs[f'{tag}laser:sigma_xy']
        laser_alpha_xy = inputs[f'{tag}laser:alpha_xy']
        laser_avg_x = inputs[f'{tag}laser:mean_x']
        laser_avg_y = inputs[f'{tag}laser:mean_y']

        r_dist = SuperGaussianRad(verbose=False, **r_params)
        rs = (r_dist.get_r_pts(count)).to(self.pvdefs[f'{tag}laser:r']['unit'])
        Pr = (dunits(str(laser_power))*r_dist.rho(rs)).to(self.pvdefs[f'{tag}laser:Pr']['unit'])

        cathode_QE = inputs[f'{tag}cathode:QE']
        cathode_MTE = inputs[f'{tag}cathode:MTE']

        hc = 1*units.h*units.c
        photon_flux = (laser_power/(hc/laser_wavelength) ).to_base_units()
        gun_current = (photon_flux*cathode_QE*(1*units.e)).to(self.pvdefs[f'{tag}gun:current']['unit'])
        #----------------------------------------------------------------------------
       

        #----------------------------------------------------------------------------
        # Create Distgen input and run generator
        #----------------------------------------------------------------------------
        distgen_input = yaml.dump(
                        {'n_particle':inputs[f'{tag}gpt:n_particle'].magnitude,
                         'random_type':'hammersley',
                         'total_charge': {'value': 0.0, 'units': 'pC'},   
                         'start': {
                             'type':'cathode',
                             'MTE': {'value': cathode_MTE.magnitude, 'units': str(cathode_MTE.units)}},

                         'r_dist': {
                             'type':'rsg',
                             'sigma_xy':{'value': laser_sigma_xy.magnitude, 'units': str(laser_sigma_xy.units)},
                             'alpha':{'value': laser_alpha_xy.magnitude, 'units': str(laser_alpha_xy.units)},},

                         'transforms':{
                             't1':{'type':'set_avg x', 'avg_x': {'value': laser_avg_x.magnitude, 'units': str(laser_avg_x.units)}},
                             't2':{'type':'set_avg y', 'avg_y': {'value': laser_avg_y.magnitude, 'units': str(laser_avg_y.units)}}
                         }})
 
        gen = Generator(distgen_input, verbose=True)     
        beam = gen.beam()   
        #----------------------------------------------------------------------------


        #----------------------------------------------------------------------------
        # Configure GPT and run
        #----------------------------------------------------------------------------
        G = GPT(input_file=os.path.join(os.getcwd(),'templates/gpt.in'), 
                initial_particles = ParticleGroup(data=beam.data()), 
                use_tempdir=True,
                workdir=os.path.join(os.getcwd(),'tmp'),
                timeout = 5,
                verbose=True)

        settings = {'gun_voltage':   inputs[f'{tag}gun:voltage'].magnitude, 
                    'sol01_current': inputs[f'{tag}sol1:current'].magnitude,
                    'sol02_current': inputs[f'{tag}sol2:current'].magnitude,
                    'npts':          inputs[f'{tag}gpt:n_screen'].magnitude+1}

        result = G.set_variables(settings)
        G.run()
        #----------------------------------------------------------------------------


        #----------------------------------------------------------------------------
        # Load all relevant data into output structure
        #----------------------------------------------------------------------------
        # laser distribution
        output = {f'{tag}laser:r':rs.magnitude, f'{tag}laser:Pr':Pr.magnitude, f'{tag}gun:current':gun_current.magnitude}

        # GPT statistical data
        stats = {'max':['r'], 'mean':['x', 'y', 'z', 'kinetic_energy'], 'sigma':['x','y']}
        for stat, variables in stats.items():
                output = {**output, **{f'{tag}beam:{stat}_{var}': self.gpt_stat_to_pv(G, f'{stat}_{var}', 'screen').magnitude for var in variables} }

        scr_numbers = [1]
        for scr_number in scr_numbers:
            z = inputs[f'{tag}scr{scr_number}:mean_z'].magnitude
            for var in ['x' ,'y']:
                output[f'{tag}scr{scr_number}:mean_{var}']  = np.interp(z, output[f'{tag}beam:mean_z'], output[f'{tag}beam:mean_{var}'])
                output[f'{tag}scr{scr_number}:sigma_{var}'] = np.interp(z, output[f'{tag}beam:mean_z'], output[f'{tag}beam:sigma_{var}'])
                
        # transmission
        output[f'{tag}beam:transmission'] = [100*len(screen['x'])/inputs[f'{tag}gpt:n_particle'].magnitude for screen in G.screen]
    
        min_clearance = np.min( (inputs[f'{tag}beampipe:radius']-self.gpt_stat_to_pv(G, f'{stat}_{var}', 'screen') ) ).to('mm')
        output[f'{tag}beam:radiation'] = output[f'{tag}gun:current']*np.max(output[f'{tag}beam:mean_kinetic_energy'])/min_clearance.magnitude
        #----------------------------------------------------------------------------


        return output