From dfe9b28a5d74b9d916984fb46327a5e900b5a510 Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Wed, 9 Oct 2024 10:37:45 -0700 Subject: [PATCH] BeamMonitor: Enable Filtering by Cycle (Turn) (#713) * Elements: Pass Cycle, Filter openPMD * Examples: Update Programmable Signature For cycle argument * Documentation * Add Example * Update Docstrings Co-authored-by: Chad Mitchell <46825199+cemitch99@users.noreply.github.com> * Rename `cycle_intervals` -> `period_sample_intervals` * Update `Marker` * Rename: `cycle` -> `period` --------- Co-authored-by: Chad Mitchell <46825199+cemitch99@users.noreply.github.com> --- docs/source/usage/examples.rst | 10 ++ docs/source/usage/parameters.rst | 5 + docs/source/usage/python.rst | 3 +- examples/CMakeLists.txt | 17 +++ examples/fodo_channel/README.rst | 74 ++++++++++ examples/fodo_channel/analysis_fodo.py | 105 ++++++++++++++ examples/fodo_channel/input_fodo.in | 60 ++++++++ examples/fodo_channel/plot_fodo.py | 132 ++++++++++++++++++ examples/fodo_channel/run_fodo.py | 71 ++++++++++ .../run_ml_surrogate_15_stage.py | 4 +- src/ImpactX.cpp | 26 ++-- src/initialization/InitElement.cpp | 4 +- src/particles/Push.H | 10 +- src/particles/Push.cpp | 13 +- src/particles/PushAll.H | 2 + src/particles/elements/Empty.H | 3 +- src/particles/elements/Marker.H | 3 +- src/particles/elements/Programmable.H | 6 +- src/particles/elements/Programmable.cpp | 7 +- src/particles/elements/diagnostics/openPMD.H | 9 +- .../elements/diagnostics/openPMD.cpp | 15 +- src/particles/elements/mixin/beamoptic.H | 12 +- src/python/elements.cpp | 17 +-- 23 files changed, 559 insertions(+), 49 deletions(-) create mode 100644 examples/fodo_channel/README.rst create mode 100755 examples/fodo_channel/analysis_fodo.py create mode 100644 examples/fodo_channel/input_fodo.in create mode 100755 examples/fodo_channel/plot_fodo.py create mode 100755 examples/fodo_channel/run_fodo.py diff --git a/docs/source/usage/examples.rst b/docs/source/usage/examples.rst index 426881b7f..100139584 100644 --- a/docs/source/usage/examples.rst +++ b/docs/source/usage/examples.rst @@ -70,6 +70,16 @@ Beam Distributions examples/distgen/README + +Channels & Rings +---------------- + +.. toctree:: + :maxdepth: 1 + + examples/fodo_channel/README.rst + + Lattice Design & Optimization ----------------------------- diff --git a/docs/source/usage/parameters.rst b/docs/source/usage/parameters.rst index d21b1f2ab..7b6a81a70 100644 --- a/docs/source/usage/parameters.rst +++ b/docs/source/usage/parameters.rst @@ -449,6 +449,11 @@ Lattice Elements openPMD `iteration encoding `__: (v)ariable based, (f)ile based, (g)roup based (default) variable based is an `experimental feature with ADIOS2 `__. + * ``.period_sample_intervals`` (``int``, default value: ``1``) + + for periodic lattice, only output every Nth period (turn). + By default, diagnostics are returned every cycle. + * ``.nonlinear_lens_invariants`` (``boolean``, default value: ``false``) Compute and output the invariants H and I within the nonlinear magnetic insert element (see: ``nonlinear_lens``). diff --git a/docs/source/usage/python.rst b/docs/source/usage/python.rst index 948de5a6b..2f004a7b8 100644 --- a/docs/source/usage/python.rst +++ b/docs/source/usage/python.rst @@ -632,7 +632,7 @@ This module provides elements for the accelerator lattice. :param rotation: rotation error in the transverse plane [degrees] :param name: an optional name for the element -.. py:class:: impactx.elements.BeamMonitor(name, backend="default", encoding="g") +.. py:class:: impactx.elements.BeamMonitor(name, backend="default", encoding="g", period_sample_intervals=1) A beam monitor, writing all beam particles at fixed ``s`` to openPMD files. @@ -649,6 +649,7 @@ This module provides elements for the accelerator lattice. :param name: name of the series :param backend: I/O backend, e.g., ``bp``, ``h5``, ``json`` :param encoding: openPMD iteration encoding: (v)ariable based, (f)ile based, (g)roup based (default) + :param period_sample_intervals: for periodic lattice, only output every Nth period (turn) .. py:property:: name diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt index d33221240..528eb8290 100644 --- a/examples/CMakeLists.txt +++ b/examples/CMakeLists.txt @@ -219,6 +219,23 @@ add_impactx_test(FODO.py.MPI examples/fodo/plot_fodo.py ) +# FODO Channel ################################################################ +# +add_impactx_test(FODO_channel + examples/fodo_channel/input_fodo.in + OFF # ImpactX MPI-parallel + examples/fodo_channel/analysis_fodo.py + examples/fodo_channel/plot_fodo.py +) +add_impactx_test(FODO_channel.py + examples/fodo_channel/run_fodo.py + OFF # ImpactX MPI-parallel + examples/fodo_channel/analysis_fodo.py + examples/fodo_channel/plot_fodo.py +) +label_impactx_test(FODO_channel slow) +label_impactx_test(FODO_channel.py slow) + # Chicane ##################################################################### # add_impactx_test(chicane diff --git a/examples/fodo_channel/README.rst b/examples/fodo_channel/README.rst new file mode 100644 index 000000000..6282e8de0 --- /dev/null +++ b/examples/fodo_channel/README.rst @@ -0,0 +1,74 @@ +.. _examples-fodo-channel: + +FODO Channel +============ + +A 300m channel of 100 stable FODO cells (3m each) with a zero-current phase advance of 67.8 degrees. + +The matched Twiss parameters at entry are: + +* :math:`\beta_\mathrm{x} = 2.82161941` m +* :math:`\alpha_\mathrm{x} = -1.59050035` +* :math:`\beta_\mathrm{y} = 2.82161941` m +* :math:`\alpha_\mathrm{y} = 1.59050035` + +We use a 2 GeV electron beam with initial unnormalized rms emittance of 2 nm. + +The second moments of the particle distribution after the FODO cell should coincide with the second moments of the particle distribution before the FODO cell, to within the level expected due to noise due to statistical sampling. + +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. +This test also demonstrates the ``period_sample_intervals`` capability of our beam monitor diagnostics, only creating output every 10th FODO cell + + +Run +--- + +This example can be run **either** as: + +* **Python** script: ``python3 run_fodo.py`` or +* ImpactX **executable** using an input file: ``impactx input_fodo.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_fodo.py + :language: python3 + :caption: You can copy this file from ``examples/fodo/run_fodo.py``. + + .. tab-item:: Executable: Input File + + .. literalinclude:: input_fodo.in + :language: ini + :caption: You can copy this file from ``examples/fodo/input_fodo.in``. + + +Analyze +------- + +We run the following script to analyze correctness: + +.. dropdown:: Script ``analysis_fodo.py`` + + .. literalinclude:: analysis_fodo.py + :language: python3 + :caption: You can copy this file from ``examples/fodo/analysis_fodo.py``. + + +Visualize +--------- + +You can run the following script to visualize the beam evolution over time: + +.. dropdown:: Script ``plot_fodo.py`` + + .. literalinclude:: plot_fodo.py + :language: python3 + :caption: You can copy this file from ``examples/fodo/plot_fodo.py``. + +.. figure:: https://gist.githubusercontent.com/ax3l/8ae7dcb9e07c361e002fa56d6b16cb16/raw/cc952670bb946cd7a62282bc7aa3f03f3d5faa16/fodo_channel.png + :alt: preserved emittance in the FODO channel. + + FODO transverse emittance evolution (preserved) diff --git a/examples/fodo_channel/analysis_fodo.py b/examples/fodo_channel/analysis_fodo.py new file mode 100755 index 000000000..127a4af0a --- /dev/null +++ b/examples/fodo_channel/analysis_fodo.py @@ -0,0 +1,105 @@ +#!/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 = series.iterations[1].particles["beam"].to_df() +final = series.iterations[last_step].particles["beam"].to_df() + +# compare number of particles +num_particles = 10000 +assert num_particles == len(initial) +assert num_particles == len(final) + +# compare beamline length: 300m +assert np.isclose( + 300.0, series.iterations[last_step].particles["beam"].get_attribute("z_ref") +) +# compare beam monitor outputs: 10 (every 10th FODO element + 1) +assert len(series.iterations) == 11 + +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 = 2.2 * 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], + [ + 7.5451170454175073e-005, + 7.5441588239210947e-005, + 9.9775878164077539e-004, + 1.9959540393751392e-009, + 2.0175015289132990e-009, + 2.0013820193294972e-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 = 2.2 * 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], + [ + 7.4790118496224206e-005, + 7.5357525169680140e-005, + 9.9775879288128088e-004, + 1.9959539836392703e-009, + 2.0175014668882125e-009, + 2.0013820380883801e-006, + ], + rtol=rtol, + atol=atol, +) diff --git a/examples/fodo_channel/input_fodo.in b/examples/fodo_channel/input_fodo.in new file mode 100644 index 000000000..e782be7f8 --- /dev/null +++ b/examples/fodo_channel/input_fodo.in @@ -0,0 +1,60 @@ +############################################################################### +# Particle Beam(s) +############################################################################### +beam.npart = 10000 +beam.units = static +beam.kin_energy = 2.0e3 +beam.charge = 1.0e-9 +beam.particle = electron +beam.distribution = waterbag_from_twiss +beam.alphaX = -1.5905003499999992 +beam.alphaY = 1.5905003499999992 +beam.alphaT = 0.0 +beam.betaX = 2.8216194100262637 +beam.betaY = 2.8216194100262637 +beam.betaT = 0.5 +beam.emittX = 2e-09 +beam.emittY = 2e-09 +beam.emittT = 2e-06 + + +############################################################################### +# Beamline: lattice elements and segments +############################################################################### +lattice.elements = monitor drift1 quad1 drift2 quad2 drift3 +lattice.nslice = 5 +lattice.periods = 101 # FODO channel of 101 periods + +monitor.type = beam_monitor +monitor.period_sample_intervals = 10 +monitor.backend = h5 + +drift1.type = drift +drift1.ds = 0.25 + +quad1.type = quad +quad1.ds = 1.0 +quad1.k = 1.0 + +drift2.type = drift +drift2.ds = 0.5 + +quad2.type = quad +quad2.ds = 1.0 +quad2.k = -1.0 + +drift3.type = drift +drift3.ds = 0.25 + + +############################################################################### +# Algorithms +############################################################################### +algo.particle_shape = 2 +algo.space_charge = false + + +############################################################################### +# Diagnostics +############################################################################### +diag.slice_step_diagnostics = false diff --git a/examples/fodo_channel/plot_fodo.py b/examples/fodo_channel/plot_fodo.py new file mode 100755 index 000000000..5ec2b6715 --- /dev/null +++ b/examples/fodo_channel/plot_fodo.py @@ -0,0 +1,132 @@ +#!/usr/bin/env python3 +# +# Copyright 2022-2023 ImpactX contributors +# Authors: Axel Huebl, Chad Mitchell +# License: BSD-3-Clause-LBNL +# + +import argparse +import glob +import re + +import matplotlib.pyplot as plt +import openpmd_api as io +import pandas as pd +from matplotlib.ticker import MaxNLocator +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) + + +def read_file(file_pattern): + for filename in glob.glob(file_pattern): + df = pd.read_csv(filename, delimiter=r"\s+") + if "step" not in df.columns: + step = int(re.findall(r"[0-9]+", filename)[0]) + df["step"] = step + yield df + + +def read_time_series(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( + read_file(file_pattern), + axis=0, + ignore_index=True, + ) # .set_index('id') + + +# options to run this script +parser = argparse.ArgumentParser(description="Plot the FODO benchmark.") +parser.add_argument( + "--save-png", action="store_true", help="non-interactive run: save to PNGs" +) +args = parser.parse_args() + + +# initial/final beam +series = io.Series("diags/openPMD/monitor.h5", io.Access.read_only) +last_step = list(series.iterations)[-1] +initial = series.iterations[1].particles["beam"].to_df() +final = series.iterations[last_step].particles["beam"].to_df() +ref_particle = read_time_series("diags/ref_particle.*") + +# scaling to units +millimeter = 1.0e3 # m->mm +mrad = 1.0e3 # ImpactX uses "static units": momenta are normalized by the magnitude of the momentum of the reference particle p0: px/p0 (rad) +# mm_mrad = 1.e6 +nm_rad = 1.0e9 + + +# select a single particle by id +# particle_42 = beam[beam["id"] == 42] +# print(particle_42) + + +# steps & corresponding z +steps = list(series.iterations) + +z = list(map(lambda step: ref_particle[ref_particle["step"] == step].z.values, steps)) +# print(f"z={z}") + + +# beam transversal size & emittance over steps +moments = list( + map( + lambda step: ( + step, + get_moments(series.iterations[step].particles["beam"].to_df()), + ), + steps, + ) +) +# print(moments) +emittance_x = list(map(lambda step_val: step_val[1][3] * nm_rad, moments)) +emittance_y = list(map(lambda step_val: step_val[1][4] * nm_rad, moments)) + +# print(sigx, sigy) + + +# print beam transversal size over steps +f = plt.figure(figsize=(9, 4.8)) +ax1 = f.gca() +im_emittance_x = ax1.plot(z, emittance_x, ".-", label=r"$\epsilon_x$") +im_emittance_y = ax1.plot(z, emittance_y, ".-", label=r"$\epsilon_y$") + +ax1.legend() +ax1.set_xlabel(r"$z$ [m]") +ax1.set_ylabel(r"$\epsilon_{x,y}$ [nm]") +ax1.set_ylim([1.98, 2.03]) +ax1.xaxis.set_major_locator(MaxNLocator(integer=True)) +plt.tight_layout() +if args.save_png: + plt.savefig("fodo_lambda.png") +else: + plt.show() diff --git a/examples/fodo_channel/run_fodo.py b/examples/fodo_channel/run_fodo.py new file mode 100755 index 000000000..9eaf171f5 --- /dev/null +++ b/examples/fodo_channel/run_fodo.py @@ -0,0 +1,71 @@ +#!/usr/bin/env python3 +# +# Copyright 2022-2024 ImpactX contributors +# Authors: Axel Huebl, Chad Mitchell, Marco Garten +# License: BSD-3-Clause-LBNL +# +# -*- coding: utf-8 -*- + +from impactx import ImpactX, distribution, elements, twiss + +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() + +# load a 2 GeV electron beam with an initial +# unnormalized rms emittance of 2 nm +kin_energy_MeV = 2.0e3 # 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_kin_energy_MeV(kin_energy_MeV) + +# particle bunch +distr = distribution.Waterbag( + **twiss( + beta_x=2.8216194100262637, + beta_y=2.8216194100262637, + beta_t=0.5, + emitt_x=2e-09, + emitt_y=2e-09, + emitt_t=2e-06, + alpha_x=-1.5905003499999992, + alpha_y=1.5905003499999992, + alpha_t=0.0, + ) +) +sim.add_particles(bunch_charge_C, distr, npart) + +# add beam diagnostics +monitor = elements.BeamMonitor("monitor", backend="h5", period_sample_intervals=10) + +# design the accelerator lattice) +ns = 5 # number of slices per ds in the element +fodo = [ + monitor, + elements.Drift(ds=0.25, nslice=ns), + elements.Quad(ds=1.0, k=1.0, nslice=ns), + elements.Drift(ds=0.5, nslice=ns), + elements.Quad(ds=1.0, k=-1.0, nslice=ns), + elements.Drift(ds=0.25, nslice=ns), +] +# assign a fodo segment +sim.lattice.extend(fodo) + +# FODO channel of 101 periods +sim.periods = 101 + +# run simulation +sim.evolve() + +# clean shutdown +sim.finalize() diff --git a/examples/pytorch_surrogate_model/run_ml_surrogate_15_stage.py b/examples/pytorch_surrogate_model/run_ml_surrogate_15_stage.py index 94c205c80..004f79bdb 100644 --- a/examples/pytorch_surrogate_model/run_ml_surrogate_15_stage.py +++ b/examples/pytorch_surrogate_model/run_ml_surrogate_15_stage.py @@ -169,7 +169,7 @@ def __init__(self, stage_i, surrogate_model, surrogate_length, stage_start): self.push = self.surrogate_push self.ds = surrogate_length - def surrogate_push(self, pc, step): + def surrogate_push(self, pc, step, period): ref_part = pc.ref_particle() ref_z_i = ref_part.z ref_z_i_LPA = ref_z_i - self.stage_start @@ -304,7 +304,7 @@ def __init__(self, sim, stage_i, lattice_index, x_or_y): self.x_or_y = x_or_y self.push = self.set_lens - def set_lens(self, pc, step): + def set_lens(self, pc, step, period): # get envelope parameters rbc = pc.reduced_beam_characteristics() alpha = rbc[f"alpha_{self.x_or_y}"] diff --git a/src/ImpactX.cpp b/src/ImpactX.cpp index 7466356af..ed08201f3 100644 --- a/src/ImpactX.cpp +++ b/src/ImpactX.cpp @@ -137,7 +137,7 @@ namespace impactx { // a global step for diagnostics including space charge slice steps in elements // before we start the evolve loop, we are in "step 0" (initial state) - int global_step = 0; + int step = 0; // check typos in inputs after step 1 bool early_params_checked = false; @@ -158,7 +158,7 @@ namespace impactx { diagnostics::DiagnosticOutput(*amr_data->m_particle_container, diagnostics::OutputType::PrintRefParticle, "diags/ref_particle", - global_step); + step); // print the initial values of reduced beam characteristics diagnostics::DiagnosticOutput(*amr_data->m_particle_container, @@ -181,10 +181,10 @@ namespace impactx { } // periods through the lattice - int periods = 1; - amrex::ParmParse("lattice").queryAdd("periods", periods); + int num_periods = 1; + amrex::ParmParse("lattice").queryAdd("periods", num_periods); - for (int cycle=0; cycle < periods; ++cycle) { + for (int period=0; period < num_periods; ++period) { // loop over all beamline elements for (auto &element_variant: m_lattice) { // update element edge of the reference particle @@ -201,9 +201,9 @@ namespace impactx { // sub-steps for space charge within the element for (int slice_step = 0; slice_step < nslice; ++slice_step) { BL_PROFILE("ImpactX::evolve::slice_step"); - global_step++; + step++; if (verbose > 0) { - amrex::Print() << " ++++ Starting global_step=" << global_step + amrex::Print() << " ++++ Starting step=" << step << " slice_step=" << slice_step << "\n"; } @@ -263,7 +263,7 @@ namespace impactx { // assuming that the distribution did not change // push all particles with external maps - Push(*amr_data->m_particle_container, element_variant, global_step); + Push(*amr_data->m_particle_container, element_variant, step, period); // move "lost" particles to another particle container collect_lost_particles(*amr_data->m_particle_container); @@ -282,14 +282,14 @@ namespace impactx { diagnostics::DiagnosticOutput(*amr_data->m_particle_container, diagnostics::OutputType::PrintRefParticle, "diags/ref_particle", - global_step, + step, true); // print slice step reduced beam characteristics to file diagnostics::DiagnosticOutput(*amr_data->m_particle_container, diagnostics::OutputType::PrintReducedBeamCharacteristics, "diags/reduced_beam_characteristics", - global_step, + step, true); } @@ -308,13 +308,13 @@ namespace impactx { diagnostics::DiagnosticOutput(*amr_data->m_particle_container, diagnostics::OutputType::PrintRefParticle, "diags/ref_particle_final", - global_step); + step); // print the final values of the reduced beam characteristics diagnostics::DiagnosticOutput(*amr_data->m_particle_container, diagnostics::OutputType::PrintReducedBeamCharacteristics, "diags/reduced_beam_characteristics_final", - global_step); + step); // output particles lost in apertures if (amr_data->m_particles_lost->TotalNumberOfParticles() > 0) @@ -323,7 +323,7 @@ namespace impactx { pp_diag.queryAdd("backend", openpmd_backend); diagnostics::BeamMonitor output_lost("particles_lost", openpmd_backend, "g"); - output_lost(*amr_data->m_particles_lost, 0); + output_lost(*amr_data->m_particles_lost, 0, 0); output_lost.finalize(); } } diff --git a/src/initialization/InitElement.cpp b/src/initialization/InitElement.cpp index dd770e51f..71ec7e4aa 100644 --- a/src/initialization/InitElement.cpp +++ b/src/initialization/InitElement.cpp @@ -395,6 +395,8 @@ namespace detail pp_element.queryAdd("backend", openpmd_backend); std::string openpmd_encoding{"g"}; pp_element.queryAdd("encoding", openpmd_encoding); + int period_sample_intervals = 1; + pp_element.queryAdd("period_sample_intervals", period_sample_intervals); // optional: add and calculate additional particle properties // property: nonlinear lens invariants @@ -412,7 +414,7 @@ namespace detail pp_element.queryAdd("cn", cn); } - m_lattice.emplace_back(diagnostics::BeamMonitor(openpmd_name, openpmd_backend, openpmd_encoding)); + m_lattice.emplace_back(diagnostics::BeamMonitor(openpmd_name, openpmd_backend, openpmd_encoding, period_sample_intervals)); } else if (element_type == "line") { // Parse the lattice elements for the sub-lattice in the line diff --git a/src/particles/Push.H b/src/particles/Push.H index 7665df7a6..268a65e28 100644 --- a/src/particles/Push.H +++ b/src/particles/Push.H @@ -23,10 +23,14 @@ namespace impactx * @param[inout] pc container of the particles to push * @param[inout] element_variant a single element to push the particles through * @param[in] step global step for diagnostics + * @param[in] period for periodic lattices, this is the current period (turn or cycle) */ - void Push (ImpactXParticleContainer & pc, - KnownElements & element_variant, - int step); + void Push ( + ImpactXParticleContainer & pc, + KnownElements & element_variant, + int step, + int period + ); } // namespace impactx diff --git a/src/particles/Push.cpp b/src/particles/Push.cpp index 5769dd5d8..54f8cd8b4 100644 --- a/src/particles/Push.cpp +++ b/src/particles/Push.cpp @@ -16,17 +16,20 @@ namespace impactx { - void Push (ImpactXParticleContainer & pc, - KnownElements & element_variant, - int step) + void Push ( + ImpactXParticleContainer & pc, + KnownElements & element_variant, + int step, + int period + ) { // here we just access the element by its respective type - std::visit([&pc, step](auto&& element) + std::visit([&pc, step, period](auto&& element) { BL_PROFILE("impactx::Push"); // push reference particle & all particles - element(pc, step); + element(pc, step, period); }, element_variant); } diff --git a/src/particles/PushAll.H b/src/particles/PushAll.H index ab0bacd65..3fec0f3bc 100644 --- a/src/particles/PushAll.H +++ b/src/particles/PushAll.H @@ -26,6 +26,7 @@ namespace impactx * @param[in,out] pc particle container to push * @param[in,out] element the beamline element * @param[in] step global step for diagnostics + * @param[in] period for periodic lattices, this is the current period (turn or cycle) * @param[in] omp_parallel allow threading via OpenMP for the particle iterator loop (note: if OMP backend is active) */ template @@ -33,6 +34,7 @@ namespace impactx ImpactXParticleContainer & pc, T_Element & element, [[maybe_unused]] int step, + [[maybe_unused]] int period, [[maybe_unused]] bool omp_parallel = true ) { diff --git a/src/particles/elements/Empty.H b/src/particles/elements/Empty.H index cf0dba130..9ce3eece8 100644 --- a/src/particles/elements/Empty.H +++ b/src/particles/elements/Empty.H @@ -37,7 +37,8 @@ namespace impactx /** Push all particles - nothing to do here */ void operator() ( ImpactXParticleContainer & /* pc */, - int /* step */ + int /* step */, + int /* period */ ) { // nothing to do } diff --git a/src/particles/elements/Marker.H b/src/particles/elements/Marker.H index 46eff72a4..3291aff79 100644 --- a/src/particles/elements/Marker.H +++ b/src/particles/elements/Marker.H @@ -41,7 +41,8 @@ namespace impactx /** Push all particles - nothing to do here */ void operator() ( ImpactXParticleContainer & /* pc */, - int /* step */ + int /* step */, + int /* period */ ) { // nothing to do } diff --git a/src/particles/elements/Programmable.H b/src/particles/elements/Programmable.H index 31e8b5fc3..89b6afa8e 100644 --- a/src/particles/elements/Programmable.H +++ b/src/particles/elements/Programmable.H @@ -46,10 +46,12 @@ namespace impactx * * @param[in,out] pc particle container to push * @param[in] step global step for diagnostics + * @param[in] period for periodic lattices, this is the current period (turn or cycle) */ void operator() ( ImpactXParticleContainer & pc, - int step + int step, + int period ) const; /** Push all particles relative to the reference particle */ @@ -103,7 +105,7 @@ namespace impactx */ bool m_threadsafe = false; - std::function m_push; //! hook for push of whole container + std::function m_push; //! hook for push of whole container (pc, step, period) std::function m_beam_particles; //! hook for beam particles std::function m_ref_particle; //! hook for reference particle std::function m_finalize; //! hook for finalize cleanup diff --git a/src/particles/elements/Programmable.cpp b/src/particles/elements/Programmable.cpp index 3f73ac413..43f37aa73 100644 --- a/src/particles/elements/Programmable.cpp +++ b/src/particles/elements/Programmable.cpp @@ -19,16 +19,17 @@ namespace impactx void Programmable::operator() ( ImpactXParticleContainer & pc, - int step + int step, + int period ) const { if (m_push == nullptr) { // TODO: print if verbose mode is set - push_all(pc, *this, step, m_threadsafe); + push_all(pc, *this, step, period, m_threadsafe); } else { BL_PROFILE("impactx::Push::Programmable"); - m_push(&pc, step); + m_push(&pc, step, period); } } diff --git a/src/particles/elements/diagnostics/openPMD.H b/src/particles/elements/diagnostics/openPMD.H index 5b6a6014f..f5ee0d4a8 100644 --- a/src/particles/elements/diagnostics/openPMD.H +++ b/src/particles/elements/diagnostics/openPMD.H @@ -76,8 +76,9 @@ namespace detail * @param series_name name of the data series, usually the element name * @param backend file format backend for openPMD, e.g., "bp" or "h5" * @param encoding openPMD iteration encoding: "v"ariable based, "f"ile based, "g"roup based (default) + * @param period_sample_intervals for periodic lattice, only output every Nth period (turn) */ - BeamMonitor (std::string series_name, std::string backend="default", std::string encoding="g"); + BeamMonitor (std::string series_name, std::string backend="default", std::string encoding="g", int period_sample_intervals=1); BeamMonitor (BeamMonitor const & other) = default; BeamMonitor (BeamMonitor && other) = default; @@ -108,10 +109,12 @@ namespace detail * * @param[in,out] pc particle container to push * @param[in] step global step for diagnostics + * @param[in] period for periodic lattices, this is the current period (turn or cycle) */ void operator() ( ImpactXParticleContainer & pc, - int step + int step, + int period ); /** Write a tile of particles @@ -156,6 +159,8 @@ namespace detail int m_file_min_digits = 6; //! minimum number of digits to iteration number in file name + int m_period_sample_intervals = 1; //! only output every Nth period (turn or cycle) of periodic lattices + /** This rank's offset in the MPI-global particle array, by level * * This MUST be updated by prepare() before the next step's output. diff --git a/src/particles/elements/diagnostics/openPMD.cpp b/src/particles/elements/diagnostics/openPMD.cpp index 163cbbe11..8b215d3e2 100644 --- a/src/particles/elements/diagnostics/openPMD.cpp +++ b/src/particles/elements/diagnostics/openPMD.cpp @@ -154,8 +154,8 @@ namespace detail #endif // ImpactX_USE_OPENPMD } - BeamMonitor::BeamMonitor (std::string series_name, std::string backend, std::string encoding) : - m_series_name(std::move(series_name)), m_OpenPMDFileType(std::move(backend)) + BeamMonitor::BeamMonitor (std::string series_name, std::string backend, std::string encoding, int period_sample_intervals) : + m_series_name(std::move(series_name)), m_OpenPMDFileType(std::move(backend)), m_period_sample_intervals(period_sample_intervals) { #ifdef ImpactX_USE_OPENPMD // pick first available backend if default is chosen @@ -179,8 +179,10 @@ namespace detail else if ( "f" == encoding ) series_encoding = openPMD::IterationEncoding::fileBased; - // legacy options from other diagnostics amrex::ParmParse pp_diag("diag"); + // turn filter + pp_diag.queryAdd("period_sample_intervals", m_period_sample_intervals); + // legacy options from other diagnostics pp_diag.queryAdd("file_min_digits", m_file_min_digits); // Ensure m_series is the same for the same names. @@ -311,9 +313,14 @@ namespace detail void BeamMonitor::operator() ( ImpactXParticleContainer & pc, - int step + int step, + int period ) { + // filter out this turn? + if (period % m_period_sample_intervals != 0) + return; + #ifdef ImpactX_USE_OPENPMD std::string profile_name = "impactx::Push::" + std::string(BeamMonitor::type); BL_PROFILE(profile_name); diff --git a/src/particles/elements/mixin/beamoptic.H b/src/particles/elements/mixin/beamoptic.H index fd58ac600..0616b5374 100644 --- a/src/particles/elements/mixin/beamoptic.H +++ b/src/particles/elements/mixin/beamoptic.H @@ -147,10 +147,16 @@ namespace detail template struct BeamOptic { - /** Push first the reference particle, then all other particles */ + /** Push first the reference particle, then all other particles + * + * @param[inout] pc container of the particles to push + * @param[in] step global step for diagnostics + * @param[in] period for periodic lattices, this is the current period (turn or cycle) + */ void operator() ( ImpactXParticleContainer & pc, - int step + int step, + int period ) { static_assert( @@ -159,7 +165,7 @@ namespace detail ); T_Element& element = *static_cast(this); - push_all(pc, element, step); + push_all(pc, element, step, period); } /** This pushes the particles on a particle iterator tile or box. diff --git a/src/python/elements.cpp b/src/python/elements.cpp index 7d4a33f76..b6403e71b 100644 --- a/src/python/elements.cpp +++ b/src/python/elements.cpp @@ -57,10 +57,10 @@ namespace using Element = typename T_PyClass::type; // py::class cl.def("push", - [](Element & el, ImpactXParticleContainer & pc, int step) { - el(pc, step); + [](Element & el, ImpactXParticleContainer & pc, int step, int period) { + el(pc, step, period); }, - py::arg("pc"), py::arg("step")=0, + py::arg("pc"), py::arg("step")=0, py::arg("period")=0, "Push first the reference particle, then all other particles." ); } @@ -214,10 +214,11 @@ void init_elements(py::module& m) py::class_ py_BeamMonitor(me, "BeamMonitor"); py_BeamMonitor - .def(py::init(), + .def(py::init(), py::arg("name"), py::arg("backend") = "default", py::arg("encoding") = "g", + py::arg("period_sample_intervals") = 1, "This element writes the particle beam out to openPMD data." ) .def_property_readonly("name", @@ -914,9 +915,9 @@ void init_elements(py::module& m) .def_property("push", [](Programmable & p) { return p.m_push; }, [](Programmable & p, - std::function new_hook + std::function new_hook ) { p.m_push = std::move(new_hook); }, - "hook for push of whole container (pc, step)" + "hook for push of whole container (pc, step, period)" ) .def_property("beam_particles", [](Programmable & p) { return p.m_beam_particles; }, @@ -1480,9 +1481,9 @@ void init_elements(py::module& m) register_beamoptics_push(py_TaperedPL); - // free-standing push function + // freestanding push function m.def("push", &Push, - py::arg("pc"), py::arg("element"), py::arg("step")=0, + py::arg("pc"), py::arg("element"), py::arg("step")=0, py::arg("period")=0, "Push particles through an element" );