Refactor python callback handling (#2703)

* added support to uninstall an external Poisson solver and return to using the default MLMG solver; also updated some callbacks.py calls to Python3

* refactor callback handling - use a map to handle all the different callbacks

* warpx_callback_py_map does not need to link to C

* Apply suggestions from code review

Co-authored-by: Axel Huebl <axel.huebl@plasma.ninja>

* further suggested changes from code review

* added function ExecutePythonCallback to reduce code duplication

* moved ExecutePythonCallback to WarpX_py

* added function IsPythonCallbackInstalled

Co-authored-by: Axel Huebl <axel.huebl@plasma.ninja>

6 files changed, 64 insertions, 143 deletions

diff --git a/Source/Evolve/WarpXEvolve.cpp b/Source/Evolve/WarpXEvolve.cpp
index c3c297b56..86bcc7fb7 100644
--- a/Source/Evolve/WarpXEvolve.cpp
+++ b/Source/Evolve/WarpXEvolve.cpp
@@ -83,10 +83,7 @@ WarpX::Evolve (int numsteps)
         if (verbose) {
             amrex::Print() << "\nSTEP " << step+1 << " starts ...\n";
         }
-        if (warpx_py_beforestep) {
-            WARPX_PROFILE("warpx_py_beforestep");
-            warpx_py_beforestep();
-        }
+        ExecutePythonCallback("beforestep");
 
         amrex::LayoutData<amrex::Real>* cost = WarpX::getCosts(0);
         if (cost) {
@@ -168,10 +165,7 @@ WarpX::Evolve (int numsteps)
         // Main PIC operation:
         // gather fields, push particles, deposit sources, update fields
 
-        if (warpx_py_particleinjection) {
-            WARPX_PROFILE("warpx_py_particleinjection");
-            warpx_py_particleinjection();
-        }
+        ExecutePythonCallback("particleinjection");
         // Electrostatic case: only gather fields and push particles,
         // deposition and calculation of fields done further below
         if (do_electrostatic != ElectrostaticSolverAlgo::None)
@@ -307,10 +301,7 @@ WarpX::Evolve (int numsteps)
         }
 
         if( do_electrostatic != ElectrostaticSolverAlgo::None ) {
-            if (warpx_py_beforeEsolve) {
-                WARPX_PROFILE("warpx_py_beforeEsolve");
-                warpx_py_beforeEsolve();
-            }
+            ExecutePythonCallback("beforeEsolve");
             // Electrostatic solver:
             // For each species: deposit charge and add the associated space-charge
             // E and B field to the grid ; this is done at the end of the PIC
@@ -319,18 +310,12 @@ WarpX::Evolve (int numsteps)
             // and so that the fields are at the correct time in the output.
             bool const reset_fields = true;
             ComputeSpaceChargeField( reset_fields );
-            if (warpx_py_afterEsolve) {
-                WARPX_PROFILE("warpx_py_afterEsolve");
-                warpx_py_afterEsolve();
-            }
+            ExecutePythonCallback("afterEsolve");
         }
 
-        // warpx_py_afterstep runs with the updated global time. It is included
+        // afterstep callback runs with the updated global time. It is included
         // in the evolve timing.
-        if (warpx_py_afterstep) {
-            WARPX_PROFILE("warpx_py_afterstep");
-            warpx_py_afterstep();
-        }
+        ExecutePythonCallback("afterstep");
 
         /// reduced diags
         if (reduced_diags->m_plot_rd != 0)
@@ -387,20 +372,13 @@ WarpX::OneStep_nosub (Real cur_time)
     //               from p^{n-1/2} to p^{n+1/2}
     // Deposit current j^{n+1/2}
     // Deposit charge density rho^{n}
-    if (warpx_py_particlescraper) {
-        WARPX_PROFILE("warpx_py_particlescraper");
-        warpx_py_particlescraper();
-    }
-    if (warpx_py_beforedeposition) {
-        WARPX_PROFILE("warpx_py_beforedeposition");
-        warpx_py_beforedeposition();
-    }
+
+    ExecutePythonCallback("particlescraper");
+    ExecutePythonCallback("beforedeposition");
+
     PushParticlesandDepose(cur_time);
 
-    if (warpx_py_afterdeposition) {
-        WARPX_PROFILE("warpx_py_afterdeposition");
-        warpx_py_afterdeposition();
-    }
+    ExecutePythonCallback("afterdeposition");
 
     // Synchronize J and rho
     SyncCurrent();
@@ -422,10 +400,7 @@ WarpX::OneStep_nosub (Real cur_time)
     if (do_pml && pml_has_particles) CopyJPML();
     if (do_pml && do_pml_j_damping) DampJPML();
 
-    if (warpx_py_beforeEsolve) {
-        WARPX_PROFILE("warpx_py_beforeEsolve");
-        warpx_py_beforeEsolve();
-    }
+    ExecutePythonCallback("beforeEsolve");
 
     // Push E and B from {n} to {n+1}
     // (And update guard cells immediately afterwards)
@@ -507,10 +482,7 @@ WarpX::OneStep_nosub (Real cur_time)
             FillBoundaryB(guard_cells.ng_alloc_EB);
     } // !PSATD
 
