Source code for assembly

"""
Assemble a series of sparse matrices using two gathering matrices. For
more information about assembling routines, see, for example,
`[F. Cuvelier, C. Japhet and G. Scarella, An efficient way to assemble
finite element matrices in vector languages, 2016]
<http://doi.org/10.1007/s10543-015-0587-4>`_.


⭕ To access the source code, click on the [source] button at the right
side or click on
:download:`[assembly.py]</contents/LIBRARY/ptc/mathischeap_ptc/assembly.py>`.
Dependence may exist. In case of error, check import and install required
packages or download required scripts. © mathischeap.com

"""

import numpy as np
from scipy import sparse as spspa




def assemble(Ms, Gr, Gc=None):
    """A (not very fast) routine for assembling matrices and
    vectors.
    
    :param Ms: The matrices/vectors to be assembled. ``Ms[i]`` refers
        to the ith (for example, in element No. i) matrix/vector. The
        matrices/vectors need to be sparse matrices/vectors. A sparse
        vector is a csc_matrix of shape :math:`(n, 1)`.
    :param Gr: The row gathering matrix.
    :type Gr: np.array
    :param Gc: (default: ``None``) The column gathering matrix. When it
        is ``None``, it means we are assembling vectors. Therefore we
        will only need the row gathering matrix, ``Gr``.
    :type Gc: None, np.array
    :return: A csc matrix representing the assembled matrix/vector.

    :example:

    >>> from crazy_mesh import CrazyMeshGlobalNumbering
    >>> from scipy import sparse as spspa
    >>> K = 2
    >>> N = 3
    >>> GM = CrazyMeshGlobalNumbering(K, N)
    >>> GM_F = GM.FP
    >>> GM_V = GM.VP
    >>> Ms = [spspa.random(N**3, 3*(N+1)*N**2, 0.1, 'csc')
    ...     for _ in range(K**3)] # generate a series of sparse matrices
    >>> M = assemble(Ms, GM_V, GM_F) # assemble matrices
    >>> M.shape
    (216, 756)
    >>> Vs = [spspa.random(N**3, 1, 0.5, 'csc')
    ...     for _ in range(K**3)] # generate a series of sparse vectors
    >>> V = assemble(Vs, GM_V)
    >>> V.shape
    (216, 1)

    """
    if Gc is None:  # we are assembling vectors
        DEP = int(np.max(Gr) + 1)
        ROW = list()
        DAT = list()
        for i, Vi in enumerate(Ms):
            assert Vi.__class__.__name__ == 'csc_matrix' and \
                   np.shape(Vi)[1] == 1, f"V[{i}] is not a sparse vector."
            indices = Vi.indices
            data = Vi.data
            ROW.extend(Gr[i][indices])
            DAT.extend(data)
        return spspa.csc_matrix((DAT, ROW, [0, len(ROW)]), shape=(DEP, 1))

    else:
        DEP = int(np.max(Gr) + 1)
        WID = int(np.max(Gc) + 1)
        ROW = list()
        COL = list()
        DAT = list()
        for i, Mi in enumerate(Ms):
            assert Mi.__class__.__name__ in ('csc_matrix', 'csr_matrix')
            indices = Mi.indices
            indptr = Mi.indptr
            data = Mi.data
            nums = np.diff(indptr)
            if Mi.__class__.__name__ == 'csc_matrix':
                for j, num in enumerate(nums):
                    idx = indices[indptr[j]:indptr[j + 1]]
                    ROW.extend(Gr[i][idx])
                    COL.extend([Gc[i][j], ] * num)
            elif Mi.__class__.__name__ == 'csr_matrix':
                for j, num in enumerate(nums):
                    idx = indices[indptr[j]:indptr[j + 1]]
                    ROW.extend([Gr[i][j], ] * num)
                    COL.extend(Gc[i][idx])
            else:
                raise Exception("I can not handle %r." % Mi)
            DAT.extend(data)

        return spspa.csc_matrix((DAT, (ROW, COL)),shape=(DEP, WID))





if __name__ == '__main__':
    import doctest
    doctest.testmod()