pass custom mpi4py communicator

pypolychord/_pypolychord.cpp

@@ -5,19 +5,29 @@
 #define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION
 #include <numpy/arrayobject.h>
 
+#ifdef USE_MPI
+#include <mpi.h>
+#include <mpi4py/mpi4py.h>
+#endif
 
 /* Initialize the module */
 #ifdef PYTHON3
 PyMODINIT_FUNC PyInit__pypolychord(void)
 {
  import_array();
+#ifdef USE_MPI
+ import_mpi4py();
+#endif
  return PyModule_Create(&_pypolychordmodule);
 }
 #else
 PyMODINIT_FUNC init_pypolychord(void)
 {
  Py_InitModule3("_pypolychord", module_methods, module_docstring);
  import_array();
+#ifdef USE_MPI
+ import_mpi4py();
+#endif
 }
 #endif
 
@@ -29,13 +39,13 @@ static PyObject *python_loglikelihood = NULL;
 double loglikelihood(double* theta, int nDims, double* phi, int nDerived)
 {
  /* Create a python version of theta */
- npy_intp theta_shape[] = {nDims}; 
+ npy_intp theta_shape[] = {nDims};
  PyObject *array_theta = PyArray_SimpleNewFromData(1, theta_shape, NPY_DOUBLE, theta);
  if (array_theta==NULL) throw PythonException();
  PyArray_CLEARFLAGS(reinterpret_cast<PyArrayObject*>(array_theta), NPY_ARRAY_WRITEABLE);
 
  /* Create a python version of phi */
- npy_intp phi_shape[] = {nDerived}; 
+ npy_intp phi_shape[] = {nDerived};
  PyObject *array_phi = PyArray_SimpleNewFromData(1, phi_shape, NPY_DOUBLE, phi);
  if (array_phi==NULL) {Py_DECREF(array_theta); throw PythonException();}
 
@@ -63,7 +73,7 @@ static PyObject *python_prior = NULL;
 void prior(double* cube, double* theta, int nDims)
 {
  /* create a python version of cube */
- npy_intp shape[] = {nDims}; 
+ npy_intp shape[] = {nDims};
  PyObject *array_cube = PyArray_SimpleNewFromData(1, shape, NPY_DOUBLE, cube);
  if (array_cube==NULL) throw PythonException();
  PyArray_CLEARFLAGS(reinterpret_cast<PyArrayObject*>(array_cube), NPY_ARRAY_WRITEABLE);
@@ -85,19 +95,19 @@ static PyObject *python_dumper = NULL;
 void dumper(int ndead, int nlive, int npars, double* live, double* dead, double* logweights, double logZ, double logZerr)
 {
  /* create a python version of live points */
- npy_intp live_shape[] = {nlive, npars}; 
+ npy_intp live_shape[] = {nlive, npars};
  PyObject *array_live = PyArray_SimpleNewFromData(2, live_shape, NPY_DOUBLE, live);
  if (array_live==NULL) throw PythonException();
  PyArray_CLEARFLAGS(reinterpret_cast<PyArrayObject*>(array_live), NPY_ARRAY_WRITEABLE);
 
  /* create a python version of dead points */
- npy_intp dead_shape[] = {ndead, npars}; 
+ npy_intp dead_shape[] = {ndead, npars};
  PyObject *array_dead = PyArray_SimpleNewFromData(2, dead_shape, NPY_DOUBLE, dead);
  if (array_dead==NULL) {Py_DECREF(array_live); throw PythonException();}
  PyArray_CLEARFLAGS(reinterpret_cast<PyArrayObject*>(array_dead), NPY_ARRAY_WRITEABLE);
 
  /* create a python version of posterior weights */
- npy_intp logweights_shape[] = {ndead}; 
+ npy_intp logweights_shape[] = {ndead};
  PyObject *array_logweights = PyArray_SimpleNewFromData(1, logweights_shape, NPY_DOUBLE, logweights);
  if (array_logweights==NULL) {Py_DECREF(array_live); Py_DECREF(array_dead); throw PythonException();}
  PyArray_CLEARFLAGS(reinterpret_cast<PyArrayObject*>(array_logweights), NPY_ARRAY_WRITEABLE);
@@ -121,12 +131,11 @@ static PyObject *run_pypolychord(PyObject *, PyObject *args)
  Settings S;
 
  PyObject *temp_logl, *temp_prior, *temp_dumper;
- PyObject* py_grade_dims, *py_grade_frac, *py_nlives;
+ PyObject* py_grade_dims, *py_grade_frac, *py_nlives, *py_comm;
  char* base_dir, *file_root;
-
 
  if (!PyArg_ParseTuple(args,
- "OOOiiiiiiiiddidiiiiiiiiiiidissO!O!O!i:run",
+ "OOOiiiiiiiiddidiiiiiiiiiiidissO!O!O!iO:run",
  &temp_logl,
  &temp_prior,
  &temp_dumper,
@@ -163,13 +172,14 @@ static PyObject *run_pypolychord(PyObject *, PyObject *args)
  &py_grade_dims,
  &PyDict_Type,
  &py_nlives,
- &S.seed
+ &S.seed,
+ &py_comm
  )
  )
  return NULL;
  S.base_dir = base_dir;
  S.file_root = file_root;
- 
+
  int nGrade = PyList_Size(py_grade_frac);
  S.grade_frac.resize(nGrade);
  S.grade_dims.resize(nGrade);
@@ -216,14 +226,20 @@ static PyObject *run_pypolychord(PyObject *, PyObject *args)
  python_dumper = temp_dumper;
 
  /* Run PolyChord */
- try{ run_polychord(loglikelihood, prior, dumper, S); }
+ try{
+#ifdef USE_MPI
+ MPI_Comm *comm = PyMPIComm_Get(py_comm);
+ run_polychord(loglikelihood, prior, dumper, S, *comm);
+#else
+ run_polychord(loglikelihood, prior, dumper, S);
+#endif
+ }
  catch (PythonException& e)
- { 
+ {
  Py_DECREF(py_grade_frac);Py_DECREF(py_grade_dims);Py_DECREF(python_loglikelihood);Py_DECREF(python_prior);
- return NULL; 
+ return NULL;
  }
 
  /* Return None */
  Py_RETURN_NONE;
 }
-

pypolychord/polychord.py

@@ -13,7 +13,7 @@ def default_dumper(live, dead, logweights, logZ, logZerr):
 
 
 def run_polychord(loglikelihood, nDims, nDerived, settings,
- prior=default_prior, dumper=default_dumper):
+ prior=default_prior, dumper=default_dumper, comm=None):
  """
  Runs PolyChord.
 
@@ -144,20 +144,22 @@ def run_polychord(loglikelihood, nDims, nDerived, settings,
  Final output evidence statistics
 
  """
-
- try:
- from mpi4py import MPI
- comm = MPI.COMM_WORLD
- rank = comm.Get_rank()
- except ImportError:
- rank = 0
+ rank = 0
+ if comm is None:
+ try:
+ from mpi4py import MPI
+ comm = MPI.COMM_WORLD
+ except ImportError:
+ pass
+ if comm is not None:
+ rank = comm.rank
 
  try:
  if rank == 0:
  os.makedirs(settings.base_dir)
  except OSError:
  pass
- 
+
  try:
  if rank == 0:
  os.makedirs(settings.cluster_dir)
@@ -213,7 +215,8 @@ def wrap_prior(cube, theta):
  settings.grade_frac,
  settings.grade_dims,
  settings.nlives,
- settings.seed)
+ settings.seed,
+ comm)
 
  if settings.cube_samples is not None:
  settings.read_resume = read_resume

