Add sqrt option to diagonal method of stencil/block matrices

.github/workflows/continuous-integration.yml

@@ -58,6 +58,7 @@ jobs:
       - name: Set default MPI and HDF5 C compilers on Ubuntu
         if: matrix.os == 'ubuntu-latest'
         run: |
+          sudo apt-get update
           sudo apt-get install --reinstall openmpi-bin libhdf5-openmpi-dev
 
       - name: Install non-Python dependencies on macOS

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

@@ -528,8 +528,9 @@ def test_in_place_operations(n1, n2, p1, p2, P1=False, P2=False):
 @pytest.mark.parametrize('n2', n2array)
 @pytest.mark.parametrize('p1', p1array)
 @pytest.mark.parametrize('p2', p2array)
+@pytest.mark.parametrize('sqrt', [False, True])
 
-def test_inverse_transpose_interaction(n1, n2, p1, p2, P1=False, P2=False):
+def test_inverse_transpose_interaction(n1, n2, p1, p2, sqrt, P1=False, P2=False):
 
     # 1. Initiate square LOs: S (V->V, StencilMatrix), S1 (W->W, StencilMatrix)
         #Initiate BlockLO: B (VxW -> VxW) and a BlockVector u element of VxW
@@ -618,6 +619,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.allclose(sqrt_diagonal_values, np.sqrt(diagonal_values))
+
     tol = 1e-5
     C = inverse(B, 'cg', tol=tol)
     P = B.diagonal(inverse=True)