Add Time-Averaged Field Diagnostics (#5285)

This PR adds time-averaged field diagnostics to the WarpX output.

To-do:
- [x] code
- [x] docs
- [x] tests
- [x] example

Follow-up PRs:
- meta-data
- make compatible with adaptive time stepping

This PR is based on work performed during the *2024 WarpX Refactoring
Hackathon* and was created together with @RevathiJambunathan.

Successfully merging this pull request may close #5165.

---------

Co-authored-by: RevathiJambunathan <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Edoardo Zoni <[email protected]>
Co-authored-by: Edoardo Zoni <[email protected]>

Author: 4 people

Docs/source/usage/parameters.rst

@@ -2633,10 +2633,11 @@ Diagnostics and output
 In-situ visualization
 ^^^^^^^^^^^^^^^^^^^^^
 
-WarpX has four types of diagnostics:
-``FullDiagnostics`` consist in dumps of fields and particles at given iterations,
-``BackTransformedDiagnostics`` are used when running a simulation in a boosted frame, to reconstruct output data to the lab frame,
-``BoundaryScrapingDiagnostics`` are used to collect the particles that are absorbed at the boundary, throughout the simulation, and
+WarpX has five types of diagnostics:
+``Full`` diagnostics consist in dumps of fields and particles at given iterations,
+``TimeAveraged`` diagnostics only allow field data, which they output after averaging over a period of time,
+``BackTransformed`` diagnostics are used when running a simulation in a boosted frame, to reconstruct output data to the lab frame,
+``BoundaryScraping`` diagnostics are used to collect the particles that are absorbed at the boundary, throughout the simulation, and
 ``ReducedDiags`` allow the user to compute some reduced quantity (particle temperature, max of a field) and write a small amount of data to text files.
 Similar to what is done for physical species, WarpX has a class Diagnostics that allows users to initialize different diagnostics, each of them with different fields, resolution and period.
 This currently applies to standard diagnostics, but should be extended to back-transformed diagnostics and reduced diagnostics (and others) in a near future.
@@ -2882,12 +2883,58 @@ In-situ capabilities can be used by turning on Sensei or Ascent (provided they a
 * ``warpx.mffile_nstreams`` (`int`) optional (default `4`)
     Limit the number of concurrent readers per file.
 
+
+.. _running-cpp-parameters-diagnostics-timeavg:
+
+Time-Averaged Diagnostics
+^^^^^^^^^^^^^^^^^^^^^^^^^
+
+``TimeAveraged`` diagnostics are a special type of ``Full`` diagnostics that allows for the output of time-averaged field data.
+This type of diagnostics can be created using ``<diag_name>.diag_type = TimeAveraged``.
+We support only field data and related options from the list at `Full Diagnostics`_.
+
+.. note::
+
+    As with ``Full`` diagnostics, ``TimeAveraged`` diagnostics output the initial **instantaneous** conditions of the selected fields on step 0 (unless more specific output intervals exclude output for step 0).
+
+In addition, ``TimeAveraged`` diagnostic options include:
+
+* ``<diag_name>.time_average_mode`` (`string`, default `none`)
+    Describes the operating mode for time averaged field output.
+
+    * ``none`` for no averaging (instantaneous fields)
+
+    * ``fixed_start`` for a diagnostic that averages all fields between the current output step and a fixed point in time
+
+    * ``dynamic_start`` for a constant averaging period and output at different points in time (non-overlapping)
+
+    .. note::
+
+        To enable time-averaged field output with intervals tightly spaced enough for overlapping averaging periods,
+        please create additional instances of ``TimeAveraged`` diagnostics.
+
+* ``<diag_name>.average_period_steps`` (`int`)
+    Configures the number of time steps in an averaging period.
+    Set this only in the ``dynamic_start`` mode and only if ``average_period_time`` has not already been set.
+    Will be ignored in the ``fixed_start`` mode (with warning).
+
+* ``<diag_name>.average_period_time`` (`float`, in seconds)
+    Configures the time (SI units) in an averaging period.
+    Set this only in the ``dynamic_start`` mode and only if ``average_period_steps`` has not already been set.
+    Will be ignored in the ``fixed_start`` mode (with warning).
+
+* ``<diag_name>.average_start_step`` (`int`)
+    Configures the time step at which time-averaging begins.
+    Set this only in the ``fixed_start`` mode.
+    Will be ignored in the ``dynamic_start`` mode (with warning).
+
 .. _running-cpp-parameters-diagnostics-btd:
 
 BackTransformed Diagnostics
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
 ``BackTransformed`` diag type are used when running a simulation in a boosted frame, to reconstruct output data to the lab frame. This option can be set using ``<diag_name>.diag_type = BackTransformed``. We support the following list of options from `Full Diagnostics`_
+
     ``<diag_name>.format``, ``<diag_name>.openpmd_backend``, ``<diag_name>.dump_rz_modes``, ``<diag_name>.file_prefix``, ``<diag_name>.diag_lo``, ``<diag_name>.diag_hi``, ``<diag_name>.write_species``, ``<diag_name>.species``.
 
     Additional options for this diagnostic include:

Docs/source/usage/python.rst

@@ -114,6 +114,8 @@ Diagnostics
 
 .. autoclass:: pywarpx.picmi.FieldDiagnostic
 
+.. autoclass:: pywarpx.picmi.TimeAveragedFieldDiagnostic
+
 .. autoclass:: pywarpx.picmi.ElectrostaticFieldDiagnostic
 
 .. autoclass:: pywarpx.picmi.Checkpoint

Examples/Physics_applications/laser_ion/CMakeLists.txt

@@ -6,8 +6,8 @@ add_warpx_test(
     2  # dims
     2  # nprocs
     inputs_test_2d_laser_ion_acc  # inputs
-    analysis_default_openpmd_regression.py  # analysis
-    diags/diag1/  # output
+    analysis_test_laser_ion.py  # analysis
+    diags/diagInst/  # output
     OFF  # dependency
 )
 
@@ -16,7 +16,7 @@ add_warpx_test(
     2  # dims
     2  # nprocs
     inputs_test_2d_laser_ion_acc_picmi.py  # inputs
-    analysis_default_openpmd_regression.py  # analysis
-    diags/diag1/  # output
+    analysis_test_laser_ion.py  # analysis
+    diags/diagInst/  # output
     OFF  # dependency
 )

Examples/Physics_applications/laser_ion/README.rst

@@ -87,7 +87,7 @@ Visualize
    :alt: Particle densities for electrons (top), protons (middle), and electrons again in logarithmic scale (bottom).
    :width: 80%
 
-    Particle densities for electrons (top), protons (middle), and electrons again in logarithmic scale (bottom).
+   Particle densities for electrons (top), protons (middle), and electrons again in logarithmic scale (bottom).
 
 Particle density output illustrates the evolution of the target in time and space.
 Logarithmic scales can help to identify where the target becomes transparent for the laser pulse (bottom panel in :numref:`fig-tnsa-densities` ).

Examples/Physics_applications/laser_ion/analysis_default_openpmd_regression.py

@@ -0,0 +1,85 @@
+#!/usr/bin/env python3
+
+import os
+import sys
+
+import numpy as np
+import openpmd_api as io
+
+sys.path.insert(1, "../../../../warpx/Regression/Checksum/")
+from checksumAPI import evaluate_checksum
+
+
+def load_field_from_iteration(
+    series, iteration: int, field: str, coord: str = None
+) -> np.ndarray:
+    """Load iteration of field data from file."""
+
+    it = series.iterations[iteration]
+    field_obj = it.meshes[f"{field}"]
+
+    if field_obj.scalar:
+        field_data = field_obj[io.Mesh_Record_Component.SCALAR].load_chunk()
+    elif coord in [item[0] for item in list(field_obj.items())]:
+        field_data = field_obj[coord].load_chunk()
+    else:
+        raise Exception(
+            f"Specified coordinate: f{coord} is not available for field: f{field}."
+        )
+    series.flush()
+
+    return field_data
+
+
+def compare_time_avg_with_instantaneous_diags(dir_inst: str, dir_avg: str):
+    """Compare instantaneous data (multiple iterations averaged in post-processing) with in-situ averaged data."""
+
+    field = "E"
+    coord = "z"
+    avg_period_steps = 5
+    avg_output_step = 100
+
+    path_tpl_inst = f"{dir_inst}/openpmd_%T.h5"
+    path_tpl_avg = f"{dir_avg}/openpmd_%T.h5"
+
+    si = io.Series(path_tpl_inst, io.Access.read_only)
+    sa = io.Series(path_tpl_avg, io.Access.read_only)
+
+    ii0 = si.iterations[0]
+    fi0 = ii0.meshes[field][coord]
+    shape = fi0.shape
+
+    data_inst = np.zeros(shape)
+
+    for i in np.arange(avg_output_step - avg_period_steps + 1, avg_output_step + 1):
+        data_inst += load_field_from_iteration(si, i, field, coord)
+
+    data_inst = data_inst / avg_period_steps
+
+    data_avg = load_field_from_iteration(sa, avg_output_step, field, coord)
+
+    # Compare the data
+    if np.allclose(data_inst, data_avg, rtol=1e-12):
+        print("Test passed: actual data is close to expected data.")
+    else:
+        print("Test failed: actual data is not close to expected data.")
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    # NOTE: works only in the example directory due to relative path import
+    # compare checksums
+    evaluate_checksum(
+        test_name=os.path.split(os.getcwd())[1],
+        output_file=sys.argv[1],
+        output_format="openpmd",
+    )
+
+    # TODO: implement intervals parser for PICMI that allows more complex output periods
+    test_name = os.path.split(os.getcwd())[1]
+    if "picmi" not in test_name:
+        # Functionality test for TimeAveragedDiagnostics
+        compare_time_avg_with_instantaneous_diags(
+            dir_inst=sys.argv[1],
+            dir_avg="diags/diagTimeAvg/",
+        )

Examples/Physics_applications/laser_ion/analysis_test_laser_ion.py

@@ -200,18 +200,32 @@ laser1.profile_focal_distance = 4.0e-6  # focal distance from the antenna [m]
 #################################
 # Diagnostics
 #
-diagnostics.diags_names = diag1 openPMDfw openPMDbw
+diagnostics.diags_names = diagInst diagTimeAvg openPMDfw openPMDbw
 
-diag1.intervals = 100
-diag1.diag_type = Full
-diag1.fields_to_plot = Ex Ey Ez Bx By Bz jx jy jz rho rho_electrons rho_hydrogen
+# instantaneous field and particle diagnostic
+diagInst.intervals = 100,96:100  # second interval only for CI testing the time-averaged diags
+diagInst.diag_type = Full
+diagInst.fields_to_plot = Ex Ey Ez Bx By Bz jx jy jz rho rho_electrons rho_hydrogen
 # reduce resolution of output fields
-diag1.coarsening_ratio = 4 4
+diagInst.coarsening_ratio = 4 4
 # demonstration of a spatial and momentum filter
-diag1.electrons.plot_filter_function(t,x,y,z,ux,uy,uz) = (uz>=0) * (x<1.0e-6) * (x>-1.0e-6)
-diag1.hydrogen.plot_filter_function(t,x,y,z,ux,uy,uz) = (uz>=0) * (x<1.0e-6) * (x>-1.0e-6)
-diag1.format = openpmd
-diag1.openpmd_backend = h5
+diagInst.electrons.plot_filter_function(t,x,y,z,ux,uy,uz) = (uz>=0) * (x<1.0e-6) * (x>-1.0e-6)
+diagInst.hydrogen.plot_filter_function(t,x,y,z,ux,uy,uz) = (uz>=0) * (x<1.0e-6) * (x>-1.0e-6)
+diagInst.format = openpmd
+diagInst.openpmd_backend = h5
+
+# time-averaged particle and field diagnostic
+diagTimeAvg.intervals = 100
+diagTimeAvg.diag_type = TimeAveraged
+diagTimeAvg.time_average_mode = dynamic_start
+#diagTimeAvg.average_period_time = 2.67e-15  # period of 800 nm light waves
+diagTimeAvg.average_period_steps = 5  # use only either `time` or `steps`
+diagTimeAvg.write_species = 0
+diagTimeAvg.fields_to_plot = Ex Ey Ez Bx By Bz jx jy jz rho rho_electrons rho_hydrogen
+# reduce resolution of output fields
+diagTimeAvg.coarsening_ratio = 4 4
+diagTimeAvg.format = openpmd
+diagTimeAvg.openpmd_backend = h5
 
 openPMDfw.intervals = 100
 openPMDfw.diag_type = Full

Examples/Physics_applications/laser_ion/inputs_test_2d_laser_ion_acc

@@ -140,7 +140,7 @@
 
 # Diagnostics
 particle_diag = picmi.ParticleDiagnostic(
-    name="diag1",
+    name="diagInst",
     period=100,
     warpx_format="openpmd",
     warpx_openpmd_backend="h5",
@@ -153,7 +153,7 @@
 for ncell_comp, cr in zip([nx, nz], coarsening_ratio):
     ncell_field.append(int(ncell_comp / cr))
 field_diag = picmi.FieldDiagnostic(
-    name="diag1",
+    name="diagInst",
     grid=grid,
     period=100,
     number_of_cells=ncell_field,
@@ -162,6 +162,18 @@
     warpx_openpmd_backend="h5",
 )
 
+field_time_avg_diag = picmi.TimeAveragedFieldDiagnostic(
+    name="diagTimeAvg",
+    grid=grid,
+    period=100,
+    number_of_cells=ncell_field,
+    data_list=["B", "E", "J", "rho", "rho_electrons", "rho_hydrogen"],
+    warpx_format="openpmd",
+    warpx_openpmd_backend="h5",
+    warpx_time_average_mode="dynamic_start",
+    warpx_average_period_time=2.67e-15,
+)
+
 particle_fw_diag = picmi.ParticleDiagnostic(
     name="openPMDfw",
     period=100,
@@ -292,6 +304,7 @@
 # Add full diagnostics
 sim.add_diagnostic(particle_diag)
 sim.add_diagnostic(field_diag)
+sim.add_diagnostic(field_time_avg_diag)
 sim.add_diagnostic(particle_fw_diag)
 sim.add_diagnostic(particle_bw_diag)
 # Add reduced diagnostics

Examples/Physics_applications/laser_ion/inputs_test_2d_laser_ion_acc_picmi.py

@@ -259,7 +259,7 @@ def visualize_particle_histogram_iteration(
         "-d",
         "--diag_dir",
         type=str,
-        default="./diags/diag1",
+        default="./diags/diagInst",
         help="Directory containing density and field diagnostics",
     )
     parser.add_argument(

Examples/Physics_applications/laser_ion/plot_2d.py

@@ -3271,6 +3271,57 @@ def diagnostic_initialize_inputs(self):
 ElectrostaticFieldDiagnostic = FieldDiagnostic
 
 
+class TimeAveragedFieldDiagnostic(FieldDiagnostic):
+    """
+    See `Input Parameters <https://warpx.readthedocs.io/en/latest/usage/parameters.html>`__ for more information.
+
+    Parameters
+    ----------
+    warpx_time_average_mode: str
+        Type of time averaging diagnostic
+        Supported values include ``"none"``, ``"fixed_start"``, and ``"dynamic_start"``
+
+            * ``"none"`` for no averaging (instantaneous fields)
+            * ``"fixed_start"`` for a diagnostic that averages all fields between the current output step and a fixed point in time
+            * ``"dynamic_start"`` for a constant averaging period and output at different points in time (non-overlapping)
+
+    warpx_average_period_steps: int, optional
+        Configures the number of time steps in an averaging period.
+        Set this only in the ``"dynamic_start"`` mode and only if ``warpx_average_period_time`` has not already been set.
+        Will be ignored in the ``"fixed_start"`` mode (with warning).
+
+    warpx_average_period_time: float, optional
+        Configures the time (SI units) in an averaging period.
+        Set this only in the ``"dynamic_start"`` mode and only if ``average_period_steps`` has not already been set.
+        Will be ignored in the ``"fixed_start"`` mode (with warning).
+
+    warpx_average_start_steps: int, optional
+        Configures the time step at which time-averaging begins.
+        Set this only in the ``"fixed_start"`` mode.
+        Will be ignored in the ``"dynamic_start"`` mode (with warning).
+    """
+
+    def init(self, kw):
+        super().init(kw)
+        self.time_average_mode = kw.pop("warpx_time_average_mode", None)
+        self.average_period_steps = kw.pop("warpx_average_period_steps", None)
+        self.average_period_time = kw.pop("warpx_average_period_time", None)
+        self.average_start_step = kw.pop("warpx_average_start_step", None)
+
+    def diagnostic_initialize_inputs(self):
+        super().diagnostic_initialize_inputs()
+
+        self.diagnostic.set_or_replace_attr("diag_type", "TimeAveraged")
+
+        if "write_species" not in self.diagnostic.argvattrs:
+            self.diagnostic.write_species = False
+
+        self.diagnostic.time_average_mode = self.time_average_mode
+        self.diagnostic.average_period_steps = self.average_period_steps
+        self.diagnostic.average_period_time = self.average_period_time
+        self.diagnostic.average_start_step = self.average_start_step
+
+
 class Checkpoint(picmistandard.base._ClassWithInit, WarpXDiagnosticBase):
     """
     Sets up checkpointing of the simulation, allowing for later restarts