From 706f7ccf8dc0998e3162ab0e37e14c8fa5b6ec94 Mon Sep 17 00:00:00 2001 From: Chad Mitchell <46825199+cemitch99@users.noreply.github.com> Date: Sun, 18 Sep 2022 22:25:54 -0700 Subject: [PATCH] Added cold expanding beam benchmark example (with space charge). (#244) * Added cold expanding beam benchmark example (with space charge). * Fix rst formatting * Cheaper resolution converges :tada: * Fix atol * Add in CMakeLists.txt & Manual * Analysis: Emittance Absolute to Zero * Expanding Beam: Python Version * Slightly Increase Tolerance * Disable Analysis for now Co-authored-by: Axel Huebl --- docs/source/usage/examples.rst | 1 + examples/CMakeLists.txt | 20 +++ examples/expanding_beam/README.rst | 48 +++++++ examples/expanding_beam/analysis_expanding.py | 124 ++++++++++++++++++ examples/expanding_beam/input_expanding.in | 38 ++++++ examples/expanding_beam/run_expanding.py | 59 +++++++++ 6 files changed, 290 insertions(+) create mode 100644 examples/expanding_beam/README.rst create mode 100755 examples/expanding_beam/analysis_expanding.py create mode 100644 examples/expanding_beam/input_expanding.in create mode 100755 examples/expanding_beam/run_expanding.py diff --git a/docs/source/usage/examples.rst b/docs/source/usage/examples.rst index 5c6c9568a..343f9dee4 100644 --- a/docs/source/usage/examples.rst +++ b/docs/source/usage/examples.rst @@ -11,6 +11,7 @@ This section allows you to **download input files** that correspond to different examples/fodo/README.rst examples/chicane/README.rst examples/cfchannel/README.rst + examples/expanding_beam/README.rst examples/kurth/README.rst examples/fodo_rf/README.rst examples/multipole/README.rst diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt index 8114a060b..9965461b7 100644 --- a/examples/CMakeLists.txt +++ b/examples/CMakeLists.txt @@ -280,6 +280,26 @@ add_impactx_test(multipole OFF # no plot script yet ) +# Expanding Beam Test ################################################################# +# +add_impactx_test(expanding_beam + examples/expanding_beam/input_expanding.in + OFF # ImpactX MPI-parallel + OFF # ImpactX Python interface + OFF # examples/expanding_beam/analysis_expanding.py + OFF # no plot script yet +) + +# Python: Expanding Beam Test ######################################################### +# +add_impactx_test(expanding_beam.py + examples/expanding_beam/run_expanding.py + OFF # ImpactX MPI-parallel + ON # ImpactX Python interface + OFF # examples/expanding_beam/analysis_expanding.py + OFF # no plot script yet +) + # Python: Chain of Multipoles Test ######################################################### # add_impactx_test(multipole.py diff --git a/examples/expanding_beam/README.rst b/examples/expanding_beam/README.rst new file mode 100644 index 000000000..7cc8caa16 --- /dev/null +++ b/examples/expanding_beam/README.rst @@ -0,0 +1,48 @@ +.. _examples-expanding: + +Expanding Beam in Free Space +============================ + +A coasting bunch expanding freely in free space under its own space charge. + +We use a cold (zero emittance) 250 MeV electron bunch whose +initial distribution is a uniformly-populated 3D ball of radius R0 = 1 mm when viewed in the bunch rest +frame. + +In the laboratory frame, the bunch is a uniformly-populated ellipsoid, which +expands to twice its original size. This is tested using the second moments of the distribution. + +In this test, the initial and final values of :math:`\sigma_x`, :math:`\sigma_y`, :math:`\sigma_t`, :math:`\epsilon_x`, :math:`\epsilon_y`, and :math:`\epsilon_t` must agree with nominal values. + + +Run +--- + +This example can be run as a Python script (``python3 run_expanding.py``) or with an app with an input file (``impactx input_expanding.in``). +Each can also be prefixed with an `MPI executor `__, such as ``mpiexec -n 4 ...`` or ``srun -n 4 ...``, depending on the system. + +.. tab-set:: + + .. tab-item:: Python Script + + .. literalinclude:: run_expanding.py + :language: python3 + :caption: You can copy this file from ``examples/expanding/run_expanding.py``. + + .. tab-item:: App Input File + + .. literalinclude:: input_expanding.in + :language: ini + :caption: You can copy this file from ``examples/expanding/input_expanding.in``. + + +Analyze +------- + +We run the following script to analyze correctness: + +.. dropdown:: Script ``analysis_expanding.py`` + + .. literalinclude:: analysis_expanding.py + :language: python3 + :caption: You can copy this file from ``examples/expanding/analysis_expanding.py``. diff --git a/examples/expanding_beam/analysis_expanding.py b/examples/expanding_beam/analysis_expanding.py new file mode 100755 index 000000000..b11008fbf --- /dev/null +++ b/examples/expanding_beam/analysis_expanding.py @@ -0,0 +1,124 @@ +#!/usr/bin/env python3 +# +# Copyright 2022 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 standard deviations of beam position & momenta + and emittance values + + Returns + ------- + sigx, sigy, sigt, emittance_x, emittance_y, emittance_t + """ + sigx = moment(beam["x"], moment=2) ** 0.5 # variance -> std dev. + sigpx = moment(beam["px"], moment=2) ** 0.5 + sigy = moment(beam["y"], moment=2) ** 0.5 + sigpy = moment(beam["py"], moment=2) ** 0.5 + sigt = moment(beam["t"], moment=2) ** 0.5 + sigpt = moment(beam["pt"], moment=2) ** 0.5 + + epstrms = beam.cov(ddof=0) + emittance_x = (sigx**2 * sigpx**2 - epstrms["x"]["px"] ** 2) ** 0.5 + emittance_y = (sigy**2 * sigpy**2 - epstrms["y"]["py"] ** 2) ** 0.5 + emittance_t = (sigt**2 * sigpt**2 - epstrms["t"]["pt"] ** 2) ** 0.5 + + return (sigx, sigy, sigt, emittance_x, emittance_y, emittance_t) + + +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/beam_000000.*") +final = read_all_files("diags/beam_final.*") + +# compare number of particles +num_particles = 10000 +assert num_particles == len(initial) +assert num_particles == len(final) + +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 = 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, emittance_t], + [ + 4.4721359550e-004, + 4.4721359550e-004, + 9.1224186858e-007, + 0.0e-006, + 0.0e-006, + 0.0e-006, + ], + rtol=rtol, + atol=atol, +) + + +print("") +print("Final Beam:") +sigx, sigy, sigt, emittance_x, emittance_y, emittance_t = get_moments(final) +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 = 1.5 * num_particles**-0.5 # from random sampling of a smooth distribution +print(f" rtol={rtol} (ignored: atol~={atol})") + +assert np.allclose( + [sigx, sigy, sigt], + [ + 8.9442719100e-004, + 8.9442719100e-004, + 1.8244837370e-006, + ], + rtol=rtol, + atol=atol, +) +atol = 1.0e-8 +rtol = 0.0 # ignored +assert np.allclose( + [emittance_x, emittance_y, emittance_t], + [ + 0.0, + 0.0, + 0.0, + ], + rtol=rtol, + atol=atol, +) diff --git a/examples/expanding_beam/input_expanding.in b/examples/expanding_beam/input_expanding.in new file mode 100644 index 000000000..e706ead43 --- /dev/null +++ b/examples/expanding_beam/input_expanding.in @@ -0,0 +1,38 @@ +############################################################################### +# Particle Beam(s) +############################################################################### +beam.npart = 10000 +beam.units = static +beam.energy = 250.0 +beam.charge = 1.0e-9 +beam.particle = electron +beam.distribution = waterbag +beam.sigmaX = 4.472135955e-4 +beam.sigmaY = 4.472135955e-4 +beam.sigmaT = 9.12241869e-7 +beam.sigmaPx = 0.0 +beam.sigmaPy = 0.0 +beam.sigmaPt = 0.0 +beam.muxpx = 0.0 +beam.muypy = 0.0 +beam.mutpt = 0.0 + + +############################################################################### +# Beamline: lattice elements and segments +############################################################################### +lattice.elements = drift1 +lattice.nslice = 30 + +drift1.type = drift +drift1.ds = 6.0 + + +############################################################################### +# Algorithms +############################################################################### +algo.particle_shape = 2 +algo.space_charge = true + +amr.n_cell = 40 40 32 +geometry.prob_relative = 1.0 diff --git a/examples/expanding_beam/run_expanding.py b/examples/expanding_beam/run_expanding.py new file mode 100755 index 000000000..72a4b3a24 --- /dev/null +++ b/examples/expanding_beam/run_expanding.py @@ -0,0 +1,59 @@ +#!/usr/bin/env python3 +# +# Copyright 2022 ImpactX contributors +# Authors: Axel Huebl, Chad Mitchell +# License: BSD-3-Clause-LBNL +# +# -*- coding: utf-8 -*- + +import amrex +from impactx import ImpactX, RefPart, distribution, elements + +pp_amr = amrex.ParmParse("amr") +pp_amr.addarr("n_cell", [40, 40, 32]) + +sim = ImpactX() + +# set numerical parameters and IO control +sim.set_particle_shape(2) # B-spline order +sim.set_space_charge(True) +sim.dynamic_size = True +sim.prob_relative = 1.0 + +# beam diagnostics +# sim.set_diagnostics(False) # benchmarking +sim.set_slice_step_diagnostics(False) + +# domain decomposition & space charge mesh +sim.init_grids() + +# load a 2 GeV electron beam with an initial +# unnormalized rms emittance of 2 nm +energy_MeV = 250 # 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(0.510998950).set_energy_MeV(energy_MeV) + +# particle bunch +distr = distribution.Waterbag( + sigmaX=4.472135955e-4, + sigmaY=4.472135955e-4, + sigmaT=9.12241869e-7, + sigmaPx=0.0, + sigmaPy=0.0, + sigmaPt=0.0, +) +sim.add_particles(bunch_charge_C, distr, npart) + +# design the accelerator lattice +sim.lattice.append(elements.Drift(ds=6.0, nslice=30)) + +# run simulation +sim.evolve() + +# clean shutdown +del sim +amrex.finalize()