Add sqrt option to diagonal method of stencil/block matrices (#434)

---------

Co-authored-by: Yaman Güçlü <yaman.guclu@gmail.com>

psydac/linalg/block.py

@@ -777,13 +777,18 @@ def __sub__(self, M):
     #--------------------------------------
     # New properties/methods
     #--------------------------------------
-    def diagonal(self, *, inverse = False, out = None):
+    def diagonal(self, *, inverse = False, sqrt = False, out = None):
         """Get the coefficients on the main diagonal as another BlockLinearOperator object.
 
         Parameters
         ----------
         inverse : bool
             If True, get the inverse of the diagonal. (Default: False).
+            Can be combined with sqrt to get the inverse square root.
+
+        sqrt : bool
+            If True, get the square root of the diagonal. (Default: False).
+            Can be combined with inverse to get the inverse square root.
 
         out : BlockLinearOperator
             If provided, write the diagonal entries into this matrix. (Default: None).
@@ -815,7 +820,7 @@ def diagonal(self, *, inverse = False, out = None):
         # Store the diagonal (or its inverse) into `out`
         for i, j in self.nonzero_block_indices:
             if i == j:
-                out[i, i] = self[i, i].diagonal(inverse = inverse, out = out[i, i])
+                out[i, i] = self[i, i].diagonal(inverse = inverse, sqrt = sqrt, out = out[i, i])
 
         return out
 

psydac/linalg/stencil.py

@@ -1372,14 +1372,19 @@ def _exchange_assembly_data_serial(self):
                 self._data[idx_to] += self._data[idx_from]
 
     # ...
-    def diagonal(self, *, inverse = False, out = None):
+    def diagonal(self, *, inverse = False, sqrt = False, out = None):
         """
         Get the coefficients on the main diagonal as a StencilDiagonalMatrix object.
 
         Parameters
         ----------
         inverse : bool
             If True, get the inverse of the diagonal. (Default: False).
+            Can be combined with sqrt to get the inverse square root.
+
+        sqrt : bool
+            If True, get the square root of the diagonal. (Default: False).
+            Can be combined with inverse to get the inverse square root.
 
         out : StencilDiagonalMatrix
             If provided, write the diagonal entries into this matrix. (Default: None).
@@ -1418,6 +1423,9 @@ def diagonal(self, *, inverse = False, out = None):
         else:
             data = diag.copy()
 
+        if sqrt:
+            np.sqrt(data, out=data)
+
         # If needed create a new StencilDiagonalMatrix object
         if out is None:
             out = StencilDiagonalMatrix(V, W, data)

psydac/linalg/tests/test_linalg.py

@@ -618,6 +618,12 @@ def test_inverse_transpose_interaction(n1, n2, p1, p2, P1=False, P2=False):
     ### -1,T & T,-1 --- -1,T,T --- -1,T,-1 --- T,-1,-1 --- T,-1,T (the combinations I test)
     ###
 
+    # Square root test
+    scaled_matrix = B * np.random.random() # Ensure the diagonal elements != 1
+    diagonal_values = scaled_matrix.diagonal(sqrt=False).toarray()
+    sqrt_diagonal_values = scaled_matrix.diagonal(sqrt=True).toarray()
+    assert np.array_equal(sqrt_diagonal_values, np.sqrt(diagonal_values))
+
     tol = 1e-5
     C = inverse(B, 'cg', tol=tol)
     P = B.diagonal(inverse=True)