Functional LUME-Astra¶

This is the functional way to run astra and return the evaluate Astra object, or simple dict of outputs

Settings is a list of settings that can appear in the input file.

In [1]:

# Useful for debugging
%load_ext autoreload
%autoreload 2

# Useful for debugging %load_ext autoreload %autoreload 2

In [2]:

from astra import run_astra, run_astra_with_generator

import matplotlib.pyplot as plt
import os

from astra import run_astra, run_astra_with_generator import matplotlib.pyplot as plt import os

Input template file and some settings

In [3]:

ASTRA_IN =  'templates/dcgun/astra.in'

ASTRA_IN = 'templates/dcgun/astra.in'

In [4]:

settings0 = {'zstop':1, 'zemit':200, 'zphase':1, 'phases':True}

settings0 = {'zstop':1, 'zemit':200, 'zphase':1, 'phases':True}

Run Astra¶

In [5]:

A1 = run_astra(settings0, astra_input_file=ASTRA_IN, verbose=True, timeout=100)

A1 = run_astra(settings0, astra_input_file=ASTRA_IN, verbose=True, timeout=100)

run_astra
zstop is in astra newrun
zstop is in astra output
zemit is in astra newrun
zphase is in astra newrun
phases is in astra newrun
loading 1 particle files
[100.0]
{'start_time': 1690854350.025616, 'run_script': '/Users/chrisonian/Code/Astra/bin/Astra astra.in', 'why_error': '', 'run_time': 0.7370400428771973, 'run_error': False}

Run Astra with Generator¶

In [6]:

GENERATOR_IN = 'templates/dcgun/generator.in'


settings0['ipart']= 2000

A2 = run_astra_with_generator(settings0, astra_input_file=ASTRA_IN,
                              generator_input_file=GENERATOR_IN, verbose=True)

GENERATOR_IN = 'templates/dcgun/generator.in' settings0['ipart']= 2000 A2 = run_astra_with_generator(settings0, astra_input_file=ASTRA_IN, generator_input_file=GENERATOR_IN, verbose=True)

run_astra_with_generator
zstop is in astra newrun
zstop is in astra output
zemit is in astra newrun
zphase is in astra newrun
phases is in astra newrun
ipart is in generator
set spacecharge mesh for n_particles: 2000 to {'nrad': 10, 'nlong_in': 20}
 --------------------------------------------------------------------------

                              generator              
                 Version 1.0 - macOS 64bit - Intel   
                        DESY,  Hamburg 2002          
                        Mon Jul 31 18:45:50 

     Working File is:    generator.in                                      
     Initializing       2000      electrons
     including 6 probe particles at standard positions
     Particles start from a cathode
     Particles are quasi randomly distributed

          WARNING: Values could not be reached:
          Energy spread too high.
                  46 times standard correction procedure



     Final check:
     Particles taken into account      N =       2000
     total charge                      Q =    -0.1000     nC
     horizontal beam position          x =     2.3378E-06 mm
     vertical beam position            y =     8.9462E-06 mm
     longitudinal beam position        z =      0.000     m
     horizontal beam size          sig x =     0.2500     mm
     vertical beam size            sig y =     0.2500     mm
     longitudinal beam size        sig z =      0.000     mm
     total emission time               t =     5.0872E-02 ns
     rms emission time             sig t =     8.3856E-03 ns
     average kinetic energy            E =     1.0197E-06 MeV
     energy spread                    dE =     4.2631E-04 keV
     average momentum                  P =     1.0208E-03 MeV/c
     transverse beam emittance     eps x =     0.1747     pi mrad mm
     correlated divergence         cor x =     6.8371E-03 mrad 
     transverse beam emittance     eps y =     0.1750     pi mrad mm
     correlated divergence         cor y =    -4.4571E-03 mrad
     longitudinal beam emittance   eps z =      0.000     pi keV mm
     correlated energy spread      cor z =      0.000     keV
     emittance ratio eps y/eps x         =     0.9982    

     phase-space distribution saved to file: generator.part                


     Generator ended with     1 warning(s)


Initial particles written to /var/folders/2f/l5_mybzs30j4qqvyj98w1_nw0000gn/T/tmp5xjn8oto/astra.particles
loading 1 particle files
[100.0]
{'start_time': 1690854350.8262599, 'run_script': '/Users/chrisonian/Code/Astra/bin/Astra astra.in', 'why_error': '', 'run_time': 0.7480919361114502, 'run_error': False}
run_astra_with_generator finished

