Add diagonal() to KroneckerStencilMatrix

psydac/fem/tests/test_dirichlet_projectors.py

@@ -40,7 +40,7 @@ class SinMapping1D(Mapping):
     _ldim        = 1
     _pdim        = 1
 
-def _test_LO_equality_using_rng(A, B):
+def _test_LO_equality_using_rng(A, B, tol=1e-15):
     """
     A simple tool to check with almost certainty that two linear operators are identical, 
     by applying them repeatedly to random vectors.
@@ -73,11 +73,10 @@ def _test_LO_equality_using_rng(A, B):
         A.dot(x, out=y1)
         B.dot(x, out=y2)
 
-        diff = y1 - y2
-        err  = diff.inner(diff) / diff.space.dimension**2
-        tol  = 1e-15
+        diff            = y1 - y2
+        scaled_err_sqr  = diff.inner(diff) / diff.space.dimension**2
 
-        assert err < tol**2
+        assert scaled_err_sqr < tol**2
 
 #===============================================================================
 @pytest.mark.parametrize('dim', [1, 2])

psydac/linalg/basic.py

@@ -801,15 +801,14 @@ def dtype(self):
         return None
 
     def tosparse(self):
-        from scipy.sparse import csr_matrix
-        out = csr_matrix(self.shape, dtype=self.dtype)
-        for a in self._addends:
-            out += a.tosparse()
+        out = self.addends[0].tosparse().tocsr()
+        for a in self.addends[1:]:
+            out += a.tosparse().tocsr()
         return out
 
     def toarray(self):
-        out = np.zeros(self.shape, dtype=self.dtype)
-        for a in self._addends:
+        out = self.addends[0].toarray()
+        for a in self.addends[1:]:
             out += a.toarray()
         return out
 

psydac/linalg/kron.py

@@ -6,7 +6,7 @@
 from scipy.sparse import coo_matrix
 
 from psydac.linalg.basic   import LinearOperator, LinearSolver
-from psydac.linalg.stencil import StencilVectorSpace, StencilVector, StencilMatrix
+from psydac.linalg.stencil import StencilVectorSpace, StencilVector, StencilMatrix, StencilDiagonalMatrix
 
 __all__ = ('KroneckerStencilMatrix',
            'KroneckerLinearSolver',
@@ -218,6 +218,57 @@ def toarray(self):
     def transpose(self, conjugate=False):
         mats_tr = [Mi.transpose(conjugate=conjugate) for Mi in self.mats]
         return KroneckerStencilMatrix(self.codomain, self.domain, *mats_tr)
+
+    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).
+
+        Returns
+        -------
+        StencilDiagonalMatrix
+            The matrix which contains the main diagonal of self (or its inverse (square root)).
+
+        """
+        # Check `inverse` and `sqrt` argument
+        assert isinstance(inverse, bool)
+        assert isinstance(sqrt, bool)
+
+        # Determine domain and codomain of the StencilDiagonalMatrix
+        V, W = self.domain, self.codomain
+        if inverse:
+            V, W = W, V
+
+        # Check `out` argument
+        if out is not None:
+            assert isinstance(out, StencilDiagonalMatrix)
+            assert out.domain is V
+            assert out.codomain is W
+
+        # Obtain nested numpy array of diagonal entries (or their inverse (square root))
+        # by using the `diagonal` method of StencilMatrices
+        diag = 1.
+        for mat, start, end in zip(self.mats[::-1], V.starts[::-1], V.ends[::-1]):
+            diag = np.array([d*diag for d in mat.diagonal(inverse=inverse, sqrt=sqrt)._data[start:end+1]])
+
+        if out is not None:
+            np.copyto(diag, out._data)
+        else:
+            out = StencilDiagonalMatrix(V, W, diag)
+
+        return out
 
 #==============================================================================
 class KroneckerDenseMatrix(LinearOperator):

psydac/linalg/stencil.py

@@ -1419,11 +1419,12 @@ def diagonal(self, *, inverse = False, sqrt = False, out = None):
         Returns
         -------
         StencilDiagonalMatrix
