diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt index e16b12350..d54adad03 100644 --- a/examples/CMakeLists.txt +++ b/examples/CMakeLists.txt @@ -640,6 +640,42 @@ add_impactx_test(hvkicker_madx.py OFF # no plot script yet ) +# IOTA s-dependent nonlinear lens test ########################################################## +# +# w/o space charge +add_impactx_test(IOTA_nll + examples/iota_lens/input_iotalens_sdep.in + ON # ImpactX MPI-parallel + OFF # ImpactX Python interface + examples/iota_lens/analysis_iotalens_sdep.py + OFF # no plot script yet +) +add_impactx_test(IOTA_nll.py + examples/iota_lens/run_iotalens_sdep.py + OFF # ImpactX MPI-parallel + ON # ImpactX Python interface + examples/iota_lens/analysis_iotalens_sdep.py + OFF # no plot script yet +) + +# IOTA nonlinear lattice test ########################################################## +# +# w/o space charge +add_impactx_test(IOTA_lattice + examples/iota_lattice/input_iotalattice_sdep.in + ON # ImpactX MPI-parallel + OFF # ImpactX Python interface + examples/iota_lattice/analysis_iotalattice_sdep.py + OFF # no plot script yet +) +add_impactx_test(IOTA_lattice.py + examples/iota_lattice/run_iotalattice_sdep.py + OFF # ImpactX MPI-parallel + ON # ImpactX Python interface + examples/iota_lattice/analysis_iotalattice_sdep.py + OFF # no plot script yet +) + # Thin dipole ######################################################## # # w/o space charge diff --git a/examples/iota_lattice/README.rst b/examples/iota_lattice/README.rst index bd92cbf61..60ce8a9a4 100644 --- a/examples/iota_lattice/README.rst +++ b/examples/iota_lattice/README.rst @@ -48,3 +48,55 @@ We run the following script to analyze correctness: .. literalinclude:: analysis_iotalattice.py :language: python3 :caption: You can copy this file from ``examples/iota_lattice/analysis_iotalattice.py``. + + +.. _examples-iotalattice-sdep: + +The full nonlinear lattice of the Fermilab IOTA storage ring +============================================================ + +The full nonlinear lattice of the IOTA storage ring, configured for operation with a 2.5 MeV proton beam. + +The special nonlinear magnetic element for integrable optics experiments is denoted ``nll``. To simplify analysis, +the lattice has been arranged so that nll appears as the first element in the sequence. + +The two functions H (Hamiltonian) and I (the second invariant) are evaluated at the entrance to the nonlinear element. +These values should be unchanged for all particles (to within acceptable tolerance), over the specified number of periods (default 5). + +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_iotalattice_sdep.py`` or +* ImpactX **executable** using an input file: ``impactx input_iotalattice_sdep.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_iotalattice_sdep.py + :language: python3 + :caption: You can copy this file from ``examples/iota_lattice/run_iotalattice_sdep.py``. + + .. tab-item:: Executable: Input File + + .. literalinclude:: input_iotalattice_sdep.in + :language: ini + :caption: You can copy this file from ``examples/iota_lattice/input_iotalattice_sdep.in``. + + +Analyze +------- + +We run the following script to analyze correctness: + +.. dropdown:: Script ``analysis_iotalattice_sdep.py`` + + .. literalinclude:: analysis_iotalattice_sdep.py + :language: python3 + :caption: You can copy this file from ``examples/iota_lattice/analysis_iotalattice_sdep.py``. diff --git a/examples/iota_lattice/analysis_iotalattice_sdep.py b/examples/iota_lattice/analysis_iotalattice_sdep.py new file mode 100755 index 000000000..13e457798 --- /dev/null +++ b/examples/iota_lattice/analysis_iotalattice_sdep.py @@ -0,0 +1,109 @@ +#!/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 +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], + [7.263202e-02, 4.454371e-02, 9.288060e-02, 8.211506e-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], + [7.263202e-02, 4.454371e-02, 9.288060e-02, 8.211506e-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 = 4.0e-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 = 5.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) diff --git a/examples/iota_lattice/input_iotalattice_sdep.in b/examples/iota_lattice/input_iotalattice_sdep.in new file mode 100644 index 000000000..815750479 --- /dev/null +++ b/examples/iota_lattice/input_iotalattice_sdep.in @@ -0,0 +1,291 @@ +############################################################################### +# 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.sigmaX = 1.865379469388e-003 +beam.sigmaY = 2.0192133150418e-003 +beam.sigmaT = 1.0e-3 +beam.sigmaPx = 1.402566720991e-003 +beam.sigmaPy = 9.57593913381e-004 +beam.sigmaPt = 0.0 +beam.muxpx = 0.482260919078473 +beam.muypy = 0.762127656873158 +beam.mutpt = 0.0 + + +############################################################################### +# Beamline: lattice elements and segments +############################################################################### +lattice.periods = 5 +lattice.elements = monitor nll after_nll qe3 after_nll_rev dnll before_nll_rev \ + before_nll monitor + +# lines +before_nll.type = line +before_nll.elements = dra1 qa1 dra2 qa2 dra3 qa3 dra4 qa4 dra5 \ + edge30 sbend30 edge30 drb1 qb1 drb2 qb2 drb2 qb3 \ + drb3 + +after_nll.type = line +after_nll.elements = drb3 qb4 drb2 qb5 drb2 qb6 drb4 \ + edge60 sbend60 edge60 drc1 qc1 drc2 qc2 drc2 qc3 drc1 \ + edge60 sbend60 edge60 drd1 qd1 drd2 qd2 drd3 qd3 drd2 qd4 drd4 \ + edge30 sbend30 edge30 dre1 qe1 dre2 qe2 dre3 + +before_nll_rev.type = line +before_nll_rev.reverse = true +before_nll_rev.elements = dra1 qa1 dra2 qa2 dra3 qa3 dra4 qa4 dra5 \ + edge30 sbend30 edge30 drb1 qb1 drb2 qb2 drb2 qb3 \ + drb3 + +after_nll_rev.type = line +after_nll_rev.reverse = true +after_nll_rev.elements = drb3 qb4 drb2 qb5 drb2 qb6 drb4 \ + edge60 sbend60 edge60 drc1 qc1 drc2 qc2 drc2 qc3 drc1 \ + edge60 sbend60 edge60 drd1 qd1 drd2 qd2 drd3 qd3 drd2 qd4 drd4 \ + edge30 sbend30 edge30 dre1 qe1 dre2 qe2 dre3 + +nll.type = line +nll.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 + +# thick element splitting for space charge +lattice.nslice = 10 + +# 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 + + +# Drift elements: + +dr.type = drift +dr.ds = 0.1 + +dr_end.type = drift +dr_end.ds = 0.05 + +dra1.type = drift +dra1.ds = 0.9125 + +dra2.type = drift +dra2.ds = 0.135 + +dra3.type = drift +dra3.ds = 0.725 + +dra4.type = drift +dra4.ds = 0.145 + +dra5.type = drift +dra5.ds = 0.3405 + +drb1.type = drift +drb1.ds = 0.3205 + +drb2.type = drift +drb2.ds = 0.14 + +drb3.type = drift +drb3.ds = 0.1525 + +drb4.type = drift +drb4.ds = 0.31437095 + +drc1.type = drift +drc1.ds = 0.42437095 + +drc2.type = drift +drc2.ds = 0.355 + +dnll.type = drift +dnll.ds = 1.8 + +drd1.type = drift +drd1.ds = 0.62437095 + +drd2.type = drift +drd2.ds = 0.42 + +drd3.type = drift +drd3.ds = 1.625 + +drd4.type = drift +drd4.ds = 0.6305 + +dre1.type = drift +dre1.ds = 0.5305 + +dre2.type = drift +dre2.ds = 1.235 + +dre3.type = drift +dre3.ds = 0.8075 + + +# Bend elements: + +sbend30.type = sbend +sbend30.ds = 0.4305191429 +sbend30.rc = 0.822230996255981 + +edge30.type = dipedge +edge30.psi = 0.0 +edge30.rc = 0.822230996255981 +edge30.g = 0.058 +edge30.K2 = 0.5 + +sbend60.type = sbend +sbend60.ds = 0.8092963858 +sbend60.rc = 0.772821121503940 + +edge60.type = dipedge +edge60.psi = 0.0 +edge60.rc = 0.772821121503940 +edge60.g = 0.058 +edge60.K2 = 0.5 + + +# Quad elements: + +qa1.type = quad +qa1.ds = 0.21 +qa1.k = -8.78017699 + +qa2.type = quad +qa2.ds = 0.21 +qa2.k = 13.24451745 + +qa3.type = quad +qa3.ds = 0.21 +qa3.k = -13.65151327 + +qa4.type = quad +qa4.ds = 0.21 +qa4.k = 19.75138652 + +qb1.type = quad +qb1.ds = 0.21 +qb1.k = -10.84199727 + +qb2.type = quad +qb2.ds = 0.21 +qb2.k = 16.24844348 + +qb3.type = quad +qb3.ds = 0.21 +qb3.k = -8.27411104 + +qb4.type = quad +qb4.ds = 0.21 +qb4.k = -7.45719247 + +qb5.type = quad +qb5.ds = 0.21 +qb5.k = 14.03362243 + +qb6.type = quad +qb6.ds = 0.21 +qb6.k = -12.23595641 + +qc1.type = quad +qc1.ds = 0.21 +qc1.k = -13.18863768 + +qc2.type = quad +qc2.ds = 0.21 +qc2.k = 11.50601829 + +qc3.type = quad +qc3.ds = 0.21 +qc3.k = -11.10445869 + +qd1.type = quad +qd1.ds = 0.21 +qd1.k = -6.78179218 + +qd2.type = quad +qd2.ds = 0.21 +qd2.k = 5.19026998 + +qd3.type = quad +qd3.ds = 0.21 +qd3.k = -5.8586173 + +qd4.type = quad +qd4.ds = 0.21 +qd4.k = 4.62460039 + +qe1.type = quad +qe1.ds = 0.21 +qe1.k = -4.49607687 + +qe2.type = quad +qe2.ds = 0.21 +qe2.k = 6.66737146 + +qe3.type = quad +qe3.ds = 0.21 +qe3.k = -6.69148177 + +# Beam Monitor: Diagnostics +monitor.type = beam_monitor +monitor.backend = h5 + + +############################################################################### +# Algorithms +############################################################################### +algo.particle_shape = 2 +algo.space_charge = false + + +############################################################################### +# Diagnostics +############################################################################### +diag.slice_step_diagnostics = true + +diag.alpha = 1.376381920471173 +diag.beta = 1.892632003628881 +diag.tn = 0.4 +diag.cn = 0.01 diff --git a/examples/iota_lattice/run_iotalattice_sdep.py b/examples/iota_lattice/run_iotalattice_sdep.py new file mode 100644 index 000000000..92e265e8b --- /dev/null +++ b/examples/iota_lattice/run_iotalattice_sdep.py @@ -0,0 +1,224 @@ +#!/usr/bin/env python3 +# +# Copyright 2022-2023 ImpactX contributors +# Authors: Chad Mitchell, Axel Huebl +# License: BSD-3-Clause-LBNL +# +# -*- coding: utf-8 -*- + +import math + +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() + +# diagnostics: IOTA nonlinear lens invariants calculation +sim.set_diag_iota_invariants( + alpha=1.376381920471173, beta=1.892632003628881, tn=0.4, cn=0.01 +) + +# init particle beam +energy_MeV = 2.5 +bunch_charge_C = 1.0e-9 # used with space charge +npart = 10000 + +# reference particle +ref = sim.particle_container().ref_particle() +ref.set_charge_qe(1.0).set_mass_MeV(938.27208816).set_kin_energy_MeV(energy_MeV) + +# particle bunch +distr = distribution.Waterbag( + sigmaX=1.865379469388e-003, + sigmaY=2.0192133150418e-003, + sigmaT=1.0e-3, + sigmaPx=1.402566720991e-003, + sigmaPy=9.57593913381e-004, + sigmaPt=0.0, + muxpx=0.482260919078473, + muypy=0.762127656873158, + mutpt=0.0, +) +sim.add_particles(bunch_charge_C, distr, npart) + +# add beam diagnostics +monitor = elements.BeamMonitor("monitor", backend="h5") + +# init accelerator lattice +ns = 10 # number of slices per ds in the element + +# Drift elements +dra1 = elements.Drift(ds=0.9125, nslice=ns) +dra2 = elements.Drift(ds=0.135, nslice=ns) +dra3 = elements.Drift(ds=0.725, nslice=ns) +dra4 = elements.Drift(ds=0.145, nslice=ns) +dra5 = elements.Drift(ds=0.3405, nslice=ns) +drb1 = elements.Drift(ds=0.3205, nslice=ns) +drb2 = elements.Drift(ds=0.14, nslice=ns) +drb3 = elements.Drift(ds=0.1525, nslice=ns) +drb4 = elements.Drift(ds=0.31437095, nslice=ns) +drc1 = elements.Drift(ds=0.42437095, nslice=ns) +drc2 = elements.Drift(ds=0.355, nslice=ns) +dnll = elements.Drift(ds=1.8, nslice=ns) +drd1 = elements.Drift(ds=0.62437095, nslice=ns) +drd2 = elements.Drift(ds=0.42, nslice=ns) +drd3 = elements.Drift(ds=1.625, nslice=ns) +drd4 = elements.Drift(ds=0.6305, nslice=ns) +dre1 = elements.Drift(ds=0.5305, nslice=ns) +dre2 = elements.Drift(ds=1.235, nslice=ns) +dre3 = elements.Drift(ds=0.8075, nslice=ns) + +# Bend elements +rc30 = 0.822230996255981 +sbend30 = elements.Sbend(ds=0.4305191429, rc=rc30) +edge30 = elements.DipEdge(psi=0.0, rc=rc30, g=0.058, K2=0.5) + +rc60 = 0.772821121503940 +sbend60 = elements.Sbend(ds=0.8092963858, rc=rc60) +edge60 = elements.DipEdge(psi=0.0, rc=rc60, g=0.058, K2=0.5) + +# Quad elements +ds_quad = 0.21 +qa1 = elements.Quad(ds=ds_quad, k=-8.78017699, nslice=ns) +qa2 = elements.Quad(ds=ds_quad, k=13.24451745, nslice=ns) +qa3 = elements.Quad(ds=ds_quad, k=-13.65151327, nslice=ns) +qa4 = elements.Quad(ds=ds_quad, k=19.75138652, nslice=ns) +qb1 = elements.Quad(ds=ds_quad, k=-10.84199727, nslice=ns) +qb2 = elements.Quad(ds=ds_quad, k=16.24844348, nslice=ns) +qb3 = elements.Quad(ds=ds_quad, k=-8.27411104, nslice=ns) +qb4 = elements.Quad(ds=ds_quad, k=-7.45719247, nslice=ns) +qb5 = elements.Quad(ds=ds_quad, k=14.03362243, nslice=ns) +qb6 = elements.Quad(ds=ds_quad, k=-12.23595641, nslice=ns) +qc1 = elements.Quad(ds=ds_quad, k=-13.18863768, nslice=ns) +qc2 = elements.Quad(ds=ds_quad, k=11.50601829, nslice=ns) +qc3 = elements.Quad(ds=ds_quad, k=-11.10445869, nslice=ns) +qd1 = elements.Quad(ds=ds_quad, k=-6.78179218, nslice=ns) +qd2 = elements.Quad(ds=ds_quad, k=5.19026998, nslice=ns) +qd3 = elements.Quad(ds=ds_quad, k=-5.8586173, nslice=ns) +qd4 = elements.Quad(ds=ds_quad, k=4.62460039, nslice=ns) +qe1 = elements.Quad(ds=ds_quad, k=-4.49607687, nslice=ns) +qe2 = elements.Quad(ds=ds_quad, k=6.66737146, nslice=ns) +qe3 = elements.Quad(ds=ds_quad, k=-6.69148177, nslice=ns) + +# Special (elliptic) nonlinear element: + +# set basic parameters of the nonlinear element +lens_length = 1.8 +num_lenses = 18 +tune_advance = 0.3 +c_parameter = 0.01 +t_strength = 0.4 +ds = lens_length / num_lenses + +# build up the nonlinear lens in segments +ds_half = ds / 2.0 +dr = elements.Drift(ds=ds_half, nslice=1) +nll = [] +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) + segment = [dr, nllens, dr] + nll.extend(segment) + +lattice_before_nll = [ + dra1, + qa1, + dra2, + qa2, + dra3, + qa3, + dra4, + qa4, + dra5, + edge30, + sbend30, + edge30, + drb1, + qb1, + drb2, + qb2, + drb2, + qb3, + drb3, +] + +lattice_after_nll = [ + drb3, + qb4, + drb2, + qb5, + drb2, + qb6, + drb4, + edge60, + sbend60, + edge60, + drc1, + qc1, + drc2, + qc2, + drc2, + qc3, + drc1, + edge60, + sbend60, + edge60, + drd1, + qd1, + drd2, + qd2, + drd3, + qd3, + drd2, + qd4, + drd4, + edge30, + sbend30, + edge30, + dre1, + qe1, + dre2, + qe2, + dre3, +] + +# build lattice: first half, qe3, then mirror +# modified to place nll as the first element +# fmt:on +sim.lattice.append(monitor) +sim.lattice.extend(nll) +sim.lattice.extend(lattice_after_nll) +sim.lattice.append(qe3) +lattice_after_nll.reverse() +sim.lattice.extend(lattice_after_nll) +sim.lattice.append(dnll) +lattice_before_nll.reverse() +sim.lattice.extend(lattice_before_nll) +lattice_before_nll.reverse() +sim.lattice.extend(lattice_before_nll) +sim.lattice.append(monitor) + +# number of turns in the ring +sim.periods = 5 + +# run simulation +sim.evolve() + +# clean shutdown +del sim +amr.finalize() diff --git a/examples/iota_lens/README.rst b/examples/iota_lens/README.rst index d1d1a5b77..d7455f4b1 100644 --- a/examples/iota_lens/README.rst +++ b/examples/iota_lens/README.rst @@ -48,3 +48,61 @@ We run the following script to analyze correctness: .. literalinclude:: analysis_iotalens.py :language: python3 :caption: You can copy this file from ``examples/iota_lens/analysis_iotalens.py``. + + +.. _examples-iotalens-sdep: + +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 :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.py`` or +* ImpactX **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. + +.. tab-set:: + + .. tab-item:: Python: Script + + .. literalinclude:: run_iotalens_sdep.py + :language: python3 + :caption: You can copy this file from ``examples/iota_lens/run_iotalens_sdep.py``. + + .. tab-item:: Executable: Input File + + .. literalinclude:: input_iotalens_sdep.in + :language: ini + :caption: You can copy this file from ``examples/iota_lens/input_iotalens_sdep.in``. + + +Analyze +------- + +We run the following script to analyze correctness: + +.. dropdown:: Script ``analysis_iotalens_sdep.py`` + + .. literalinclude:: analysis_iotalens_sdep.py + :language: python3 + :caption: You can copy this file from ``examples/iota_lens/analysis_iotalens_sdep.py``. diff --git a/examples/iota_lens/analysis_iotalens_sdep.py b/examples/iota_lens/analysis_iotalens_sdep.py new file mode 100755 index 000000000..46a3b7755 --- /dev/null +++ b/examples/iota_lens/analysis_iotalens_sdep.py @@ -0,0 +1,111 @@ +#!/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 +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 = 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) diff --git a/examples/iota_lens/input_iotalens_sdep.in b/examples/iota_lens/input_iotalens_sdep.in new file mode 100644 index 000000000..55ed2d454 --- /dev/null +++ b/examples/iota_lens/input_iotalens_sdep.in @@ -0,0 +1,109 @@ +############################################################################### +# 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.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 +beam.sigmaX = 1.397456296195e-003 +beam.sigmaY = 1.397456296195e-003 +beam.sigmaT = 1.0e-4 +beam.sigmaPx = 1.256184325020e-003 +beam.sigmaPy = 1.256184325020e-003 +beam.sigmaPt = 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 + +############################################################################### +# Algorithms +############################################################################### +algo.particle_shape = 2 +algo.space_charge = false + + +############################################################################### +# Diagnostics +############################################################################### +diag.alpha = 1.376381920471173 +diag.beta = 1.892632003628881 +diag.tn = 0.4 +diag.cn = 0.01 diff --git a/examples/iota_lens/run_iotalens_sdep.py b/examples/iota_lens/run_iotalens_sdep.py new file mode 100644 index 000000000..ec9289220 --- /dev/null +++ b/examples/iota_lens/run_iotalens_sdep.py @@ -0,0 +1,98 @@ +#!/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 + +import amrex.space3d as amr +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() + +# diagnostics: IOTA nonlinear lens invariants calculation +sim.set_diag_iota_invariants( + alpha=1.376381920471173, beta=1.892632003628881, tn=0.4, cn=0.01 +) + +# 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( + sigmaX=1.397456296195e-003, + sigmaY=1.397456296195e-003, + sigmaT=1.0e-4, + sigmaPx=1.256184325020e-003, + sigmaPy=1.256184325020e-003, + sigmaPt=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") +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 +del sim +amr.finalize() diff --git a/src/particles/diagnostics/NonlinearLensInvariants.H b/src/particles/diagnostics/NonlinearLensInvariants.H index 8713b4c73..7f05d0ff3 100644 --- a/src/particles/diagnostics/NonlinearLensInvariants.H +++ b/src/particles/diagnostics/NonlinearLensInvariants.H @@ -82,8 +82,8 @@ namespace impactx::diagnostics // convert transverse phase space coordinates to normalized units amrex::ParticleReal const xn = x/(m_cn*sqrt(m_beta)); amrex::ParticleReal const yn = y/(m_cn*sqrt(m_beta)); - amrex::ParticleReal const pxn = px*sqrt(m_beta)/m_cn + m_alpha*x; - amrex::ParticleReal const pyn = py*sqrt(m_beta)/m_cn + m_alpha*y; + amrex::ParticleReal const pxn = px*sqrt(m_beta)/m_cn + m_alpha*xn; + amrex::ParticleReal const pyn = py*sqrt(m_beta)/m_cn + m_alpha*yn; // assign complex position zeta = x + iy Complex const zeta(xn, yn);