Run Astra with Distgen¶

In [7]:

from astra.astra_distgen import run_astra_with_distgen
DISTGEN_IN =  'templates/dcgun/distgen.yaml'

settings0 = {'zstop':1, 'zemit':200, 'zphase':1, 'phases':True}
settings0['distgen:n_particle'] = 2000

A3 = run_astra_with_distgen(settings0, astra_input_file=ASTRA_IN,
                              distgen_input_file=DISTGEN_IN, verbose=True)

from astra.astra_distgen import run_astra_with_distgen DISTGEN_IN = 'templates/dcgun/distgen.yaml' settings0 = {'zstop':1, 'zemit':200, 'zphase':1, 'phases':True} settings0['distgen:n_particle'] = 2000 A3 = run_astra_with_distgen(settings0, astra_input_file=ASTRA_IN, distgen_input_file=DISTGEN_IN, verbose=True)

run_astra_with_generator
zstop is in astra newrun
zstop is in astra output
zemit is in astra newrun
zphase is in astra newrun
phases is in astra newrun
Setting distgen n_particle = 2000
Distribution format: None
   Warning: no output file specified, defaulting to "None".
Output file: None

Creating beam distribution....
   Beam starting from: cathode
   Total charge: 100 pC.
   Number of macroparticles: 2000.
   Assuming cylindrical symmetry...
   r distribution: radial uniform
      min_r = 0 mm, max_r = 0.5 mm
   theta distribution: uniform theta
      min_theta = 0 rad, max_theta = 6.28319 rad
   t distribution: Gaussian
      avg_t = 0 ps, sigma_t = 8.500 ps
      Left n_sigma_cutoff = 3, Right n_sigma_cutoff = -3
   px distribution: Gaussian
      avg_px = 0 eV/c, sigma_px = 357.421 eV/c
   py distribution: Gaussian
      avg_py = 0 eV/c, sigma_py = 357.421 eV/c
   pz distribution: Gaussian
      avg_pz = 0 eV/c, sigma_pz = 357.421 eV/c
   Shifting avg_x = -0.000182722 mm -> 0 mm
   Scaling sigma_x = 0.249836 mm -> 0.25 mm
   Shifting avg_y = -5.96585E-05 mm -> 0 mm
   Scaling sigma_y = 0.249997 mm -> 0.25 mm
   Shifting avg_px = -1.43404 eV/c -> 0 eV/c
   Scaling sigma_px = 356.94 eV/c -> 357.421 eV/c
   Shifting avg_py = -0.99808 eV/c -> 0 eV/c
   Scaling sigma_py = 356.407 eV/c -> 357.421 eV/c
   Shifting avg_pz = -2.93972 eV/c -> 0 eV/c
   Scaling sigma_pz = 356.195 eV/c -> 357.421 eV/c
   Shifting avg_t = -0.0272122 ps -> 0 ps
   Scaling sigma_t = 8.36979 ps -> 8.38592 ps
   Cathode start: fixing pz momenta to forward hemisphere
      avg_pz -> 285.455 eV/c, sigma_pz -> 215.094 eV/c
...done. Time Elapsed: 20.7162 ms.

   Created particles in .particles: 
   ParticleGroup with 2000 particles with total charge 9.999999999999999e-11 Cset spacecharge mesh for n_particles: 2000 to {'nrad': 10, 'nlong_in': 20}
Initial particles written to /var/folders/2f/l5_mybzs30j4qqvyj98w1_nw0000gn/T/tmppi727x6e/astra.particles
loading 1 particle files
[100.0]
{'start_time': 1690854351.635436, 'run_script': '/Users/chrisonian/Code/Astra/bin/Astra astra.in', 'why_error': '', 'run_time': 0.7446029186248779, 'run_error': False}

These Generator and Distgen inputs are set up to produce nearly the same initial beams

In [8]:

plt.plot(A2.stat('sigma_x'))
plt.plot(A3.stat('sigma_x'))

plt.plot(A2.stat('sigma_x')) plt.plot(A3.stat('sigma_x'))

Out[8]:

[<matplotlib.lines.Line2D at 0x151a1e880>]

Evaluate functions¶

For scans an optimizations, a user often wants to run a simulation many times and examine some particular output. The package provides several evaluate_ functions that are similar to the run_ functions above, but apply some merit function to the output and returns that. Additionally, if an archive_path is provided, the .archive method will be called to save the complete object output

This is the default merit function that is applied. The user can supply a different function if needed.

In [9]:

from astra.evaluate import default_astra_merit

default_astra_merit(A2)

from astra.evaluate import default_astra_merit default_astra_merit(A2)

Out[9]:

{'error': False,
 'end_mean_z': 1.0,
 'end_mean_t': 4.0153e-09,
 'end_mean_x': 6.1453e-09,
 'end_sigma_x': 0.00046093,
 'end_sigma_xp': 0.00051108,
 'end_norm_emit_x': 1.7455e-07,
 'end_cov_x__xp': 2.1221217199999998e-07,
 'end_mean_y': 3.9381999999999995e-08,
 'end_sigma_y': 0.00046181,
 'end_sigma_yp': 0.0005110700000000001,
 'end_norm_emit_y': 1.7527e-07,
 'end_cov_y__yp': 2.125018715e-07,
 'end_mean_kinetic_energy': 499790.0,
 'end_sigma_z': 0.0021705,
 'end_sigma_energy': 1.5773,
 'end_norm_emit_z': 0.0034225,
 'end_cov_z__energy': 7.868713650000001e-05,
 'end_n_particle_loss': 0,
 'end_total_charge': 9.964999999999999e-11,
 'end_higher_order_energy_spread': 1.5777476779634658}

evaluate_astra_with_generator¶

An even simpler run, returns a simple dict of outputs. For use in optimization

In [10]:

from astra import evaluate_astra_with_generator

MY_GENERATOR_TEMPLATE = 'templates/dcgun/generator.in'

settings0 = {'zstop':1, 'zemit':200, 'zphase':1, 'phases':True}
settings0['ipart'] = 2000

outputs1 = evaluate_astra_with_generator(settings0, 
                                        astra_input_file=ASTRA_IN,
                                        generator_input_file=GENERATOR_IN,
                                        archive_path = '.')
outputs1

from astra import evaluate_astra_with_generator MY_GENERATOR_TEMPLATE = 'templates/dcgun/generator.in' settings0 = {'zstop':1, 'zemit':200, 'zphase':1, 'phases':True} settings0['ipart'] = 2000 outputs1 = evaluate_astra_with_generator(settings0, astra_input_file=ASTRA_IN, generator_input_file=GENERATOR_IN, archive_path = '.') outputs1

Out[10]:

{'error': False,
 'end_mean_z': 1.0,
 'end_mean_t': 4.0153e-09,
 'end_mean_x': 6.1453e-09,
 'end_sigma_x': 0.00046093,
 'end_sigma_xp': 0.00051108,
 'end_norm_emit_x': 1.7455e-07,
 'end_cov_x__xp': 2.1221217199999998e-07,
 'end_mean_y': 3.9381999999999995e-08,
 'end_sigma_y': 0.00046181,
 'end_sigma_yp': 0.0005110700000000001,
 'end_norm_emit_y': 1.7527e-07,
 'end_cov_y__yp': 2.125018715e-07,
 'end_mean_kinetic_energy': 499790.0,
 'end_sigma_z': 0.0021705,
 'end_sigma_energy': 1.5773,
 'end_norm_emit_z': 0.0034225,
 'end_cov_z__energy': 7.868713650000001e-05,
 'end_n_particle_loss': 0,
 'end_total_charge': 9.964999999999999e-11,
 'end_higher_order_energy_spread': 1.5777476779634658,
 'fingerprint': '342ec49887deb72ecbf7505f4bda8e25',
 'archive': '/Users/chrisonian/Code/GitHub/lume-astra/docs/examples/342ec49887deb72ecbf7505f4bda8e25.h5'}

The archive can be loaded back:

In [11]:

from astra import Astra
AX = Astra.from_archive(outputs1['archive'])

from astra import Astra AX = Astra.from_archive(outputs1['archive'])

Check the merit:

In [12]:

default_astra_merit(AX)

default_astra_merit(AX)

Out[12]:

