""" dict_sbmat module

Helper functions to deal with sparse block matrices using a dictionary
with index tuples as keys. This is beneficial e.g. in assembling coupled
circuit/FEM systems in the matrix level.

E.g. if you have sparse matrices A B C D and you want to create a sparse
block matrix like
[[A, 0, 0],
 [0, B, C],
 [0,-D, 0]]

you can do the following:
> sm = {}
> sm[(0,0)] = A
> sm[(1,1)] = B
> sm[(1,2)] = C
> sm[(2,1)] = -D

Inspect the block structure with print_blocks
> dict_tools.print_blocks(sm)

Create a scipy bmat with tolist
> S = scipy.sparse.bmat(dict_tools.tolist(sm))

Pick subblocks corresponding to the block indices of the resulting sparse 
matrix with 'submat' and 'mk_selector_builder'. 

To e.g. pick blocks

S11 = [[A]]
S12 = [[0,0]]
S21 = [[0],
       [0]]
S22 = [[B,C],
       [-D,0]]

use

> builder = mk_selector_builder(sm)
> P11,Q11 = builder([0],[0])
> S11 = P11*S*Q11
> P12,Q12 = builder([0], [1,2])
> S12 = P12*S*P12
> P21,Q21 = builder([1,2], [0])
> S21 = P21*S*Q21
> P22,Q22 = builder([1,2], [1,2])
> S22 = P22*S*Q22

At first this seems terribly inefficient, but it really isn't. Using the 
sparse linear algebra * to pick rows and columnt sacrifices some memory but 
is extremely simple to use and e.g. utilizes the sparsity patterns of all 
matrices efficiently. """


import numpy as np
import scipy.sparse as sps

from itertools import product

def tolist(dmat):
    """ Convert dmat to a list format [[A,None,...], ...] where 
    empty blocks are filled with None. This can be given as an input to 
    scipy.sparse.bmat """

    inds = np.array(list(dmat.keys()))        
    nrows = np.max(inds[:,0])+1
    ncols = np.max(inds[:,1])+1
    return [[dmat.get((row,col),None) 
             for col in range(0,ncols)]
             for row in range(0,nrows)]
    
def print_blocks(dmat):
    
    inds = np.array(list(dmat.keys()))
    
    xdim = np.max(inds[:,0])+1
    ydim = np.max(inds[:,1])+1
    
    strs = []
    for i in range(0, xdim):
        rowstr = []
        for j in range(0, ydim):
            m = dmat.get((i,j), None)
            if m is None:
                rowstr.append("{:10s}".format(" "))
            else:
                try:
                    s = m.shape
                except AttributeError:
                    s = "(s)"
                rowstr.append("{:10s}".format(s))
        strs.append(rowstr)
        
    print ("\n".join([ "|" + ", ".join(ss) + "|" for ss in strs]))

    
def submat(X,ii,jj):
    """ Pick a block submatrix from block matrix X by rows ii and cols jj """
    return {(ix,jx): X.get((i,j), None) 
            for (ix,i),(jx,j) in product(enumerate(ii), 
                                         enumerate(jj))}
    
