Add python interface

RaulPPelaez · RaulPPelaez · commit 39d6b126a36d · 2022-01-25T20:01:47.000+01:00
diff --git a/.gitmodules b/.gitmodules
@@ -0,0 +1,3 @@
+[submodule "extern/pybind11"]
+	path = extern/pybind11
+	url = https://github.com/pybind/pybind11
diff --git a/Makefile b/Makefile
@@ -2,9 +2,13 @@
 
 NVCC=nvcc
 CXX=g++
+#To compile the python wrapper
+PYTHON3=python3
+#Pybind is cloned as a submodule to this location
+PYBIND_INCLUDE=extern/pybind11/include
 
 #Uncomment for a GPU enabled library
-CUDA_ENABLED=-DCUDA_ENABLED
+#CUDA_ENABLED=-DCUDA_ENABLED
 
 #Uncomment to compile in double precision mode
 #DOUBLE_PRECISION=-DDOUBLE_PRECISION
@@ -15,31 +19,44 @@ NVCCLDFLAGS= -lcublas
 LDFLAGS= -llapacke -lcblas
 
 LIBNAME=liblanczos.so
+PYTHON_MODULE_NAME=Lanczos
 
 
-CXXFLAGS=-fPIC -w -O3 -g -std=c++14 $(INCLUDEFLAGS) $(DOUBLE_PRECISION) $(CUDA_ENABLED)
+CXXFLAGS=-fPIC -w -O3 -g -std=c++14 $(INCLUDEFLAGS) $(DOUBLE_PRECISION)
 NVCCFLAGS=-ccbin=$(CXX) -Xcompiler "$(CXXFLAGS)" -std=c++14 -O3 $(INCLUDEFLAGS) $(DOUBLE_PRECISION) $(CUDA_ENABLED)
 
+PYTHON_LIBRARY_NAME=python/$(PYTHON_MODULE_NAME)$(shell $(PYTHON3)-config --extension-suffix)
+
 ifndef CUDA_ENABLED
 COMPILER=$(CXX)
-CXXFLAGS:=$(CXXFLAGS)  -xc++
+CXXFLAGS_BOTH:=$(CXXFLAGS)  -xc++
+LDFLAGS_BOTH:=$(LDFLAGS)
 else
 COMPILER=$(NVCC) 
-LDFLAGS:=$(LDFLAGS) $(NVCCLDFLAGS)
-CXXFLAGS:=$(NVCCFLAGS) -I$(CUDA_ROOT)/include
+LDFLAGS_BOTH:=$(LDFLAGS) $(NVCCLDFLAGS)
+CXXFLAGS_BOTH:=$(NVCCFLAGS) -I$(CUDA_ROOT)/include $(CUDA_ENABLED)
 endif
 
