Linalg.py 文件源码

def block_diag(*ms):
    '''
    Create a block diagonal matrix from provided ones.

    Parameters
    ----------
    ms : list of 2d ndarray
        The input matrices.

    Returns
    -------
    2d ndarray
        The constructed block diagonal matrix.
    '''
    if len(ms)==0: ms=[np.zeros((0,0))]
    shapes=np.array([a.shape for a in ms])
    dtype=np.find_common_type([m.dtype for m in ms],[])
    result=np.zeros(np.sum(shapes,axis=0),dtype=dtype)
    r,c=0,0
    for i,(cr,cc) in enumerate(shapes):
        result[r:r+cr,c:c+cc]=ms[i]
        r+=cr
        c+=cc
    return result