Skip to content

Commit

Permalink
Example
Browse files Browse the repository at this point in the history
  • Loading branch information
ax3l committed Feb 11, 2025
1 parent ffd408d commit 1f763ad
Show file tree
Hide file tree
Showing 6 changed files with 261 additions and 1 deletion.
2 changes: 1 addition & 1 deletion docs/source/usage/examples.rst
Original file line number Diff line number Diff line change
Expand Up @@ -83,7 +83,7 @@ Channels & Rings
:maxdepth: 1

examples/fodo_channel/README.rst
examples/fodo_channel_restart/README.rst
examples/solenoid_restart/README.rst


Lattice Design & Optimization
Expand Down
19 changes: 19 additions & 0 deletions examples/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -624,6 +624,25 @@ add_impactx_test(solenoid.MADX.py
)


# Ideal, Hard-Edge Solenoid (Restart + 270 degrees) ###########################
#
# w/o space charge
add_impactx_test(solenoid.restart
examples/solenoid_restart/input_solenoid.in
OFF # ImpactX MPI-parallel
examples/solenoid_restart/analysis_solenoid.py
OFF # no plot script yet
)
add_impactx_test(solenoid.restart.py
examples/solenoid_restart/run_solenoid.py
OFF # ImpactX MPI-parallel
examples/solenoid_restart/analysis_solenoid.py
OFF # no plot script yet
)
set_property(TEST solenoid.restart.run APPEND PROPERTY DEPENDS "solenoid.run")
set_property(TEST solenoid.restart.py.run APPEND PROPERTY DEPENDS "solenoid.py.run")


# Pole-face rotation ##########################################################
#
# w/o space charge
Expand Down
50 changes: 50 additions & 0 deletions examples/solenoid_restart/README.rst
Original file line number Diff line number Diff line change
@@ -0,0 +1,50 @@
.. _examples-solenoid-restart:

Solenoid channel (Restart)
==========================

This is the same example as the :ref:`proton beam undergoing 90 deg X-Y rotation in an ideal solenoid channel <examples-solenoid>`, but it restarts the resulting beam and rotates it another 3 channel periods to the initial X-Y conditions.

The solenoid magnetic field corresponds to B = 2 T.

The second moments of the particle distribution after the solenoid channel are rotated by 90 (start) + 270 = 360 degrees: the final horizontal moments should coincide with the
initial vertical moments, and vice-versa, to within the level expected due to noise due to statistical sampling.

In this test, the initial and final values of :math:`\sigma_x`, :math:`\sigma_y`, :math:`\sigma_t`, :math:`\epsilon_x`, :math:`\epsilon_y`, and :math:`\epsilon_t` must agree with nominal values.


Run
---

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

* **Python** script: ``python3 run_solenoid.py`` or
* ImpactX **executable** using an input file: ``impactx input_solenoid.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_solenoid.py
:language: python3
:caption: You can copy this file from ``examples/solenoid_restart/run_solenoid.py``.

.. tab-item:: Executable: Input File

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


Analyze
-------

We run the following script to analyze correctness:

.. dropdown:: Script ``analysis_solenoid.py``

.. literalinclude:: analysis_solenoid.py
:language: python3
:caption: You can copy this file from ``examples/solenoid_restart/analysis_solenoid.py``.
96 changes: 96 additions & 0 deletions examples/solenoid_restart/analysis_solenoid.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,96 @@
#!/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)
first_step = list(series.iterations)[0]
last_step = list(series.iterations)[-1]
initial = series.iterations[first_step].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)

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 = 1.3 * 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],
[
2.205510139392e-3,
1.559531175539e-3,
6.404930308742e-3,
2.0e-6,
1.0e-6,
1.0e-6,
],
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.3 * num_particles**-0.5 # from random sampling of a smooth distribution
print(f" rtol={rtol} (ignored: atol~={atol})")

assert np.allclose(
[sigx, sigy, emittance_x, emittance_y, emittance_t],
[
1.559531175539e-3,
2.205510139392e-3,
1.0e-6,
2.0e-6,
1.0e-6,
],
rtol=rtol,
atol=atol,
)
41 changes: 41 additions & 0 deletions examples/solenoid_restart/input_solenoid.in
Original file line number Diff line number Diff line change
@@ -0,0 +1,41 @@
###############################################################################
# Reference Particle
###############################################################################
beam.kin_energy = 250.0
beam.particle = proton

# ignored
beam.charge = 1.0e-9
beam.units = static
beam.distribution = empty


###############################################################################
# Beamline: lattice elements and segments
###############################################################################
lattice.elements = source monitor sol1 sol1 sol1 monitor
lattice.nslice = 1

source.type = source
source.distribution = openPMD
source.openpmd_path = "../solenoid/diags/openPMD/monitor.h5"

monitor.type = beam_monitor
monitor.backend = h5

sol1.type = solenoid
sol1.ds = 3.820395
sol1.ks = 0.8223219329893234


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


###############################################################################
# Diagnostics
###############################################################################
diag.slice_step_diagnostics = true
54 changes: 54 additions & 0 deletions examples/solenoid_restart/run_solenoid.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,54 @@
#!/usr/bin/env python3
#
# Copyright 2022-2023 ImpactX contributors
# Authors: Marco Garten, Axel Huebl, Chad Mitchell
# License: BSD-3-Clause-LBNL
#
# -*- coding: utf-8 -*-

from impactx import ImpactX, elements, push

sim = ImpactX()

# set numerical parameters and IO control
sim.particle_shape = 2 # B-spline order
sim.space_charge = False
# sim.diagnostics = False # benchmarking
sim.slice_step_diagnostics = True

# domain decomposition & space charge mesh
sim.init_grids()

# load a 250 MeV proton beam with an initial
# horizontal rms emittance of 1 um and an
# initial vertical rms emittance of 2 um
kin_energy_MeV = 250.0 # reference energy

# reference particle
pc = sim.particle_container()
ref = pc.ref_particle()
ref.set_charge_qe(1.0).set_mass_MeV(938.27208816).set_kin_energy_MeV(kin_energy_MeV)

# load particle bunch from file
push(pc, elements.Source("openPMD", "../solenoid.py/diags/openPMD/monitor.h5"))

# add beam diagnostics
monitor = elements.BeamMonitor("monitor", backend="h5")

# design the accelerator lattice
sol1 = elements.Sol(name="sol1", ds=3.820395, ks=0.8223219329893234)
sim.lattice.extend(
[
monitor,
sol1,
sol1,
sol1,
monitor,
]
)

# run simulation
sim.track_particles()

# clean shutdown
sim.finalize()

0 comments on commit 1f763ad

Please sign in to comment.