Backported mpi4py

.github/workflows/ci.yml

@@ -119,6 +119,8 @@ jobs:
         echo "OMPI_MCA_btl_base_warn_component_unused=0" >> $GITHUB_ENV
         echo "OMPI_MCA_btl_base_verbose=0" >> $GITHUB_ENV
         echo "MPIEXEC=mpirun --oversubscribe" >> $GITHUB_ENV
+        pip install --user mpi4py
+        python -c "import mpi4py"
     - name: Build PLUMED
       run: |
         ccache -s -M 100M

python/cplumed.pxd

@@ -39,6 +39,7 @@ cdef extern from "Plumed.h" namespace "PLMD":
 # see https://stackoverflow.com/questions/42610108/is-overloading-broken-in-cppclass-cython-c-definitions
          void cmd_int "cmd" (const char*key, int val) except +exceptions_handler
          void cmd_float "cmd" (const char*key, double val) except +exceptions_handler
+         void cmd_mpi "cmd" (const char*key, const void*val) except +exceptions_handler
          void cmd(const char*key) except +exceptions_handler
          bool valid() except +exceptions_handler
          @staticmethod

python/plumed.pyx

@@ -126,17 +126,24 @@ cdef class Plumed:
          self.c_plumed.cmd_float( ckey, val)
      cdef cmd_int(self, ckey, int val):
          self.c_plumed.cmd_int( ckey, val)
+     cdef cmd_mpi(self, ckey, val):
+         import mpi4py.MPI as MPI
+         cdef size_t comm_addr = MPI._addressof(val)
+         self.c_plumed.cmd_mpi( ckey, <const void*>comm_addr)
      def cmd( self, key, val=None ):
          cdef bytes py_bytes = key.encode()
          cdef char* ckey = py_bytes
          cdef char* cval
          if val is None :
             self.c_plumed.cmd( ckey)
-         elif isinstance(val, (int,long) ):
+            return
+         if isinstance(val, (int,long) ):
             self.cmd_int(ckey, val)
-         elif isinstance(val, float ) :
+            return
+         if isinstance(val, float ) :
             self.cmd_float(ckey, val)
-         elif HAS_NUMPY and isinstance(val, np.ndarray) :
+            return
+         if HAS_NUMPY and isinstance(val, np.ndarray) :
             # See https://numpy.org/doc/stable/user/basics.types.html
             if( val.dtype==np.double):
                self.cmd_ndarray_double(ckey, val)
@@ -146,7 +153,8 @@ cdef class Plumed:
                self.cmd_ndarray_long(ckey, val)
             else :
                raise ValueError("ndarrys should be np.double, np.intc, or np.int_")
-         elif isinstance(val, array.array) :
+            return
+         if isinstance(val, array.array) :
             if( (val.typecode=="d" or val.typecode=="f") and val.itemsize==8):
                self.cmd_array_double(ckey, val)
             elif( (val.typecode=="i" or val.typecode=="I") ) :
@@ -155,13 +163,18 @@ cdef class Plumed:
                self.cmd_array_long(ckey, val)
             else :
                raise ValueError("ndarrays should be double (size=8), int, or long")
-         elif isinstance(val, str ) :
-              py_bytes = val.encode()
-              cval = py_bytes
-              self.c_plumed.cmd( ckey, <const char*>cval )
-         else :
-            raise ValueError("Unknown value type ({})".format(str(type(val))))
-
+            return
+         if isinstance(val, str ) :
+            py_bytes = val.encode()
+            cval = py_bytes
+            self.c_plumed.cmd( ckey, <const char*>cval )
+            return
+         if 'mpi4py' in sys.modules:
+            import mpi4py.MPI as MPI
+            if isinstance(val, MPI.Comm):
+              self.cmd_mpi(ckey, val)
+              return
+         raise ValueError("Unknown value type ({})".format(str(type(val))))
 
 class FormatError(Exception):
     """Custom error reported by read_as_pandas.

python/test/test_mpi.py

@@ -0,0 +1,24 @@
+# simple test to check setMPIComm and mpi4py
+import plumed
+import os
+import unittest
+try:
+  from mpi4py import MPI
+  HAS_MPI4PY=True
+except ImportError:
+  HAS_MPI4PY=False
+
+class Test(unittest.TestCase):
+  if HAS_MPI4PY:
+    def test(self):
+
+      comm = MPI.COMM_WORLD
+      p = plumed.Plumed()
+      p.cmd("setNatoms",2)
+      p.cmd("setLogFile","test.log")
+      p.cmd("setMPIComm",comm)
+      p.cmd("init")
+
+
+if __name__ == "__main__":
+    unittest.main()