From bd72ef7996c345098a1327c5a28c4538ae04d049 Mon Sep 17 00:00:00 2001
From: Axel Kohlmeyer <akohlmey@gmail.com>
Date: Wed, 9 Sep 2020 22:03:11 -0400
Subject: [PATCH] add API to library/python interface to extract the MPI
 communicator

---
 doc/src/pg_lib_properties.rst |  5 +++++
 python/lammps.py              | 24 ++++++++++++++++++++++++
 src/library.cpp               | 34 ++++++++++++++++++++++++++++++++++
 src/library.h                 |  1 +
 4 files changed, 64 insertions(+)

diff --git a/doc/src/pg_lib_properties.rst b/doc/src/pg_lib_properties.rst
index 789c7e2441..cdc7617d91 100644
--- a/doc/src/pg_lib_properties.rst
+++ b/doc/src/pg_lib_properties.rst
@@ -21,6 +21,11 @@ event as atoms are migrating between sub-domains.
 
 -----------------------
 
+.. doxygenfunction:: lammps_get_mpi_comm
+   :project: progguide
+
+-----------------------
+
 .. doxygenfunction:: lammps_get_natoms
    :project: progguide
 
diff --git a/python/lammps.py b/python/lammps.py
index b76608af7d..6e910b49bf 100644
--- a/python/lammps.py
+++ b/python/lammps.py
@@ -332,6 +332,8 @@ class lammps(object):
 
     self.lib.lammps_version.argtypes = [c_void_p]
 
+    self.lib.lammps_get_mpi_comm.argtypes = [c_void_p]
+
     self.lib.lammps_decode_image_flags.argtypes = [self.c_imageint, POINTER(c_int*3)]
 
     self.lib.lammps_extract_atom.argtypes = [c_void_p, c_char_p]
@@ -601,6 +603,28 @@ class lammps(object):
 
   # -------------------------------------------------------------------------
 
+  def get_mpi_comm(self):
+    """Get the MPI communicator in use by the current LAMMPS instance
+
+    This is a wrapper around the :cpp:func:`lammps_get_mpi_comm` function
+    of the C-library interface.  It will return ``None`` if either the
+    LAMMPS library was compiled without MPI support or the mpi4py
+    Python module is not available.
+
+    :return: MPI communicator
+    :rtype:  MPI_Comm
+    """
+
+    if self.has_mpi4py and self.has_mpi_support:
+        from mpi4py import MPI
+        f_comm = self.lib.lammps_get_mpi_comm(self.lmp)
+        c_comm = MPI.Comm.f2py(f_comm)
+        return c_comm
+    else:
+        return None
+
+  # -------------------------------------------------------------------------
+
   def file(self, path):
     """Read LAMMPS commands from a file.
 
diff --git a/src/library.cpp b/src/library.cpp
index 72105bdd15..0e0c4070ea 100644
--- a/src/library.cpp
+++ b/src/library.cpp
@@ -532,6 +532,40 @@ int lammps_version(void *handle)
   return atoi(lmp->universe->num_ver);
 }
 
+
+/* ---------------------------------------------------------------------- */
+
+/** Return current LAMMPS world communicator as integer
+ *
+\verbatim embed:rst
+
+This will take the LAMMPS "world" communicator and convert it to an
+integer using ``MPI_Comm_c2f()``, so it is equivalent to the
+corresponding MPI communicator in Fortran. This way it can be safely
+passed around between different programming languages.  To convert it
+to the C language representation use ``MPI_Comm_f2c()``.
+
+If LAMMPS was compiled with MPI_STUBS, this function returns -1.
+
+.. versionadded:: 15Sep2020
+
+\endverbatim
+ * \sa lammps_open_fortran
+ *
+ * \param  handle  pointer to a previously created LAMMPS instance
+ * \return         Fortran representation of the LAMMPS world communicator */
+
+int lammps_get_mpi_comm(void *handle)
+{
+#ifdef MPI_STUBS
+  return -1;
+#else
+  LAMMPS *lmp = (LAMMPS *) handle;
+  MPI_Fint f_comm = MPI_Comm_c2f(lmp->world);
+  return f_comm;
+#endif
+}
+
 /* ---------------------------------------------------------------------- */
 
 /** Return the total number of atoms in the system.
diff --git a/src/library.h b/src/library.h
index 0ffd111d7b..f5a02502e7 100644
--- a/src/library.h
+++ b/src/library.h
@@ -98,6 +98,7 @@ void  lammps_commands_string(void *handle, const char *str);
  * ----------------------------------------------------------------------- */
 
 int    lammps_version(void *handle);
+int    lammps_get_mpi_comm(void* handle);
 double lammps_get_natoms(void *handle);
 double lammps_get_thermo(void *handle, char *keyword);