add examples and tests for python wrapper

python/examples/example.py

+#!/usr/bin/env python
+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)
+a.data[:, :] = np.random.rand(a.na_rows, a.na_cols).astype(np.float64)
+
+print("Call ELPA eigenvectors")
+sys.stdout.flush()
+# now compute nev of na eigenvectors and eigenvalues
+e.eigenvectors(a.data, eigenvalues, eigenvectors.data)
+print("Done")
+
+# 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
+
+print("Call ELPA eigenvalues")
+sys.stdout.flush()
+# now compute nev of na eigenvalues
+e.eigenvalues(a.data, eigenvalues)
+print("Done")
+
+# now eigenvalues contains all computed eigenvalues on all cores

python/tests/test_elpa_import.py

+def test_pyelpa_import():
+    import pyelpa
+
+def test_elpa_supported_api_version():
+    from pyelpa import Elpa
+    e = Elpa()

python/tests/test_mpi4py.py

+def test_import():
+    from mpi4py import MPI
+    comm = MPI.COMM_WORLD
+
+
+def test_barrier():
+    from mpi4py import MPI
+    comm = MPI.COMM_WORLD
+    comm.Barrier()

python/tests/test_numroc.py

+def test_numroc():
+    from pyelpa import DistributedMatrix
+    n = 100
+    nb = 16
+    assert(DistributedMatrix.numroc(n, nb, 0, 0, 3) == 36)
+    assert(DistributedMatrix.numroc(n, nb, 1, 0, 3) == 32)
+    assert(DistributedMatrix.numroc(n, nb, 1, 1, 3) == 36)

python/tests/test_with_mpi.py

