diff --git a/python/examples/example.py b/python/examples/example.py index 6c5b0d799c78703a16a6cd541f7ad5c5f37bfa4b..510b75f395d345c803660573b021dec822392dfe 100644 --- a/python/examples/example.py +++ b/python/examples/example.py @@ -1,7 +1,6 @@ #!/usr/bin/env python import numpy as np -from pyelpa import ProcessorLayout, DistributedMatrix, Elpa -from mpi4py import MPI +from pyelpa import ProcessorLayout, DistributedMatrix import sys # set some parameters for matrix layout @@ -9,37 +8,37 @@ na = 1000 nev = 200 nblk = 16 -# initialize processor layout, needed for calling ELPA -comm = MPI.COMM_WORLD -layout_p = ProcessorLayout(comm) +# create distributed matrix +a = DistributedMatrix.from_comm_world(na, nev, nblk) -# 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) +def set_matrix(a): + for global_row, global_col in a.global_indices(): + a.set_data_for_global_index(global_row, global_col, + global_row*global_col) print("Call ELPA eigenvectors") sys.stdout.flush() + +set_matrix(a) # now compute nev of na eigenvectors and eigenvalues -e.eigenvectors(a.data, eigenvalues, eigenvectors.data) +data = a.compute_eigenvectors() +eigenvalues = data['eigenvalues'] +eigenvectors = data['eigenvectors'] + 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 +# which is stored in a block-cyclic distributed layout and eigenvalues contains +# all computed eigenvalues on all cores print("Call ELPA eigenvalues") sys.stdout.flush() + +set_matrix(a) # now compute nev of na eigenvalues -e.eigenvalues(a.data, eigenvalues) +set_matrix(a) +eigenvalues = a.compute_eigenvalues() + print("Done") # now eigenvalues contains all computed eigenvalues on all cores