diff --git a/docs/source/usage/examples.rst b/docs/source/usage/examples.rst index 8e8ef3dac..62b91f5e4 100644 --- a/docs/source/usage/examples.rst +++ b/docs/source/usage/examples.rst @@ -31,6 +31,7 @@ This section allows you to **download input files** that correspond to different examples/thin_dipole/README.rst examples/aperture/README.rst examples/pytorch_surrogate_model/README.rst + examples/apochromatic/README.rst Unit tests diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt index d54adad03..1b5bebf6d 100644 --- a/examples/CMakeLists.txt +++ b/examples/CMakeLists.txt @@ -711,3 +711,21 @@ add_impactx_test(aperture.py examples/aperture/analysis_aperture.py OFF # no plot script yet ) + +# Apochromat example ######################################################## +# +# w/o space charge +add_impactx_test(apochromat + examples/apochromatic/input_apochromatic.in + ON # ImpactX MPI-parallel + OFF # ImpactX Python interface + examples/apochromatic/analysis_apochromatic.py + OFF # no plot script yet +) +add_impactx_test(apochromat.py + examples/apochromatic/run_apochromatic.py + OFF # ImpactX MPI-parallel + ON # ImpactX Python interface + examples/apochromatic/analysis_apochromatic.py + OFF # no plot script yet +) diff --git a/examples/apochromatic/README.rst b/examples/apochromatic/README.rst new file mode 100644 index 000000000..b8798693d --- /dev/null +++ b/examples/apochromatic/README.rst @@ -0,0 +1,58 @@ +.. _examples-apochromat: + +Apochromatic Drift-Quadrupole Beamline +====================================== + +Electron beam matched to the 1st-order apochromatic drift-quadrupole beamline appearing in Fig. 4a of: +C. A. Lindstrom and E. Adli, "Design of general apochromatic drift-quadrupole beam lines," Phys. Rev. Accel. Beams 19, 071002 (2016). + +The matched Twiss parameters at entry are: + +* :math:`\beta_\mathrm{x} = 0.325` m +* :math:`\alpha_\mathrm{x} = 0` +* :math:`\beta_\mathrm{y} = 0.325` m +* :math:`\alpha_\mathrm{y} = 0` + +We use a 100 GeV electron beam with an initially 6D Gaussian distribution of normalized rms emittance 1 micron and relative energy spread of 1%. + +The second moments of the particle distribution after the focusing beamline should coincide with the second moments of the particle distribution before the beamline, to within the level expected due to noise due to statistical sampling. +The emittance growth due to chromatic effects remain below 1%. In the absence of chromatic correction, the projected emittance growth is near 10%. + +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_apochromatic.py`` or +* ImpactX **executable** using an input file: ``impactx input_apochromatic.in`` + +For `MPI-parallel `__ runs, prefix these lines with ``mpiexec -n 4 ...`` or ``srun -n 4 ...``, depending on the system. + +.. tab-set:: + + .. tab-item:: Python: Script + + .. literalinclude:: run_apochromatic.py + :language: python3 + :caption: You can copy this file from ``examples/apochromatic/run_apochromatic.py``. + + .. tab-item:: Executable: Input File + + .. literalinclude:: input_apochromatic.in + :language: ini + :caption: You can copy this file from ``examples/apochromatic/input_apochromatic.in``. + + +Analyze +------- + +We run the following script to analyze correctness: + +.. dropdown:: Script ``analysis_apochromatic.py`` + + .. literalinclude:: analysis_apochromatic.py + :language: python3 + :caption: You can copy this file from ``examples/apochromatic/analysis_apochromatic.py``. diff --git a/examples/apochromatic/analysis_apochromatic.py b/examples/apochromatic/analysis_apochromatic.py new file mode 100755 index 000000000..78f98e707 --- /dev/null +++ b/examples/apochromatic/analysis_apochromatic.py @@ -0,0 +1,126 @@ +#!/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["divergence_x"], moment=2) ** 0.5 + sigy = moment(beam["position_y"], moment=2) ** 0.5 + sigpy = moment(beam["divergence_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 = 100000 +assert num_particles == len(initial) +assert num_particles == len(final) + +scale = ( + (1.0 - initial.momentum_t) ** 2 + + (initial.momentum_x) ** 2 + + (initial.momentum_y) ** 2 +) +xp = initial.momentum_x / np.sqrt(scale) +initial["divergence_x"] = xp +yp = initial.momentum_y / np.sqrt(scale) +initial["divergence_y"] = yp + +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 = 3.0 * 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], + [ + 1.288697604e-6, + 1.288697604e-6, + 1.0e-6, + 5.10997388810014764e-12, + 5.10997388810014764e-12, + 1.0e-8, + ], + rtol=rtol, + atol=atol, +) + + +scale = ( + (1.0 - final.momentum_t) ** 2 + (final.momentum_x) ** 2 + (final.momentum_y) ** 2 +) +xp = final.momentum_x / np.sqrt(scale) +final["divergence_x"] = xp +yp = final.momentum_y / np.sqrt(scale) +final["divergence_y"] = yp + +print("") +print("Final Beam:") +sigx, sigy, sigt, emittance_xf, emittance_yf, emittance_tf = get_moments(final) +demittance_x = 100 * (emittance_xf - emittance_x) / emittance_x +demittance_y = 100 * (emittance_yf - emittance_y) / emittance_y +demittance_t = 100 * (emittance_tf - emittance_t) / emittance_t + +print(f" sigx={sigx:e} sigy={sigy:e} sigt={sigt:e}") +print( + f" emittance change x (%)={demittance_x:e} emittance change y (%)={demittance_y:e} emittance change t (%)={demittance_t:e}" +) + +atol = 0.0 # ignored +rtol = 15.0 * num_particles**-0.5 # from random sampling of a smooth distribution +print(f" rtol={rtol} (ignored: atol~={atol})") + +assert np.allclose( + [sigx, sigy, sigt, demittance_x, demittance_y, emittance_t], + [ + 1.245e-6, + 1.245e-6, + 1.0e-6, + 0.94, + 0.94, + 1.0e-8, + ], + rtol=rtol, + atol=atol, +) diff --git a/examples/apochromatic/input_apochromatic.in b/examples/apochromatic/input_apochromatic.in new file mode 100644 index 000000000..1074f7b40 --- /dev/null +++ b/examples/apochromatic/input_apochromatic.in @@ -0,0 +1,77 @@ +############################################################################### +# Particle Beam(s) +############################################################################### +beam.npart = 100000 +beam.units = static +beam.kin_energy = 100.0e3 # 100 GeV nominal energy +beam.charge = 1.0e-9 +beam.particle = electron +beam.distribution = gaussian +beam.sigmaX = 1.288697604e-6 +beam.sigmaY = 1.288697604e-6 +beam.sigmaT = 1.0e-6 +beam.sigmaPx = 3.965223396e-6 +beam.sigmaPy = 3.965223396e-6 +beam.sigmaPt = 0.01 #1% energy spread +beam.muxpx = 0.0 +beam.muypy = 0.0 +beam.mutpt = 0.0 + +############################################################################### +# Beamline: lattice elements and segments +############################################################################### +lattice.elements = monitor dr1 q1 q2 q3 dr2 q4 q5 dr2 q6 q7 q8 dr1 monitor +lattice.nslice = 1 + +monitor.type = beam_monitor +monitor.backend = h5 + +dr1.type = drift_chromatic +dr1.ds = 1.0 + +dr2.type = drift_chromatic +dr2.ds = 10.0 + +q1.type = quad_chromatic +q1.ds = 1.2258333333 +q1.k = 0.5884 + +q2.type = quad_chromatic +q2.ds = 1.5677083333 +q2.k = -0.7525 + +q3.type = quad_chromatic +q3.ds = 1.205625 +q3.k = 0.5787 + +q4.type = quad_chromatic +q4.ds = 1.2502083333 +q4.k = -0.6001 + +q5.type = quad_chromatic +q5.ds = 1.2502083333 +q5.k = 0.6001 + +q6.type = quad_chromatic +q6.ds = 1.205625 +q6.k = -0.5787 + +q7.type = quad_chromatic +q7.ds = 1.5677083333 +q7.k = 0.7525 + +q8.type = quad_chromatic +q8.ds = 1.2258333333 +q8.k = -0.5884 + +############################################################################### +# Algorithms +############################################################################### +algo.particle_shape = 2 +algo.space_charge = false + + +############################################################################### +# Diagnostics +############################################################################### +diag.slice_step_diagnostics = true diff --git a/examples/apochromatic/run_apochromatic.py b/examples/apochromatic/run_apochromatic.py new file mode 100644 index 000000000..022300827 --- /dev/null +++ b/examples/apochromatic/run_apochromatic.py @@ -0,0 +1,77 @@ +#!/usr/bin/env python3 +# +# Copyright 2022-2023 ImpactX contributors +# Authors: Axel Huebl, Chad Mitchell +# License: BSD-3-Clause-LBNL +# +# -*- coding: utf-8 -*- + +import amrex.space3d as amr +from impactx import ImpactX, RefPart, distribution, elements + +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 2 GeV electron beam with an initial +# unnormalized rms emittance of 2 nm +kin_energy_MeV = 100.0e3 # reference energy +bunch_charge_C = 1.0e-9 # used with space charge +npart = 100000 # number of macro particles + +# reference particle +ref = sim.particle_container().ref_particle() +ref.set_charge_qe(-1.0).set_mass_MeV(0.510998950).set_kin_energy_MeV(kin_energy_MeV) + +# particle bunch +distr = distribution.Gaussian( + sigmaX=1.288697604e-6, + sigmaY=1.288697604e-6, + sigmaT=1.0e-6, + sigmaPx=3.965223396e-6, + sigmaPy=3.965223396e-6, + sigmaPt=0.01, # 1% energy spread + muxpx=0.0, + muypy=0.0, + mutpt=0.0, +) +sim.add_particles(bunch_charge_C, distr, npart) + +# add beam diagnostics +monitor = elements.BeamMonitor("monitor", backend="h5") + +# design the accelerator lattice) +ns = 25 # number of slices per ds in the element + +# Drift elements +dr1 = elements.ChrDrift(ds=1.0, nslice=ns) +dr2 = elements.ChrDrift(ds=10.0, nslice=ns) + +# Quad elements +q1 = elements.ChrQuad(ds=1.2258333333, k=0.5884, nslice=ns) +q2 = elements.ChrQuad(ds=1.5677083333, k=-0.7525, nslice=ns) +q3 = elements.ChrQuad(ds=1.205625, k=0.5787, nslice=ns) +q4 = elements.ChrQuad(ds=1.2502083333, k=-0.6001, nslice=ns) +q5 = elements.ChrQuad(ds=1.2502083333, k=0.6001, nslice=ns) +q6 = elements.ChrQuad(ds=1.205625, k=-0.5787, nslice=ns) +q7 = elements.ChrQuad(ds=1.5677083333, k=0.7525, nslice=ns) +q8 = elements.ChrQuad(ds=1.2258333333, k=-0.5884, nslice=ns) + +lattice_line = [monitor, dr1, q1, q2, q3, dr2, q4, q5, dr2, q6, q7, q8, dr1, monitor] + +# define the lattice +sim.lattice.extend(lattice_line) + +# run simulation +sim.evolve() + +# clean shutdown +del sim +amr.finalize()