Skip to content

Commit

Permalink
Add Example
Browse files Browse the repository at this point in the history
  • Loading branch information
ax3l committed Sep 25, 2024
1 parent 30b53b2 commit 5eaa4c8
Show file tree
Hide file tree
Showing 7 changed files with 469 additions and 0 deletions.
10 changes: 10 additions & 0 deletions docs/source/usage/examples.rst
Original file line number Diff line number Diff line change
Expand Up @@ -70,6 +70,16 @@ Beam Distributions

examples/distgen/README


Channels & Rings
----------------

.. toctree::
:maxdepth: 1

examples/fodo_channel/README.rst


Lattice Design & Optimization
-----------------------------

Expand Down
17 changes: 17 additions & 0 deletions examples/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -219,6 +219,23 @@ add_impactx_test(FODO.py.MPI
examples/fodo/plot_fodo.py
)

# FODO Channel ################################################################
#
add_impactx_test(FODO_channel
examples/fodo_channel/input_fodo.in
OFF # ImpactX MPI-parallel
examples/fodo_channel/analysis_fodo.py
examples/fodo_channel/plot_fodo.py
)
add_impactx_test(FODO_channel.py
examples/fodo_channel/run_fodo.py
OFF # ImpactX MPI-parallel
examples/fodo_channel/analysis_fodo.py
examples/fodo_channel/plot_fodo.py
)
label_impactx_test(FODO_channel slow)
label_impactx_test(FODO_channel.py slow)