-    if (warpx_py_afterEsolve) {
-        WARPX_PROFILE("warpx_py_afterEsolve");
-        warpx_py_afterEsolve();
-    }
+    ExecutePythonCallback("afterEsolve");
 }
 
 void
diff --git a/Source/FieldSolver/ElectrostaticSolver.cpp b/Source/FieldSolver/ElectrostaticSolver.cpp
index 18f90431f..e1d70be16 100644
--- a/Source/FieldSolver/ElectrostaticSolver.cpp
+++ b/Source/FieldSolver/ElectrostaticSolver.cpp
@@ -172,10 +172,7 @@ WarpX::AddSpaceChargeFieldLabFrame ()
     std::array<Real, 3> beta = {0._rt};
 
     // Compute the potential phi, by solving the Poisson equation
-    if (warpx_py_poissonsolver) {
-        WARPX_PROFILE("warpx_py_poissonsolver");
-        warpx_py_poissonsolver();
-    }
+    if ( IsPythonCallBackInstalled("poissonsolver") ) ExecutePythonCallback("poissonsolver");
     else computePhi( rho_fp, phi_fp, beta, self_fields_required_precision,
                      self_fields_absolute_tolerance, self_fields_max_iters,
                      self_fields_verbosity );
@@ -185,7 +182,7 @@ WarpX::AddSpaceChargeFieldLabFrame ()
 #ifndef AMREX_USE_EB
     computeE( Efield_fp, phi_fp, beta );
 #else
-    if (warpx_py_poissonsolver) computeE( Efield_fp, phi_fp, beta );
+    if ( IsPythonCallBackInstalled("poissonsolver") ) computeE( Efield_fp, phi_fp, beta );
 #endif
 
     // Compute the magnetic field
diff --git a/Source/Python/WarpXWrappers.H b/Source/Python/WarpXWrappers.H
index 42066b672..8d2ee4364 100644
--- a/Source/Python/WarpXWrappers.H
+++ b/Source/Python/WarpXWrappers.H
@@ -48,19 +48,9 @@ extern "C" {
 
     typedef void(*WARPX_CALLBACK_PY_FUNC_0)();
 
-    void warpx_set_callback_py_afterinit (WARPX_CALLBACK_PY_FUNC_0);
-    void warpx_set_callback_py_beforeEsolve (WARPX_CALLBACK_PY_FUNC_0);
-    void warpx_set_callback_py_poissonsolver (WARPX_CALLBACK_PY_FUNC_0);
-    void warpx_set_callback_py_afterEsolve (WARPX_CALLBACK_PY_FUNC_0);
-    void warpx_set_callback_py_beforedeposition (WARPX_CALLBACK_PY_FUNC_0);
-    void warpx_set_callback_py_afterdeposition (WARPX_CALLBACK_PY_FUNC_0);
-    void warpx_set_callback_py_particlescraper (WARPX_CALLBACK_PY_FUNC_0);
-    void warpx_set_callback_py_particleloader (WARPX_CALLBACK_PY_FUNC_0);
-    void warpx_set_callback_py_beforestep (WARPX_CALLBACK_PY_FUNC_0);
-    void warpx_set_callback_py_afterstep (WARPX_CALLBACK_PY_FUNC_0);
-    void warpx_set_callback_py_afterrestart (WARPX_CALLBACK_PY_FUNC_0);
-    void warpx_set_callback_py_particleinjection (WARPX_CALLBACK_PY_FUNC_0);
-    void warpx_set_callback_py_appliedfields (WARPX_CALLBACK_PY_FUNC_0);
+    void warpx_set_callback_py (const char* char_callback_name,
+                                WARPX_CALLBACK_PY_FUNC_0 callback);
+    void warpx_clear_callback_py (const char* char_callback_name);
 
     void warpx_evolve (int numsteps);  // -1 means the inputs parameter will be used.
 
diff --git a/Source/Python/WarpXWrappers.cpp b/Source/Python/WarpXWrappers.cpp
index d4589ec8a..7350bbbb6 100644
--- a/Source/Python/WarpXWrappers.cpp
+++ b/Source/Python/WarpXWrappers.cpp
@@ -164,14 +164,8 @@ namespace
     {
         WarpX& warpx = WarpX::GetInstance();
         warpx.InitData();
-        if (warpx_py_afterinit) {
-            WARPX_PROFILE("warpx_py_afterinit");
-            warpx_py_afterinit();
-        }
-        if (warpx_py_particleloader) {
-            WARPX_PROFILE("warpx_py_particleloader");
-            warpx_py_particleloader();
-        }
+        ExecutePythonCallback("afterinit");
+        ExecutePythonCallback("particleloader");
     }
 
     void warpx_finalize ()
@@ -179,57 +173,17 @@ namespace
         WarpX::ResetInstance();
     }
 
