From e538399d99121d37ece3d23f4444342117a9a723 Mon Sep 17 00:00:00 2001
From: Roelof Groenewald <40245517+roelof-groenewald@users.noreply.github.com>
Date: Tue, 19 Apr 2022 17:40:53 -0700
Subject: [PATCH] add after diagnostic python callback (#3043)

---
 Python/pywarpx/callbacks.py   | 17 +++++++++++++++++
 Source/Evolve/WarpXEvolve.cpp |  3 +++
 Source/Python/WarpX_py.H      |  4 ++--
 3 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/Python/pywarpx/callbacks.py b/Python/pywarpx/callbacks.py
index 8fa08e6ef87..cb9293c5a8b 100644
--- a/Python/pywarpx/callbacks.py
+++ b/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')
diff --git a/Source/Evolve/WarpXEvolve.cpp b/Source/Evolve/WarpXEvolve.cpp
index 39fc1489676..de2b5969187 100644
--- a/Source/Evolve/WarpXEvolve.cpp
+++ b/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
diff --git a/Source/Python/WarpX_py.H b/Source/Python/WarpX_py.H
index 15076ea35f7..a0392a20332 100644
--- a/Source/Python/WarpX_py.H
+++ b/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;