def mk_selector_builder(dmat):
    """ Builds a function which generated pairs of matrices (Pr, Pc) which 
    can be used to choose subblocks of dmat. Dmat needs to be in a dict
    block form. The matrices generated by the returned function act in the 
    resulting sparse matrix 
    Input: dmat - a dictionary {(row, col): sparse matrix / None} representing
                  a sparse block matrix
    Output: fun: (rows: [Int], cols: [Int]) -> (Pr, Pc) , where Pr picks 
            the rows and Pc picks the columns of a sparse matrix

    """
    
    # find out how many block rows and columns we have
    inds = np.array(list(dmat.keys()))   
    nrows = np.max(inds[:,0])+1
    ncols = np.max(inds[:,1])+1

    rowshapes = [None]*nrows
    colshapes = [None]*ncols
    
    # find out the shapes of rows and columns
    for i in range(0, nrows):
        for j in range(0, ncols):
            m = dmat.get((i,j), None)
            if m is None:
                # an empty block doesn't give us any information
                continue
            else:
                shp = m.shape
                # double check if we already have a dimension that it matches
                # the discovered one
                # If we don't have a dimension, add it to the lists
                existing_rows = rowshapes[i]
                existing_cols = colshapes[j]
                if existing_rows is not None:
                    assert shp[0] == existing_rows, "Shapes don't match"
                else:
                    rowshapes[i] = shp[0]
                if existing_cols is not None:
                    assert shp[1] == existing_cols, "Shapes don't match"
                else:
                    colshapes[j] = shp[1]

    rows_in_matrix = sum(rowshapes)
    cols_in_matrix = sum(colshapes)
    
    rowcsum = np.cumsum([0]+rowshapes)
    colcsum = np.cumsum([0]+colshapes)

    def fun(rows, cols):
        """ Rows, cols are lists of indices """
        try:
            # get the amount of rows in each block in the slice
            nrows = [rowshapes[row] for row in rows]
        except TypeError:
            # If rows is not a list, try to convert it into a singleton
            nrows = [rowshapes[rows]]
            rows = [rows]
        try:               
            ncols = [colshapes[col] for col in cols]
        except TypeError:
            ncols = [colshapes[cols]]
            cols = [cols]
            
        # total number of rows to be in the slice
        rows_in_slice = np.sum(nrows)
        cols_in_slice = np.sum(ncols)
        # the ones to be inserted into Pr and Pc
        rowi = np.ones((rows_in_slice,))
        coli = np.ones((cols_in_slice,))
        # The indices from where the blocks in the slice start 
        rowstarts = [rowcsum[row] for row in rows]
        colstarts = [colcsum[col] for col in cols]
        # the row indices from the whole slice
        rowvector = np.concatenate([np.arange(rowstart, rowstart+rows) 
                                   for rows, rowstart in zip(nrows, rowstarts)])
        # row permutation matrix is (rows, totalrows)                 
        Pr = sps.csc_matrix((rowi, (np.arange(0,rows_in_slice),
                                    rowvector)), 
                            shape=(rows_in_slice, rows_in_matrix))
        
        colvector = np.concatenate([np.arange(colstart, colstart+cols) 
                                   for cols, colstart in zip(ncols, colstarts)])
        # col permutation matrix
        Pc = sps.csc_matrix((coli, (colvector,
                                    np.arange(0,cols_in_slice))), 
                            shape=(cols_in_matrix, cols_in_slice))
        return (Pr, Pc)
    
    return fun

    
if __name__ == "__main__":
    
#    import scipy.sparse as sps
    
    A = dict()
     
    a = np.array([[1,2,3], [4,5,6]])
    b = np.array([[1,2], [3,4]])
    c = np.ones((2,3))

    A[(0,0)] = b
    A[(1,2)] = a
    A[(2,1)] = 10*a
    A[(2,0)] = 10*b
    A[(2,2)] = c
    
    
    spmat = sps.bmat(tolist(A)).tocsc()
    
    f = mk_selector_builder(A)
    
    def match(a1, a2):
        return np.all(a1 == a2)
    
    # Small and cute unit test suite
    
    Pr2, Pc1 = f(2,1)
    assert match(10*a, (Pr2 @ spmat @ Pc1).toarray()), "Block (2,1)" 
    Pr0, Pc0 = f(0,0)
    assert match(b, (Pr0 @ spmat @ Pc0).toarray()), "Block (0,0)" 
    Pr1, Pc2 = f(1,2)
    assert match(a, (Pr1 @ spmat @ Pc2).toarray()), "Block (1,2)" 
    assert match(10*b, (Pr2 @ spmat @ Pc0).toarray()), "Block (2,1)" 
    
    assert np.all(np.zeros((2,3)) == (Pr1 @ spmat @ Pc1).toarray()), "Block(1,1) is zero"
    
    Pr12, Pc12 = f([1,2], [1,2])
    m = sps.bmat([[None,    a],
                  [10*a, c]]).toarray()
    assert np.all(m == (Pr12 @ spmat @ Pc12).toarray()), "Block([1,2], [1,2])"
    
    Pr02,Pc02 = f([0,2], [0,2])
    m = sps.bmat([[b, None],
                  [10*b, c]]).toarray()
    assert np.all(m == (Pr02 @ spmat @ Pc02).toarray()), "Block([0,2], [0,2])"
    

