Callback modernization (ECP-WarpX#4)

* pyAMReX: Build System

* refactor callbacks.py

* added binding code in `pyWarpX.cpp` to add or remove keys from the callback dictionary

* minor PR cleanups

* CMake: Undo Unintentional Change

Likely from rebase/merge.

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

Source/FieldSolver/ElectrostaticSolver.cpp

@@ -242,7 +242,7 @@ WarpX::AddSpaceChargeFieldLabFrame ()
     setPhiBC(phi_fp);
 
     // Compute the potential phi, by solving the Poisson equation
-    if (IsPythonCallBackInstalled("poissonsolver")) {
+    if (IsPythonCallbackInstalled("poissonsolver")) {
 
         // Use the Python level solver (user specified)
         ExecutePythonCallback("poissonsolver");
@@ -266,7 +266,7 @@ WarpX::AddSpaceChargeFieldLabFrame ()
 #ifndef AMREX_USE_EB
     computeE( Efield_fp, phi_fp, beta );
 #else
-    if ( IsPythonCallBackInstalled("poissonsolver") ) computeE( Efield_fp, phi_fp, beta );
+    if ( IsPythonCallbackInstalled("poissonsolver") ) computeE( Efield_fp, phi_fp, beta );
 #endif
 
     // Compute the magnetic field

Source/Python/CMakeLists.txt

@@ -2,8 +2,11 @@
 # These are the files equivalent to the primary C++ implementation - but here
 # we define how they will appear in our Python module (aka Python bindings).
 #
+target_sources(WarpX
+  PRIVATE
+    WarpX_py.cpp
+)
 target_sources(pyWarpX
   PRIVATE
     WarpX.cpp
-    WarpX_py.cpp  # TODO: modernize via https://pybind11.readthedocs.io/en/stable/advanced/cast/functional.html
 )

Source/Python/WarpX_py.H

@@ -1,14 +1,14 @@
-/* Copyright 2019 David Grote, Maxence Thevenet, Weiqun Zhang
- *
+/* Copyright 2019-2022 The WarpX Community
  *
  * This file is part of WarpX.
  *
+ * Authors: David Grote, Maxence Thevenet, Weiqun Zhang, Roelof Groenewald
+ *
  * License: BSD-3-Clause-LBNL
  */
 #ifndef WARPX_PY_H_
 #define WARPX_PY_H_
 
-#include "WarpXWrappers.H"
 #include "Utils/WarpXProfilerWrapper.H"
 
 #include <map>
@@ -18,22 +18,29 @@
  * 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, beforecollisions, aftercollisions, beforeEsolve, poissonsolver,
- * afterEsolve, beforedeposition, afterdeposition, particlescraper,
- * particleloader, beforestep, afterstep, afterdiagnostics, afterrestart,
- * oncheckpointsignal, particleinjection and appliedfields.
+ * functions will be called. See ``WarpX/Python/pywarpx/callbacks.py`` for a
+ * list of currently supported callback names.
 */
-extern std::map< std::string, WARPX_CALLBACK_PY_FUNC_0 > warpx_callback_py_map;
+extern std::map< std::string, std::function<void()> > warpx_callback_py_map;
+
+/**
+ * \brief Function to install the given name and function in warpx_callback_py_map
+ */
+void InstallPythonCallback ( std::string name, std::function<void()> callback );
 
 /**
  * \brief Function to check if the given name is a key in warpx_callback_py_map
  */
-bool IsPythonCallBackInstalled ( std::string name );
+bool IsPythonCallbackInstalled ( std::string name );
 
 /**
  * \brief Function to look for and execute Python callbacks
  */
 void ExecutePythonCallback ( std::string name );
 
-#endif
+/**
+ * \brief Function to clear the given callback name from warpx_callback_py_map
+ */
+void ClearPythonCallback ( std::string name );
+
+#endif

Source/Python/WarpX_py.cpp

@@ -1,24 +1,35 @@
-/* Copyright 2019 David Grote, Maxence Thevenet, Weiqun Zhang
- *
+/* Copyright 2019-2022 The WarpX Community
  *
  * This file is part of WarpX.
  *
+ * Authors: David Grote, Maxence Thevenet, Weiqun Zhang, Roelof Groenewald
+ *
  * License: BSD-3-Clause-LBNL
  */
 #include "WarpX_py.H"
 
-std::map< std::string, WARPX_CALLBACK_PY_FUNC_0 > warpx_callback_py_map;
+std::map< std::string, std::function<void()> > warpx_callback_py_map;
+
+void InstallPythonCallback ( std::string name, std::function<void()> callback )
+{
+    warpx_callback_py_map[name] = callback;
+}
 
-bool IsPythonCallBackInstalled ( std::string name )
+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) ) {
+    if ( IsPythonCallbackInstalled(name) ) {
         WARPX_PROFILE("warpx_py_"+name);
         warpx_callback_py_map[name]();
     }
 }
+
+void ClearPythonCallback ( std::string name )
+{
+    warpx_callback_py_map.erase(name);
+}

Source/Python/pyWarpX.H

@@ -6,6 +6,8 @@
  * Authors: Axel Huebl
  * License: BSD-3-Clause-LBNL
  */
+#include "WarpX_py.H"
+
 #include <pybind11/pybind11.h>
 #include <pybind11/stl.h>
 #include <pybind11/stl_bind.h>

Source/Python/pyWarpX.cpp

@@ -69,4 +69,8 @@ PYBIND11_MODULE(PYWARPX_MODULE_NAME, m) {
     // auto numpy = py::module::import("numpy");
     // auto npversion = numpy.attr("__version__");
     // std::cout << "numpy version: " << py::str(npversion) << std::endl;
+
+    // Expose the python callback function installation and removal functions
+    m.def("add_python_callback", &InstallPythonCallback);
+    m.def("remove_python_callback", &ClearPythonCallback);
 }