diff --git a/python/examples/example.py b/python/examples/example.py index 03d7d6589a38aadf5c9e5229fb27f65a452eb7aa..a8e3cec491b0d4c14d4ab6ca2ec35f3b3c63c026 100644 --- a/python/examples/example.py +++ b/python/examples/example.py @@ -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() diff --git a/python/pyelpa/distributedmatrix.py b/python/pyelpa/distributedmatrix.py index 4b65284d379b8d297e5d884db06bd0e4e30473b4..a9787f2eec0a083ea95d968820b676ee0b0f0e8c 100644 --- a/python/pyelpa/distributedmatrix.py +++ b/python/pyelpa/distributedmatrix.py @@ -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)) diff --git a/python/tests/test_with_mpi.py b/python/tests/test_with_mpi.py index f9f0e4601163831b1b63164ef2624423860ca86c..d5e9c2c9d3df0591754c112b4e614311b11added 100644 --- a/python/tests/test_with_mpi.py +++ b/python/tests/test_with_mpi.py @@ -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))