{'error': False,
 'end_cov_x__xp': 2.1221217199999998e-07,
 'end_cov_y__yp': 2.125018715e-07,
 'end_cov_z__energy': 7.868713650000001e-05,
 'end_mean_kinetic_energy': 499790.0,
 'end_mean_t': 4.0153e-09,
 'end_mean_x': 6.1453e-09,
 'end_mean_y': 3.9381999999999995e-08,
 'end_mean_z': 1.0,
 'end_norm_emit_x': 1.7455e-07,
 'end_norm_emit_y': 1.7527e-07,
 'end_norm_emit_z': 0.0034225,
 'end_sigma_energy': 1.5773,
 'end_sigma_x': 0.00046093,
 'end_sigma_xp': 0.00051108,
 'end_sigma_y': 0.00046181,
 'end_sigma_yp': 0.0005110700000000001,
 'end_sigma_z': 0.0021705,
 'end_n_particle_loss': 0,
 'end_total_charge': 9.964999999999999e-11,
 'end_higher_order_energy_spread': 1.5777476779634658}

The Generator can also be loaded from the same archive:

In [13]:

from astra import AstraGenerator
G = AstraGenerator.from_archive(outputs1['archive'])
G.input

from astra import AstraGenerator G = AstraGenerator.from_archive(outputs1['archive']) G.input

Out[13]:

{'add': False,
 'c_sig_clock': 3,
 'c_sig_ekin': 5,
 'c_sig_px': 5,
 'c_sig_x': 5,
 'c_sig_y': 5,
 'cathode': True,
 'cor_px': 0,
 'cor_py': 0,
 'dist_px': 'g',
 'dist_py': 'g',
 'dist_pz': 'g',
 'dist_x': 'r',
 'dist_y': 'r',
 'dist_z': 'g',
 'fname': 'generator.part',
 'high_res': True,
 'ipart': 2000,
 'n_add': 0,
 'noise_reduc': True,
 'probe': True,
 'q_total': 0.1,
 'ref_clock': 0,
 'ref_ekin': 1e-06,
 'ref_zpos': 0,
 'sig_clock': 0.00849257,
 'sig_ekin': 0.0005,
 'sig_px': 357.7,
 'sig_py': 357.7,
 'sig_x': 0.25,
 'sig_y': 0.25,
 'species': 'electrons'}

In [14]:

# Cleanup
os.remove(outputs1['archive'])

# Cleanup os.remove(outputs1['archive'])

evaluate_astra_with_distgen¶

In [15]:

from astra import evaluate_astra_with_distgen

?evaluate_astra_with_distgen

from astra import evaluate_astra_with_distgen ?evaluate_astra_with_distgen

Signature:
evaluate_astra_with_distgen(
    settings,
    astra_input_file=None,
    distgen_input_file=None,
    workdir=None,
    astra_bin='$ASTRA_BIN',
    timeout=2500,
    verbose=False,
    auto_set_spacecharge_mesh=True,
    archive_path=None,
    merit_f=None,
)
Docstring:
Similar to run_astra_with_distgen, but returns a flat dict of outputs as processed by merit_f. 

If no merit_f is given, a default one will be used. See:
    astra.evaluate.default_astra_merit

Will raise an exception if there is an error. 
File:      ~/Code/GitHub/lume-astra/astra/astra_distgen.py
Type:      function

In [16]:

settings5 = {'zstop':1, 'zemit':200, 'zphase':1, 'phases':True}
settings5['distgen:n_particle'] = 2000

outputs2 = evaluate_astra_with_distgen(settings5, astra_input_file=ASTRA_IN,
                              distgen_input_file=DISTGEN_IN, archive_path='.')
outputs2

settings5 = {'zstop':1, 'zemit':200, 'zphase':1, 'phases':True} settings5['distgen:n_particle'] = 2000 outputs2 = evaluate_astra_with_distgen(settings5, astra_input_file=ASTRA_IN, distgen_input_file=DISTGEN_IN, archive_path='.') outputs2

Out[16]:

{'error': False,
 'end_mean_z': 1.0,
 'end_mean_t': 4.0156e-09,
 'end_mean_x': 7.553300000000001e-11,
 'end_sigma_x': 0.0004619,
 'end_sigma_xp': 0.00051138,
 'end_norm_emit_x': 1.7514e-07,
 'end_cov_x__xp': 2.1275114000000003e-07,
 'end_mean_y': -3.6688e-09,
 'end_sigma_y': 0.00046157,
 'end_sigma_yp': 0.00051168,
 'end_norm_emit_y': 1.7457e-07,
 'end_cov_y__yp': 2.12876084e-07,
 'end_mean_kinetic_energy': 499790.0,
 'end_sigma_z': 0.0021695000000000004,
 'end_sigma_energy': 1.5373999999999999,
 'end_norm_emit_z': 0.0033353,
 'end_cov_z__energy': -2.1510375550000005e-05,
 'end_n_particle_loss': 0,
 'end_total_charge': 9.999999999999999e-11,
 'end_higher_order_energy_spread': 1.5373494617932955,
 'fingerprint': '87298760ea8c8ab1a8d95865f84b07bd',
 'archive': '/Users/chrisonian/Code/GitHub/lume-astra/docs/examples/87298760ea8c8ab1a8d95865f84b07bd.h5'}

In [17]:

from distgen import Generator

from distgen import Generator

In [18]:

G = Generator()
G.load_archive(outputs2['archive'])
#

G = Generator() G.load_archive(outputs2['archive']) #

In [19]:

# Note that there are no particles, this is just the input. 
# Particles are 
G.particles == None

# Note that there are no particles, this is just the input. # Particles are G.particles == None

Out[19]:

True

In [20]:

# This will essentially recreate evaluate_astra_with_distgen
A = Astra(verbose=False)

A.load_archive(outputs2['archive'])

# Make particles
G.run()

A.initial_particles = G.particles
A.configure()
A.run()
default_astra_merit(A)

# This will essentially recreate evaluate_astra_with_distgen A = Astra(verbose=False) A.load_archive(outputs2['archive']) # Make particles G.run() A.initial_particles = G.particles A.configure() A.run() default_astra_merit(A)

No input data specified.
Run Aborted Traceback (most recent call last):
  File "/Users/chrisonian/Code/GitHub/lume-astra/astra/astra.py", line 313, in run_astra
    self.load_output()
  File "/Users/chrisonian/Code/GitHub/lume-astra/astra/astra.py", line 243, in load_output
    assert len(nlist) == 1, f'Stat keys do not all have the same length: {[len(stats[k]) for k in stats]}'
AssertionError: Stat keys do not all have the same length: []

Out[20]:

{'error': True}

In [21]:

# Cleanup
os.remove(outputs2['archive'])

# Cleanup os.remove(outputs2['archive'])

In [22]:

# Compare Generator, Distgen. Thee should be similar, but not exactly the same. 
for k in outputs1:
    print(k, outputs1[k], outputs2[k])

# Compare Generator, Distgen. Thee should be similar, but not exactly the same. for k in outputs1: print(k, outputs1[k], outputs2[k])

error False False
end_mean_z 1.0 1.0
end_mean_t 4.0153e-09 4.0156e-09
end_mean_x 6.1453e-09 7.553300000000001e-11
end_sigma_x 0.00046093 0.0004619
end_sigma_xp 0.00051108 0.00051138
end_norm_emit_x 1.7455e-07 1.7514e-07
end_cov_x__xp 2.1221217199999998e-07 2.1275114000000003e-07
end_mean_y 3.9381999999999995e-08 -3.6688e-09
end_sigma_y 0.00046181 0.00046157
end_sigma_yp 0.0005110700000000001 0.00051168
end_norm_emit_y 1.7527e-07 1.7457e-07
end_cov_y__yp 2.125018715e-07 2.12876084e-07
end_mean_kinetic_energy 499790.0 499790.0
end_sigma_z 0.0021705 0.0021695000000000004
end_sigma_energy 1.5773 1.5373999999999999
end_norm_emit_z 0.0034225 0.0033353
end_cov_z__energy 7.868713650000001e-05 -2.1510375550000005e-05
end_n_particle_loss 0 0
end_total_charge 9.964999999999999e-11 9.999999999999999e-11
end_higher_order_energy_spread 1.5777476779634658 1.5373494617932955
fingerprint 342ec49887deb72ecbf7505f4bda8e25 87298760ea8c8ab1a8d95865f84b07bd
archive /Users/chrisonian/Code/GitHub/lume-astra/docs/examples/342ec49887deb72ecbf7505f4bda8e25.h5 /Users/chrisonian/Code/GitHub/lume-astra/docs/examples/87298760ea8c8ab1a8d95865f84b07bd.h5