+import pytest
+
+# combinations of na, nev, nblk to run all the tests with
+parameter_list = [
+    (200, 20, 16),
+    (200, 200, 16),
+    (200, 20, 64),
+    (200, 200, 64),
+    (200, 20, 4),
+    (200, 200, 4),
+    (50, 20, 16),
+    (100, 20, 16),
+]
+
+def get_random_vector(size):
+    """generate random vector with given size that is equal on all cores"""
+    import numpy as np
+    from mpi4py import MPI
+    comm = MPI.COMM_WORLD
+    am_i_root = comm.Get_rank() == 0
+    vector = np.empty(size)
+    if am_i_root:
+        vector[:] = np.random.rand(size)
+    comm.Bcast(vector)
+    return vector
+
+
+def test_processor_layout():
+    from pyelpa import ProcessorLayout
+    from mpi4py import MPI
+    comm = MPI.COMM_WORLD
+    layout_p = ProcessorLayout(comm)
+    assert(comm.Get_size() == layout_p.np_cols*layout_p.np_rows)
+    assert(layout_p.my_prow >= 0)
+    assert(layout_p.my_pcol >= 0)
+    assert(layout_p.my_prow <= comm.Get_size())
+    assert(layout_p.my_pcol <= comm.Get_size())
+
+
+@pytest.mark.parametrize("na,nev,nblk", parameter_list)
+def test_distributed_matrix_from_processor_layout(na, nev, nblk):
+    import numpy as np
+    from pyelpa import ProcessorLayout, DistributedMatrix
+    from mpi4py import MPI
+    comm = MPI.COMM_WORLD
+    layout_p = ProcessorLayout(comm)
+
+    for dtype in [np.float64, np.float32, np.complex64, np.complex128]:
+        a = DistributedMatrix(layout_p, na, nev, nblk, dtype=dtype)
+        assert(a.data.dtype == dtype)
+        assert(a.data.shape == (a.na_rows, a.na_cols))
+
+
+@pytest.mark.parametrize("na,nev,nblk", parameter_list)
+def test_distributed_matrix_from_communicator(na, nev, nblk):
+    import numpy as np
+    from pyelpa import ProcessorLayout, DistributedMatrix
+    from mpi4py import MPI
+    comm = MPI.COMM_WORLD
+
+    for dtype in [np.float64, np.float32, np.complex64, np.complex128]:
+        a = DistributedMatrix.from_communicator(comm, na, nev, nblk,
+                                                dtype=dtype)
+        assert(a.data.dtype == dtype)
+        assert(a.data.shape == (a.na_rows, a.na_cols))
+
+
+@pytest.mark.parametrize("na,nev,nblk", parameter_list)
+def test_distributed_matrix_from_world(na, nev, nblk):
+    import numpy as np
+    from pyelpa import ProcessorLayout, DistributedMatrix
+
+    for dtype in [np.float64, np.float32, np.complex64, np.complex128]:
+        a = DistributedMatrix.from_comm_world(na, nev, nblk, dtype=dtype)
+        assert(a.data.dtype == dtype)
+        assert(a.data.shape == (a.na_rows, a.na_cols))
+
+
+@pytest.mark.parametrize("na,nev,nblk", parameter_list)
+def test_distributed_matrix_like_other_matrix(na, nev, nblk):
+    import numpy as np
+    from pyelpa import ProcessorLayout, DistributedMatrix
+    from mpi4py import MPI
+    comm = MPI.COMM_WORLD
+    layout_p = ProcessorLayout(comm)
+
+    for dtype in [np.float64, np.float32, np.complex64, np.complex128]:
+        a = DistributedMatrix(layout_p, na, nev, nblk, dtype=dtype)
+        b = DistributedMatrix.like(a)
+        assert(a.na == b.na)
+        assert(a.nev == b.nev)
+        assert(a.nblk == b.nblk)
+        assert(a.data.dtype == b.data.dtype)
+        assert(a.data.shape == b.data.shape)
+
+
+@pytest.mark.parametrize("na,nev,nblk", parameter_list)
+def test_call_eigenvectors(na, nev, nblk):
+    import numpy as np
+    from pyelpa import ProcessorLayout, DistributedMatrix, Elpa
+    from mpi4py import MPI
+
+    comm = MPI.COMM_WORLD
+    layout_p = ProcessorLayout(comm)
+
+    for dtype in [np.float64, np.complex128]:
+        # create arrays
+        a = DistributedMatrix(layout_p, na, nev, nblk, dtype=dtype)
+        a.data[:, :] = np.random.rand(a.na_rows, a.na_cols).astype(dtype)
+        q = DistributedMatrix(layout_p, na, nev, nblk, dtype=dtype)
+        ev = np.zeros(na, dtype=np.float64)
+
+        e = Elpa.from_distributed_matrix(a)
+        e.eigenvectors(a.data, ev, q.data)
+
+
+@pytest.mark.parametrize("na,nev,nblk", parameter_list)
+def test_call_eigenvalues(na, nev, nblk):
+    import numpy as np
+    from pyelpa import ProcessorLayout, DistributedMatrix, Elpa
+    from mpi4py import MPI
+
+    comm = MPI.COMM_WORLD
+    layout_p = ProcessorLayout(comm)
+
+    for dtype in [np.float64, np.complex128]:
+        # create arrays
+        a = DistributedMatrix(layout_p, na, nev, nblk, dtype=dtype)
+        a.data[:, :] = np.random.rand(a.na_rows, a.na_cols).astype(dtype)
+        ev = np.zeros(na, dtype=np.float64)
+
+        e = Elpa.from_distributed_matrix(a)
+        e.eigenvalues(a.data, ev)
+
+
+@pytest.mark.parametrize("na,nev,nblk", parameter_list)
+def test_compare_eigenvalues_to_those_from_eigenvectors(na, nev, nblk):
+    import numpy as np
+    from pyelpa import ProcessorLayout, DistributedMatrix, Elpa
+    from mpi4py import MPI
+
+    comm = MPI.COMM_WORLD
+    layout_p = ProcessorLayout(comm)
+
+    for dtype in [np.float64, np.complex128]:
+        # create arrays
+        a = DistributedMatrix(layout_p, na, nev, nblk, dtype=dtype)
+        random_matrix = np.random.rand(a.na_rows, a.na_cols).astype(dtype)
+        a.data[:, :] = random_matrix
+        q = DistributedMatrix(layout_p, na, nev, nblk, dtype=dtype)
+        ev = np.zeros(na, dtype=np.float64)
+        ev2 = np.zeros(na, dtype=np.float64)
+
+        e = Elpa.from_distributed_matrix(a)
+        e.eigenvectors(a.data, ev, q.data)
+
+        a.data[:, :] = random_matrix
+        e.eigenvalues(a.data, ev2)
+
+        assert(np.allclose(ev, ev2))
+
+
+@pytest.mark.parametrize("na,nev,nblk", parameter_list)
+def test_compare_eigenvalues_to_those_from_eigenvectors_self_functions(
+        na, nev, nblk):
+    import numpy as np
+    from pyelpa import ProcessorLayout, DistributedMatrix
+
+    for dtype in [np.float64, np.complex128]:
+        # create arrays
+        a = DistributedMatrix.from_comm_world(na, nev, nblk, dtype=dtype)
+        random_matrix = np.random.rand(a.na_rows, a.na_cols).astype(dtype)
+        a.data[:, :] = random_matrix
+        data = a.compute_eigenvectors()
+
+        a.data[:, :] = random_matrix
+        eigenvalues = a.compute_eigenvalues()
+
+        assert(np.allclose(data['eigenvalues'], eigenvalues))
+
+
+@pytest.mark.parametrize("na,nev,nblk", parameter_list)
+def test_distributed_matrix_global_index(na, nev, nblk):
+    import numpy as np
+    from pyelpa import ProcessorLayout, DistributedMatrix
+    from mpi4py import MPI
+    comm = MPI.COMM_WORLD
+    layout_p = ProcessorLayout(comm)
+
+    for dtype in [np.float64, np.complex128]:
+        a = DistributedMatrix(layout_p, na, nev, nblk, dtype=dtype)
+        for local_row in range(a.na_rows):
+            for local_col in range(a.na_cols):
+                global_row, global_col = a.get_global_index(local_row,
+                                                            local_col)
+                l_row, l_col = a.get_local_index(global_row, global_col)
+                assert(global_row >= 0 and global_row < a.na)
+                assert(global_col >= 0 and global_col < a.na)
+                assert(local_row == l_row and local_col == l_col)
+
+
+@pytest.mark.parametrize("na,nev,nblk", parameter_list)
+def test_distributed_matrix_local_index(na, nev, nblk):
+    import numpy as np
+    from pyelpa import ProcessorLayout, DistributedMatrix
+    from mpi4py import MPI
+    comm = MPI.COMM_WORLD
+    layout_p = ProcessorLayout(comm)
+
+    for dtype in [np.float64, np.complex128]:
+        a = DistributedMatrix(layout_p, na, nev, nblk, dtype=dtype)
+        for global_row in range(a.na):
+            for global_col in range(a.na):
+                if not a.is_local_index(global_row, global_col):
+                    continue
+                local_row, local_col = a.get_local_index(global_row,
+                                                         global_col)
+                g_row, g_col = a.get_global_index(local_row, local_col)
+                assert(local_row >= 0 and local_row < a.na_rows)
+                assert(local_col >= 0 and local_col < a.na_cols)
+                assert(global_row == g_row and global_col == g_col)
+
+
+@pytest.mark.parametrize("na,nev,nblk", parameter_list)
+def test_distributed_matrix_indexing_loop(na, nev, nblk):
+    import numpy as np
+    from pyelpa import ProcessorLayout, DistributedMatrix
+    from mpi4py import MPI
+    comm = MPI.COMM_WORLD
+    layout_p = ProcessorLayout(comm)
+
+    for dtype in [np.float64, np.complex128]:
+        a = DistributedMatrix(layout_p, na, nev, nblk, dtype=dtype)
+        for local_row in range(a.na_rows):
+            for local_col in range(a.na_cols):
+                global_row, global_col = a.get_global_index(local_row,
+                                                            local_col)
+                a.data[local_row, local_col] = global_row*10 + global_col
+
+        for global_row in range(a.na):
+            for global_col in range(a.na):
+                if not a.is_local_index(global_row, global_col):
+                    continue
+                local_row, local_col = a.get_local_index(global_row,
+                                                         global_col)
+                assert(a.data[local_row, local_col] ==
+                       global_row*10 + global_col)
+
+
+@pytest.mark.parametrize("na,nev,nblk", parameter_list)
+def test_setting_global_matrix(na, nev, nblk):
+    import numpy as np
+    from pyelpa import ProcessorLayout, DistributedMatrix
+    from mpi4py import MPI
+    comm = MPI.COMM_WORLD
+    layout_p = ProcessorLayout(comm)
+
+    for dtype in [np.float64, np.complex128]:
+        a = DistributedMatrix(layout_p, na, nev, nblk, dtype=dtype)
+        # get global matrix that is equal on all cores
+        matrix = get_random_vector(na*na).reshape(na, na).astype(dtype)
+        a.set_data_from_global_matrix(matrix)
+
+        # check data
+        for global_row in range(a.na):
+            for global_col in range(a.na):
+                if not a.is_local_index(global_row, global_col):
+                    continue
+                local_row, local_col = a.get_local_index(global_row,
+                                                         global_col)
+                assert(a.data[local_row, local_col] ==
+                       matrix[global_row, global_col])
+
+
+@pytest.mark.parametrize("na,nev,nblk", parameter_list)
+def test_dot_product(na, nev, nblk):
+    import numpy as np
+    from pyelpa import ProcessorLayout, DistributedMatrix
+
+    for dtype in [np.float64, np.complex128]:
+        a = DistributedMatrix.from_comm_world(na, nev, nblk, dtype=dtype)
+        # get global matrix and vector that is equal on all cores
+        matrix = get_random_vector(na*na).reshape(na, na).astype(dtype)
+        vector = get_random_vector(na).astype(dtype)
+
+        a.set_data_from_global_matrix(matrix)
+
+        product_distributed = a.dot(vector)
+        product_naive = a._dot_naive(vector)
+        product_serial = np.dot(matrix, vector)
+
+        assert(np.allclose(product_distributed, product_serial))
+        assert(np.allclose(product_distributed, product_naive))
+
+@pytest.mark.parametrize("na,nev,nblk", parameter_list)
+def test_dot_product_incompatible_size(na, nev, nblk):
+    import numpy as np
+    from pyelpa import ProcessorLayout, DistributedMatrix
+
+    for dtype in [np.float64, np.complex128]:
+        a = DistributedMatrix.from_comm_world(na, nev, nblk, dtype=dtype)
+        # get global matrix and vector that is equal on all cores
+        matrix = get_random_vector(na*na).reshape(na, na).astype(dtype)
+        vector = get_random_vector(na*2).astype(dtype)
+
+        a.set_data_from_global_matrix(matrix)
+
+        with pytest.raises(ValueError):
+            product_distributed = a.dot(vector)
+
+
+@pytest.mark.parametrize("na,nev,nblk", parameter_list)
+def test_validate_eigenvectors(na, nev, nblk):
+    import numpy as np
+    from pyelpa import ProcessorLayout, DistributedMatrix
+
+    for dtype in [np.float64, np.complex128]:
+        a = DistributedMatrix.from_comm_world(na, nev, nblk, dtype=dtype)
+        # get a symmetric/hermitian matrix
+        matrix = get_random_vector(na*na).reshape(na, na).astype(dtype)
+        matrix = 0.5*(matrix + np.conj(matrix.T))
+        a.set_data_from_global_matrix(matrix)
+
+        data = a.compute_eigenvectors()
+        eigenvalues = data['eigenvalues']
+        eigenvectors = data['eigenvectors']
+        # reset data of a
+        a.set_data_from_global_matrix(matrix)
+        for index in range(a.nev):
+            eigenvector = eigenvectors.get_column(index)
+            scaled_eigenvector = eigenvalues[index]*eigenvector
+            # test solution
+            assert(np.allclose(a.dot(eigenvector),
+                               scaled_eigenvector))
+
+
+@pytest.mark.parametrize("na,nev,nblk", parameter_list)
+def test_validate_eigenvectors_to_numpy(na, nev, nblk):
+    import numpy as np
+    from numpy import linalg
+    from pyelpa import ProcessorLayout, DistributedMatrix
+
+    for dtype in [np.float64, np.complex128]:
+        a = DistributedMatrix.from_comm_world(na, nev, nblk, dtype=dtype)
+        # get a symmetric/hermitian matrix
+        matrix = get_random_vector(na*na).reshape(na, na).astype(dtype)
+        matrix = 0.5*(matrix + np.conj(matrix.T))
+        a.set_data_from_global_matrix(matrix)
+
+        data = a.compute_eigenvectors()
+        eigenvalues = data['eigenvalues']
+        eigenvectors = data['eigenvectors']
+
+        # get numpy solution
+        eigenvalues_np, eigenvectors_np = linalg.eigh(matrix)
+
+        assert(np.allclose(eigenvalues, eigenvalues_np))
+        for index in range(a.nev):
+            eigenvector = eigenvectors.get_column(index)
+            assert(np.allclose(eigenvector, eigenvectors_np[:, index]) or
+                   np.allclose(eigenvector, -eigenvectors_np[:, index]))
+
+
+@pytest.mark.parametrize("na,nev,nblk", parameter_list)
+def test_accessing_matrix(na, nev, nblk):
+    import numpy as np
+    from pyelpa import ProcessorLayout, DistributedMatrix
+
+    for dtype in [np.float64, np.complex128]:
+        a = DistributedMatrix.from_comm_world(na, nev, nblk, dtype=dtype)
+        matrix = get_random_vector(na*na).reshape(na, na).astype(dtype)
+        a.set_data_from_global_matrix(matrix)
+
+        for index in range(a.na):
+            column = a.get_column(index)
+            assert(np.allclose(column, matrix[:, index]))
+            row = a.get_row(index)
+            assert(np.allclose(row, matrix[index, :]))