# Useful for debugging
%load_ext autoreload
%autoreload 2
%config InlineBackend.figure_format = 'retina'
Transformation utilties¶
from distgen import Generator
from distgen.plot import plot_dist2d
from distgen.physical_constants import unit_registry as unit
from matplotlib import pyplot as plt
import yaml
import numpy as np
units = {
"x": "mm",
"y": "mm",
"r": "mm",
"z": "mm",
"px": "keV/c",
"py": "keV/c",
"pz": "MeV/c",
"t": "ps",
"q": "pC",
"thetax": "mrad",
"thetay": "mrad",
}
A set of transformations can be applied to a beam when it is created by adding a dictionary called 'transforms' at the top level of the Distgen input structure (in either dictionary or YAML format).
Transformations are added to this dictionary with a unique identifier key. The key 'order'
is protected and should not be used. Each transformation definition must be a dictionary containing a 'type'
key associated with a known transform function and the variable(s) it acts on. A simple example is given by a translation in x: the type key is then "type: translate x". All input parameters for the transformation should be input at the same level as the 'type' key. Physical quantities with units should be dictionaries with a 'value' and 'units' key/value pair supplied.
Unless specified, transformations are applied in the order they are input in the transforms dictionary. Because transformation operations often do not commute, the user may specify a desired order by adding a list of transformation ID's to the tranform dictionary using the key 'order'
. Because python and yaml do not universally guarantee order is preserved, it is recommended to always specify the order list when the number of transforms is more than one.
In the example below, two transforms are include: scaling a round beam so that it is oval, and then rotated the beam by 45 deg. In this example the full input to Distgen is printed, in subsequent examples, only the transform details are printed.
gen = Generator("data/beer.can.in.yaml", verbose=0)
initial_beam = gen.beam()
setstdx = {"type": "set_std x", "sigma_x": {"value": 3, "units": "mm"}}
rot2dxy = {"type": "rotate2d x:y", "angle": {"value": 45, "units": "deg"}}
transy = {"type": "translate y", "delta": {"value": 1, "units": "mm"}}
gen["transforms"] = {
"t1": setstdx,
"t2": rot2dxy,
"t3": transy,
"order": ["t1", "t2", "t3"],
}
final_beam = gen.beam()
print("Input:\n", yaml.dump(gen.input["transforms"]))
fig, ax = plt.subplots(1, 2, sharex="col", constrained_layout=True)
plot_dist2d(initial_beam, "x", units["x"], "y", units["y"], axis="equal", ax=ax[0])
ax[0].set_title("Before transforms")
plot_dist2d(final_beam, "x", units["x"], "y", units["y"], axis="equal", ax=ax[1])
ax[1].set_title("After transforms");
Input: order: - t1 - t2 - t3 t1: sigma_x: units: millimeter value: 3.0 type: set_std x t2: angle: units: degree value: 45.0 type: rotate2d x:y t3: delta: units: millimeter value: 1.0 type: translate y
Distgen provides a set of basic transformation utilities that can be applied the particle coordinates of a beam object. Here the basic examples are discused: the transformation functions are defined as well as how to use them from the standard Distgen input structure. Note the these transformation operations do not in general commute.
The primary example used is that of a uniform radial distribution.
Translations¶
Translations of the coordinate $u$ are defined by: $u\rightarrow u + \Delta u$.
# Translations: a translation of a single coordinate are handled by transforms.translate
gen = Generator("data/beer.can.in.yaml", verbose=0)
beam1 = gen.beam()
transx = {"type": "translate x", "delta": {"value": +3, "units": "mm"}}
transy = {"type": "translate y", "delta": "-1 mm"}
gen["transforms"] = {"tx": transx, "ty": transy, "order": ["tx", "ty"]}
beam2 = gen.beam()
print("Input:\n", yaml.dump(gen.input["transforms"]))
fig, ax = plt.subplots(1, 2, sharex="col", constrained_layout=True)
plot_dist2d(
beam1, "x", units["x"], "y", units["y"], axis="equal", ax=ax[0], title_on=True
)
ax[0].set_title(f"Before\n{ax[0].get_title()}")
plot_dist2d(
beam2, "x", units["x"], "y", units["y"], axis="equal", ax=ax[1], title_on=True
)
ax[1].set_title(f"After\n{ax[1].get_title()}");
Input: order: - tx - ty tx: delta: units: millimeter value: 3.0 type: translate x ty: delta: units: millimeter value: -1 type: translate y
gen = Generator("data/beer.can.in.yaml", verbose=0)
beam1 = gen.beam()
setavgx = {"type": "set_avg x", "avg_x": {"value": +3, "units": "mm"}}
setavgy = {"type": "set_avg y", "avg_y": {"value": -1, "units": "mm"}}
gen["transforms"] = {"tx": setavgx, "ty": setavgy, "order": ["tx", "ty"]}
beam2 = gen.beam()
print("Input:\n", yaml.dump(gen.input["transforms"]))
fig, ax = plt.subplots(1, 2, sharex="col", constrained_layout=True)
plot_dist2d(
beam1, "x", units["x"], "y", units["y"], axis="equal", ax=ax[0], title_on=True
)
ax[0].set_title(f"Before\n{ax[0].get_title()}")
plot_dist2d(
beam2, "x", units["x"], "y", units["y"], axis="equal", ax=ax[1], title_on=True
)
ax[1].set_title(f"After\n{ax[1].get_title()}");
Input: order: - tx - ty tx: avg_x: units: millimeter value: 3.0 type: set_avg x ty: avg_y: units: millimeter value: -1.0 type: set_avg y
Scaling¶
Basic scaling is handled using transforms.scale. To scale the $x$ coordinate of the beam by $\alpha$ use:
scale(beam,'x',$\alpha$)
where $\alpha$ is a dimensionless quantity or float. Note that if the $<x>\neq0$ then $<x>\rightarrow\alpha<x>$. It is possible to fix the average value under scaling using:
scale(beam, 'x', $\alpha$, fix_average='True')
gen = Generator("data/beer.can.in.yaml", verbose=0)
beam1 = gen.beam()
transx = {"type": "translate x", "delta": {"value": 3, "units": "mm"}}
scalex = {
"type": "scale x",
"scale": 2,
}
gen["transforms"] = {"t1": transx, "t2": scalex, "order": ["t1", "t2"]}
beam1 = gen.beam()
scalex["fix_average"] = True
gen["transforms"] = {"t": transx, "s": scalex, "order": ["t", "s"]}
beam2 = gen.beam()
print("Input:\n", yaml.dump(gen.input["transforms"]))
fig, ax = plt.subplots(1, 2, constrained_layout=True)
plot_dist2d(beam1, "x", units["x"], "y", units["y"], axis="equal", ax=ax[0])
ax[0].set_title("Default scaling (average scaled)")
plot_dist2d(beam2, "x", units["x"], "y", units["y"], axis="equal", ax=ax[1])
ax[1].set_title("Scaling with average fixed");
Input: order: - t - s s: fix_average: true scale: 2 type: scale x t: delta: units: millimeter value: 3.0 type: translate x
gen = Generator("data/beer.can.in.yaml", verbose=0)
beam1 = gen.beam()
scalex = {
"type": "set_stdxy x:y",
"sigma_xy": {"value": 3.4, "units": "mm"},
}
gen["transforms"] = {"t1": scalex}
beam2 = gen.beam()
plot_dist2d(beam2, "x", units["x"], "y", units["y"], axis="equal", title_on=True);
Shift and Scale a single coordinate¶
It is possible to combine a translation with a scaling operation to shift and rescale a coordinate to have a new standard deviation while keeping the form of the underlying coordinate distribution unchanged. This is accomplished using transforms.set_avg_and_std(...)
gen = Generator("data/beer.can.in.yaml", verbose=0)
setx = {
"type": "set_avg_and_std x",
"avg_x": {"value": -3, "units": "mm"},
"sigma_x": {"value": 5, "units": "mm"},
}
gen["transforms"] = {"sx": setx}
beam = gen.beam()
print("Input:\n", yaml.dump(gen.input["transforms"]))
plot_dist2d(beam, "x", units["x"], "y", units["y"], axis="equal", title_on=True);
Input: sx: avg_x: units: millimeter value: -3.0 sigma_x: units: millimeter value: 5.0 type: set_avg_and_std x
Rotating coordinates¶
Rotation between two coordinates is accomplished using transforms.rotate2d. In addition to the beam object, the user must specify the variables to rotate. This can be done in a string of the form 'var1:var2' or as a list of strs ['var1','var2']. The variables supplied must have the same type of units. The user must also specify an angle in radians or degrees to rotate by.
Note, the default behavior is to rotate around the coordinate origins (not the coordinate averages) as seen in the plot on the left below. The rotation can be performed about a different origin by setting the keyword arguement 'origin'. Often it is desirable to set the origin of rotation to be the coordinate centroids. This is done by setting origin='centroid' (as seen in the plot on the right below).
gen = Generator("data/beer.can.in.yaml", verbose=0)
scalex = {"type": "scale x", "scale": 3}
shiftx = {"type": "translate x", "delta": {"value": 5, "units": "mm"}}
rotxy = {"type": "rotate2d x:y", "angle": {"value": 45, "units": "deg"}}
gen["transforms"] = {
"sx": scalex,
"shx": shiftx,
"rxy": rotxy,
"order": ["sx", "shx", "rxy"],
}
obeam = gen.beam()
# do with rotation around centroid
gen["transforms:rxy"] = {
"type": "rotate2d x:y",
"angle": {"value": 45, "units": "deg"},
"origin": "centroid",
}
cbeam = gen.beam()
print("YAML input:\n", yaml.dump(gen.input["transforms"]))
fig, ax = plt.subplots(1, 2, sharex="col", constrained_layout=True)
plot_dist2d(obeam, "x", units["x"], "y", units["y"], ax=ax[0], axis="equal")
ax[0].set_title("Rotated around origin")
plot_dist2d(cbeam, "x", units["x"], "y", units["y"], ax=ax[1], axis="equal")
ax[1].set_title("Rotated around centroids");
YAML input: order: - sx - shx - rxy rxy: angle: units: degree value: 45.0 origin: centroid type: rotate2d x:y shx: delta: units: millimeter value: 5.0 type: translate x sx: scale: 3 type: scale x
Shear¶
The shear operation allows one to apply a sheer in a 2D subspace according of the form:
$v\rightarrow v + \alpha u$
This can be useful in a variety of cases such as drifting particles.
gen = Generator("data/x.y.uniform.in.yaml", verbose=0)
beam1 = gen.beam()
gen["transforms"] = {
"s": {"type": "shear x:y", "shear_coefficient": {"value": 0.5, "units": ""}}
}
beam2 = gen.beam()
print("YAML:\n", yaml.dump(gen.input["transforms"]))
fig, ax = plt.subplots(1, 2, constrained_layout=True)
plot_dist2d(beam1, "x", units["x"], "y", units["y"], axis="equal", ax=ax[0])
ax[0].set_title("Before")
plot_dist2d(beam2, "x", units["x"], "y", units["y"], axis="equal", ax=ax[1])
ax[1].set_title("After");
YAML: s: shear_coefficient: units: dimensionless value: 0.5 type: shear x:y
Magnetizing a cylindrical beam provides a more physically relevant application of the sheer function. Here the magnetization $\mathcal{L}$ is added to the particle momentum in the form $p_x\rightarrow p_x + \frac{\mathcal{L}}{\sigma_{x,y}^2}y$ and $p_y\rightarrow p_y - \frac{\mathcal{L}}{\sigma_{x,y}^2}x$. This results in a transverse emittance of $\sqrt{\epsilon_{n,x,uncor}^2 + \mathcal{L}^2}$. Note that using the definitions of cylindrical variables it is possible to show this is equivalent to a sheer of $p_{\theta}\rightarrow p_{\theta}-\frac{\mathcal{L}}{\sigma_{x,y}^2}r$.
For symplicity, a magnetization function has been defined to perform the above transform given the magnetization $\mathcal{L}$. Currently this assumes a cylindrically symmetric bunch.
from distgen import PHYSICAL_CONSTANTS
c = PHYSICAL_CONSTANTS["speed of light in vacuum"]
MC2 = PHYSICAL_CONSTANTS.species("electron")["mc2"]
gen = Generator("data/beer.can.in.yaml", verbose=0)
fig, ax = plt.subplots(1, 2, constrained_layout=True, figsize=(10, 4))
ibeam = gen.beam()
plot_dist2d(ibeam, "r", units["r"], "ptheta", units["px"], ax=ax[0], title_on=True)
ax[0].set_title(f"Before\n{ax[0].get_title()}")
Lmag = -50 * unit("micrometer") * MC2.magnitude * unit("eV/c")
print(f"{Lmag.units:~P}")
gen["transforms"] = {
"m": {
"type": "magnetize r:ptheta",
"magnetization": {"value": Lmag.magnitude, "units": str(Lmag.units)},
}
}
print("YAML:\n", yaml.dump(gen.input["transforms"]))
fbeam = gen.beam()
plot_dist2d(fbeam, "r", units["r"], "ptheta", units["px"], ax=ax[1], title_on=True)
ax[1].set_title(f"After\n{ax[1].get_title()}")
eni = 0.5 * (ibeam.emitt("x", "normalized") + ibeam.emitt("y", "normalized"))
enf = 0.5 * (fbeam.emitt("x", "normalized") + fbeam.emitt("y", "normalized"))
enf0 = np.sqrt(eni**2 + (Lmag / (MC2 / c).to("eV/c")) ** 2)
err = enf / enf0 - 1
print(f'Initial emittance: {eni.to("um"):G~P}')
print(f'Final emittance: {enf.to("um"):G~P}')
print(f"Error: {100*err.magnitude} %")
eV·µm/c YAML: m: magnetization: units: electron_volt * micrometer / speed_of_light value: -25549947.499999996 type: magnetize r:ptheta
Initial emittance: 0.541796 µm Final emittance: 50.0029 µm Error: -5.464299750457258e-05 %
Polynomial¶
It is possible to apply a polynomial map in one beam variable to a second using the polynomial function:
$p \rightarrow p + \sum_{n=0}^N a_n (x-\mathcal{O})^n$.
Here $p_i$ is the dependent variable before the transformation (this term can be left out by specifying the keyword zero_dependent_var=True), the $a_n$ specify the polynomial coefficients, and $\mathcal{O}$ is the expansion origin. In the example used below, a quartic polynomial in $z$ is applied to $p_z$ mimicing the effect of an RF cavity:
import numpy as np
gen = Generator("data/gaussian.in.yaml", verbose=0)
V0 = 1000 * unit("MeV/c")
w = 2 * PHYSICAL_CONSTANTS.pi * 1.3 * unit("GHz")
k = w / c
phi = -2.5 * unit("deg")
c3 = -0.5 * V0 * k**2
polytrans = {
"type": "polynomial z:pz",
"coefficients": [
{"value": V0.magnitude, "units": str(V0.units)},
{"value": 0.0, "units": "eV/c/meter"},
{"value": c3.magnitude, "units": str(c3.units)},
],
}
gen["transforms"] = {"pt": polytrans}
beam = gen.beam()
print("YAML input:\n", yaml.dump(gen.input["transforms"]))
plot_dist2d(beam, "z", units["z"], "pz", units["pz"]);
YAML input: pt: coefficients: - units: megaelectron_volt / speed_of_light value: 1000.0 - units: electron_volt / meter / speed_of_light value: 0.0 - units: gigahertz ** 2 * megaelectron_volt * radian ** 2 * second ** 2 / meter ** 2 / speed_of_light value: -3.7117185708535e-13 type: polynomial z:pz
Cosine¶
It is often convenient to impart a cosine like energy spread to beam's longitudinal momentum. This can be accomplished using the cosine transform:
$p \rightarrow p + A\cos(\omega v + \phi)$.
import numpy as np
gen = Generator("data/gaussian.in.yaml", verbose=0)
w = 2 * PHYSICAL_CONSTANTS.pi * 1.3 * unit("GHz")
k = w / c
tcos = {
"type": "cosine z:pz",
"amplitude": {"value": 1000, "units": "MeV/c"},
"phase": {"value": -2.5, "units": "deg"},
"omega": {"value": k.magnitude, "units": str(k.units)},
}
gen["transforms"] = {"tc": tcos}
beam = gen.beam()
print("YAML input:\n", yaml.dump(gen.input["transforms"]))
plot_dist2d(beam, "z", units["z"], "pz", units["pz"]);
YAML input: tc: amplitude: units: megaelectron_volt / speed_of_light value: 1000.0 omega: units: gigahertz * radian * second / meter value: 2.724598528537186e-08 phase: units: degree value: -2.5 type: cosine z:pz
Setting Twiss parameters¶
Often for beams at energy the user may wish to set beam Twiss parameters $\beta$, $\alpha$, and $\epsilon$ for a desired 2D phase space.
filename = "data/gaussian.in.yaml"
gen = Generator(filename, verbose=0)
boost_pz = {
"type": "translate pz",
"delta": {
"value": 1,
"units": "GeV/c",
},
}
gen["transforms"] = {"boost": boost_pz}
beam1 = gen.beam()
twiss_x = {
"type": "set_twiss x",
"beta": {
"value": 12.5,
"units": "m",
},
"alpha": {"value": -1, "units": ""},
"emittance": {"value": 2, "units": "nm"},
}
gen["transforms"] = {"boost": boost_pz, "twiss": twiss_x, "order": ["boost", "twiss"]}
beam2 = gen.beam()
print("YAML input:\n", yaml.dump(gen.input["transforms"]))
print("Initial Horizontal Twiss params:")
print(
f'beta: {beam1.Beta("x"):G~P}, alpha: {beam1.Alpha("x"):G~P}, eps: {beam1.emitt("x","geometric").to("nm"):0.3f~P}'
)
fig, ax = plt.subplots(1, 2, constrained_layout=True)
plot_dist2d(beam1, "x", units["x"], "thetax", units["thetax"], ax=ax[0])
ax[0].set_title("Before")
plot_dist2d(beam2, "x", units["x"], "thetax", units["thetax"], ax=ax[1])
ax[1].set_title("After")
print("\nFinal Horizontal Twiss params:")
print(
f'beta: {beam2.Beta("x").to("m"):G~P}, alpha: {beam2.Alpha("x"):G~P}, eps: {beam2.emitt("x","geometric").to("nm"):0.3f~P}'
)
YAML input: boost: delta: units: gigaelectron_volt / speed_of_light value: 1.0 type: translate pz order: - boost - twiss twiss: alpha: units: dimensionless value: -1.0 beta: units: meter value: 12.5 emittance: units: nanometer value: 2.0 type: set_twiss x Initial Horizontal Twiss params: beta: 999999 mm, alpha: 1.94566E-05, eps: 1.000 nm Final Horizontal Twiss params: beta: 12.5 m, alpha: -1, eps: 2.000 nm
Example: offsetting beam on cathode to avoid a dead zone¶
gen = Generator("data/beer.can.in.yaml", verbose=0)
R = 4.0 # mm
Theta = 45 # deg
set_x = {"type": "set_avg x", "avg_x": {"value": R, "units": "mm"}}
rotxy = {"type": "rotate2d x:y", "angle": {"value": Theta, "units": "deg"}}
gen["transforms"] = {"r": set_x, "theta": rotxy}
P = gen.run()
P.plot("x", "y", figsize=(5, 5))
(
1e3 * P["mean_x"],
R * np.cos(np.pi / 180 * Theta),
1e3 * np.sqrt(P["mean_x"] ** 2 + P["mean_y"] ** 2),
np.arctan2(P["mean_y"], P["mean_x"]) * 180 / np.pi,
)
(np.float64(2.8284271247461894), np.float64(2.8284271247461903), np.float64(3.9999999999999982), np.float64(45.0))