Add PICMI functionality

Author: n01r

Docs/source/usage/parameters.rst

@@ -2888,7 +2888,11 @@ Time-Averaged Diagnostics
 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`_.
 
-    In addition, `TimeAveraged` diagnostic options include:
+.. note::
+
+    As with ``FullDiagnostics``, ``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.

Examples/Physics_applications/laser_ion/inputs_test_2d_laser_ion_acc

@@ -213,7 +213,7 @@ diagInst.hydrogen.plot_filter_function(t,x,y,z,ux,uy,uz) = (uz>=0) * (x<1.0e-6)
 diagInst.format = openpmd
 diagInst.openpmd_backend = h5
 
-diagTimeAvg.intervals = 100::100  # TODO just write 100 and make step 0 an instantaneous diagnostic and do not forget to document
+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

Examples/Physics_applications/laser_ion/inputs_test_2d_laser_ion_acc_picmi.py

@@ -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_steps=10,
+)
+
 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

Python/pywarpx/picmi.py

@@ -3270,6 +3270,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):
     """

Source/Diagnostics/FullDiagnostics.H

@@ -36,7 +36,7 @@ private:
      */
     bool m_solver_deposits_current = true;
     /** Whether the diagnostics are averaging data over time or not */
-    TimeAverageType m_time_average_type = TimeAverageType::None;
+    TimeAverageType m_time_average_mode = TimeAverageType::None;
     /** Period to average fields over: in steps */
     int m_average_period_steps = -1;
     /** Period to average fields over: in seconds */

Source/Diagnostics/FullDiagnostics.cpp

@@ -59,7 +59,7 @@ void
 FullDiagnostics::DerivedInitData() {
     if (m_diag_type == DiagTypes::TimeAveraged) {
         auto & warpx = WarpX::GetInstance();
-        if (m_time_average_type == TimeAverageType::Dynamic) {
+        if (m_time_average_mode == TimeAverageType::Dynamic) {
 
             // already checked in ReadParameters that only one of the parameters is set
             // calculate the other averaging period parameter from the other one, respectively
@@ -128,11 +128,11 @@ FullDiagnostics::ReadParameters ()
         pp_diag_name.get("time_average_mode", m_time_average_mode_str);
 
         if (m_time_average_mode_str == "fixed_start") {
-                m_time_average_type = TimeAverageType::Static;
+            m_time_average_mode = TimeAverageType::Static;
             } else if (m_time_average_mode_str == "dynamic_start") {
-                m_time_average_type = TimeAverageType::Dynamic;
+            m_time_average_mode = TimeAverageType::Dynamic;
             } else if (m_time_average_mode_str == "none") {
-                m_time_average_type = TimeAverageType::None;
+            m_time_average_mode = TimeAverageType::None;
             } else {
             WARPX_ABORT_WITH_MESSAGE(
                     "Unknown time averaging mode. Valid entries are: none, fixed_start, dynamic_start"
@@ -146,7 +146,7 @@ FullDiagnostics::ReadParameters ()
                 "average_period_time", m_average_period_time
         );
 
-        if (m_time_average_type == TimeAverageType::Static) {
+        if (m_time_average_mode == TimeAverageType::Static) {
             // This fails if users do not specify a start.
             pp_diag_name.get("average_start_step", m_average_start_step);
             if (m_average_start_step == 0) {
@@ -170,7 +170,7 @@ FullDiagnostics::ReadParameters ()
 
         }
 
-        if (m_time_average_type == TimeAverageType::Dynamic) {
+        if (m_time_average_mode == TimeAverageType::Dynamic) {
             // one of the two averaging period options must be set but neither none nor both
             if (
                     (averaging_period_steps_specified && averaging_period_time_specified)
@@ -238,7 +238,7 @@ FullDiagnostics::Flush ( int i_buffer, bool /* force_flush */ )
     // to accommodate a user workflow that only uses that type of diagnostic.
     // This allows for quicker turnaround in setup by avoiding having to set an additional instantaneous diagnostic.
     if (m_diag_type == DiagTypes::TimeAveraged && step > 0) {
-        if (m_time_average_type == TimeAverageType::Static || m_time_average_type == TimeAverageType::Dynamic) {
+        if (m_time_average_mode == TimeAverageType::Static || m_time_average_mode == TimeAverageType::Dynamic) {
             // Loop over the output levels and divide by the number of steps in the averaging period
             for (int lev = 0; lev < nlev_output; ++lev) {
                 m_sum_mf_output.at(i_buffer).at(lev).mult(1._rt/static_cast<amrex::Real>(m_average_period_steps));
@@ -251,7 +251,7 @@ FullDiagnostics::Flush ( int i_buffer, bool /* force_flush */ )
                     m_file_min_digits, m_plot_raw_fields, m_plot_raw_fields_guards);
 
             // Reset the values in the dynamic start time-averaged diagnostics after flush
-            if (m_time_average_type == TimeAverageType::Dynamic) {
+            if (m_time_average_mode == TimeAverageType::Dynamic) {
                 for (int lev = 0; lev < nlev_output; ++lev) {
                     m_sum_mf_output.at(i_buffer).at(lev).setVal(0.);
                 }
@@ -298,12 +298,12 @@ FullDiagnostics::DoComputeAndPack (int step, bool force_flush)
 
         if (step > 0) {
 
-            if (m_time_average_type == TimeAverageType::Dynamic) {
+            if (m_time_average_mode == TimeAverageType::Dynamic) {
                 m_average_start_step = m_intervals.nextContains(step) - m_average_period_steps;
                 // check that the periods do not overlap and that the start step is not negative
                 if (m_average_start_step > 0) {
                     // The start step cannot be on an interval step because then we would begin a new period and also output the old one
-                    if (m_average_start_step <= m_intervals.previousContains(step)) {
+                    if (m_average_start_step < m_intervals.previousContains(step)) {
                         WARPX_ABORT_WITH_MESSAGE(
                                 "Averaging periods may not overlap within a single diagnostic. "
                                 "Please create a second diagnostic for overlapping time averaging periods "
@@ -321,7 +321,7 @@ FullDiagnostics::DoComputeAndPack (int step, bool force_flush)
                 if (step >= m_average_start_step && step <= m_intervals.nextContains(step)) {
                     in_averaging_period = true;
 
-                    if (m_time_average_type == TimeAverageType::Static) {
+                    if (m_time_average_mode == TimeAverageType::Static) {
                         // Update time averaging period to current step
                         m_average_period_steps = step - m_average_start_step;
                     }