diff --git a/psydac/linalg/block.py b/psydac/linalg/block.py index d54e63626..49e29b60d 100644 --- a/psydac/linalg/block.py +++ b/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 diff --git a/psydac/linalg/stencil.py b/psydac/linalg/stencil.py index 0c818af26..882d2c854 100644 --- a/psydac/linalg/stencil.py +++ b/psydac/linalg/stencil.py @@ -1372,7 +1372,7 @@ 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. @@ -1380,6 +1380,11 @@ def diagonal(self, *, inverse = False, out = None): ---------- 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) diff --git a/psydac/linalg/tests/test_linalg.py b/psydac/linalg/tests/test_linalg.py index b34e222c1..c9962c616 100644 --- a/psydac/linalg/tests/test_linalg.py +++ b/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)