""" dict_sbmat module

Helper functions to deal with sparse block matrices using a dictionary
with index tuples as keys. This is beneficial e.g. in assembling coupled
circuit/FEM systems in the matrix level.

E.g. if you have sparse matrices A B C D and you want to create a sparse
block matrix like
[[A, 0, 0],
 [0, B, C],
 [0,-D, 0]]

you can do the following:
> sm = {}
> sm[(0,0)] = A
> sm[(1,1)] = B
> sm[(1,2)] = C
> sm[(2,1)] = -D

Inspect the block structure with print_blocks
> dict_tools.print_blocks(sm)

Create a scipy bmat with tolist
> S = scipy.sparse.bmat(dict_tools.tolist(sm))

Pick subblocks corresponding to the block indices of the resulting sparse 
matrix with 'submat' and 'mk_selector_builder'. 

To e.g. pick blocks

S11 = [[A]]
S12 = [[0,0]]
S21 = [[0],
       [0]]
S22 = [[B,C],
       [-D,0]]

use

> builder = mk_selector_builder(sm)
> P11,Q11 = builder([0],[0])
> S11 = P11*S*Q11
> P12,Q12 = builder([0], [1,2])
> S12 = P12*S*P12
> P21,Q21 = builder([1,2], [0])
> S21 = P21*S*Q21
> P22,Q22 = builder([1,2], [1,2])
> S22 = P22*S*Q22

At first this seems terribly inefficient, but it really isn't. Using the 
sparse linear algebra * to pick rows and columnt sacrifices some memory but 
is extremely simple to use and e.g. utilizes the sparsity patterns of all 
matrices efficiently. """


import numpy as np
import scipy.sparse as sps

from itertools import product

def tolist(dmat):
    """ Convert dmat to a list format [[A,None,...], ...] where 
    empty blocks are filled with None. This can be given as an input to 
    scipy.sparse.bmat """

    inds = np.array(list(dmat.keys()))        
    nrows = np.max(inds[:,0])+1
    ncols = np.max(inds[:,1])+1
    return [[dmat.get((row,col),None) 
             for col in range(0,ncols)]
             for row in range(0,nrows)]
    
def print_blocks(dmat):
    
    inds = np.array(list(dmat.keys()))
    
    xdim = np.max(inds[:,0])+1
    ydim = np.max(inds[:,1])+1
    
    strs = []
    for i in range(0, xdim):
        rowstr = []
        for j in range(0, ydim):
            m = dmat.get((i,j), None)
            if m is None:
                rowstr.append("{:10s}".format(" "))
            else:
                try:
                    s = m.shape
                except AttributeError:
                    s = "(s)"
                rowstr.append("{:10s}".format(s))
        strs.append(rowstr)
        
    print ("\n".join([ "|" + ", ".join(ss) + "|" for ss in strs]))

    
def submat(X,ii,jj):
    """ Pick a block submatrix from block matrix X by rows ii and cols jj """
    return {(ix,jx): X.get((i,j), None) 
            for (ix,i),(jx,j) in product(enumerate(ii), 
                                         enumerate(jj))}
    
