Commit eafc2a15 authored by Sebastian Ohlmann's avatar Sebastian Ohlmann

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
parent 9e3bfb2e
......@@ -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
......
......@@ -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
......
"""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']
"""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
"""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)
elif a.dtype == np.float32:
self.eigenvectors_f(a, ev, q)
elif a.dtype == np.complex128:
self.eigenvectors_dc(a, ev, q)
elif a.dtype == np.complex64:
self.eigenvectors_fc(a, ev, q)
else:
raise TypeError("Type not known.")
def eigenvalues_d(self,
np.ndarray[np.float64_t, ndim=2] a,
np.ndarray[np.float64_t, ndim=1] ev):
cdef int error
elpa_eigenvalues_d(<elpa_t>self.handle, <np.float64_t *>a.data,
<np.float64_t *>ev.data, <int*>&error)
if error != ELPA_OK:
raise RuntimeError("ELPA returned error value {:d}.".format(error))
def eigenvalues_f(self,
np.ndarray[np.float32_t, ndim=2] a,
np.ndarray[np.float32_t, ndim=1] ev):
cdef int error
elpa_eigenvalues_f(<elpa_t>self.handle, <np.float32_t *>a.data,
<np.float32_t *>ev.data, <int*>&error)
if error != ELPA_OK:
raise RuntimeError("ELPA returned error value {:d}.".format(error))
def eigenvalues_dc(self,
np.ndarray[np.complex128_t, ndim=2] a,
np.ndarray[np.float64_t, ndim=1] ev):
cdef int error
elpa_eigenvalues_dc(<elpa_t>self.handle, <np.complex128_t *>a.data,
<np.float64_t *>ev.data, <int*>&error)
if error != ELPA_OK:
raise RuntimeError("ELPA returned error value {:d}.".format(error))
def eigenvalues_fc(self,
np.ndarray[np.complex64_t, ndim=2] a,
np.ndarray[np.float32_t, ndim=1] ev):
cdef int error
elpa_eigenvalues_fc(<elpa_t>self.handle, <np.complex64_t *>a.data,
<np.float32_t *>ev.data, <int*>&error)
if error != ELPA_OK:
raise RuntimeError("ELPA returned error value {:d}.".format(error))
def eigenvalues(self, a, ev):
"""Compute eigenvalues.
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
"""
if a.dtype == np.float64:
self.eigenvalues_d(a, ev)
elif a.dtype == np.float32:
self.eigenvalues_f(a, ev)
elif a.dtype == np.complex128:
self.eigenvalues_dc(a, ev)
elif a.dtype == np.complex64:
self.eigenvalues_fc(a, ev)
else:
raise TypeError("Type not known.")