Update analysis script.

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
 )

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,
+)
+