-            The matrix which contains the main diagonal of self (or its inverse).
+            The matrix which contains the main diagonal of self (or its inverse (square root)).
 
         """
-        # Check `inverse` argument
+        # Check `inverse` and `sqrt` argument
         assert isinstance(inverse, bool)
+        assert isinstance(sqrt, bool)
 
         # Determine domain and codomain of the StencilDiagonalMatrix
         V, W = self.domain, self.codomain
@@ -1436,7 +1437,6 @@ def diagonal(self, *, inverse = False, sqrt = False, out = None):
             assert out.domain is V
             assert out.codomain is W
 
-
         # Extract diagonal data from self and identify output array
         diagonal_indices = self._get_diagonal_indices()
         diag = self._data[diagonal_indices]
@@ -2049,28 +2049,34 @@ def copy(self, *, out=None):
 
         return out
 
-    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.
 
-        In the default case (inverse=False, out=None) self is returned.
+        In the default case (inverse=False, sqrt=False, out=None) self is returned.
 
         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).
 
         Returns
         -------
         StencilDiagonalMatrix
-            Either self, or another StencilDiagonalMatrix with the diagonal inverse.
+            Either self, or another StencilDiagonalMatrix with the diagonal (or its inverse (square root)).
 
         """
-        # Check `inverse` argument
+        # Check `inverse` and `sqrt` argument
         assert isinstance(inverse, bool)
+        assert isinstance(sqrt, bool)
 
         # Determine domain and codomain of the `out` matrix
         V, W = self.domain, self.codomain
@@ -2086,10 +2092,16 @@ def diagonal(self, *, inverse = False, out = None):
             assert out.codomain is W
             data = out._data
 
+        diag = self._data
+
         # Calculate entries, or set `out=self` in default case
         if inverse:
             data = np.divide(1, diag, out=data)
-        elif out:
+            if sqrt:
+                data = np.sqrt(data, out=data)
+        elif sqrt:
+            data = np.sqrt(diag, out=data)
+        elif out not in (None, self):
             np.copyto(data, diag)
         else:
             out = self

psydac/linalg/tests/test_kron_stencil_matrix.py

@@ -1,14 +1,23 @@
-from functools import reduce
+from    functools import reduce
+
+import  pytest
+import  numpy                   as np
+from    mpi4py                  import MPI
+from    scipy.sparse            import kron
+
+from    sympde.calculus         import inner
+from    sympde.expr             import integral, BilinearForm
+from    sympde.topology         import elements_of, Square, Line, Derham
+
+from    psydac.api.discretization import discretize
+from    psydac.api.settings     import PSYDAC_BACKEND_GPYCCEL
+from    psydac.ddm.cart         import DomainDecomposition, CartDecomposition
+from    psydac.linalg.kron      import KroneckerStencilMatrix
+from    psydac.linalg.block     import BlockLinearOperator
+from    psydac.linalg.stencil   import StencilVectorSpace
+from    psydac.linalg.stencil   import StencilVector
+from    psydac.linalg.stencil   import StencilMatrix
 
-import pytest
-import numpy as np
-from scipy.sparse import kron
-
-from psydac.ddm.cart       import DomainDecomposition, CartDecomposition
-from psydac.linalg.stencil import StencilVectorSpace
-from psydac.linalg.stencil import StencilVector
-from psydac.linalg.stencil import StencilMatrix
-from psydac.linalg.kron    import KroneckerStencilMatrix
 #===============================================================================
 def compute_global_starts_ends(domain_decomposition, npts):
     ndims         = len(npts)
@@ -112,3 +121,96 @@ def test_KroneckerStencilMatrix(dtype, npts, pads, periodic):
 
     # Test dot product
     assert np.array_equal(M_sp.dot(w.toarray()), M.dot(w).toarray())