-all: shared $(patsubst %.cu, %, $(wildcard *.cu)) $(patsubst %.cpp, %, $(wildcard *.cpp))
+all: shared  python $(patsubst %.cu, %, $(wildcard *.cu)) $(patsubst %.cpp, %, $(wildcard *.cpp)) Makefile
 
 $(LIBNAME): $(wildcard include/*.cu)
-	$(COMPILER) -DSHARED_LIBRARY_COMPILATION -shared $(CXXFLAGS) $^ -o $@ $(LDFLAGS)
+	$(COMPILER) -DSHARED_LIBRARY_COMPILATION -shared $(CXXFLAGS_BOTH) $^ -o $@ $(LDFLAGS_BOTH)
+
+
+shared: $(LIBNAME) Makefile
+
 
+python: $(PYTHON_LIBRARY_NAME) Makefile
+#	-DLANCZOS_PYTHON_NAME=$(PYTHON_MODULE_NAME)
 
-shared: $(LIBNAME)
+$(PYTHON_LIBRARY_NAME): python/python_wrapper.cpp python/lanczos_trampoline.o
+	$(CXX) $(CXXFLAGS) `$(PYTHON3)-config --includes` -I $(PYBIND_INCLUDE) -shared  $^ -o $@ $(LDFLAGS)
 
+python/lanczos_trampoline.o: python/lanczos_trampoline.cpp Makefile
+	$(CXX) $(CXXFLAGS) -c $<  -o $@
 
 %: %.cu Makefile
-	$(COMPILER) $(CXXFLAGS) $<  -o $@ $(LDFLAGS)
+	$(COMPILER) $(CXXFLAGS_BOTH) $<  -o $@ $(LDFLAGS_BOTH)
 
 
 
@@ -48,4 +65,4 @@ shared: $(LIBNAME)
 #%.clean:
 #rm -f $(@:.clean=.so)
 clean:
-	rm -rf include/*.o $(LIBNAME) example
+	rm -rf include/*.o python/*.o $(LIBNAME) example $(PYTHON_LIBRARY_NAME)
diff --git a/README.md b/README.md
@@ -78,6 +78,14 @@ Note, however, that the heavy-weight of this solver comes from the Matrix-vector
 
 See the Makefile for further instructions.  
 
+## Python interface
+
+The python/ folder contains a python wrapper to the solver. A class defining the matrix vector product can be written directly in python an provided to the solver.
+See python/example.py for more information.
+
+The root folder's Makefile will try to compile the python library as well. It expects pybind11 to be placed under the extern/ folder. Pybind11 is included as a submodule, so make sure to clone this repository with --recursive.
+Note that the python wrapper can only be compiled in CPU mode.
+
 ## References:  
 
   [1] Krylov subspace methods for computing hydrodynamic interactions in Brownian dynamics simulations  J. Chem. Phys. 137, 064106 (2012); doi: 10.1063/1.4742347  
diff --git a/example.cpp b/example.cpp
@@ -23,10 +23,10 @@ struct DiagonalMatrix: public lanczos::MatrixDot{
   int size;
   DiagonalMatrix(int size): size(size){}
   
-  void operator()(real* v, real* Mv){
+  void dot(real* v, real* Mv) override{
     //an example diagonal matrix
     for(int i=0; i<size; i++){
-      Mv[i] = (2+i/10.0)*v[i];
+      Mv[i] = (2+i/10.0)*v[i]*2;
     }
   }
 
diff --git a/example.cu b/example.cu
@@ -23,7 +23,7 @@ struct DiagonalMatrix: public lanczos::MatrixDot{
   int size;
   DiagonalMatrix(int size): size(size){}
   
-  void operator()(real* v, real* Mv){
+  void dot(real* v, real* Mv) override{
     //An example diagonal matrix
     for(int i=0; i<size; i++){
       Mv[i] = (2+i/10.0)*v[i];
diff --git a/extern/pybind11 b/extern/pybind11
@@ -0,0 +1 @@
+Subproject commit 9ec1128c7aac3d069a4ec2bd1dfc7f57c6526d1c
diff --git a/include/LanczosAlgorithm.cu b/include/LanczosAlgorithm.cu
@@ -9,6 +9,10 @@ References:
 #include<string.h>
 #include"utils/lapack_and_blas_defines.h"
 #include<stdexcept>
+#ifdef CUDA_ENABLED
+#include"utils/debugTools.h"
+#endif
+
 namespace lanczos{
 
   Solver::Solver(real tolerance):
@@ -23,6 +27,18 @@ namespace lanczos{
 #endif    
   }
 
+  Solver::~Solver(){
+#ifdef CUDA_ENABLED
+    CublasSafeCall(cublasDestroy(cublas_handle));
+#endif
+  }
+
+  real* Solver::getV(int N){
+    if(N != this->N) numElementsChanged(N);
+    return detail::getRawPointer(V);
+  }
+
+  
   void Solver::numElementsChanged(int newN){
     this-> N = newN;
     try{
@@ -45,7 +61,7 @@ namespace lanczos{
     this->max_iter += inc;
   }
 
-    int Solver::solve(MatrixDot *dot, real *Bz, real*z, int N){
+    int Solver::solve(MatrixDot *dot, real *Bz, const real*z, int N){
     //Handles the case of the number of elements changing since last call
     if(N != this->N){
       real * d_V = detail::getRawPointer(V);
@@ -103,7 +119,8 @@ namespace lanczos{
     real* d_V =  detail::getRawPointer(V);
     real * d_w = detail::getRawPointer(w);
     /*w = D·vi*/
-    dot->operator()(d_V+N*i, d_w);
+    dot->setSize(N);
+    dot->dot(d_V+N*i, d_w);
     if(i>0){
       /*w = w-h[i-1][i]·vi*/
       real alpha = -hsup[i-1];
diff --git a/include/LanczosAlgorithm.h b/include/LanczosAlgorithm.h
@@ -24,49 +24,33 @@ Some notes:
 #include<vector>
 #include<memory>
 #include"utils/device_container.h"
-
-#ifdef CUDA_ENABLED
-#include"utils/debugTools.h"
-#endif
+#include"utils/MatrixDot.h"
 namespace lanczos{
 
-  struct MatrixDot{
-    
-    virtual void operator()(real* Mv, real*v) = 0;
-    
-  };
-  
   struct Solver{
     Solver(real tolerance = 1e-3);
 
-    ~Solver(){
-#ifdef CUDA_ENABLED
-      CublasSafeCall(cublasDestroy(cublas_handle));
-#endif
-    }
-
+    ~Solver();
+    
     //Given a Dotctor that computes a product M·v (where M is handled by Dotctor ), computes Bv = sqrt(M)·v
     //Returns the number of iterations performed
     //B = sqrt(M)
-    int solve(MatrixDot *dot, real *Bv, real* v, int N);
+    int solve(MatrixDot *dot, real *Bv, const real* v, int N);
     
     //Overload for a shared_ptr
-    int solve(std::shared_ptr<MatrixDot> dot, real *Bv, real* v, int N){
+    int solve(std::shared_ptr<MatrixDot> dot, real *Bv, const real* v, int N){
       return this->solve(dot.get(), Bv, v, N);
     }
 
     //Overload for an instance
     template<class SomeDot>
-    int solve(SomeDot &dot, real *Bv, real* v, int N){
+    int solve(SomeDot &dot, real *Bv, const real* v, int N){
       MatrixDot* ptr = static_cast<MatrixDot*>(&dot);
       return this->solve(ptr, Bv, v, N);
     }
 
     //You can use this array as input to the solve operation, which will save some memory
-    real * getV(int N){
-      if(N != this->N) numElementsChanged(N);
-      return detail::getRawPointer(V);
-    }
+    real * getV(int N);
 
 #ifdef CUDA_ENABLED
     //The solver will use this cuda stream when possible
diff --git a/include/utils/MatrixDot.h b/include/utils/MatrixDot.h
@@ -0,0 +1,13 @@
+#ifndef LANCZOS_MATRIX_DOT_H
+#define LANCZOS_MATRIX_DOT_H
+#include"defines.h"
+namespace lanczos{
+  
+  struct MatrixDot{
+    void setSize(int newsize){this->m_size = newsize;}
+    virtual void dot(real* v, real*Mv) = 0;
+  protected:
+    int m_size;
+  };
+}
+#endif
diff --git a/python/example.py b/python/example.py
@@ -0,0 +1,38 @@
+#Raul P. Pelaez 2022. Usage example for the Lanczos solver's Python interface
+#A class that computes the dot product of a matrix, M,  and an arbitrary vector, v, must be written to use the solver (see DiagonalMatrix below).
+#The class must inherit from Lanczos.MatrixDot and provide a function called "dot" that given an arbitrary vector, v, returns the product Mv.
+#When provided with an instance of this class, the function "solve" in Lanczos.Solver will return the product sqrt(M)v
+#IMPORTANT: Remeber to use the same numerical precision here and when compiling the library (see the Makefile for more info)
+#Try help(Lanczos)
+
+import Lanczos 
+import numpy as np
+
+#Lanczos provides the precision it was compiled in via this function.
+precision = np.float32 if Lanczos.getPrecision() else np.float64;
+
+
+# A simple class that computes the product of a diagonal matrix (2*I) by the input vector
+class DiagonalMatrix(Lanczos.MatrixDot):
+
+    def dot(self, v):
+#        size=v.size()
+        Mv = v*2.0
+        return Mv
+
+#Create the solver and provide a tolerance
+solver = Lanczos.Solver(tolerance=1e-3)
+
+#Let us compute the result of sqrt(2*I)*v, where v=[1,1,1....1] and I the identity matrix
+#The result vector will be filled with sqrt(2)
+size = 1000000
+result = np.zeros(size, precision);
+v = np.ones(size, precision);
+
+dotProduct = DiagonalMatrix()
+#The solve function fills the result vector with sqrt(M)*v and returns the number of iterations required to do so.
+numiter = solver.solve(dotProduct, result,v, size)
+
+print("Done after "+ str(numiter) + " iterations.")
+print("Result vector (should be filled with ~sqrt(2)="+str(np.sqrt(2))+"):")
+print(result)
diff --git a/python/lanczos_trampoline.cpp b/python/lanczos_trampoline.cpp
@@ -0,0 +1,15 @@
+/* Raul P. Pelaez 2022. Implementation for the trampoline lanczos class.
+Simply defer the calls to the actual lanczos solver code.
+*/
+#include"lanczos_trampoline.h"
+#include"../include/LanczosAlgorithm.h"
+
+using namespace lanczos;
+
+LanczosTrampoline::LanczosTrampoline(real tolerance):
+  solver(std::make_shared<lanczos::Solver>(tolerance)){}
+
+int LanczosTrampoline::solve(MatrixDot *dot, real* result, const real* v, int size){
+  int numberIterations = solver->solve(dot, result, v, size);
+  return numberIterations;
+}
diff --git a/python/lanczos_trampoline.h b/python/lanczos_trampoline.h
@@ -0,0 +1,24 @@
+/* Raul P. Pelaez 2022. A class that bridges with lanczos::Solver without actually including any of its code.
+This allows to compile the lanczos::Solver library separatedly from the python wrapper.
+
+ */
+#ifndef LANCZOS_TRAMPOLINE_H
+#define LANCZOS_TRAMPOLINE_H
+#include<memory>
+#include"../include/utils/MatrixDot.h"
+
+namespace lanczos{
+  class Solver;
+}
+  
+
+class LanczosTrampoline{
+  std::shared_ptr<lanczos::Solver> solver;
+public:
+  LanczosTrampoline(lanczos::real tolerance);
+
+  int solve(lanczos::MatrixDot *dot, lanczos::real* result, const lanczos::real* v, int size);
+  
+};
+
+#endif
diff --git a/python/python_wrapper.cpp b/python/python_wrapper.cpp
@@ -0,0 +1,69 @@
+/* Raul P. Pelaez 2022. Pybind11 python wrappers for the Lanczos solver library
+ */
+#include "../include/utils/MatrixDot.h"
+#include"lanczos_trampoline.h"
+#include<pybind11/pybind11.h>
+#include<pybind11/numpy.h>
+#include<iostream>
+
+using real = lanczos::real;
+
+namespace py = pybind11;
+using namespace pybind11::literals;
+
+
+//This class allows to inherit lanczos::MatrixDot from python
+struct MatrixDotTrampoline: public lanczos::MatrixDot{  
+  using MatrixDot::MatrixDot;
+  
+  void dot(real* v, real* Mv) override{
+    pybind11::gil_scoped_acquire gil;  // Acquire the GIL while in this scope.
+    // Try to look up the overridden method on the Python side.
+    pybind11::function overridef = pybind11::get_override(this, "dot");
+    if (overridef) {  // method is found
+      //Pybind needs a dummy object to know about the lifetime of the py::arrays.
+      //without it it will copy the memory.
+      py::str dummy;
+      // Call the Python function.
+      py::array_t<real> Mvp = overridef(py::array_t<real, py::array::c_style>(this->m_size, v, dummy));
+      std::copy(Mvp.data(), Mvp.data()+this->m_size, Mv);
+    }
+  }
+};
+
+//Python wrapper class for the Lanczos solver.
+class PyLanczos{
+  std::shared_ptr<LanczosTrampoline> solver;
+public:
+  PyLanczos(real tolerance):
+    solver(std::make_shared<LanczosTrampoline>(tolerance)){}
+
+  int solve(lanczos::MatrixDot* dot, py::array_t<real> &result,  py::array_t<real> &v, int size){
+    dot->setSize(size);
+    return solver->solve(dot, result.mutable_data(), v.data(), size);
+  }
+
+};
+
+#ifndef LANCZOS_PYTHON_NAME
+#define LANCZOS_PYTHON_NAME Lanczos
+#endif
+
+
+PYBIND11_MODULE(LANCZOS_PYTHON_NAME, m){
+  py::class_<lanczos::MatrixDot, MatrixDotTrampoline>(m, "MatrixDot", "The virtual class required by the Lanczos solver").
+    def(py::init<>()).
+    def("dot", &lanczos::MatrixDot::dot,
+	"Given a result (Mv) and a vector (v), this method must write in Mv the result of multiplying the target matrix and v.",
+	"v"_a, "The input vector", "Mv"_a, "The output result vector");
+
+  py::class_<PyLanczos>(m, "Solver", "A Lanczos iterative solver. Computes sqrt(M)*v in O(N^2) operations, being N the size of the matrix, which must be square.").
+    def(py::init<real>(), "tolerance"_a, "The tolerance of the algorithm").
+    //MatrixDot *dot, real* result, real* v, int size
+    def("solve", &PyLanczos::solve,
+	"dot"_a, "A pointer to the matrix multiplication functor",
+	"result"_a, "Output array storing the result of sqrt(M)*v",
+	"v"_a, "The input array to be multiplied by sqrt(M)",
+	"size"_a, "The size of the input array (and the matrix, size x size)");
+  m.def("getPrecision", [](){return (std::is_same<lanczos::real, float>::value)?"float":"double";});
+}

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,3 @@`
	`1`	`+[submodule "extern/pybind11"]`
	`2`	`+ path = extern/pybind11`
	`3`	`+ url = https://github.com/pybind/pybind11`
Original file line number	Diff line number	Diff line change
`@@ -23,10 +23,10 @@ struct DiagonalMatrix: public lanczos::MatrixDot{`
`23`	`23`	`int size;`
`24`	`24`	`DiagonalMatrix(int size): size(size){}`
`25`	`25`
`26`		`- void operator()(real* v, real* Mv){`
	`26`	`+ void dot(real* v, real* Mv) override{`
`27`	`27`	`//an example diagonal matrix`
`28`	`28`	`for(int i=0; i<size; i++){`
`29`		`- Mv[i] = (2+i/10.0)*v[i];`
	`29`	`+ Mv[i] = (2+i/10.0)v[i]2;`
`30`	`30`	`}`
`31`	`31`	`}`
`32`	`32`