add after diagnostic python callback (#3043)

Python/pywarpx/callbacks.py

@@ -29,6 +29,7 @@
  - afterdeposition <installafterdeposition>: after particle deposition (for charge and/or current)
  - beforestep <installbeforestep>: before the time step
  - afterstep <installafterstep>: after the time step
+ - afterdiagnostics <installafterdiagnostics>: after diagnostic output
  - particlescraper <installparticlescraper>: just after the particle boundary conditions are applied
                                              but before lost particles are processed
  - particleloader <installparticleloader>: at the time that the standard particle loader is called
@@ -261,6 +262,7 @@ def callfuncsinlist(self,*args,**kw):
 _particleloader = CallbackFunctions('particleloader')
 _beforestep = CallbackFunctions('beforestep')
 _afterstep = CallbackFunctions('afterstep')
+_afterdiagnostics = CallbackFunctions('afterdiagnostics')
 _afterrestart = CallbackFunctions('afterrestart',lcallonce=1)
 _particleinjection = CallbackFunctions('particleinjection')
 _appliedfields = CallbackFunctions('appliedfields')
@@ -279,6 +281,7 @@ def printcallbacktimers(tmin=1.,lminmax=False,ff=None):
               _particlescraper,
               _particleloader,
               _beforestep,_afterstep,
+              _afterdiagnostics,
               _afterrestart,
               _particleinjection,
               _appliedfields]:
@@ -471,6 +474,20 @@ def isinstalledafterstep(f):
     "Checks if the function is called after a step"
     return _afterstep.isinstalledfuncinlist(f)
 
+# ----------------------------------------------------------------------------
+def callfromafterdiagnostics(f):
+    installafterdiagnostics(f)
+    return f
+def installafterdiagnostics(f):
+    "Adds a function to the list of functions called after diagnostic output"
+    _afterdiagnostics.installfuncinlist(f)
+def uninstallafterdiagnostics(f):
+    "Removes the function from the list of functions called after diagnostic output"
+    _afterdiagnostics.uninstallfuncinlist(f)
+def isinstalledafterdiagnostics(f):
+    "Checks if the function is called after diagnostic output"
+    return _afterdiagnostics.isinstalledfuncinlist(f)
+
 # ----------------------------------------------------------------------------
 def callfromafterrestart(f):
     raise Exception('restart call back not implemented yet')

Source/Evolve/WarpXEvolve.cpp

@@ -335,6 +335,9 @@ WarpX::Evolve (int numsteps)
         }
         multi_diags->FilterComputePackFlush( step );
 
+        // execute afterdiagnostic callbacks
+        ExecutePythonCallback("afterdiagnostics");
+
         // inputs: unused parameters (e.g. typos) check after step 1 has finished
         if (!early_params_checked) {
             amrex::Print() << "\n"; // better: conditional \n based on return value

Source/Python/WarpX_py.H

@@ -21,8 +21,8 @@
  * functions will be called. Currently supported keys (callback points) are
  * afterinit, beforecollisions, aftercollisions, beforeEsolve, poissonsolver,
  * afterEsolve, beforedeposition, afterdeposition, particlescraper,
- * particleloader, beforestep, afterstep, afterrestart, particleinjection and
- * appliedfields.
+ * particleloader, beforestep, afterstep, afterdiagnostics, afterrestart,
+ * particleinjection and appliedfields.
 */
 extern std::map< std::string, WARPX_CALLBACK_PY_FUNC_0 > warpx_callback_py_map;