Merge pull request #1247 from germasch/python_fixes

WIP: Python fixes
ornladios · Mar 3, 2019 · 9930706 · 9930706
2 parents 9bd95e8 + aafda0a
commit 9930706
Show file tree

Hide file tree

Showing 6 changed files with 93 additions and 6 deletions.
diff --git a/bindings/Python/py11IO.cpp b/bindings/Python/py11IO.cpp
@@ -13,6 +13,10 @@
 #include "adios2/ADIOSMacros.h"
 #include "adios2/helper/adiosFunctions.h" //GetType<T>
 
+#ifdef ADIOS2_HAVE_MPI
+#include <mpi4py/mpi4py.h>
+#endif
+
 #include "py11types.h"
 
 namespace adios2
@@ -243,12 +247,20 @@ Engine IO::Open(const std::string &name, const int mode)
 }
 
 #ifdef ADIOS2_HAVE_MPI
-Engine IO::Open(const std::string &name, const Mode mode, MPI_Comm comm)
+Engine IO::Open(const std::string &name, const int mode, pybind11::object &comm)
 {
+    if (import_mpi4py() < 0)
+    {
+        throw std::runtime_error("ERROR: could not import mpi4py "
+                                 "communicator\n");
+    }
+
     helper::CheckForNullptr(m_IO,
                             "for engine " + name + ", in call to IO::Open");
+
+    MPI_Comm *mpiCommPtr = PyMPIComm_Get(comm.ptr());
     return Engine(
-        &m_IO->Open(name, static_cast<adios2::Mode>(mode), m_IO->m_MPIComm));
+        &m_IO->Open(name, static_cast<adios2::Mode>(mode), *mpiCommPtr));
 }
 #endif
 

diff --git a/bindings/Python/py11IO.h b/bindings/Python/py11IO.h
@@ -87,7 +87,8 @@ class IO
     Engine Open(const std::string &name, const int openMode);
 
 #ifdef ADIOS2_HAVE_MPI
-    Engine Open(const std::string &name, const Mode mode, MPI_Comm comm);
+    Engine Open(const std::string &name, const int openMode,
+                pybind11::object &comm);
 #endif
 
     void FlushAll();

diff --git a/bindings/Python/py11glue.cpp b/bindings/Python/py11glue.cpp
@@ -366,6 +366,10 @@ PYBIND11_MODULE(adios2, m)
         .def("Open", (adios2::py11::Engine (adios2::py11::IO::*)(
                          const std::string &, const int)) &
                          adios2::py11::IO::Open)
+        .def("Open",
+             (adios2::py11::Engine (adios2::py11::IO::*)(
+                 const std::string &, const int, pybind11::object &comm)) &
+                 adios2::py11::IO::Open)
         .def("AvailableVariables", &adios2::py11::IO::AvailableVariables)
         .def("AvailableAttributes", &adios2::py11::IO::AvailableAttributes)
         .def("FlushAll", &adios2::py11::IO::FlushAll)

diff --git a/examples/hello/bpReader/helloBPReaderHeatMap2D.py b/examples/hello/bpReader/helloBPReaderHeatMap2D.py
@@ -41,9 +41,8 @@
 adios = adios2.ADIOS(comm)
 ioWrite = adios.DeclareIO("ioWriter")
 
-varTemperature = ioWrite.DefineVariable("temperature2D", shape,
-                                        start, count, adios2.ConstantDims,
-                                        temperatures)
+varTemperature = ioWrite.DefineVariable("temperature2D", temperatures, shape,
+                                        start, count, adios2.ConstantDims)
 
 obpStream = ioWrite.Open('HeatMap2D_py.bp', adios2.Mode.Write)
 obpStream.Put(varTemperature, temperatures)

diff --git a/testing/adios2/bindings/python/CMakeLists.txt b/testing/adios2/bindings/python/CMakeLists.txt
@@ -17,6 +17,8 @@ if(ADIOS2_HAVE_MPI)
                     SCRIPT TestBPWriteTypesHighLevelAPI.py)
     python_add_test(NAME PythonBPReadMultisteps ${test_parameters} 
                     SCRIPT TestBPReadMultisteps.py)
+    python_add_test(NAME PythonBPWriteRead2D ${test_parameters} 
+                    SCRIPT TestBPWriteRead2D.py)
     if (ADIOS2_HAVE_HDF5)
         python_add_test(NAME PythonBPWriteReadHighLevelAPI_HDF5 
                         SCRIPT TestBPWriteTypesHighLevelAPI_HDF5.py)

diff --git a/testing/adios2/bindings/python/TestBPWriteRead2D.py b/testing/adios2/bindings/python/TestBPWriteRead2D.py
@@ -0,0 +1,69 @@
+#
+# Distributed under the OSI-approved Apache License, Version 2.0.  See
+# accompanying file Copyright.txt for details.
+#
+# TestBPWriteRead2D.py
+#
+#
+#  Created on: Mar 3rd, 2019
+#      Author: Kai Germaschewski <[email protected]>
+#              William F Godoy [email protected]
+#
+
+from mpi4py import MPI
+import numpy as np
+import adios2
+
+# MPI
+comm = MPI.COMM_WORLD
+rank = comm.Get_rank()
+size = comm.Get_size()
+
+# User data
+Nx = 10
+Ny = 10
+
+count = [Nx, Ny]
+start = [rank * Nx, 0]
+shape = [size * Nx, Ny]
+
+temperatures = np.zeros(count, dtype=np.int)
+
+for i in range(0, Nx):
+    for j in range(0, Ny):
+        temperatures[i, j] = (start[0] + i) * shape[1] + (j + start[1])
+
+# print(temperatures)
+# ADIOS2 read
+adios = adios2.ADIOS(comm)
+ioWrite = adios.DeclareIO("ioWriter")
+
+varTemperature = ioWrite.DefineVariable("temperature2D", temperatures, shape,
+                                        start, count, adios2.ConstantDims)
+
+obpStream = ioWrite.Open('HeatMap2D_py.bp', adios2.Mode.Write)
+obpStream.Put(varTemperature, temperatures)
+obpStream.Close()
+
+
+if rank == 0:
+    # ADIOS2 read
+    ioRead = adios.DeclareIO("ioReader")
+    ibpStream = ioRead.Open('HeatMap2D_py.bp', adios2.Mode.Read, MPI.COMM_SELF)
+    var_inTemperature = ioRead.InquireVariable("temperature2D")
+
+    assert var_inTemperature is not None
+    readOffset = [2, 2]
+    readSize = [4, 4]
+
+    var_inTemperature.SetSelection([readOffset, readSize])
+    inTemperatures = np.zeros(readSize, dtype=np.int)
+    ibpStream.Get(var_inTemperature, inTemperatures, adios2.Mode.Sync)
+    ibpStream.Close()
+
+    # print('Incoming temperature map\n', inTemperatures)
+    expected = np.array([[22, 23, 24, 25],
+                         [32, 33, 34, 35],
+                         [42, 43, 44, 45],
+                         [52, 53, 54, 55]], np.int)
+    assert np.array_equal(inTemperatures, expected)