Skip to content

Commit 0cc637c

Browse files
ax3lcemitch99
ax3l
and
cemitch99
authored
Transverse Alignment Support (#490)
* Transversal Alignment: Mixin Class Create a mixin class that can be used for transversal alignment offsets and tilts. Co-authored-by: Chad Mitchell <[email protected]> * Update Elements for `Alignment` Support Co-authored-by: Chad Mitchell <[email protected]> * Inputs: `ParmParse` w/ Alignment * Inputs: Python w/ Alignment * Docs: Alignment Input * Reword Rotation Error Docstring Now: `rotation error in the transverse plane` * Add misalignment example. * Alignment Example: Docs & CMake * Fix Python Example: Protons --------- Co-authored-by: Chad Mitchell <[email protected]>
1 parent 8d64edb commit 0cc637c

35 files changed

+1657
-396
lines changed

‎docs/source/index.rst

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -65,9 +65,9 @@ Usage
6565
:hidden:
6666

6767
usage/how_to_run
68+
usage/examples
6869
usage/python
6970
usage/parameters
70-
usage/examples
7171
usage/workflows
7272

7373
Data Analysis

‎docs/source/usage/examples.rst

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ This section allows you to **download input files** that correspond to different
1313
examples/cfchannel/README.rst
1414
examples/expanding_beam/README.rst
1515
examples/kurth/README.rst
16+
examples/alignment/README.rst
1617
examples/rfcavity/README.rst
1718
examples/fodo_rf/README.rst
1819
examples/fodo_chromatic/README.rst

‎docs/source/usage/parameters.rst

Lines changed: 66 additions & 67 deletions
Large diffs are not rendered by default.

‎docs/source/usage/python.rst

Lines changed: 93 additions & 28 deletions
Large diffs are not rendered by default.

‎examples/CMakeLists.txt

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -780,3 +780,21 @@ add_impactx_test(apochromat.py
780780
examples/apochromatic/analysis_apochromatic.py
781781
OFF # no plot script yet
782782
)
783+
784+
# Misalignment Quad ###########################################################
785+
#
786+
# w/o space charge
787+
add_impactx_test(alignment
788+
examples/alignment/input_alignment.in
789+
ON # ImpactX MPI-parallel
790+
OFF # ImpactX Python interface
791+
examples/alignment/analysis_alignment.py
792+
OFF # no plot script yet
793+
)
794+
add_impactx_test(alignment.py
795+
examples/alignment/run_alignment.py
796+
OFF # ImpactX MPI-parallel
797+
ON # ImpactX Python interface
798+
examples/alignment/analysis_alignment.py
799+
OFF # no plot script yet
800+
)

‎examples/alignment/README.rst

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
1+
.. _examples-alignment:
2+
3+
Quadrupole with Alignment Errors
4+
================================
5+
6+
A 2 GeV proton beam propagates through a single quadrupole with 3 mm horizontal misalignment and 30 degree rotation error.
7+
8+
The first and second moments of the particle distribution before and after the quadrupole should coincide with
9+
analytical predictions, to within the level expected due to noise due to statistical sampling.
10+
11+
In this test, the initial and final values of :math:`\mu_x`, :math:`\mu_y`, :math:`\sigma_x`, :math:`\sigma_y`,
12+
:math:`\sigma_t`, :math:`\epsilon_x`, :math:`\epsilon_y`, and :math:`\epsilon_t` must agree with nominal values.
13+
14+
15+
Run
16+
---
17+
18+
This example can be run **either** as:
19+
20+
* **Python** script: ``python3 run_alignment.py`` or
21+
* ImpactX **executable** using an input file: ``impactx input_alignment.in``
22+
23+
For `MPI-parallel <https://www.mpi-forum.org>`__ runs, prefix these lines with ``mpiexec -n 4 ...`` or ``srun -n 4 ...``, depending on the system.
24+
25+
.. tab-set::
26+
27+
.. tab-item:: Python: Script
28+
29+
.. literalinclude:: run_alignment.py
30+
:language: python3
31+
:caption: You can copy this file from ``examples/alignment/run_alignment.py``.
32+
33+
.. tab-item:: Executable: Input File
34+
35+
.. literalinclude:: input_alignment.in
36+
:language: ini
37+
:caption: You can copy this file from ``examples/alignment/input_alignment.in``.
38+
39+
40+
Analyze
41+
-------
42+
43+
We run the following script to analyze correctness:
44+
45+
.. dropdown:: Script ``analysis_alignment.py``
46+
47+
.. literalinclude:: analysis_alignment.py
48+
:language: python3
49+
:caption: You can copy this file from ``examples/alignment/analysis_alignment.py``.

‎examples/alignment/analysis_alignment.py

Lines changed: 137 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,137 @@
1+
#!/usr/bin/env python3
2+
#
3+
# Copyright 2022-2023 ImpactX contributors
4+
# Authors: Axel Huebl, Chad Mitchell
5+
# License: BSD-3-Clause-LBNL
6+
#
7+
8+
9+
import numpy as np
10+
import openpmd_api as io
11+
from scipy.stats import moment, tmean
12+
13+
14+
def get_moments(beam):
15+
"""Calculate standard deviations of beam position & momenta
16+
and emittance values
17+
18+
Returns
19+
-------
20+
meanx, meany, sigx, sigy, sigt, emittance_x, emittance_y, emittance_t
21+
"""
22+
meanx = tmean(beam["position_x"])
23+
meany = tmean(beam["position_y"])
24+
sigx = moment(beam["position_x"], moment=2) ** 0.5 # variance -> std dev.
25+
sigpx = moment(beam["momentum_x"], moment=2) ** 0.5
26+
sigy = moment(beam["position_y"], moment=2) ** 0.5
27+
sigpy = moment(beam["momentum_y"], moment=2) ** 0.5
28+
sigt = moment(beam["position_t"], moment=2) ** 0.5
29+
sigpt = moment(beam["momentum_t"], moment=2) ** 0.5
30+
31+
epstrms = beam.cov(ddof=0)
32+
emittance_x = (
33+
sigx**2 * sigpx**2 - epstrms["position_x"]["momentum_x"] ** 2
34+
) ** 0.5
35+
emittance_y = (
36+
sigy**2 * sigpy**2 - epstrms["position_y"]["momentum_y"] ** 2
37+
) ** 0.5
38+
emittance_t = (
39+
sigt**2 * sigpt**2 - epstrms["position_t"]["momentum_t"] ** 2
40+
) ** 0.5
41+
42+
return (meanx, meany, sigx, sigy, sigt, emittance_x, emittance_y, emittance_t)
43+
44+
45+
# initial/final beam
46+
series = io.Series("diags/openPMD/monitor.h5", io.Access.read_only)
47+
last_step = list(series.iterations)[-1]
48+
initial = series.iterations[1].particles["beam"].to_df()
49+
final = series.iterations[last_step].particles["beam"].to_df()
50+
51+
# compare number of particles
52+
num_particles = 100000
53+
assert num_particles == len(initial)
54+
assert num_particles == len(final)
55+
56+
print("Initial Beam:")
57+
meanx, meany, sigx, sigy, sigt, emittance_x, emittance_y, emittance_t = get_moments(
58+
initial
59+
)
60+
print(f" meanx={meanx:e} meany={meany:e}")
61+
print(f" sigx={sigx:e} sigy={sigy:e} sigt={sigt:e}")
62+
print(
63+
f" emittance_x={emittance_x:e} emittance_y={emittance_y:e} emittance_t={emittance_t:e}"
64+
)
65+
66+
atol = 3.0 * num_particles ** (-0.5) * sigx
67+
print(f" atol~={atol}")
68+
69+
assert np.allclose(
70+
[meanx, meany],
71+
[
72+
0.0,
73+
0.0,
74+
],
75+
atol=atol,
76+
)
77+
78+
atol = 0.0 # ignored
79+
rtol = 1.8 * num_particles**-0.5 # from random sampling of a smooth distribution
80+
print(f" rtol={rtol} (ignored: atol~={atol})")
81+
82+
assert np.allclose(
83+
[sigx, sigy, sigt, emittance_x, emittance_y, emittance_t],
84+
[
85+
1.160982600086e-3,
86+
1.160982600086e-3,
87+
1.0e-3,
88+
6.73940299e-7,
89+
6.73940299e-7,
90+
2.0e-6,
91+
],
92+
rtol=rtol,
93+
atol=atol,
94+
)
95+
96+
97+
print("")
98+
print("Final Beam:")
99+
meanx, meany, sigx, sigy, sigt, emittance_x, emittance_y, emittance_t = get_moments(
100+
final
101+
)
102+
print(f" meanx={meanx:e} meany={meany:e}")
103+
print(f" sigx={sigx:e} sigy={sigy:e} sigt={sigt:e}")
104+
print(
105+
f" emittance_x={emittance_x:e} emittance_y={emittance_y:e} emittance_t={emittance_t:e}"
106+
)
107+
108+
109+
atol = 3.0 * num_particles ** (-0.5) * sigx
110+
print(f" atol~={atol}")
111+
112+
assert np.allclose(
113+
[meanx, meany],
114+
[
115+
1.79719761842e-4,
116+
3.24815908981e-4,
117+
],
118+
atol=atol,
119+
)
120+
121+
atol = 0.0 # ignored
122+
rtol = 3.0 * num_particles**-0.5 # from random sampling of a smooth distribution
123+
print(f" rtol={rtol} (ignored: atol~={atol})")
124+
125+
assert np.allclose(
126+
[sigx, sigy, sigt, emittance_x, emittance_y, emittance_t],
127+
[
128+
1.2372883901369e-3,
129+
1.3772750218080e-3,
130+
1.027364e-03,
131+
7.39388142e-7,
132+
7.39388142e-7,
133+
2.0e-6,
134+
],
135+
rtol=rtol,
136+
atol=atol,
137+
)

‎examples/alignment/input_alignment.in

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
1+
###############################################################################
2+
# Particle Beam(s)
3+
###############################################################################
4+
beam.npart = 100000
5+
beam.units = static
6+
beam.kin_energy = 2.0e3
7+
beam.charge = 1.0e-9
8+
beam.particle = proton
9+
beam.distribution = waterbag
10+
beam.sigmaX = 1.16098260008648811e-3
11+
beam.sigmaY = 1.16098260008648811e-3
12+
beam.sigmaT = 1.0e-3
13+
beam.sigmaPx = 0.580491300043e-3
14+
beam.sigmaPy = 0.580491300043e-3
15+
beam.sigmaPt = 2.0e-3
16+
beam.muxpx = 0.0
17+
beam.muypy = 0.0
18+
beam.mutpt = 0.0
19+
20+
21+
###############################################################################
22+
# Beamline: lattice elements and segments
23+
###############################################################################
24+
lattice.elements = monitor quad_err monitor
25+
lattice.nslice = 1
26+
27+
monitor.type = beam_monitor
28+
monitor.backend = h5
29+
30+
quad_err.type = quad
31+
quad_err.ds = 1.0
32+
quad_err.k = 0.25
33+
quad_err.dx = 0.003
34+
quad_err.rotation = 30.0
35+
36+
37+
###############################################################################
38+
# Algorithms
39+
###############################################################################
40+
algo.particle_shape = 2
41+
algo.space_charge = false
42+
43+
44+
###############################################################################
45+
# Diagnostics
46+
###############################################################################
47+
diag.slice_step_diagnostics = true

‎examples/alignment/run_alignment.py

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,63 @@
1+
#!/usr/bin/env python3
2+
#
3+
# Copyright 2022-2023 ImpactX contributors
4+
# Authors: Axel Huebl, Chad Mitchell
5+
# License: BSD-3-Clause-LBNL
6+
#
7+
# -*- coding: utf-8 -*-
8+
9+
import amrex.space3d as amr
10+
from impactx import ImpactX, RefPart, distribution, elements
11+
12+
sim = ImpactX()
13+
14+
# set numerical parameters and IO control
15+
sim.particle_shape = 2 # B-spline order
16+
sim.space_charge = False
17+
# sim.diagnostics = False # benchmarking
18+
sim.slice_step_diagnostics = True
19+
20+
# domain decomposition & space charge mesh
21+
sim.init_grids()
22+
23+
# load a 2 GeV proton beam
24+
kin_energy_MeV = 2.0e3 # reference energy
25+
bunch_charge_C = 1.0e-9 # used with space charge
26+
npart = 100000 # number of macro particles
27+
28+
# reference particle
29+
ref = sim.particle_container().ref_particle()
30+
ref.set_charge_qe(1.0).set_mass_MeV(938.27208816).set_kin_energy_MeV(kin_energy_MeV)
31+
32+
# particle bunch
33+
distr = distribution.Waterbag(
34+
sigmaX=1.16098260008648811e-3,
35+
sigmaY=1.16098260008648811e-3,
36+
sigmaT=1.0e-3,
37+
sigmaPx=0.580491300043e-3,
38+
sigmaPy=0.580491300043e-3,
39+
sigmaPt=2.0e-3,
40+
muxpx=0.0,
41+
muypy=0.0,
42+
mutpt=0.0,
43+
)
44+
sim.add_particles(bunch_charge_C, distr, npart)
45+
46+
# add beam diagnostics
47+
monitor = elements.BeamMonitor("monitor", backend="h5")
48+
49+
# design the accelerator lattice)
50+
ns = 1 # number of slices per ds in the element
51+
lattice = [
52+
monitor,
53+
elements.Quad(ds=1.0, k=0.25, dx=0.003, dy=0.0, rotation=30.0, nslice=ns),
54+
monitor,
55+
]
56+
sim.lattice.extend(lattice)
57+
58+
# run simulation
59+
sim.evolve()
60+
61+
# clean shutdown
62+
del sim
63+
amr.finalize()

0 commit comments

Comments
 (0)