Distgen basic usage¶
# Useful for debugging
%load_ext autoreload
%autoreload 2
%config InlineBackend.figure_format = 'retina'
from matplotlib import pyplot as plt
import numpy as np
import os
import yaml
Distgen supports nearly arbitrary generation of 6D + time particle coordinates.
Generator¶
Generator class defines the main object that creates a beam distribution. Typical usage is to supply the Generator.__init__(input, verbose) with an input file name and a verbose integer value to control the level of output printed to the user. Input file types can be either json or YAML. The user may also supply a dictionary as the input argument.
from distgen import Generator
gen = Generator("data/rad.gaussian.in.yaml", verbose=0)
Printing a Generator object displays the basic input data read from a distgen input file:
print(gen)
This data is stored internally in fully parsed form. To get this structure it, use the convenience property Generator.input. Note: there is no setter.
gen.input
The input structure allows for four top levels: generator, beam, {variable}_dists, transforms, and output. Any other input will throw an exception when attempting to create the beam:
try:
gen["foo"] = "bar"
except Exception as ex:
print(ex)
Distgen generator.input is a nested dictionary by design. While this logically organizes the input, it can be somewhat cumbersome for changing input. Thus the generator.input dictionary can be accessed with pure strings via the [...] method with a flattened dictionary (nested keys separated by ':'):
gen["start"]
gen["start:MTE"]
Brackets can also be used to set parameters. When setting a parameter that has associated units, one must either specify whether setting the value or units, or pass in only the value:
gen["start:MTE:value"] = 120
print("New MTE setting:", gen["start:MTE"])
gen["start:MTE:units"] = "eV"
print("New MTE setting:", gen["start:MTE"])
gen["start:MTE"] = 100
print("New MTE setting:", gen["start:MTE"])
To create a distribution, run Generator.beam():
gen = Generator("data/rad.gaussian.in.yaml", verbose=1)
beam = gen.beam()
Alternatively, call .run(), which creates an openPMD-beamphysics style ParticleGroup
gen.verbose = False
gen.run()
gen.particles
Internally, Distgen uses the Pint module to handle/check units automatically, see PINT.
Physical Constants¶
All physical constants used by Distgen are accessed via the main PHYSICAL_CONSTANTS object. Access to specific constants is done via [...], and returns the associated quantity (value with units) from the SciPy physical_constants module.
from distgen import PHYSICAL_CONSTANTS
PHYSICAL_CONSTANTS["elementary charge"]
Particle Species¶
As of v2.0.0, additional particle species can be generated. The supported list is:
PHYSICAL_CONSTANTS.species_list
Species specific data can be accessed via the .species function, which returns a dictionary of species specific parameters:
PHYSICAL_CONSTANTS.species("electron")
Beam Object¶
The particle coordinates created by the Generator class are stored internally a beam obect. The main underlying coordinates are the 6D phase space coordinates $x$, $y$, $z$, $p_x$, $p_y$, $p_z$, and time $t$. These can be accessed via the [...] operator. The bunch charge is given by beam.q.
gen = Generator("data/rad.gaussian.in.yaml", verbose=0)
beam = gen.beam()
print("X coordinates:", beam["x"])
print("Bunch charge:", beam["q"])
Currently, the code assumes one particle species (default: electrons) per beam object. The bunch contains an array of particle weights (currently uniform) for use with averaging. The particle weights $w$ are normalized $\sum_i{w_i}=1$ and are used for computing expectation values over the particle ensemble.
Simple examples include beam.avg: $\langle\mathcal{O}\rangle = \sum_i w_i\mathcal{O}_i$ and beam.std(): $\sigma_{\mathcal{O}} = \sqrt{\sum_i{w_i(\mathcal{O}_i-\langle\mathcal{O}\rangle)^2}}$.
Other examples include the normalized and geometric emittance and the twiss parameters $\beta$ and $\alpha$ (See the distgen.beam.py).
print("Beam species:", beam.species)
print("Weights:", beam["w"])
print("Normalization sum(weights): ", np.sum(beam["w"]))
print(f'Avgerage of X: {beam.avg("x"):G~P}')
print(f'Standard Deviation of X: {beam.std("x",desired_units="mm"):G~P}')
Generating OpenPMD-beamphysics Particle Group objects¶
Distgen supports direct creation of OpenPMD beam physics particle group objects, a very useful object for handling beam distribitions and writing them to various codes. Please see https://github.com/ChristopherMayes/openPMD-beamphysics for details. To generate a OpenPMD-beamphysics particle group, use the run() command. The resulting ParticleGroup is saved as Generator.particles:
gen.run()
pg = gen.particles
pg.plot("x", "px")
Distribution Objects¶
Distgen handles distirbutions using classes controlled by Generator. Currently there are three types supported: 1D distributions (Dist1d), radial distributions for the $r$ coordinate, and 2D distributions. Each distribution type keeps an internal list of the required and optional parameters that must be passed to the object. When a distribution object is initialized, the input parameters are checked against this list. Unexpected inputs will throw exceptions:
gen = Generator("data/rad.gaussian.in.yaml", verbose=0)
try:
gen["r_dist:foo"] = "bar"
gen.beam()
except Exception as ex:
print(ex)
As will missing required parameters:
gen = Generator("data/beer.can.in.yaml", verbose=0)
try:
gen["r_dist"] = {}
gen.beam()
except Exception as ex:
print(ex)
Output writer functionality¶
Distgen supports writer functions for various output formats/codes. Currently these include GPT and ASTRA, and the openPMD specification. To write to a file, use distgen.writers.writer with appropriate code type specified as well as the beam object and desired output file. The writer function accepts additional parameters required for any code dependent output.
from distgen.writers import writer
gen = Generator("data/beer.can.in.yaml", verbose=0)
beam = gen.beam()
# Typically the user can just call the general write function and specify the desired format:
writer("gpt", beam, "gpt.out.txt", verbose=1)
writer("astra", beam, "astra.out.txt", verbose=1)
writer("openPMD", beam, "openPMD.out.h5", verbose=1)
Plotting¶
Distgen provides the user with some basic plotting routines for help in visualizing the beam coordinates. The most common plot types include plot_dist2d, plot_dist1d, plot_radial_dist, and plot_current_profile. Below shows an example of how to plot 2 coordinates, in this case $x$ and $y$. The coloring scheme is a scatter histogram.
# The distgen module has some basic plotting routines
from distgen.plot import (
plot_dist1d,
plot_dist2d,
plot_radial_dist,
plot_current_profile,
)
units = {
"x": "mm",
"y": "mm",
"z": "mm",
"r": "mm",
"px": "keV/c",
"py": "keV/c",
"pz": "keV/c",
"t": "ps",
"q": "pC",
"thetax": "mrad",
"I": "A",
}
gen = Generator("data/beer.can.in.yaml", verbose=0)
beam = gen.beam()
fig = plt.figure(1)
# X-Y Plot
plot_dist2d(
beam,
"x",
units["x"],
"y",
units["y"],
style="scatter_hist2d",
nbins=100,
axis="equal",
title_on=True,
);
For plotting the histogram of the radial coordinate $r$, use plot_radial_dist:
gen = Generator("data/beer.can.in.yaml", verbose=0)
beam = gen.beam()
plot_radial_dist(beam, units["r"], scale="charge", nbins=100, title_on=True);
For plotting 1D projections of the beam distribution, use plot_dist1d:
gen = Generator("data/beer.can.in.yaml", verbose=0)
beam = gen.beam()
plot_dist1d(beam, "x", units["x"], scale="number", nbins=50, title_on=True)
For plotting the current profile, use plot_current_profile:
gen = Generator("data/beer.can.in.yaml", verbose=0)
beam = gen.beam()
plot_current_profile(beam, units["t"], units["I"], title_on=True, nbins=100);
Running Distgen with driver functions¶
The main in python driver function for distgen is distgen.drivers.run_distgen(). The function creates its own Generator object and returns a beam object. Inputs can either be an inputfile or a parsed distgen supported input dictionary. Here the function is run with an input dictionary created above.
from distgen.drivers import run_distgen
with open("data/rad.gaussian.in.yaml") as fid:
p = yaml.safe_load(fid)
p
beam = run_distgen(inputs=p, verbose=1)
Below the function is called with 'inputs' pointing to an input file. The function can also take accept a flattened settings dict, which updates inputs:
new_settings = {"start:MTE:value": 0.1}
beam = run_distgen(settings=new_settings, inputs=p, verbose=1)
Finally, is also possible to run distgen as an executible script using /bin/DistGen.py
import subprocess
output = subprocess.check_output(
["Distgen", "-f", "data/rad.gaussian.in.yaml", "-v", "2"]
)
output = str(output.decode("utf-8")).split("\n")
for line in output:
print(line)
Scaling/Rotating/Shifting Coordinates¶
The distgen module allows the user with several options for applying coordinate transformations. Some of the basic operations include: shifting, scaling, and rotating coordinates. See the transform.ipynb for a detailed description of this functionality.
Cathode Distributions¶
Currently Distgen supports several cathode emission models. Currently these all assume the particles are emitted on a flat surface at z = 0 m. The particle time coordinates in this case effectively respresents the emission time of the particle from the cathode, and thus the time distribution should be set by the user. For photocathodes, the time coordinate can be related to the laser pulse longitudinal intensity. To designate emission from a cathode, please set the input key start:type = cathode
WARNING: the user is ultimately responsible to make sure the cathode model is physically consistent with the particle species generated.
1. Maxwell-Boltzmann (default): thermalized particle momenta¶
Particles emitted from this type of cathode are assumed to have a Maxwell-Boltzmann distribution for their total mometnum $|p|$ parameterized by an energy scale denoted $MTE$:
$\rho_p(|p|) = \left(\frac{1}{2\pi mMTE}\right)^{3/2} 4\pi |p|^2 \exp\left(-\frac{|p|^2}{2m MTE}\right)$, with $\int\rho_p(|p|)d|p| = 1$.
Assuming spherical symmetry, it is easy to show the corresponding PDF for the components of the momenta is given by:
$\rho(p_x,p_y,p_z) = \left(\frac{1}{2\pi mMTE}\right)^{3/2}\exp\left(-\frac{p_x^2+p_y^2+p_z^2}{2m MTE}\right)$, with $\int\rho(p_x,p_y,p_z) d^3p=1$.
From the equipartition theorem the average energy $\frac{3}{2}MTE$ is divided equally among each direction and thus the energy scale $MTE$ can be identified as the Mean Transverse Energy: $MTE=\frac{\langle p_x^2\rangle}{2m} + \frac{\langle p_y^2\rangle}{2m}$.
This expression can be used to relate the MTE to the initial cathode emittance:
$\epsilon_{n,x} = \sigma_{x}\sqrt{ \frac{MTE}{mc^2} }$
assuming there is no distinction between the $p_x$ and $p_y$ distributions inherent to the emission process. To account for the fact that particles with $p_z<0$ are not emitted from the cathode, distgen takes $p_z = |p_z|$. Becauset this model serves as the default, simpy setting the input key start:MTE to a value (with units) will instruct the Generator to assume thermalized momenta. Note the absolute value of the $p_z$ components is taken.
from distgen import Generator
gen = Generator(input_file="data/rad.gaussian.in.yaml", verbose=1)
beam = gen.beam()
fig, ax = plt.subplots(1, 2, sharex="col", constrained_layout=True)
plot_dist2d(beam, "x", units["x"], "px", units["px"], ax=ax[0])
ax[0].set_title("x-Px Phase Space")
plot_dist2d(beam, "t", units["t"], "pz", units["pz"], ax=ax[1])
ax[1].set_title("t-Pz Phase Space");
Fermi-Dirac 3 Step Barrier Photocathode¶
Calculates the PDF for electrons emitted from a photocathode following the model described in [1].
The input parameters for this are the photon (laser) wavelength, as well as the cathode temperature, work function, and Fermi energy.
Shortly, electrons are initially populated in momentum space as in the Sommerfeld model. They escape over the work function barrier when their longitudinal energy is high enough and lose momentum along the direction of the surface normal to satisfy energy conservation. This sampling concept is the same described in [2].
[1] Dowell, D. H., & Schmerge, J. F. (2009). Quantum efficiency and thermal emittance of
metal photocathodes. Physical Review Special Topics - Accelerators and Beams, 12(7).
https://doi.org/10.1103/PhysRevSTAB.12.074201
[2] Pierce, C. M., Durham, D. B., Riminucci, F., Dhuey, S., Bazarov, I., Maxson, J.,
Minor, A. M., & Filippetto, D. (2023). Experimental Characterization of Photoemission
from Plasmonic Nanogroove Arrays. Physical Review Applied, 19(3), 034034.
https://doi.org/10.1103/PhysRevApplied.19.034034gen = Generator("data/fermi_dirac_3step_barrier_photocathode.in.yaml", verbose=1)
beam = gen.beam()
fig, ax = plt.subplots(1, 2, sharex="col", constrained_layout=True, figsize=(8, 5))
plot_dist2d(beam, "px", units["px"], "pz", units["pz"], ax=ax[0])
ax[0].set_title("px-pz")
plot_dist2d(beam, "t", units["t"], "x", units["x"], ax=ax[1])
ax[1].set_title("t-x");
from distgen.physical_constants import PHYSICAL_CONSTANTS
c = PHYSICAL_CONSTANTS["speed of light in vacuum"]
MC2 = PHYSICAL_CONSTANTS.species("electron")["mc2"]
MTE = (c * beam["px"].std()) ** 2 / MC2
print("estimated MTE:", MTE.to("meV"))
3. General Surface Emission defined by |P| or KE and $\theta_p$, $\phi_p$ distribution(s)¶
In addition to the two cathode emission models described above, the user may also specify a distribution for $|p|$ or the kinetic energy, as well as the momenta angle coordinates (azimuthal and polar momentum angles). If no angular distribution parameters are set, then particles are emitted uniformly into the forward hemisphere. In the example below, an initial KE distribution is created using a bi-exponential.
KE = np.linspace(0, 300, 10000)
E0 = 0
A1, m1 = 0.8, 8
A2, m2 = 0.1, 90
PKE = A1 * np.exp(-np.abs(KE - E0) / m1) + A2 * np.exp(-np.abs(KE - E0) / m2)
PKE = PKE / np.trapezoid(PKE, KE) # Numerically intergate to normalize
plt.plot(KE, PKE)
plt.xlabel("KE (eV)")
plt.ylabel("$\\rho(KE)$ (1/eV)");
input_yaml = """
n_particle: 100000
species: electron
start:
type: cathode
random:
type: hammersley
total_charge:
units: C
value: 1.60217663e-17
"""
input = yaml.safe_load(input_yaml)
input["KE_dist"] = {"type": "dist1d", "KE": KE, "PKE": PKE, "units": "eV"}
gen = Generator(input)
P = gen.run()
P.plot("kinetic_energy")
fig = plt.figure()
ax = fig.add_subplot(projection="3d")
x = P["px"] / P.p
y = P["py"] / P.p
z = P["pz"] / P.p
ax.scatter(x[::100], y[::100], z[::100], ".")
ax.set_xlabel(r"$\hat{p}_x$")
ax.set_ylabel(r"$\hat{p}_y$")
ax.set_zlabel(r"$\hat{p}_z$")
Time Start¶
Distgen also allows the user to start a beam at an arbitrary time. In this case, distgen will ignore any user distribution settings for the time coordinate of the beam particles. The beam has some initial Gaussian momentum. To make this more interesting, the beam is boosted to 1 GeV, the Twiss Parameters set, and a cosine like energy spread applied:
from distgen.physical_constants import unit_registry, PHYSICAL_CONSTANTS
pi, c = PHYSICAL_CONSTANTS.pi, PHYSICAL_CONSTANTS["speed of light in vacuum"]
gen = Generator("data/gaussian.in.yaml", verbose=0)
f = 1.3 * unit_registry("GHz")
w = 2 * pi * f / c
linac_params = {
"type": "cosine z:pz",
"amplitude": {"value": 1, "units": "GeV/c"},
"omega": {"value": w.magnitude, "units": str(w.units)},
"phase": {"value": -1, "units": "deg"},
}
xbeta_params = {
"type": "set_twiss x",
"beta": {"value": 12.5, "units": "m"},
"alpha": {"value": -1, "units": ""},
"emittance": {"value": 2, "units": "nm"},
}
ybeta_params = {
"type": "set_twiss y",
"beta": {"value": 12.5, "units": "m"},
"alpha": {"value": -1, "units": ""},
"emittance": {"value": 2, "units": "nm"},
}
gen["transforms"] = {
"t1": linac_params,
"t2": xbeta_params,
"t3": ybeta_params,
"order": ["t1", "t2", "t3"],
}
beam = gen.beam()
print("YAML transforms input\n", yaml.dump(gen.input["transforms"]))
print("\nFinal Horizontal Twiss params:")
print(
f'beta: {beam.Beta("x"):G~P}, alpha: {beam.Alpha("x"):G~P}, eps: {beam.emitt("x","geometric").to("nm"):G~P}'
)
# print(f'beta: {beam.Beta("x"):G~P}, alpha: {beam.Alpha("x"):G~P}, eps: {beam.emitt("x","geometric").to("nm"):Gf~P}')
print("\nFinal Vertical Twiss params:")
print(
f'beta: {beam.Beta("y"):G~P}, alpha: {beam.Alpha("y"):G~P}, eps: {beam.emitt("y","geometric").to("nm"):G~P}'
)
fig, ax = plt.subplots(1, 3, constrained_layout=True)
plot_dist2d(beam, "x", units["x"], "y", units["y"], axis="equal", ax=ax[0])
plot_dist2d(beam, "x", units["x"], "thetax", units["thetax"], ax=ax[1])
plot_dist2d(beam, "z", units["z"], "pz", "GeV/c", ax=ax[2]);
Combining Beams¶
Currently distgen supports combining beams in to ways. The first makes use of the OpenPMD-beamphysics Particle Group's addition property. This is achieved by making two separate beams:
filename = "data/gaussian.in.yaml"
gen = Generator(filename, verbose=0)
gen["pz_dist:avg_pz"] = 1 # Set beam 1 at pz = 1 GeV
gen["z_dist:avg_z"] = -10 # Set beam 1 at z = -10 mmm
xbeta_params = {
"type": "set_twiss x",
"beta": {"value": 12.5, "units": "m"},
"alpha": {"value": -1, "units": ""},
"emittance": {"value": 2, "units": "nm"},
}
ybeta_params = {
"type": "set_twiss y",
"beta": {"value": 12.5, "units": "m"},
"alpha": {"value": -1, "units": ""},
"emittance": {"value": 2, "units": "nm"},
}
gen["transforms"] = {"t1": xbeta_params, "t2": ybeta_params, "order": ["t1", "t2"]}
pg1 = gen.run()
pg1.plot("x", "px")
pg1.plot("z")
Next, make second beam. Here shift the z position of the beam and change the $\alpha$ Twiss parameter's sign to create a different momenta distribution:
gen["z_dist:avg_z"] = 10 # Set beam 1 at z = +10 mmm
xbeta_params["alpha"]["value"], ybeta_params["alpha"]["value"] = +1, +1
gen["transforms"] = {"t1": xbeta_params, "t2": ybeta_params, "order": ["t1", "t2"]}
pg2 = gen.run()
The particle group objects support addition:
pg = pg1 + pg2
pg.plot("x", "px")
pg.plot("z")
It is also possible to superimpose built in distgen particle distributions. Please see the examples_dists.ipynb for details on how to use this functionality.
Archiving¶
All input and output can be saved and loaded to an hdf5 file using .archive() and .load_archive functions.
Input is archived as a flattended dict, with keys separated by :
Particlces are archived as openPMD-beamphysics, but only if they are created with the run() command.
gen = Generator("data/gaussian.in.yaml", verbose=0)
gen.run()
If no filename is given, a unique one will be written based on .input
afile = gen.archive()
afile
Loading into a new object
G2 = Generator(verbose=True)
G2.load_archive(afile)
open h5 handles can also be written to, using the same routine.
import h5py
with h5py.File("archive.h5", "w") as h5:
G2.archive(h5)
Clean up¶
os.remove("rad.gaussian.out.txt")
os.remove("gpt.out.txt")
os.remove("astra.out.txt")
os.remove("openPMD.out.h5")
# os.remove('beer.can.out.txt')
os.remove(afile)
os.remove("archive.h5")