-    void warpx_set_callback_py_afterinit (WARPX_CALLBACK_PY_FUNC_0 callback)
-    {
-        warpx_py_afterinit = callback;
-    }
-    void warpx_set_callback_py_beforeEsolve (WARPX_CALLBACK_PY_FUNC_0 callback)
+    void warpx_set_callback_py (
+        const char* char_callback_name, WARPX_CALLBACK_PY_FUNC_0 callback)
     {
-        warpx_py_beforeEsolve = callback;
+        const std::string callback_name(char_callback_name);
+        warpx_callback_py_map[callback_name] = callback;
     }
-    void warpx_set_callback_py_poissonsolver (WARPX_CALLBACK_PY_FUNC_0 callback)
-    {
-        warpx_py_poissonsolver = callback;
-    }
-    void warpx_set_callback_py_afterEsolve (WARPX_CALLBACK_PY_FUNC_0 callback)
-    {
-        warpx_py_afterEsolve = callback;
-    }
-    void warpx_set_callback_py_beforedeposition (WARPX_CALLBACK_PY_FUNC_0 callback)
-    {
-        warpx_py_beforedeposition = callback;
-    }
-    void warpx_set_callback_py_afterdeposition (WARPX_CALLBACK_PY_FUNC_0 callback)
-    {
-        warpx_py_afterdeposition = callback;
-    }
-    void warpx_set_callback_py_particlescraper (WARPX_CALLBACK_PY_FUNC_0 callback)
-    {
-        warpx_py_particlescraper = callback;
-    }
-    void warpx_set_callback_py_particleloader (WARPX_CALLBACK_PY_FUNC_0 callback)
-    {
-        warpx_py_particleloader = callback;
-    }
-    void warpx_set_callback_py_beforestep (WARPX_CALLBACK_PY_FUNC_0 callback)
-    {
-        warpx_py_beforestep = callback;
-    }
-    void warpx_set_callback_py_afterstep (WARPX_CALLBACK_PY_FUNC_0 callback)
-    {
-        warpx_py_afterstep = callback;
-    }
-    void warpx_set_callback_py_afterrestart (WARPX_CALLBACK_PY_FUNC_0 callback)
-    {
-        warpx_py_afterrestart = callback;
-    }
-    void warpx_set_callback_py_particleinjection (WARPX_CALLBACK_PY_FUNC_0 callback)
-    {
-        warpx_py_particleinjection = callback;
-    }
-    void warpx_set_callback_py_appliedfields (WARPX_CALLBACK_PY_FUNC_0 callback)
+
+    void warpx_clear_callback_py (const char* char_callback_name)
     {
-        warpx_py_appliedfields = callback;
+        const std::string callback_name(char_callback_name);
+        warpx_callback_py_map.erase(callback_name);
     }
 
     void warpx_evolve (int numsteps)
diff --git a/Source/Python/WarpX_py.H b/Source/Python/WarpX_py.H
index 42697fbac..efc248591 100644
--- a/Source/Python/WarpX_py.H
+++ b/Source/Python/WarpX_py.H
@@ -9,24 +9,30 @@
 #define WARPX_PY_H_
 
 #include "WarpXWrappers.H"
+#include "Utils/WarpXProfilerWrapper.H"
 
