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 either as:
Python script:
python3 run_iotalens.pyorImpactX executable using an input file:
impactx input_iotalens.in
For MPI-parallel runs, prefix these lines with 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 -*-
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
kin_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_kin_energy_MeV(kin_energy_MeV)
# particle bunch
distr = distribution.Waterbag(
lambdaX=2.0e-3,
lambdaY=2.0e-3,
lambdaT=1.0e-3,
lambdaPx=3.0e-4,
lambdaPy=3.0e-4,
lambdaPt=0.0,
)
sim.add_particles(bunch_charge_C, distr, npart)
# add beam diagnostics
monitor = elements.BeamMonitor("monitor", backend="h5")
monitor.nonlinear_lens_invariants = True
monitor.alpha = 0.0
monitor.beta = 1.0
monitor.tn = 0.4
monitor.cn = 0.01
# design the accelerator lattice
constEnd = elements.ConstF(ds=0.05, kx=1.0, ky=1.0, kt=1.0e-12)
nllens = elements.NonlinearLens(knll=4.0e-6, cnll=0.01)
const = elements.ConstF(ds=0.1, kx=1.0, ky=1.0, kt=1.0e-12)
num_lenses = 18
nllens_lattice = (
[monitor, constEnd]
+ [nllens, const] * (num_lenses - 1)
+ [nllens, constEnd, monitor]
)
# add elements to the lattice segment
sim.lattice.extend(nllens_lattice)
# run simulation
sim.evolve()
# clean shutdown
sim.finalize()
###############################################################################
# Particle Beam(s)
###############################################################################
beam.npart = 10000
beam.units = static
beam.kin_energy = 2.5
beam.charge = 1.0e-9
beam.particle = proton
beam.distribution = waterbag
beam.lambdaX = 2.0e-3
beam.lambdaY = 2.0e-3
beam.lambdaT = 1.0e-3
beam.lambdaPx = 3.0e-4
beam.lambdaPy = 3.0e-4
beam.lambdaPt = 0.0
beam.muxpx = 0.0
beam.muypy = 0.0
beam.mutpt = 0.0
###############################################################################
# Beamline: lattice elements and segments
###############################################################################
lattice.elements = monitor const_end nllens foclens const_end monitor
foclens.type = line
foclens.elements = const nllens
foclens.repeat = 17
nllens.type = nonlinear_lens
nllens.knll = 4.0e-6
nllens.cnll = 0.01
const_end.type = constf
const_end.ds = 0.05
const_end.kx = 1.0
const_end.ky = 1.0
const_end.kt = 1.0e-12
const.type = constf
const.ds = 0.1
const.kx = 1.0
const.ky = 1.0
const.kt = 1.0e-12
monitor.type = beam_monitor
monitor.backend = h5
monitor.nonlinear_lens_invariants = true
monitor.alpha = 0.0
monitor.beta = 1.0
monitor.tn = 0.4
monitor.cn = 0.01
###############################################################################
# Algorithms
###############################################################################
algo.particle_shape = 2
algo.space_charge = false
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 numpy as np
import openpmd_api as io
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)
# 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)
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)
A nonlinear focusing channel based on the physical IOTA nonlinear magnet
A representation of the physical IOTA nonlinear magnetic element with realistic s-dependence, obtained using a sequence of nonlinear lenses and drifts equivalent to the use of a second-order symplectic integrator.
A thin linear lens is added at the exit of the nonlinear element, representing the ideal map associated with the remainder of the lattice.
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) are evaluated at the entrance to the nonlinear element, and then again after the thin lens (representing a single period). These values should be unchanged for all particles (to within acceptable tolerance).
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 either as:
Python script:
python3 run_iotalens_sdep.pyorImpactX executable using an input file:
impactx input_iotalens_sdep.in
For MPI-parallel runs, prefix these lines with 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 math
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
kin_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_kin_energy_MeV(kin_energy_MeV)
# particle bunch
distr = distribution.Waterbag(
lambdaX=1.397456296195e-003,
lambdaY=1.397456296195e-003,
lambdaT=1.0e-4,
lambdaPx=1.256184325020e-003,
lambdaPy=1.256184325020e-003,
lambdaPt=0.0,
muxpx=0.8090169943749474,
muypy=0.8090169943749474,
mutpt=0.0,
)
sim.add_particles(bunch_charge_C, distr, npart)
# add beam diagnostics
monitor = elements.BeamMonitor("monitor", backend="h5")
monitor.nonlinear_lens_invariants = True
monitor.alpha = 1.376381920471173
monitor.beta = 1.892632003628881
monitor.tn = 0.4
monitor.cn = 0.01
sim.lattice.append(monitor)
# defining parameters of the nonlinear lens
lens_length = 1.8
num_lenses = 18
tune_advance = 0.3
c_parameter = 0.01
t_strength = 0.4
ds = lens_length / num_lenses
# drift elements
ds_half = ds / 2.0
dr = elements.Drift(ds=ds_half)
# define the nonlinear lens segments
for j in range(0, num_lenses):
s = -lens_length / 2.0 + ds_half + j * ds
beta_star = lens_length / 2.0 * 1.0 / math.tan(math.pi * tune_advance)
beta = beta_star * (
1.0 + (2.0 * s * math.tan(math.pi * tune_advance) / lens_length) ** 2
)
knll_s = t_strength * c_parameter**2 * ds / beta
cnll_s = c_parameter * math.sqrt(beta)
nllens = elements.NonlinearLens(knll=knll_s, cnll=cnll_s)
segments = [dr, nllens, dr]
sim.lattice.extend(segments)
# focusing lens
const = elements.ConstF(
ds=1.0e-8, kx=12060.113295833, ky=12060.113295833, kt=1.0e-12, nslice=1
)
sim.lattice.append(const)
sim.lattice.append(monitor)
# number of periods
sim.periods = 1
# run simulation
sim.evolve()
# clean shutdown
sim.finalize()
###############################################################################
# Particle Beam(s)
###############################################################################
beam.npart = 10000
beam.units = static
beam.kin_energy = 2.5
beam.charge = 1.0e-9
beam.particle = proton
beam.distribution = waterbag
#beam.lambdaX = 2.0e-3
#beam.lambdaY = 2.0e-3
#beam.lambdaT = 1.0e-3
#beam.lambdaPx = 3.0e-4
#beam.lambdaPy = 3.0e-4
#beam.lambdaPt = 0.0
#beam.muxpx = 0.0
#beam.muypy = 0.0
#beam.mutpt = 0.0
beam.lambdaX = 1.397456296195e-003
beam.lambdaY = 1.397456296195e-003
beam.lambdaT = 1.0e-4
beam.lambdaPx = 1.256184325020e-003
beam.lambdaPy = 1.256184325020e-003
beam.lambdaPt = 0.0
beam.muxpx = 0.8090169943749474
beam.muypy = 0.8090169943749474
beam.mutpt = 0.0
###############################################################################
# Beamline: lattice elements and segments
###############################################################################
lattice.elements = monitor lens const monitor
lens.type = line
lens.elements = = dr_end nll1 dr nll2 dr nll3 dr nll4 dr nll5 dr nll6 \
dr nll7 dr nll8 dr nll9 dr nll9 dr nll8 dr nll7 \
dr nll6 dr nll5 dr nll4 dr nll3 dr nll2 dr nll1 \
dr_end
# Nonlinear lens segments
nll1.type = nonlinear_lens
nll1.knll = 2.2742558121e-6
nll1.cnll = 0.013262040169952
nll2.type = nonlinear_lens
nll2.knll = 2.641786366e-6
nll2.cnll = 0.012304986694423
nll3.type = nonlinear_lens
nll3.knll = 3.076868353e-6
nll3.cnll = 0.011401855643727
nll4.type = nonlinear_lens
nll4.knll = 3.582606522e-6
nll4.cnll = 0.010566482535866
nll5.type = nonlinear_lens
nll5.knll = 4.151211157e-6
nll5.cnll = 0.009816181575902
nll6.type = nonlinear_lens
nll6.knll = 4.754946985e-6
nll6.cnll = 0.0091718544892154
nll7.type = nonlinear_lens
nll7.knll = 5.337102374e-6
nll7.cnll = 0.008657195579489
nll8.type = nonlinear_lens
nll8.knll = 5.811437818e-6
nll8.cnll = 0.008296371635942
nll9.type = nonlinear_lens
nll9.knll = 6.081693334e-6
nll9.cnll = 0.008109941789663
dr.type = drift
dr.ds = 0.1
dr_end.type = drift
dr_end.ds = 0.05
# Focusing of the external lattice
const.type = constf
const.ds = 1.0e-8
const.kx = 12060.113295833
const.ky = 12060.113295833
const.kt = 1.0e-12
const.nslice = 1
# Beam Monitor: Diagnostics
monitor.type = beam_monitor
monitor.backend = h5
monitor.nonlinear_lens_invariants = true
monitor.alpha = 1.376381920471173
monitor.beta = 1.892632003628881
monitor.tn = 0.4
monitor.cn = 0.01
###############################################################################
# Algorithms
###############################################################################
algo.particle_shape = 2
algo.space_charge = false
Analyze
We run the following script to analyze correctness:
Script analysis_iotalens_sdep.py
#!/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 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)
# 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)
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 = 3.0 * 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],
# [5.993291e-02, 3.426664e-02, 8.513875e-02, 7.022481e-02],
[6.016450e-02, 3.502064e-02, 8.560226e-02, 7.148169e-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 = 3.0 * 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],
# [5.993291e-02, 3.426664e-02, 8.513875e-02, 7.022481e-02],
[6.016450e-02, 3.502064e-02, 8.560226e-02, 7.148169e-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
Hrms = np.sqrt(sigH**2 + meanH**2)
Irms = np.sqrt(sigI**2 + meanI**2)
atol = 2.5e-3 * Hrms
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.5e-3 * Irms
rtol = 0.0
print()
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)