add looping over blocks for python wrapper

This is more efficient because the python loop processes only (number of
blocks)^2 elements.

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()

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))

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))