run_pypolychord.py

@@ -73,4 +73,4 @@ def dumper(live, dead, logweights, logZ, logZerr):
  except ImportError:
  print("Install matplotlib and getdist for plotting examples")
 
- print("Install anesthetic or getdist  for for plotting examples")
+ print("Install anesthetic or getdist for plotting examples")

setup.py

@@ -3,21 +3,37 @@
 
 Polychord is a tool to solve high dimensional problems.
 """
+import os
+import sys
+import subprocess
+import shutil
 
 from setuptools import setup, Extension, find_packages, Distribution
 from setuptools.command.build_py import build_py as _build_py
+from setuptools.command.develop import develop as _develop
 from distutils.command.clean import clean as _clean
 
-import os, sys, subprocess, shutil
-
 import numpy
 
-def check_compiler(default_CC="gcc"):
+
+try:
+ import mpi4py
+except ImportError:
+ mpi4py_get_include = None
+else:
+ mpi4py_get_include = mpi4py.get_include()
+
+
+def check_compiler(default_CC=None):
  """Checks what compiler is being used (clang, intel, or gcc)."""
 
- CC = default_CC if "CC" not in os.environ else os.environ["CC"]
+ default_CC = 'gcc'
+ if mpi4py_get_include:
+ default_CC = 'mpicc'
+ CC = os.getenv('CC', default_CC)
+ os.environ['CC'] = CC
  CC_version = subprocess.check_output([CC, "-v"], stderr=subprocess.STDOUT).decode("utf-8").lower()
- 
+
  if "clang" in CC_version:
  CC_family = "clang"
  elif "icc" in CC_version:
@@ -63,64 +79,79 @@ def get_version(short=False):
  return line[44:50]
 
 
-class DistributionWithOption(Distribution, object):
+class DistributionWithOption(Distribution):
+
  def __init__(self, *args, **kwargs):
- self.global_options = self.global_options \
- + [("no-mpi", None, "Don't compile with MPI support."),
- ("debug-flags", None, "Compile in debug mode.")]
+ self.global_options = self.global_options + [("no-mpi", None, "Compile without MPI support."), ("debug-flags", None, "Compile in debug mode.")]
  self.no_mpi = None
  self.debug_flags = None
  super(DistributionWithOption, self).__init__(*args, **kwargs)
 
-class CustomBuildPy(_build_py, object):
+
+class CustomBuildPy(_build_py):
+
  def run(self):
- env = {}
- env["PATH"] = os.environ["PATH"]
  if self.distribution.no_mpi is None:
- env["MPI"] = "1"
+ os.environ["MPI"] = "1"
  # These need to be set so that build_ext uses the right compilers
- cc_compiler = subprocess.check_output(["make", "print_CC"]).decode('utf-8').strip()
- os.environ["CC"] = cc_compiler
+ #cc_compiler = subprocess.check_output(["make", "print_CC"]).decode('utf-8').strip()
+ #os.environ"CC"] = cc_compiler
 
- cxx_compiler = subprocess.check_output(["make", "print_CXX"]).decode('utf-8').strip()
- os.environ["CXX"] = cxx_compiler
+ #cxx_compiler = subprocess.check_output(["make", "print_CXX"]).decode('utf-8').strip()
+ #os.environ["CXX"] = cxx_compiler
  else:
- env["MPI"] = "0"
+ os.environ["MPI"] = "0"
 
  if self.distribution.debug_flags is not None:
  self.distribution.ext_modules[0].extra_compile_args += ["-g", "-O0"]
- env["DEBUG"] = "1"
- 
+ os.environ["DEBUG"] = "1"
+
  BASE_PATH = os.path.dirname(os.path.abspath(__file__))
- env["PWD"] = BASE_PATH
- env.update({k : os.environ[k] for k in ["CC", "CXX", "FC"] if k in os.environ})
- subprocess.check_call(["make", "-e", "libchord.so"], env=env, cwd=BASE_PATH)
+ os.environ["PWD"] = BASE_PATH
+ subprocess.check_call(["make", "-e", "libchord.so"], cwd=BASE_PATH)
  if not os.path.isdir("pypolychord/lib/"):
  os.makedirs(os.path.join(BASE_PATH, "pypolychord/lib/"))
- shutil.copy(os.path.join(BASE_PATH, "lib/libchord.so"), 
+ shutil.copy(os.path.join(BASE_PATH, "lib/libchord.so"),
  os.path.join(BASE_PATH, "pypolychord/lib/"))
  self.run_command("build_ext")
  return super(CustomBuildPy, self).run()
 
+
+class CustomDevelop(_develop):
+
+ def run(self):
+ self.run_command('build_py')
+ super(CustomDevelop, self).run()
+
+
 class CustomClean(_clean):
+
  def run(self):
  subprocess.run(["make", "veryclean"], check=True, env=os.environ)
  return super().run()
 
+
+include_dirs = ['src/polychord', numpy.get_include()]
+
 if "--no-mpi" in sys.argv:
  NAME += '_nompi'
  DOCLINES[1] = DOCLINES[1] + ' (cannot be used with MPI)'
 
-pypolychord_module = Extension(
- name='_pypolychord',
- library_dirs=['lib'],
- include_dirs=['src/polychord', numpy.get_include()],
- libraries=['chord',],
- extra_link_args=RPATH_FLAG + CPPRUNTIMELIB_FLAG,
- extra_compile_args= ["-std=c++11"] + RPATH_FLAG + CPPRUNTIMELIB_FLAG,
- runtime_library_dirs=['lib'],
- sources=['pypolychord/_pypolychord.cpp']
- )
+elif mpi4py_get_include:
+ CPPRUNTIMELIB_FLAG += ["-DUSE_MPI"]
+ print(mpi4py_get_include)
+ include_dirs += [mpi4py_get_include]
+
+
+pypolychord_module = Extension(name='_pypolychord',
+ library_dirs=['lib'],
+ include_dirs=include_dirs,
+ libraries=['chord'],
+ extra_link_args=RPATH_FLAG + CPPRUNTIMELIB_FLAG,
+ extra_compile_args=["-std=c++11"] + RPATH_FLAG + CPPRUNTIMELIB_FLAG,
+ runtime_library_dirs=['lib'],
+ sources=['pypolychord/_pypolychord.cpp'])
+
 
 setup(name=NAME,
  version=get_version(),
@@ -130,12 +161,13 @@ def run(self):
  author_email='wh260@cam.ac.uk',
  license='PolyChord',
  packages=find_packages(),
- install_requires=['numpy','scipy'],
+ install_requires=['numpy', 'scipy'],
  extras_require={'plotting': 'getdist'},
  distclass=DistributionWithOption,
  ext_modules=[pypolychord_module],
- cmdclass={'build_py' : CustomBuildPy,
- 'clean' : CustomClean},
- package_data={"" : ["lib/libchord.so"]},
+ cmdclass={'build_py': CustomBuildPy,
+ 'develop': CustomDevelop,
+ 'clean': CustomClean},
+ package_data={NAME: ["lib/libchord.so"]},
  include_package_data=True,
  zip_safe=False)

src/drivers/polychord_CC.cpp

@@ -36,7 +36,7 @@ int main()
  settings.boost_posterior= 5.0;
 
  setup_loglikelihood();
- run_polychord(loglikelihood,prior,dumper,settings) ;
+ run_polychord(loglikelihood,prior,dumper,settings);
 
 
 }

src/drivers/polychord_CC_ini.cpp

@@ -7,12 +7,11 @@ int main(int argc, char *argv[])
  if (argc > 1) {
  std::string input_file = argv[1];
  set_ini(input_file);
- run_polychord(loglikelihood,setup_loglikelihood,input_file) ;
+ run_polychord(loglikelihood,setup_loglikelihood,input_file);
  return 0;
  }
  else{
  std::cerr << "PolyChord should be called with at most one argument, the input file" << std::endl;
  return 1;
  }
 }
-