diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt index 0759ea70f..3333c0165 100644 --- a/examples/CMakeLists.txt +++ b/examples/CMakeLists.txt @@ -1160,12 +1160,12 @@ label_impactx_test(pytorch_surrogate_model slow) 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 + 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.py + examples/iota_lens/analysis_iotalens_sdep_aperture.py OFF # no plot script yet ) diff --git a/examples/iota_lens/analysis_iotalens_sdep_aperture.py b/examples/iota_lens/analysis_iotalens_sdep_aperture.py new file mode 100755 index 000000000..1b232acec --- /dev/null +++ b/examples/iota_lens/analysis_iotalens_sdep_aperture.py @@ -0,0 +1,117 @@ +#!/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.400000e+00, + ], + 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, +) +