+
+#==============================================================================
+def test_KroneckerStencilMatrix_diagonal(comm=None):
+    """We create three mass matrices (Stencil/Block and Kronecker) belonging to a 2D de Rham sequence, and compare their diagonals."""
+
+    ncells   = [6, 7]
+    degree   = [3, 2]
+    mult     = [1, 2]
+    periodic = [False, True]
+
+    backend = PSYDAC_BACKEND_GPYCCEL
+
+    # 1. Obtain StencilMatrix / BlockLinearOperator (of StencilMatrices) mass matrices
+
+    domain = Square('S', bounds1=(0,1), bounds2=(0,2))
+    derham = Derham(domain, sequence=['h1', 'hcurl', 'l2'])
+
+    domain_h = discretize(domain, ncells=ncells, periodic=periodic, comm=comm)
+    derham_h = discretize(derham, domain_h, degree=degree, multiplicity=mult)
+
+    V0, V1, V2       = derham.spaces
+    V0h, V1h, V2h    = derham_h.spaces
+    V0cs, V1cs, V2cs = [Vh.coeff_space for Vh in derham_h.spaces]
+
+    u0, v0 = elements_of(V0, names='u0, v0')
+    u1, v1 = elements_of(V1, names='u1, v1')
+    u2, v2 = elements_of(V2, names='u2, v2')
+
+    a0 = BilinearForm((u0, v0), integral(domain, u0*v0))
+    a1 = BilinearForm((u1, v1), integral(domain, inner(u1, v1)))
+    a2 = BilinearForm((u2, v2), integral(domain, u2*v2))
+
+    a0h = discretize(a0, domain_h, (V0h, V0h), backend=backend)
+    a1h = discretize(a1, domain_h, (V1h, V1h), backend=backend)
+    a2h = discretize(a2, domain_h, (V2h, V2h), backend=backend)
+
+    M0 = a0h.assemble()
+    M1 = a1h.assemble()
+    M2 = a2h.assemble()
+
+    # 2. Obtain KroneckerStencilMatrix / BlockLinearOperator (of KroneckerStencilMatrices) mass matrices
+
+    domain_1d_x = Line('L', bounds=domain.bounds1)
+    domain_1d_y = Line('L', bounds=domain.bounds2)
+    domains_1d  = (domain_1d_x, domain_1d_y)
+
+    M0s_1d = []
+    M1s_1d = []
+
+    for n, d, m, p, domain_1d in zip(ncells, degree, mult, periodic, domains_1d):
+        derham_1d = Derham(domain_1d)
+
+        domain_1d_h = discretize(domain_1d, ncells=[n, ], periodic=[p, ])
+        derham_1d_h = discretize(derham_1d, domain_1d_h, degree=[d, ], multiplicity=[m, ])
+
+        V0_1d,  V1_1d  = derham_1d.spaces
+        V0h_1d, V1h_1d = derham_1d_h.spaces
+
+        u0_1d, v0_1d = elements_of(V0_1d, names='u0, v0')
+        u1_1d, v1_1d = elements_of(V1_1d, names='u1, v1')
+
+        a0_1d = BilinearForm((u0_1d, v0_1d), integral(domain_1d, u0_1d*v0_1d))
+        a1_1d = BilinearForm((u1_1d, v1_1d), integral(domain_1d, u1_1d*v1_1d))
+
+        a0h_1d = discretize(a0_1d, domain_1d_h, (V0h_1d, V0h_1d))
+        a1h_1d = discretize(a1_1d, domain_1d_h, (V1h_1d, V1h_1d))
+
+        M0s_1d.append(a0h_1d.assemble())
+        M1s_1d.append(a1h_1d.assemble())
+
+    M0_kron = KroneckerStencilMatrix(V0cs, V0cs, *M0s_1d)
+    M1_kron = BlockLinearOperator(V1cs, V1cs, [[KroneckerStencilMatrix(V1cs[0], V1cs[0], M1s_1d[0], M0s_1d[1]), None],
+                                               [None, KroneckerStencilMatrix(V1cs[1], V1cs[1], M0s_1d[0], M1s_1d[1])]])
+    M2_kron = KroneckerStencilMatrix(V2cs, V2cs, *M1s_1d)
+
+    # 3. Test whether M0/1/2.diagonal() is equal to M0/1/2_kron.diagonal() for all possible kwargs
+
+    for M, M_kron in zip((M0, M1, M2), (M0_kron, M1_kron, M2_kron)):
+        options = [True, False]
+
+        for inverse in options:
+            for sqrt in options:
+                M_diag = M.diagonal(inverse=inverse, sqrt=sqrt)
+                M_kron_diag = M_kron.diagonal(inverse=inverse, sqrt=sqrt)
+
+                from psydac.fem.tests.test_dirichlet_projectors import _test_LO_equality_using_rng
+                _test_LO_equality_using_rng(M_diag, M_kron_diag, tol=1e-13)
+
+#==============================================================================
+@pytest.mark.parallel
+def test_KroneckerStencilMatrix_diagonal_parallel():
+    comm = MPI.COMM_WORLD
+    test_KroneckerStencilMatrix_diagonal(comm=comm)