First draft of test of nonlinear lens with aperture.

examples/CMakeLists.txt

@@ -1152,3 +1152,20 @@ add_impactx_test(pytorch_surrogate_model
     examples/pytorch_surrogate_model/visualize_ml_surrogate_15_stage.py
 )
 label_impactx_test(pytorch_surrogate_model slow)
+
+
+# IOTA s-dependent nonlinear lens with aperture ##############################
+#
+# w/o space charge
+add_impactx_test(IOTA_nll_aperture
+    examples/iota_lens/input_iotalens_sdep_aperture.in
+      ON   # ImpactX MPI-parallel
+#    examples/iota_lens/analysis_iotalens_sdep.py
+    OFF  # no plot script yet
+)
+add_impactx_test(IOTA_nll_aperture.py
+    examples/iota_lens/run_iotalens_sdep_aperture.py
+      OFF   # ImpactX MPI-parallel
+#    examples/iota_lens/analysis_iotalens_sdep.py
+    OFF  # no plot script yet
+)

examples/iota_lens/README.rst

@@ -106,3 +106,52 @@ We run the following script to analyze correctness:
    .. literalinclude:: analysis_iotalens_sdep.py
       :language: python3
       :caption: You can copy this file from ``examples/iota_lens/analysis_iotalens_sdep.py``.
+
+
+.. _examples-iotalens-sdep-aperture:
+
+A nonlinear focusing channel based on the physical IOTA nonlinear magnet with aperture
+=======================================================================================
+
+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.  Transverse aperture restrictions
+are included.  The elliptical aperture dimensions are based on the physical magnet
+design.
+
+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 :math:`\mu_H`, :math:`\sigma_H`, :math:`\mu_I`, :math:`\sigma_I` must agree with nominal values.
+
+
+Run
+---
+
+This example can be run **either** as:
+
+* **Python** script: ``python3 run_iotalens_sdep_aperture.py`` or
+* ImpactX **executable** using an input file: ``impactx input_iotalens_sdep_aperture.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_iotalens_sdep_aperture.py
+          :language: python3
+          :caption: You can copy this file from ``examples/iota_lens/run_iotalens_sdep_aperture.py``.
+
+   .. tab-item:: Executable: Input File
+
+       .. literalinclude:: input_iotalens_sdep_aperture.in
+          :language: ini
+          :caption: You can copy this file from ``examples/iota_lens/input_iotalens_sdep_aperture.in``.
+

examples/iota_lens/input_iotalens_sdep_aperture.in

@@ -0,0 +1,152 @@
+###############################################################################
+# 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 = beam.lambdaX
+#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 = beam.lambdaX
+beam.lambdaT = 1.0e-4
+beam.lambdaPx = 1.256184325020e-003
+beam.lambdaPy = beam.lambdaPx
+beam.lambdaPt = 0.0
+beam.muxpx = 0.8090169943749474
+beam.muypy = beam.muxpx
+beam.mutpt = 0.0
+
+
+###############################################################################
+# Beamline: lattice elements and segments
+###############################################################################
+lattice.elements = monitor lens const monitor
+
+lens.type = line
+lens.elements = = dr_end ap1 nll1 dr ap2 nll2 dr ap3 nll3 dr ap4 nll4 dr ap5 nll5 dr ap6 nll6  \
+                       dr ap7 nll7 dr ap8 nll8 dr ap9 nll9 dr ap9 nll9 dr ap8 nll8 dr ap7 nll7  \
+                       dr ap6 nll6 dr ap5 nll5 dr ap4 nll4 dr ap3 nll3 dr ap2 nll2 dr ap1 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
+
+# Aperture collimation
+ap1.type = aperture
+ap1.shape = elliptical
+ap1.aperture_x = 0.0069138074
+ap1.aperture_y = 0.00921840994
+
+ap2.type = aperture
+ap2.shape = elliptical
+ap2.aperture_x = 0.0063622465
+ap2.aperture_y = 0.00848299541
+
+ap3.type = aperture
+ap3.shape = elliptical
+ap3.aperture_x = 0.0058417668
+ap3.aperture_y = 0.00778902251
+
+ap4.type = aperture
+ap4.shape = elliptical
+ap4.aperture_x = 0.0053603421
+ap4.aperture_y = 0.0071471228
+
+ap5.type = aperture
+ap5.shape = elliptical
+ap5.aperture_x = 0.0049279501
+ap5.aperture_y = 0.00657060016
+
+ap6.type = aperture
+ap6.shape = elliptical
+ap6.aperture_x = 0.0045566354
+ap6.aperture_y = 0.00607551398
+
+ap7.type = aperture
+ap7.shape = elliptical
+ap7.aperture_x = 0.0042600509
+ap7.aperture_y = 0.00568006786
+
+ap8.type = aperture
+ap8.shape = elliptical
+ap8.aperture_x = 0.0040521202
+ap8.aperture_y = 0.00540282702
+
+ap9.type = aperture
+ap9.shape = elliptical
+ap9.aperture_x = 0.0039446881
+ap9.aperture_y = 0.00525958413
+
+# 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

examples/iota_lens/run_iotalens_sdep_aperture.py

@@ -0,0 +1,109 @@
+#!/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
+
+# aperture restrictions
+ap_x = [0.0069138074, 0.0063622465, 0.0058417668, 0.0053603421, 0.0049279501, 0.0045566354, 0.0042600509, 0.0040521202, 0.0039446881,
+        0.0039446881, 0.0040521202, 0.0042600509, 0.0045566354, 0.0049279501, 0.0053603421, 0.0058417668, 0.0063622465, 0.0069138074]
+ap_y = [0.00921840994, 0.00848299541, 0.00778902251, 0.0071471228, 0.00657060016, 0.00607551398, 0.00568006786, 0.00540282702, 0.00525958413,
+        0.00525958413, 0.00540282702, 0.00568006786, 0.00607551398, 0.00657060016, 0.0071471228, 0.00778902251, 0.00848299541, 0.00921840994]
+
+# drift elements
+ds_half = ds / 2.0
+dr = elements.Drift(name="dr", 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)
+    ap = elements.Aperture(name="ap" + str(j), type="elliptical", aperture_x=ap_x[j], aperture_y=ap_y[j])
+    nllens = elements.NonlinearLens(name="nllens" + str(j), knll=knll_s, cnll=cnll_s)
+    segments = [dr, ap, nllens, dr]
+    sim.lattice.extend(segments)
+
+# focusing lens
+const = elements.ConstF(
+    name="const",
+    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.track_particles()
+
+# clean shutdown
+sim.finalize()