def mk_selector_builder(dmat):
    """ Builds a function which generated pairs of matrices (Pr, Pc) which 
    can be used to choose subblocks of dmat. Dmat needs to be in a dict
    block form. The matrices generated by the returned function act in the 
    resulting sparse matrix 
    Input: dmat - a dictionary {(row, col): sparse matrix / None} representing
                  a sparse block matrix
    Output: fun: (rows: [Int], cols: [Int]) -> (Pr, Pc) , where Pr picks 
            the rows and Pc picks the columns of a sparse matrix

    """
    
    # find out how many block rows and columns we have
    inds = np.array(list(dmat.keys()))   
    nrows = np.max(inds[:,0])+1
    ncols = np.max(inds[:,1])+1

    rowshapes = [None]*nrows
    colshapes = [None]*ncols
    
    # find out the shapes of rows and columns
    for i in range(0, nrows):
        for j in range(0, ncols):
            m = dmat.get((i,j), None)
            if m is None:
                # an empty block doesn't give us any information
                continue
            else:
                shp = m.shape
                # double check if we already have a dimension that it matches
                # the discovered one
                # If we don't have a dimension, add it to the lists
                existing_rows = rowshapes[i]
                existing_cols = colshapes[j]
                if existing_rows is not None:
                    assert shp[0] == existing_rows, "Shapes don't match"
                else:
                    rowshapes[i] = shp[0]
                if existing_cols is not None:
                    assert shp[1] == existing_cols, "Shapes don't match"
                else:
                    colshapes[j] = shp[1]

    rows_in_matrix = sum(rowshapes)
    cols_in_matrix = sum(colshapes)
    
    rowcsum = np.cumsum([0]+rowshapes)
    colcsum = np.cumsum([0]+colshapes)

    def fun(rows, cols):
        """ Rows, cols are lists of indices """
        try:
            # get the amount of rows in each block in the slice
            nrows = [rowshapes[row] for row in rows]
        except TypeError:
            # If rows is not a list, try to convert it into a singleton
            nrows = [rowshapes[rows]]
            rows = [rows]
        try:               
            ncols = [colshapes[col] for col in cols]
        except TypeError:
            ncols = [colshapes[cols]]
            cols = [cols]
            
        # total number of rows to be in the slice
        rows_in_slice = np.sum(nrows)
        cols_in_slice = np.sum(ncols)
        # the ones to be inserted into Pr and Pc
        rowi = np.ones((rows_in_slice,))
        coli = np.ones((cols_in_slice,))
        # The indices from where the blocks in the slice start 
        rowstarts = [rowcsum[row] for row in rows]
        colstarts = [colcsum[col] for col in cols]
        # the row indices from the whole slice
        rowvector = np.concatenate([np.arange(rowstart, rowstart+rows) 
                                   for rows, rowstart in zip(nrows, rowstarts)])
        # row permutation matrix is (rows, totalrows)                 
        Pr = sps.csc_matrix((rowi, (np.arange(0,rows_in_slice),
                                    rowvector)), 
                            shape=(rows_in_slice, rows_in_matrix))
        
        colvector = np.concatenate([np.arange(colstart, colstart+cols) 
                                   for cols, colstart in zip(ncols, colstarts)])
        # col permutation matrix
        Pc = sps.csc_matrix((coli, (colvector,
                                    np.arange(0,cols_in_slice))), 
                            shape=(cols_in_matrix, cols_in_slice))
        return (Pr, Pc)
    
    return fun

    
if __name__ == "__main__":
    