# Chicane #####################################################################
#
add_impactx_test(chicane
Expand Down
74 changes: 74 additions & 0 deletions examples/fodo_channel/README.rst
Original file line number Diff line number Diff line change
@@ -0,0 +1,74 @@
.. _examples-fodo-channel:

FODO Channel
============

A 300m channel of 100 stable FODO cells (3m each) with a zero-current phase advance of 67.8 degrees.

The matched Twiss parameters at entry are:

* :math:`\beta_\mathrm{x} = 2.82161941` m
* :math:`\alpha_\mathrm{x} = -1.59050035`
* :math:`\beta_\mathrm{y} = 2.82161941` m
* :math:`\alpha_\mathrm{y} = 1.59050035`

We use a 2 GeV electron beam with initial unnormalized rms emittance of 2 nm.

The second moments of the particle distribution after the FODO cell should coincide with the second moments of the particle distribution before the FODO cell, to within the level expected due to noise due to statistical sampling.

In this test, the initial and final values of :math:`\lambda_x`, :math:`\lambda_y`, :math:`\lambda_t`, :math:`\epsilon_x`, :math:`\epsilon_y`, and :math:`\epsilon_t` must agree with nominal values.
This test also demonstrates the ``cycle_intervals`` capability of our beam monitor diagnostics, only creating output every 10th FODO cell


Run
---

This example can be run **either** as:

* **Python** script: ``python3 run_fodo.py`` or
* ImpactX **executable** using an input file: ``impactx input_fodo.in``

For `MPI-parallel <https://www.mpi-forum.org>`__ runs, prefix these lines with ``mpiexec -n 4 ...`` or ``srun -n 4 ...``, depending on the system.

.. tab-set::

.. tab-item:: Python: Script

.. literalinclude:: run_fodo.py
:language: python3
:caption: You can copy this file from ``examples/fodo/run_fodo.py``.

.. tab-item:: Executable: Input File

.. literalinclude:: input_fodo.in
:language: ini
:caption: You can copy this file from ``examples/fodo/input_fodo.in``.


Analyze
-------

We run the following script to analyze correctness:

.. dropdown:: Script ``analysis_fodo.py``

.. literalinclude:: analysis_fodo.py
:language: python3
:caption: You can copy this file from ``examples/fodo/analysis_fodo.py``.


Visualize
---------

You can run the following script to visualize the beam evolution over time:

.. dropdown:: Script ``plot_fodo.py``

.. literalinclude:: plot_fodo.py
:language: python3
:caption: You can copy this file from ``examples/fodo/plot_fodo.py``.

.. figure:: https://gist.githubusercontent.com/ax3l/8ae7dcb9e07c361e002fa56d6b16cb16/raw/cc952670bb946cd7a62282bc7aa3f03f3d5faa16/fodo_channel.png
:alt: preserved emittance in the FODO channel.

FODO transversal emittance evolution (preserved)
105 changes: 105 additions & 0 deletions examples/fodo_channel/analysis_fodo.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,105 @@
#!/usr/bin/env python3
#
# Copyright 2022-2023 ImpactX contributors
# Authors: Axel Huebl, Chad Mitchell
# License: BSD-3-Clause-LBNL
#


import numpy as np
import openpmd_api as io
from scipy.stats import moment


def get_moments(beam):
"""Calculate standard deviations of beam position & momenta
and emittance values
Returns
-------
sigx, sigy, sigt, emittance_x, emittance_y, emittance_t
"""
sigx = moment(beam["position_x"], moment=2) ** 0.5 # variance -> std dev.
sigpx = moment(beam["momentum_x"], moment=2) ** 0.5
sigy = moment(beam["position_y"], moment=2) ** 0.5
sigpy = moment(beam["momentum_y"], moment=2) ** 0.5
sigt = moment(beam["position_t"], moment=2) ** 0.5
sigpt = moment(beam["momentum_t"], moment=2) ** 0.5

epstrms = beam.cov(ddof=0)
emittance_x = (sigx**2 * sigpx**2 - epstrms["position_x"]["momentum_x"] ** 2) ** 0.5
emittance_y = (sigy**2 * sigpy**2 - epstrms["position_y"]["momentum_y"] ** 2) ** 0.5
emittance_t = (sigt**2 * sigpt**2 - epstrms["position_t"]["momentum_t"] ** 2) ** 0.5

return (sigx, sigy, sigt, emittance_x, emittance_y, emittance_t)


# initial/final beam
series = io.Series("diags/openPMD/monitor.h5", io.Access.read_only)
last_step = list(series.iterations)[-1]
initial = series.iterations[1].particles["beam"].to_df()
final = series.iterations[last_step].particles["beam"].to_df()

# compare number of particles
num_particles = 10000
assert num_particles == len(initial)
assert num_particles == len(final)

# compare beamline length: 300m
assert np.isclose(
300.0, series.iterations[last_step].particles["beam"].get_attribute("z_ref")
)
# compare beam monitor outputs: 10 (every 10th FODO element + 1)
assert len(series.iterations) == 11

print("Initial Beam:")
sigx, sigy, sigt, emittance_x, emittance_y, emittance_t = get_moments(initial)
print(f" sigx={sigx:e} sigy={sigy:e} sigt={sigt:e}")
print(
f" emittance_x={emittance_x:e} emittance_y={emittance_y:e} emittance_t={emittance_t:e}"
)

atol = 0.0 # ignored
rtol = 2.2 * num_particles**-0.5 # from random sampling of a smooth distribution
print(f" rtol={rtol} (ignored: atol~={atol})")

assert np.allclose(
[sigx, sigy, sigt, emittance_x, emittance_y, emittance_t],
[
7.5451170454175073e-005,
7.5441588239210947e-005,
9.9775878164077539e-004,
1.9959540393751392e-009,
2.0175015289132990e-009,
2.0013820193294972e-006,
],
rtol=rtol,
atol=atol,
)


print("")
print("Final Beam:")
sigx, sigy, sigt, emittance_x, emittance_y, emittance_t = get_moments(final)
print(f" sigx={sigx:e} sigy={sigy:e} sigt={sigt:e}")
print(
f" emittance_x={emittance_x:e} emittance_y={emittance_y:e} emittance_t={emittance_t:e}"
)

atol = 0.0 # ignored
rtol = 2.2 * num_particles**-0.5 # from random sampling of a smooth distribution
print(f" rtol={rtol} (ignored: atol~={atol})")

assert np.allclose(
[sigx, sigy, sigt, emittance_x, emittance_y, emittance_t],
[
7.4790118496224206e-005,
7.5357525169680140e-005,
9.9775879288128088e-004,
1.9959539836392703e-009,
2.0175014668882125e-009,
2.0013820380883801e-006,
],
rtol=rtol,
atol=atol,
)
60 changes: 60 additions & 0 deletions examples/fodo_channel/input_fodo.in
Original file line number Diff line number Diff line change
@@ -0,0 +1,60 @@
###############################################################################
# Particle Beam(s)
###############################################################################
beam.npart = 10000
beam.units = static
beam.kin_energy = 2.0e3
beam.charge = 1.0e-9
beam.particle = electron
beam.distribution = waterbag_from_twiss
beam.alphaX = -1.5905003499999992
beam.alphaY = 1.5905003499999992
beam.alphaT = 0.0
beam.betaX = 2.8216194100262637
beam.betaY = 2.8216194100262637
beam.betaT = 0.5
beam.emittX = 2e-09
beam.emittY = 2e-09
beam.emittT = 2e-06


###############################################################################
# Beamline: lattice elements and segments
###############################################################################
lattice.elements = monitor drift1 quad1 drift2 quad2 drift3
lattice.nslice = 5
lattice.periods = 101 # FODO channel of 101 periods

monitor.type = beam_monitor
monitor.cycle_intervals = 10
monitor.backend = h5

drift1.type = drift
drift1.ds = 0.25

quad1.type = quad
quad1.ds = 1.0
quad1.k = 1.0

drift2.type = drift
drift2.ds = 0.5

quad2.type = quad
quad2.ds = 1.0
quad2.k = -1.0

drift3.type = drift
drift3.ds = 0.25


###############################################################################
# Algorithms
###############################################################################
algo.particle_shape = 2
algo.space_charge = false


###############################################################################
# Diagnostics
###############################################################################
diag.slice_step_diagnostics = false
Loading

0 comments on commit 5eaa4c8

Please sign in to comment.