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

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,15 +11,11 @@ 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():
a.set_data_for_global_index(global_row, global_col,
global_row*global_col)
# set a
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()
......@@ -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