Test of IOTA nonlinear magnet model with physical aperture. (#802)

* First draft of test of nonlinear lens with aperture. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Use existing analysis script. * Correct type -> shape input * Modify app input - Np and turns * Modify Python input - Np and turns * Update analysis script. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update examples/iota_lens/README.rst Co-authored-by: Axel Huebl <[email protected]> * Update README.rst Change objective description. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Axel Huebl <[email protected]>
ECP-WarpX · Jan 18, 2025 · e78a9fc · e78a9fc
1 parent fe6e789
commit e78a9fc
Show file tree

Hide file tree

Showing 5 changed files with 477 additions and 0 deletions.
diff --git a/examples/CMakeLists.txt b/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_aperture.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_aperture.py
+    OFF  # no plot script yet
+)
diff --git a/examples/iota_lens/README.rst b/examples/iota_lens/README.rst
@@ -106,3 +106,49 @@ 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 beam is transported over 3 periods.
+
+In this test, the initial and final values of :math:`\lambda_x`, :math:`\lambda_y`, :math:`\lambda_t`, :math:`\epsilon_x`, :math:`\epsilon_y`, and :math:`\epsilon_t` must agree with nominal values.
+
+The fraction of charge lost after 3 periods of transport must agree with the nominal value obtained from high-resolution simulation.
+
+
+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``.
diff --git a/examples/iota_lens/analysis_iotalens_sdep_aperture.py b/examples/iota_lens/analysis_iotalens_sdep_aperture.py
@@ -0,0 +1,114 @@
+#!/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)
+last_step = list(series.iterations)[-1]
+initial_beam = series.iterations[1].particles["beam"]
+final_beam = series.iterations[last_step].particles["beam"]
+initial = initial_beam.to_df()
+final = final_beam.to_df()
+
+# compare number of particles
+num_particles = 100000
+assert num_particles == len(initial)
+
+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],
+    [
+        2.378921e-03,
+        2.376389e-03,
+        1.001072e-04,
+        2.984246e-06,
+        2.982321e-06,
+    ],
+    rtol=rtol,
+    atol=atol,
+)
+
+
+print("")
+print("Final Beam:")
+sigx, sigy, sigt, emittance_x, emittance_y, emittance_t = get_moments(final)
+s_ref = final_beam.get_attribute("s_ref")
+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}\n"
+)
+
+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, s_ref],
+    [
+        1.926162e-03,
+        2.781398e-03,
+        1.020826e-04,
+        2.868312e-06,
+        3.163618e-06,
+        5.400000e00,
+    ],
+    rtol=rtol,
+    atol=atol,
+)
+
+charge_i = initial_beam.get_attribute("charge_C")
+charge_f = final_beam.get_attribute("charge_C")
+
+loss_pct = 100.0 * (charge_i - charge_f) / charge_i
+
+print(f" fractional loss (%) = {loss_pct}")
+
+atol = 0.2  # tolerance 0.2%
+print(f"  atol={atol}")
+assert np.allclose(
+    [loss_pct],
+    [6.0824],
+    atol=atol,
+)
diff --git a/examples/iota_lens/input_iotalens_sdep_aperture.in b/examples/iota_lens/input_iotalens_sdep_aperture.in
@@ -0,0 +1,153 @@
+###############################################################################
+# Particle Beam(s)
+###############################################################################
+beam.npart = 100000  #note: we use 100K particles to get reasonable loss statistics
+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.periods = 3
+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