A nonlinear focusing channel based on the IOTA nonlinear lens
A constant focusing channel with nonlinear focusing, using a string of thin IOTA nonlinear lens elements alternating with constant focusing elements.
We use a 2.5 MeV proton beam, corresponding to the nominal IOTA proton energy.
The two functions H (Hamiltonian) and I (the second invariant) should remain unchanged for all particles.
In this test, the initial and final values of \(\mu_H\), \(\sigma_H\), \(\mu_I\), \(\sigma_I\) must agree with nominal values.
Run
This example can be run as a Python script (python3 run_iotalens.py) or with an app with an input file (impactx input_iotalens.in).
Each can also be prefixed with an MPI executor, such as mpiexec -n 4 ... or srun -n 4 ..., depending on the system.
#!/usr/bin/env python3
#
# Copyright 2022-2023 ImpactX contributors
# Authors: Ryan Sandberg, Axel Huebl, Chad Mitchell
# License: BSD-3-Clause-LBNL
#
# -*- coding: utf-8 -*-
import amrex
from impactx import ImpactX, 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.5 MeV proton beam
energy_MeV = 2.5 # reference energy
bunch_charge_C = 1.0e-9 # used with space charge
npart = 10000 # number of macro particles
# reference particle
ref = sim.particle_container().ref_particle()
ref.set_charge_qe(1.0).set_mass_MeV(938.27208816).set_energy_MeV(energy_MeV)
# particle bunch
distr = distribution.Waterbag(
sigmaX=2.0e-3,
sigmaY=2.0e-3,
sigmaT=1.0e-3,
sigmaPx=3.0e-4,
sigmaPy=3.0e-4,
sigmaPt=0.0,
)
sim.add_particles(bunch_charge_C, distr, npart)
constEnd = elements.ConstF(ds=0.0025, kx=1.0, ky=1.0, kt=1.0e-12)
nllens = elements.NonlinearLens(knll=2.0e-7, cnll=0.01)
const = elements.ConstF(ds=0.005, kx=1.0, ky=1.0, kt=1.0e-12)
# design the accelerator lattice
num_lenses = 10
nllens_lattice = [constEnd] + [nllens, const] * (num_lenses - 1) + [nllens, constEnd]
# add elements to the lattice segment
sim.lattice.extend(nllens_lattice)
# run simulation
sim.evolve()
# clean shutdown
del sim
amrex.finalize()
###############################################################################
# Particle Beam(s)
###############################################################################
beam.npart = 10000
beam.units = static
beam.energy = 2.5
beam.charge = 1.0e-9
beam.particle = proton
beam.distribution = waterbag
beam.sigmaX = 2.0e-3
beam.sigmaY = 2.0e-3
beam.sigmaT = 1.0e-3
beam.sigmaPx = 3.0e-4
beam.sigmaPy = 3.0e-4
beam.sigmaPt = 0.0
beam.muxpx = 0.0
beam.muypy = 0.0
beam.mutpt = 0.0
###############################################################################
# Beamline: lattice elements and segments
###############################################################################
lattice.elements = const_end nllens const nllens const nllens const nllens const
nllens const nllens const nllens const nllens const nllens const nllens
const_end
nllens.type = nonlinear_lens
nllens.knll = 2.0e-7
nllens.cnll = 0.01
const_end.type = constf
const_end.ds = 0.0025
const_end.kx = 1.0
const_end.ky = 1.0
const_end.kt = 1.0e-12
const.type = constf
const.ds = 0.005
const.kx = 1.0
const.ky = 1.0
const.kt = 1.0e-12
###############################################################################
# Algorithms
###############################################################################
algo.particle_shape = 2
algo.space_charge = false
###############################################################################
# Diagnostics
###############################################################################
diag.alpha = 0.0
diag.beta = 1.0
diag.tn = 0.4
diag.cn = 0.01
Analyze
We run the following script to analyze correctness:
Script analysis_iotalens.py
#!/usr/bin/env python3
#
# Copyright 2022-2023 ImpactX contributors
# Authors: Axel Huebl, Chad Mitchell
# License: BSD-3-Clause-LBNL
#
import glob
import numpy as np
import pandas as pd
from scipy.stats import moment
def get_moments(beam):
"""Calculate mean and std dev of functions defining the IOTA invariants
Returns
-------
meanH, sigH, meanI, sigI
"""
meanH = np.mean(beam["H"])
sigH = moment(beam["H"], moment=2) ** 0.5
meanI = np.mean(beam["I"])
sigI = moment(beam["I"], moment=2) ** 0.5
return (meanH, sigH, meanI, sigI)
def read_all_files(file_pattern):
"""Read in all CSV files from each MPI rank (and potentially OpenMP
thread). Concatenate into one Pandas dataframe.
Returns
-------
pandas.DataFrame
"""
return pd.concat(
(
pd.read_csv(filename, delimiter=r"\s+")
for filename in glob.glob(file_pattern)
),
axis=0,
ignore_index=True,
).set_index("id")
# initial/final beam on rank zero
initial = read_all_files("diags/nonlinear_lens_invariants_000000.*")
final = read_all_files("diags/nonlinear_lens_invariants_final.*")
# compare number of particles
num_particles = 10000
assert num_particles == len(initial)
assert num_particles == len(final)
print("Initial Beam:")
meanH, sigH, meanI, sigI = get_moments(initial)
print(f" meanH={meanH:e} sigH={sigH:e} meanI={meanI:e} sigI={sigI:e}")
atol = 0.0 # a big number
rtol = 1.5 * num_particles**-0.5 # from random sampling of a smooth distribution
print(f" rtol={rtol} (ignored: atol~={atol})")
assert np.allclose(
[meanH, sigH, meanI, sigI],
[4.122650e-02, 4.235181e-02, 7.356057e-02, 8.793753e-02],
rtol=rtol,
atol=atol,
)
print("")
print("Final Beam:")
meanH, sigH, meanI, sigI = get_moments(final)
print(f" meanH={meanH:e} sigH={sigH:e} meanI={meanI:e} sigI={sigI:e}")
atol = 0.0 # a big number
rtol = 1.5 * num_particles**-0.5 # from random sampling of a smooth distribution
print(f" rtol={rtol} (ignored: atol~={atol})")
assert np.allclose(
[meanH, sigH, meanI, sigI],
[4.122704e-02, 4.230576e-02, 7.348275e-02, 8.783157e-02],
rtol=rtol,
atol=atol,
)
# join tables on particle ID, so we can compare the same particle initial->final
beam_joined = final.join(initial, lsuffix="_final", rsuffix="_initial")
# add new columns: dH and dI
beam_joined["dH"] = (beam_joined["H_initial"] - beam_joined["H_final"]).abs()
beam_joined["dI"] = (beam_joined["I_initial"] - beam_joined["I_final"]).abs()
# print(beam_joined)
# particle-wise comparison of H & I initial to final
atol = 2.0e-3
rtol = 0.0 # large number
print()
print(f" atol={atol} (ignored: rtol~={rtol})")
print(f" dH_max={beam_joined['dH'].max()}")
assert np.allclose(beam_joined["dH"], 0.0, rtol=rtol, atol=atol)
atol = 3.0e-3
print(f" atol={atol} (ignored: rtol~={rtol})")
print(f" dI_max={beam_joined['dI'].max()}")
assert np.allclose(beam_joined["dI"], 0.0, rtol=rtol, atol=atol)