example.py 1.6 KB
Newer Older
1 2
#!/usr/bin/env python
import numpy as np
3
from pyelpa import DistributedMatrix
4 5 6 7 8 9 10
import sys

# set some parameters for matrix layout
na = 1000
nev = 200
nblk = 16

11 12
# create distributed matrix
a = DistributedMatrix.from_comm_world(na, nev, nblk)
13

14
# set matrix a by looping over indices
15
# this is the easiest but also slowest way
16 17 18
for global_row, global_col in a.global_indices():
    a.set_data_for_global_index(global_row, global_col,
                                global_row*global_col)
19

20 21
print("Call ELPA eigenvectors")
sys.stdout.flush()
22

23
# now compute nev of na eigenvectors and eigenvalues
24 25 26 27
data = a.compute_eigenvectors()
eigenvalues = data['eigenvalues']
eigenvectors = data['eigenvectors']

28 29 30
print("Done")

# now eigenvectors.data contains the local part of the eigenvector matrix
31 32
# which is stored in a block-cyclic distributed layout and eigenvalues contains
# all computed eigenvalues on all cores
33

34
# set a again because it has changed after calling elpa
35 36 37 38 39 40 41 42
# this time set it by looping over blocks, this is more efficient
for global_row, global_col, row_block_size, col_block_size in \
        a.global_block_indices():
    # set block with product of indices
    x = np.arange(global_row, global_row + row_block_size)[:, None] * \
        np.arange(global_col, global_col + col_block_size)[None, :]
    a.set_block_for_global_index(global_row, global_col,
                                 row_block_size, col_block_size, x)
43

44 45
print("Call ELPA eigenvalues")
sys.stdout.flush()
46

47
# now compute nev of na eigenvalues
48 49
eigenvalues = a.compute_eigenvalues()

50 51 52
print("Done")

# now eigenvalues contains all computed eigenvalues on all cores