Commit 47e2cb5f by 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,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!