add first version of python wrapper

- depends on python, cython, mpi4py and numpy
- currently only dependency check on python
- should be done optionally
- converts wrapper.pyx to wrapper.c using cython, then compiles this to
  a python extension
- installs the files such that they can be imported with `import pyelpa`
  in python

Makefile.am

@@ -612,6 +612,17 @@ endif
 #test_c_cannon@SUFFIX@_LDADD = $(test_program_ldadd) $(FCLIBS)
 #test_c_cannon@SUFFIX@_CFLAGS = $(test_program_cflags)
 
+# python wrapper
+pyelpadir = $(pythondir)/pyelpa
+pyelpa_PYTHON = python/__init__.py python/distributedmatrix.py
+pyelpa_LTLIBRARIES = wrapper.la
+wrapper_la_SOURCES = python/wrapper.c
+wrapper_la_LDFLAGS = -module -avoid-version -shared $(AM_LDFLAGS)
+wrapper_la_LIBADD = libelpa@SUFFIX@.la
+wrapper_la_CFLAGS = -I$(PYTHON_INCLUDE) $(AM_CFLAGS)
+
+python/wrapper.c: python/wrapper.pyx
+	cython $<
 
 
 # test scripts

configure.ac

@@ -1275,6 +1275,22 @@ else
         AC_MSG_RESULT([no])
 fi
 
+AM_PATH_PYTHON([3.6])
+AC_ARG_VAR([PYTHON_INCLUDE], [Include flags for python, bypassing python-config])
+AC_ARG_VAR([PYTHON_CONFIG], [Path to python-config])
+AS_IF([test -z "$PYTHON_INCLUDE"], [
+  AS_IF([test -z "$PYTHON_CONFIG"], [
+    AC_PATH_PROGS([PYTHON_CONFIG],
+                  [python$PYTHON_VERSION-config python-config],
+                  [no],
+                  [`dirname $PYTHON`])
+    AS_IF([test "$PYTHON_CONFIG" = no], [AC_MSG_ERROR([cannot find python-config for $PYTHON.])])
+  ])
+  AC_MSG_CHECKING([python include flags])
+  PYTHON_INCLUDE=`$PYTHON_CONFIG --includes`
+  AC_MSG_RESULT([$PYTHON_INCLUDE])
+])
+
 
 AC_OUTPUT
 

python/__init__.py

+"""pyelpa -- python wrapper for ELPA
+
+This wrapper uses cython to wrap the C API of ELPA (Eigenvalue SoLvers for
+Petaflop-Applications) so that it can be called from python.
+
+Examples:
+
+1. Use the Elpa object to access the eigenvectors/eigenvalues wrapper:
+
+>>> import numpy as np
+... from pyelpa import ProcessorLayout, DistributedMatrix, Elpa
+... from mpi4py import MPI
+... import sys
+... 
+... # set some parameters for matrix layout
+... na = 1000
+... nev = 200
+... nblk = 16
+... 
+... # initialize processor layout, needed for calling ELPA
+... comm = MPI.COMM_WORLD
+... layout_p = ProcessorLayout(comm)
+... 
+... # create arrays
+... a = DistributedMatrix(layout_p, na, nev, nblk)
+... eigenvectors = DistributedMatrix(layout_p, na, nev, nblk)
+... eigenvalues = np.zeros(na, dtype=np.float64)
+... 
+... # initialize elpa
+... e = Elpa.from_distributed_matrix(a)
+... 
+... # set input matrix (a.data) on this core (a is stored in a block-cyclic
+... # distributed layout; local size: a.na_rows x a.na_cols)
+... # Caution: using this, the global matrix will not be symmetric; this is just
+... # and example to show how to access the data
+... a.data[:, :] = np.random.rand(a.na_rows, a.na_cols).astype(np.float64)
+... 
+... # now compute nev of na eigenvectors and eigenvalues
+... e.eigenvectors(a.data, eigenvalues, eigenvectors.data)
+... 
+... # now eigenvectors.data contains the local part of the eigenvector matrix
+... # which is stored in a block-cyclic distributed layout
+... 
+... # now eigenvalues contains all computed eigenvalues on all cores
+... 
+... # now compute nev of na eigenvalues
+... e.eigenvalues(a.data, eigenvalues)
+... 
+... # now eigenvalues contains all computed eigenvalues on all cores
+
+
+2. Use the functions provided by the DistributedMatrix object:
+
+>>> import numpy as np
+... from pyelpa import DistributedMatrix
+... 
+... # set some parameters for matrix layout
+... na = 1000
+... nev = 200
+... nblk = 16
+... 
+... a = DistributedMatrix.from_comm_world(na, nev, nblk)
+... # use a diagonal matrix as input
+... matrix = np.diagflat(np.arange(na)**2)
+... # set from global matrix
+... a.set_data_from_global_matrix(matrix)
+... 
+... data = a.compute_eigenvectors()
+... eigenvalues = data['eigenvalues']
+... eigenvectors = data['eigenvectors']
+... # now eigenvectors.data contains the local part of the eigenvector matrix
+... # which is stored in a block-cyclic distributed layout
+... 
+... # now eigenvalues contains all computed eigenvalues on all cores
+"""
+from .wrapper import Elpa
+from .distributedmatrix import ProcessorLayout, DistributedMatrix
+
+__all__ = ['ProcessorLayout', 'DistributedMatrix', 'Elpa']