#    import scipy.sparse as sps
    
    A = dict()
     
    a = np.array([[1,2,3], [4,5,6]])
    b = np.array([[1,2], [3,4]])
    c = np.ones((2,3))

    A[(0,0)] = b
    A[(1,2)] = a
    A[(2,1)] = 10*a
    A[(2,0)] = 10*b
    A[(2,2)] = c
    
    
    spmat = sps.bmat(tolist(A)).tocsc()
    
    f = mk_selector_builder(A)
    
    def match(a1, a2):
        return np.all(a1 == a2)
    
    # Small and cute unit test suite
    
    Pr2, Pc1 = f(2,1)
    assert match(10*a, (Pr2 @ spmat @ Pc1).toarray()), "Block (2,1)" 
    Pr0, Pc0 = f(0,0)
    assert match(b, (Pr0 @ spmat @ Pc0).toarray()), "Block (0,0)" 
    Pr1, Pc2 = f(1,2)
    assert match(a, (Pr1 @ spmat @ Pc2).toarray()), "Block (1,2)" 
    assert match(10*b, (Pr2 @ spmat @ Pc0).toarray()), "Block (2,1)" 
    
    assert np.all(np.zeros((2,3)) == (Pr1 @ spmat @ Pc1).toarray()), "Block(1,1) is zero"
    
    Pr12, Pc12 = f([1,2], [1,2])
    m = sps.bmat([[None,    a],
                  [10*a, c]]).toarray()
    assert np.all(m == (Pr12 @ spmat @ Pc12).toarray()), "Block([1,2], [1,2])"
    
    Pr02,Pc02 = f([0,2], [0,2])
    m = sps.bmat([[b, None],
                  [10*b, c]]).toarray()
    assert np.all(m == (Pr02 @ spmat @ Pc02).toarray()), "Block([0,2], [0,2])"
    

dict_sbmat module  Helper functions to deal with sparse block matrices using a dictionary with index tuples as keys. This is beneficial e.g. in assembling coupled circuit/FEM systems in the matrix level.  E.g. if you have sparse matrices A B C D and you want to create a sparse block matrix like [[A, 0, 0],  [0, B, C],  [0,-D, 0]]  you can do the following: > sm = {} > sm[(0,0)] = A > sm[(1,1)] = B > sm[(1,2)] = C > sm[(2,1)] = -D  Inspect the block structure with print_blocks > dict_tools.print_blocks(sm)  Create a scipy bmat with tolist > S = scipy.sparse.bmat(dict_tools.tolist(sm))  Pick subblocks corresponding to the block indices of the resulting sparse  matrix with 'submat' and 'mk_selector_builder'.   To e.g. pick blocks  S11 = [[A]] S12 = [[0,0]] S21 = [[0],        [0]] S22 = [[B,C],        [-D,0]]  use  > builder = mk_selector_builder(sm) > P11,Q11 = builder([0],[0]) > S11 = P11*S*Q11 > P12,Q12 = builder([0], [1,2]) > S12 = P12*S*P12 > P21,Q21 = builder([1,2], [0]) > S21 = P21*S*Q21 > P22,Q22 = builder([1,2], [1,2]) > S22 = P22*S*Q22  At first this seems terribly inefficient, but it really isn't. Using the  sparse linear algebra * to pick rows and columnt sacrifices some memory but  is extremely simple to use and e.g. utilizes the sparsity patterns of all  matrices efficiently. Convert dmat to a list format [[A,None,...], ...] where      empty blocks are filled with None. This can be given as an input to      scipy.sparse.bmat Pick a block submatrix from block matrix X by rows ii and cols jj Builds a function which generated pairs of matrices (Pr, Pc) which      can be used to choose subblocks of dmat. Dmat needs to be in a dict     block form. The matrices generated by the returned function act in the      resulting sparse matrix      Input: dmat - a dictionary {(row, col): sparse matrix / None} representing                   a sparse block matrix     Output: fun: (rows: [Int], cols: [Int]) -> (Pr, Pc) , where Pr picks              the rows and Pc picks the columns of a sparse matrix # find out how many block rows and columns we have # find out the shapes of rows and columns # an empty block doesn't give us any information # double check if we already have a dimension that it matches # the discovered one # If we don't have a dimension, add it to the lists Rows, cols are lists of indices # get the amount of rows in each block in the slice # If rows is not a list, try to convert it into a singleton # total number of rows to be in the slice # the ones to be inserted into Pr and Pc # The indices from where the blocks in the slice start # the row indices from the whole slice # row permutation matrix is (rows, totalrows) # col permutation matrix #    import scipy.sparse as sps # Small and cute unit test suite