+#include <map>
+#include <string>
 
-extern "C" {
+/**
+ * Declare global map to hold python callback functions.
+ *
+ * The keys of the map describe at what point in the simulation the python
+ * functions will be called. Currently supported keys (callback points) are
+ * afterinit, beforeEsolve, poissonsolver, afterEsolve, beforedeposition,
+ * afterdeposition, particlescraper, particleloader, beforestep, afterstep,
+ * afterrestart, particleinjection and appliedfields.
+*/
+extern std::map< std::string, WARPX_CALLBACK_PY_FUNC_0 > warpx_callback_py_map;
 
-    extern WARPX_CALLBACK_PY_FUNC_0 warpx_py_afterinit;
-    extern WARPX_CALLBACK_PY_FUNC_0 warpx_py_beforeEsolve;
-    extern WARPX_CALLBACK_PY_FUNC_0 warpx_py_poissonsolver;
-    extern WARPX_CALLBACK_PY_FUNC_0 warpx_py_afterEsolve;
-    extern WARPX_CALLBACK_PY_FUNC_0 warpx_py_beforedeposition;
-    extern WARPX_CALLBACK_PY_FUNC_0 warpx_py_afterdeposition;
-    extern WARPX_CALLBACK_PY_FUNC_0 warpx_py_particlescraper;
-    extern WARPX_CALLBACK_PY_FUNC_0 warpx_py_particleloader;
-    extern WARPX_CALLBACK_PY_FUNC_0 warpx_py_beforestep;
-    extern WARPX_CALLBACK_PY_FUNC_0 warpx_py_afterstep;
-    extern WARPX_CALLBACK_PY_FUNC_0 warpx_py_afterrestart;
-    extern WARPX_CALLBACK_PY_FUNC_0 warpx_py_particleinjection;
-    extern WARPX_CALLBACK_PY_FUNC_0 warpx_py_appliedfields;
+/**
+ * \brief Function to check if the given name is a key in warpx_callback_py_map
+ */
+bool IsPythonCallBackInstalled ( std::string name );
 
-}
+/**
+ * \brief Function to look for and execute Python callbacks
+ */
+void ExecutePythonCallback ( std::string name );
 
 #endif
diff --git a/Source/Python/WarpX_py.cpp b/Source/Python/WarpX_py.cpp
index 1d69ea8f3..7e0e491f8 100644
--- a/Source/Python/WarpX_py.cpp
+++ b/Source/Python/WarpX_py.cpp
@@ -7,16 +7,18 @@
  */
 #include "WarpX_py.H"
 
-WARPX_CALLBACK_PY_FUNC_0 warpx_py_afterinit = nullptr;
-WARPX_CALLBACK_PY_FUNC_0 warpx_py_beforeEsolve = nullptr;
-WARPX_CALLBACK_PY_FUNC_0 warpx_py_poissonsolver = nullptr;
-WARPX_CALLBACK_PY_FUNC_0 warpx_py_afterEsolve = nullptr;
-WARPX_CALLBACK_PY_FUNC_0 warpx_py_beforedeposition = nullptr;
-WARPX_CALLBACK_PY_FUNC_0 warpx_py_afterdeposition = nullptr;
-WARPX_CALLBACK_PY_FUNC_0 warpx_py_particlescraper = nullptr;
-WARPX_CALLBACK_PY_FUNC_0 warpx_py_particleloader = nullptr;
-WARPX_CALLBACK_PY_FUNC_0 warpx_py_beforestep = nullptr;
-WARPX_CALLBACK_PY_FUNC_0 warpx_py_afterstep = nullptr;
-WARPX_CALLBACK_PY_FUNC_0 warpx_py_afterrestart = nullptr;
-WARPX_CALLBACK_PY_FUNC_0 warpx_py_particleinjection = nullptr;
-WARPX_CALLBACK_PY_FUNC_0 warpx_py_appliedfields = nullptr;
+std::map< std::string, WARPX_CALLBACK_PY_FUNC_0 > warpx_callback_py_map;
+
+bool IsPythonCallBackInstalled ( std::string name )
+{
+    return (warpx_callback_py_map.count(name) == 1u);
+}
+
+// Execute Python callbacks of the type given by the input string
+void ExecutePythonCallback ( std::string name )
+{
+    if ( IsPythonCallBackInstalled(name) ) {
+        WARPX_PROFILE("warpx_py_"+name);
+        warpx_callback_py_map[name]();
+    }
+}