python/distributedmatrix.py

+"""distributedmatrix.py -- classes for distributed matrices
+
+This file contains the python classes to use with the wrapper.
+"""
+import numpy as np
+from functools import wraps
+from .wrapper import Elpa
+
+class ProcessorLayout:
+    """Create rectangular processor layout for use with distributed matrices"""
+    def __init__(self, comm):
+        """Initialize processor layout.
+
+        Args:
+            comm: MPI communicator from mpi4py
+        """
+        nprocs = comm.Get_size()
+        rank = comm.Get_rank()
+        for np_cols in range(int(np.sqrt(nprocs)), 0, -1):
+            if nprocs % np_cols == 0:
+              break
+        #if nprocs == 1:
+        #    np_cols = 1
+        np_rows = nprocs//np_cols
+        # column major distribution of processors
+        my_pcol = rank // np_rows
+        my_prow = rank % np_rows
+        self.np_cols, self.np_rows = np_cols, np_rows
+        self.my_pcol, self.my_prow = my_pcol, my_prow
+        self.comm = comm
+        self.comm_f = comm.py2f()
+
+
+class DistributedMatrix:
+    """Class for generating a distributed block-cyclic matrix
+
+    The data attribute contains the array in the correct size for the local
+    processor.
+    """
+    def __init__(self, processor_layout, na, nev, nblk, dtype=np.float64):
+        """Initialize distributed matrix for a given processor layout.
+
+        Args:
+            processor_layout (ProcessorLayout): has to be created from MPI
+                communicator
+            na (int): dimension of matrix
+            nev (int): number of eigenvectors/eigenvalues to be computed
+            nblk (int): block size of distributed matrix
+            dtype: data type of matrix
+        """
+        self.na = na
+        self.nev = nev
+        self.nblk = nblk
+        self.processor_layout = processor_layout
+
+        # get local size
+        self.na_rows = self.numroc(na, nblk, processor_layout.my_prow, 0,
+                                   processor_layout.np_rows)
+        self.na_cols = self.numroc(na, nblk, processor_layout.my_pcol, 0,
+                                   processor_layout.np_cols)
+        # create array
+        self.data = np.empty((self.na_rows, self.na_cols),
+                             dtype=dtype, order='F')
+
+        self.elpa = None
+
+    @classmethod
+    def from_communicator(cls, comm, na, nev, nblk, dtype=np.float64):
+        """Initialize distributed matrix from a MPI communicator.
+
+        Args:
+            comm: MPI communicator from mpi4py
+            na (int): dimension of matrix
+            nev (int): number of eigenvectors/eigenvalues to be computed
+            nblk (int): block size of distributed matrix
+            dtype: data type of matrix
+        """
+        processor_layout = ProcessorLayout(comm)
+        return cls(processor_layout, na, nev, nblk, dtype)
+
+    @classmethod
+    def from_comm_world(cls, na, nev, nblk, dtype=np.float64):
+        """Initialize distributed matrix from the MPI_COMM_WORLD communicator.
+
+        Args:
+            na (int): dimension of matrix
+            nev (int): number of eigenvectors/eigenvalues to be computed
+            nblk (int): block size of distributed matrix
+            dtype: data type of matrix
+        """
+        from mpi4py import MPI
+        comm = MPI.COMM_WORLD
+        processor_layout = ProcessorLayout(comm)
+        return cls(processor_layout, na, nev, nblk, dtype)
+
+    @classmethod
+    def like(cls, matrix):
+        """Get a DistributedMatrix with the same parameters as matrix"""
+        return cls(matrix.processor_layout, matrix.na, matrix.nev, matrix.nblk,
+                   matrix.data.dtype)
+
+    def get_local_index(self, global_row, global_col):
+        """compute local row and column indices from global ones
+
+        Returns a tuple of the local row and column indices
+        """
+        local_row = self.indxg2l(global_row, self.nblk,
+                                 self.processor_layout.my_prow, 0,
+                                 self.processor_layout.np_rows)
+        local_col = self.indxg2l(global_col, self.nblk,
+                                 self.processor_layout.my_pcol, 0,
+                                 self.processor_layout.np_cols)
+        return local_row, local_col
+
+    def get_global_index(self, local_row, local_col):
+        """compute global row and column indices from local ones
+
+        Returns a tuple of the global row and column indices
+        """
+        global_row = self.indxl2g(local_row, self.nblk,
+                                  self.processor_layout.my_prow, 0,
+                                  self.processor_layout.np_rows)
+        global_col = self.indxl2g(local_col, self.nblk,
+                                  self.processor_layout.my_pcol, 0,
+                                  self.processor_layout.np_cols)
+        return global_row, global_col
+
+    def is_local_index(self, global_row, global_col):
+        """check if global index is stored on current processor"""
+        return self.is_local_row(global_row) and self.is_local_col(global_col)
+
+    def is_local_row(self, global_row):
+        """check if global row is stored on this processor"""
+        process_row = self.indxg2p(global_row, self.nblk,
+                                   self.processor_layout.my_prow, 0,
+                                   self.processor_layout.np_rows)
+        return process_row == self.processor_layout.my_prow
+
+    def is_local_col(self, global_col):
+        process_col = self.indxg2p(global_col, self.nblk,
+                                   self.processor_layout.my_pcol, 0,
+                                   self.processor_layout.np_cols)
+        return process_col == self.processor_layout.my_pcol
+
+    @staticmethod
+    def indxg2l(indxglob, nb, iproc, isrcproc, nprocs):
+        """compute local index from global index indxglob
+
+        original netlib scalapack source:
+
+        .. code-block:: fortran
+
+            INDXG2L = NB*((INDXGLOB-1)/(NB*NPROCS))+MOD(INDXGLOB-1,NB)+1
+        """
+        # adapt to python 0-based indexing
+        return nb*(indxglob//(nb*nprocs)) + indxglob%nb
+
+    @staticmethod
+    def indxl2g(indxloc, nb, iproc, isrcproc, nprocs):
+        """compute global index from local index indxloc
+
+        original netlib scalapack source:
+
+        .. code-block:: fortran
+
+            INDXL2G = NPROCS*NB*((INDXLOC-1)/NB) + MOD(INDXLOC-1,NB) +
+                      MOD(NPROCS+IPROC-ISRCPROC, NPROCS)*NB + 1
+        """
+        # adapt to python 0-based indexing
+        return nprocs*nb*(indxloc//nb) + indxloc%nb + \
+            ((nprocs+iproc-isrcproc)%nprocs)*nb
+
+    @staticmethod
+    def indxg2p(indxglob, nb, iproc, isrcproc, nprocs):
+        """compute process coordinate for global index
+
+        original netlib scalapack source:
+
+        .. code-block:: fortran
+
+            INDXG2P = MOD( ISRCPROC + (INDXGLOB - 1) / NB, NPROCS )
+        """
+        # adapt to python 0-based indexing
+        return (isrcproc + indxglob // nb) % nprocs
+
+    @staticmethod
+    def numroc(n, nb, iproc, isrcproc, nprocs):
+        """Get local dimensions of distributed block-cyclic matrix.
+
+        Programmed after scalapack source (tools/numroc.f on netlib).
+        """
+        mydist = (nprocs + iproc - isrcproc) % nprocs
+        nblocks = n // nb
+        result = (nblocks // nprocs) * nb
+        extrablks = nblocks % nprocs
+        if mydist < extrablks:
+            result += nb
+        elif mydist == extrablks:
+            result += n % nb
+        return int(result)
+
+    def _initialized_elpa(function):
+        # wrapper to ensure one-time initialization of Elpa object
+        @wraps(function)
+        def wrapped_function(self):
+            if self.elpa is None:
+                self.elpa = Elpa.from_distributed_matrix(self)
+            return function(self)
+        return wrapped_function
+
+    @_initialized_elpa
+    def compute_eigenvectors(self):
+        """Compute eigenvalues and eigenvectors
+
+        The eigenvectors are stored in columns.
+        This function returns a dictionary with entries 'eigenvalues' and
+        'eigenvectors'.
+
+        After computing the eigenvectors, the original content of the matrix is
+        lost.
+        """
+        eigenvectors = DistributedMatrix.like(self)
+        eigenvalues = np.zeros(self.na, dtype=np.float64)
+        # call ELPA
+        self.elpa.eigenvectors(self.data, eigenvalues, eigenvectors.data)
+        return {'eigenvalues': eigenvalues, 'eigenvectors': eigenvectors}
+
+    @_initialized_elpa
+    def compute_eigenvalues(self):
+        """Compute only the eigenvalues.
+
+        This function returns the eigenvalues as an array.
+
+        After computing the eigenvalues, the original content of the matrix is
+        lost.
+        """
+        eigenvalues = np.zeros(self.na, dtype=np.float64)
+        # call ELPA
+        self.elpa.eigenvalues(self.data, eigenvalues)
+        return eigenvalues
+
+    def set_data_from_global_matrix(self, matrix):
+        """Set local part of the global matrix"""
+        for local_row in range(self.na_rows):
+            for local_col in range(self.na_cols):
+                global_row, global_col = self.get_global_index(local_row,
+                                                               local_col)
+                self.data[local_row, local_col] = matrix[global_row,
+                                                         global_col]
+
+    def dot(self, vector):
+        """Compute dot product of matrix with vector.
+
+        This blocked implementation is much faster than the naive
+        implementation.
+        """
+        if len(vector.shape) > 1 or vector.shape[0] != self.na:
+            raise ValueError("Error: shape of vector {} incompatible to "
+                             "matrix of size {:d}x{:d}.".format(
+                                 vector.shape, self.na, self.na))
+        from mpi4py import MPI
+        summation = np.zeros_like(vector)
+        # loop only over blocks here
+        for local_row in range(0, self.na_rows, self.nblk):
+            for local_col in range(0, self.na_cols, self.nblk):
+                # do not go beyond the end of the matrix
+                row_block_size = min(local_row + self.nblk,
+                                     self.na_rows) - local_row
+                col_block_size = min(local_col + self.nblk,
+                                     self.na_cols) - local_col
+                global_row, global_col = self.get_global_index(local_row,
+                                                               local_col)
+                # use numpy for faster dot product of local block
+                summation[global_row:global_row+row_block_size] += \
+                    np.dot(self.data[local_row:local_row + row_block_size,
+                                     local_col:local_col + col_block_size],
+                           vector[global_col:global_col+col_block_size])
+        result = np.zeros_like(vector)
+        self.processor_layout.comm.Allreduce(summation, result, op=MPI.SUM)
+        return result
+
+    def _dot_naive(self, vector):
+        """Compute naive dot product of matrix with vector.
+
+        Still in here as an example and for testing purposes.
+        """
+        from mpi4py import MPI
+        summation = np.zeros_like(vector)
+        for local_row in range(self.na_rows):
+            for local_col in range(self.na_cols):
+                global_row, global_col = self.get_global_index(local_row,
+                                                               local_col)
+                summation[global_row] += self.data[local_row, local_col] *\
+                    vector[global_col]
+        result = np.zeros_like(vector)
+        self.processor_layout.comm.Allreduce(summation, result, op=MPI.SUM)
+        return result
+
+    def get_column(self, global_col):
+        """Return global column"""
+        from mpi4py import MPI
+        column = np.zeros(self.na, dtype=self.data.dtype)
+        temporary = np.zeros_like(column)
+        if self.is_local_col(global_col):
+            for global_row in range(self.na):
+                if not self.is_local_row(global_row):
+                    continue
+                local_row, local_col = self.get_local_index(global_row,
+                                                            global_col)
+                temporary[global_row] = self.data[local_row, local_col]
+        # this could be done more efficiently with a gather
+        self.processor_layout.comm.Allreduce(temporary, column, op=MPI.SUM)
+        return column
+
+    def get_row(self, global_row):
+        """Return global row"""
+        from mpi4py import MPI
+        row = np.zeros(self.na, dtype=self.data.dtype)
+        temporary = np.zeros_like(row)
+        if self.is_local_row(global_row):
+            for global_col in range(self.na):
+                if not self.is_local_col(global_col):
+                    continue
+                local_row, local_col = self.get_local_index(global_row,
+                                                            global_col)
+                temporary[global_col] = self.data[local_row, local_col]
+        # this could be done more efficiently with a gather
+        self.processor_layout.comm.Allreduce(temporary, row, op=MPI.SUM)
+        return row

python/wrapper.pyx

+"""wrapper.pyx -- python wrapper for ELPA
+
+This file contains the cython part of the wrapper.
+"""
+cimport numpy as np
+import numpy as np
+import sys
+
+if 'mpi4py.MPI' in sys.modules.keys():
+    raise NotImplementedError('Please load the pyelpa module before mpi4py, '
+                              'otherwise there will be MPI problems.')
+
+# import the function definitions from the ELPA header
+cdef import from "<elpa/elpa.h>":
+    cdef struct elpa_struct:
+        pass
+    ctypedef elpa_struct *elpa_t
+    int elpa_init(int api_version)
+    void elpa_uninit()
+    elpa_t elpa_allocate(int *error)
+    void elpa_deallocate(elpa_t handle)
+    int elpa_setup(elpa_t handle)
+    void elpa_set_integer(elpa_t handle, const char *name, int value, int *error)
+    void elpa_get_integer(elpa_t handle, const char *name, int *value, int *error)
+    void elpa_set_double(elpa_t handle, const char *name, double value, int *error)
+    void elpa_get_double(elpa_t handle, const char *name, double *value, int *error)
+    void elpa_eigenvectors_d(elpa_t handle, double *a, double *ev, double *q, int *error)
+    void elpa_eigenvectors_f(elpa_t handle, float *a, float *ev, float *q, int *error)
+    void elpa_eigenvectors_dc(elpa_t handle, double complex *a, double *ev, double complex *q, int *error)
+    void elpa_eigenvectors_fc(elpa_t handle, float complex *a, float *ev, float complex *q, int *error)
+    void elpa_eigenvalues_d(elpa_t handle, double *a, double *ev, int *error)
+    void elpa_eigenvalues_f(elpa_t handle, float *a, float *ev, int *error)
+    void elpa_eigenvalues_dc(elpa_t handle, double complex *a, double *ev, int *error)
+    void elpa_eigenvalues_fc(elpa_t handle, float complex *a, float *ev, int *error)
+    int ELPA_OK
+    int ELPA_SOLVER_2STAGE
+
+
+cdef class Elpa:
+    """Wrapper for ELPA C interface.
+
+    Provides routines for initialization, deinitialization, setting and getting
+    properties and for calling the eigenvectors and eigenvalues routines.
+    The routines eigenvectors and eigenvalues select the right ELPA routine to
+    call depending on the argument type.
+    """
+    cdef elpa_t handle
+
+    def __init__(self):
+        """Run initialization and allocation of handle"""
+        if elpa_init(20171201) != ELPA_OK:
+            raise RuntimeError("ELPA API version not supported")
+        cdef int error
+        handle = elpa_allocate(&error)
+        self.handle = handle
+
+    def set_integer(self, description, int value):
+        """Wraps elpa_set_integer"""
+        cdef int error
+        if isinstance(description, unicode):
+            # encode to ascii for passing to C
+            description = (<unicode>description).encode('ascii')
+        cdef const char* c_string = description
+        elpa_set_integer(<elpa_t>self.handle, description, value, &error)
+
+    def get_integer(self, description):
+        """Wraps elpa_get_integer"""
+        cdef int error
+        if isinstance(description, unicode):
+            # encode to ascii for passing to C
+            description = (<unicode>description).encode('ascii')
+        cdef const char* c_string = description
+        cdef int tmp
+        elpa_get_integer(<elpa_t>self.handle, c_string, &tmp, &error)
+        return tmp
+
+    def set_double(self, description, double value):
+        """Wraps elpa_set_double"""
+        cdef int error
+        if isinstance(description, unicode):
+            # encode to ascii for passing to C
+            description = (<unicode>description).encode('ascii')
+        cdef const char* c_string = description
+        elpa_set_double(<elpa_t>self.handle, description, value, &error)
+
+    def get_double(self, description):
+        """Wraps elpa_get_double"""
+        cdef int error
+        if isinstance(description, unicode):
+            # encode to ascii for passing to C
+            description = (<unicode>description).encode('ascii')
+        cdef const char* c_string = description
+        cdef double tmp
+        elpa_get_double(<elpa_t>self.handle, c_string, &tmp, &error)
+        return tmp
+
+    def setup(self):
+        """call setup function"""
+        elpa_setup(<elpa_t>self.handle)
+
+    def __del__(self):
+        """Deallocation of handle and deinitialization"""
+        elpa_deallocate(<elpa_t>self.handle)
+        elpa_uninit()
+
+    def eigenvectors_d(self,
+                       np.ndarray[np.float64_t, ndim=2] a,
+                       np.ndarray[np.float64_t, ndim=1] ev,
+                       np.ndarray[np.float64_t, ndim=2] q):
+        cdef int error
+        elpa_eigenvectors_d(<elpa_t>self.handle, <np.float64_t *>a.data,
+                            <np.float64_t *>ev.data, <np.float64_t *>q.data,
+                            <int*>&error)
+        if error != ELPA_OK:
+            raise RuntimeError("ELPA returned error value {:d}.".format(error))
+
+    def eigenvectors_f(self,
+                       np.ndarray[np.float32_t, ndim=2] a,
+                       np.ndarray[np.float32_t, ndim=1] ev,
+                       np.ndarray[np.float32_t, ndim=2] q):
+        cdef int error
+        elpa_eigenvectors_f(<elpa_t>self.handle, <np.float32_t *>a.data,
+                            <np.float32_t *>ev.data, <np.float32_t *>q.data,
+                            <int*>&error)
+        if error != ELPA_OK:
+            raise RuntimeError("ELPA returned error value {:d}.".format(error))
+
+    def eigenvectors_dc(self,
+                        np.ndarray[np.complex128_t, ndim=2] a,
+                        np.ndarray[np.float64_t, ndim=1] ev,
+                        np.ndarray[np.complex128_t, ndim=2] q):
+        cdef int error
+        elpa_eigenvectors_dc(<elpa_t>self.handle, <np.complex128_t *>a.data,
+                             <np.float64_t *>ev.data, <np.complex128_t *>q.data,
+                             <int*>&error)
+        if error != ELPA_OK:
+            raise RuntimeError("ELPA returned error value {:d}.".format(error))
+
+    def eigenvectors_fc(self,
+                        np.ndarray[np.complex64_t, ndim=2] a,
+                        np.ndarray[np.float32_t, ndim=1] ev,
+                        np.ndarray[np.complex64_t, ndim=2] q):
+        cdef int error
+        elpa_eigenvectors_fc(<elpa_t>self.handle, <np.complex64_t *>a.data,
+                             <np.float32_t *>ev.data, <np.complex64_t *>q.data,
+                             <int*>&error)
+        if error != ELPA_OK:
+            raise RuntimeError("ELPA returned error value {:d}.".format(error))
+
+    def eigenvectors(self, a, ev, q):
+        """Compute eigenvalues and eigenvectors.
+
+        The data type of a is tested and the corresponding ELPA routine called
+
+        Args:
+            a (DistributedMatrix): problem matrix
+            ev (numpy.ndarray): array of size a.na to store eigenvalues
+            q (DistributedMatrix): store eigenvectors
+        """
+        if a.dtype == np.float64:
+            self.eigenvectors_d(a, ev, q)