Commit 47e2cb5f authored by Sebastian Ohlmann's avatar Sebastian Ohlmann
Browse files

add looping over blocks for python wrapper

This is more efficient because the python loop processes only (number of
blocks)^2 elements.
parent 954eb7f2
......@@ -11,16 +11,12 @@ nblk = 16
# create distributed matrix
a = DistributedMatrix.from_comm_world(na, nev, nblk)
# function for setting the matrix
# set matrix a by looping over indices
# this is the easiest but also slowest way
def set_matrix(a):
for global_row, global_col in a.global_indices():
for global_row, global_col in a.global_indices():
a.set_data_for_global_index(global_row, global_col,
global_row*global_col)
# set a
set_matrix(a)
print("Call ELPA eigenvectors")
sys.stdout.flush()
......@@ -36,7 +32,14 @@ print("Done")
# all computed eigenvalues on all cores
# set a again because it has changed after calling elpa
set_matrix(a)
# 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)
print("Call ELPA eigenvalues")
sys.stdout.flush()
......
......@@ -351,3 +351,49 @@ class DistributedMatrix:
else:
raise ValueError('Index out of bounds: global row {:d}, '
'global col {:d}'.format(global_row, global_col))
def global_block_indices(self):
"""Return iterator over global indices of matrix blocks.
Use together with set_block_global_index and get_block_global_index
for more efficient loops.
"""
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)
yield global_row, global_col, row_block_size, col_block_size
def set_block_for_global_index(self, global_row, global_col,
row_block_size, col_block_size, value):
"""Set value of block of matrix at global coordinates"""
if self.is_local_index(global_row, global_col):
local_row, local_col = self.get_local_index(global_row, global_col)
if value.shape != (row_block_size, col_block_size):
raise ValueError("value has the wrong shape. "
"Expected: {}, found: {}."
.format((row_block_size, col_block_size),
value.shape)
)
self.data[local_row:local_row+row_block_size,
local_col:local_col+col_block_size] = value
def get_block_for_global_index(self, global_row, global_col,
row_block_size, col_block_size):
"""Get value of block of matrix at global coordinates"""
if self.is_local_index(global_row, global_col):
local_row, local_col = self.get_local_index(global_row, global_col)
if local_row+row_block_size > self.na_rows or \
local_col+col_block_size > self.na_cols:
raise ValueError("Block size wrong: exceeds dimensions of"
" matrix.")
return self.data[local_row:local_row+row_block_size,
local_col:local_col+col_block_size]
else:
raise ValueError('Index out of bounds: global row {:d}, '
'global col {:d}'.format(global_row, global_col))
......@@ -402,3 +402,37 @@ def test_global_index_access(na, nev, nblk):
for i, j in a.global_indices():
x = a.get_data_for_global_index(i, j)
assert(np.isclose(x, i*j))
@pytest.mark.parametrize("na,nev,nblk", parameter_list)
def test_global_block_iterator(na, nev, nblk):
import numpy as np
from pyelpa import DistributedMatrix
for dtype in [np.float64, np.complex128]:
a = DistributedMatrix.from_comm_world(na, nev, nblk, dtype=dtype)
for i, j, blk_i, blk_j in a.global_block_indices():
assert(a.is_local_index(i, j))
assert(blk_i <= nblk)
assert(blk_j <= nblk)
assert(i+blk_i <= na)
assert(j+blk_j <= na)
@pytest.mark.parametrize("na,nev,nblk", parameter_list)
def test_global_block_access(na, nev, nblk):
import numpy as np
from pyelpa import DistributedMatrix
for dtype in [np.float64, np.complex128]:
a = DistributedMatrix.from_comm_world(na, nev, nblk, dtype=dtype)
for i, j, blk_i, blk_j in a.global_block_indices():
x = np.arange(i, i+blk_i)[:, None] * np.arange(j, j+blk_j)[None, :]
a.set_block_for_global_index(i, j, blk_i, blk_j, x)
for i, j, blk_i, blk_j in a.global_block_indices():
original = np.arange(i, i+blk_i)[:, None] * np.arange(j, j+blk_j)[None, :]
x = a.get_block_for_global_index(i, j, blk_i, blk_j)
assert(np.allclose(x, original))
for i, j in a.global_indices():
x = a.get_data_for_global_index(i, j)
assert(np.isclose(x, i*j))
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment