Python sputils.upcast函数代码示例

    def _binopt(self, other, op):
        """apply the binary operation fn to two sparse matrices"""
        other = self.__class__(other)

        # e.g. csr_plus_csr, csr_minus_csr, etc.
        fn = getattr(sparsetools, self.format + op + self.format)

        maxnnz  = self.nnz + other.nnz
        indptr  = np.empty_like(self.indptr)
        indices = np.empty(maxnnz, dtype=np.intc)
        data    = np.empty(maxnnz, dtype=upcast(self.dtype,other.dtype))

        fn(self.shape[0], self.shape[1], \
                self.indptr,  self.indices,  self.data,
                other.indptr, other.indices, other.data,
                indptr, indices, data)

        actual_nnz = indptr[-1]
        indices = indices[:actual_nnz]
        data    = data[:actual_nnz]
        if actual_nnz < maxnnz // 2:
            #too much waste, trim arrays
            indices = indices.copy()
            data    = data.copy()

        A = self.__class__((data, indices, indptr), shape=self.shape)

        return A

    def _binopt(self, other, op, in_shape=None, out_shape=None):
        """apply the binary operation fn to two sparse matrices"""

        # ideally we'd take the GCDs of the blocksize dimensions
        # and explode self and other to match
        other = self.__class__(other, blocksize=self.blocksize)

        # e.g. bsr_plus_bsr, etc.
        fn = getattr(sparsetools, self.format + op + self.format)

        R,C = self.blocksize

        max_bnnz = len(self.data) + len(other.data)
        indptr  = np.empty_like(self.indptr)
        indices = np.empty(max_bnnz, dtype=np.intc)
        data    = np.empty(R*C*max_bnnz, dtype=upcast(self.dtype,other.dtype))

        fn(self.shape[0]//R, self.shape[1]//C, R, C,
                self.indptr,  self.indices,  np.ravel(self.data),
                other.indptr, other.indices, np.ravel(other.data),
                indptr,       indices,       data)

        actual_bnnz = indptr[-1]
        indices = indices[:actual_bnnz]
        data    = data[:R*C*actual_bnnz]

        if actual_bnnz < max_bnnz/2:
            indices = indices.copy()
            data    = data.copy()

        data = data.reshape(-1,R,C)

        return self.__class__((data, indices, indptr), shape=self.shape)

 def diagonal(self):
     """Returns the main diagonal of the matrix
     """
     #TODO support k-th diagonal
     fn = getattr(sparsetools, self.format + "_diagonal")
     y = np.empty( min(self.shape), dtype=upcast(self.dtype) )
     fn(self.shape[0], self.shape[1], self.indptr, self.indices, self.data, y)
     return y

 def diagonal(self):
     """Returns the main diagonal of the matrix
     """
     M, N = self.shape
     R, C = self.blocksize
     y = np.empty(min(M, N), dtype=upcast(self.dtype))
     sparsetools.bsr_diagonal(M // R, N // C, R, C, self.indptr, self.indices, np.ravel(self.data), y)
     return y

 def _mul_multivector(self, other):
     #matrix * multivector
     M,N = self.shape
     n_vecs = other.shape[1] #number of column vectors
     result = np.zeros( (M,n_vecs), dtype=upcast(self.dtype,other.dtype) )
     for (i,j),v in self.iteritems():
         result[i,:] += v * other[j,:]
     return result

    def _mul_vector(self, other):
        M, N = self.shape
        R, C = self.blocksize

        result = np.zeros(self.shape[0], dtype=upcast(self.dtype, other.dtype))

        bsr_matvec(M // R, N // C, R, C, self.indptr, self.indices, self.data.ravel(), other, result)

        return result

    def _mul_vector(self, other):
        M,N = self.shape

        #output array
        result = np.zeros( self.shape[0], dtype=upcast(self.dtype,other.dtype) )

        # csr_matvec or csc_matvec
        fn = getattr(sparsetools,self.format + '_matvec')
        fn(M, N, self.indptr, self.indices, self.data, other, result)

        return result

    def _mul_multivector(self, other):
        M,N = self.shape
        n_vecs = other.shape[1] #number of column vectors

        result = np.zeros( (M,n_vecs), dtype=upcast(self.dtype,other.dtype) )

        # csr_matvecs or csc_matvecs
        fn = getattr(sparsetools,self.format + '_matvecs')
        fn(M, N, n_vecs, self.indptr, self.indices, self.data, other.ravel(), result.ravel())

        return result

    def _mul_multivector(self,other):
        R,C = self.blocksize
        M,N = self.shape
        n_vecs = other.shape[1] #number of column vectors

        result = np.zeros((M,n_vecs), dtype=upcast(self.dtype,other.dtype))

        bsr_matvecs(M//R, N//C, n_vecs, R, C, \
                self.indptr, self.indices, self.data.ravel(), \
                other.ravel(), result.ravel())

        return result

    def _mul_vector(self, other):
        x = other

        y = np.zeros( self.shape[0], dtype=upcast(self.dtype,x.dtype))

        L = self.data.shape[1]

        M,N = self.shape

        dia_matvec(M,N, len(self.offsets), L, self.offsets, self.data, x.ravel(), y.ravel())

        return y

    def tocsc(self):
        indptr = np.empty(self.shape[1] + 1, dtype=np.intc)
        indices = np.empty(self.nnz, dtype=np.intc)
        data = np.empty(self.nnz, dtype=upcast(self.dtype))

        csr_tocsc(self.shape[0], self.shape[1], self.indptr, self.indices, self.data, indptr, indices, data)

        from csc import csc_matrix

        A = csc_matrix((data, indices, indptr), shape=self.shape)
        A.has_sorted_indices = True
        return A

    def tocsr(self):
        M,N = self.shape
        indptr  = np.empty(M + 1,    dtype=np.intc)
        indices = np.empty(self.nnz, dtype=np.intc)
        data    = np.empty(self.nnz, dtype=upcast(self.dtype))

        csc_tocsr(M, N, \
                 self.indptr, self.indices, self.data, \
                 indptr, indices, data)

        from csr import csr_matrix
        A = csr_matrix((data, indices, indptr), shape=self.shape)
        A.has_sorted_indices = True
        return A

    def _mul_sparse_matrix(self, other):
        M, K1 = self.shape
        K2, N = other.shape

        indptr = np.empty_like(self.indptr)

        R, n = self.blocksize

        # convert to this format
        if isspmatrix_bsr(other):
            C = other.blocksize[1]
        else:
            C = 1

        from csr import isspmatrix_csr

        if isspmatrix_csr(other) and n == 1:
            other = other.tobsr(blocksize=(n, C), copy=False)  # lightweight conversion
        else:
            other = other.tobsr(blocksize=(n, C))

        csr_matmat_pass1(M // R, N // C, self.indptr, self.indices, other.indptr, other.indices, indptr)

        bnnz = indptr[-1]
        indices = np.empty(bnnz, dtype=np.intc)
        data = np.empty(R * C * bnnz, dtype=upcast(self.dtype, other.dtype))

        bsr_matmat_pass2(
            M // R,
            N // C,
            R,
            C,
            n,
            self.indptr,
            self.indices,
            np.ravel(self.data),
            other.indptr,
            other.indices,
            np.ravel(other.data),
            indptr,
            indices,
            data,
        )

        data = data.reshape(-1, R, C)

        # TODO eliminate zeros

        return bsr_matrix((data, indices, indptr), shape=(M, N), blocksize=(R, C))

    def _mul_sparse_matrix(self, other):
        M, K1 = self.shape
        K2, N = other.shape

        major_axis = self._swap((M, N))[0]
        indptr = np.empty(major_axis + 1, dtype=np.intc)

        other = self.__class__(other)  # convert to this format
        fn = getattr(sparsetools, self.format + "_matmat_pass1")
        fn(M, N, self.indptr, self.indices, other.indptr, other.indices, indptr)

        nnz = indptr[-1]
        indices = np.empty(nnz, dtype=np.intc)
        data = np.empty(nnz, dtype=upcast(self.dtype, other.dtype))

        fn = getattr(sparsetools, self.format + "_matmat_pass2")
        fn(M, N, self.indptr, self.indices, self.data, other.indptr, other.indices, other.data, indptr, indices, data)

        return self.__class__((data, indices, indptr), shape=(M, N))

def kronsum(A, B, format=None):
    """kronecker sum of sparse matrices A and B

    Kronecker sum of two sparse matrices is a sum of two Kronecker
    products kron(I_n,A) + kron(B,I_m) where A has shape (m,m)
    and B has shape (n,n) and I_m and I_n are identity matrices
    of shape (m,m) and (n,n) respectively.

    Parameters
    ----------
    A
        square matrix
    B
        square matrix
    format : string
        format of the result (e.g. "csr")

    Returns
    -------
    kronecker sum in a sparse matrix format

    Examples
    --------


    """
    A = coo_matrix(A)
    B = coo_matrix(B)

    if A.shape[0] != A.shape[1]:
        raise ValueError('A is not square')

    if B.shape[0] != B.shape[1]:
        raise ValueError('B is not square')

    dtype = upcast(A.dtype, B.dtype)

    L = kron(identity(B.shape[0],dtype=dtype), A, format=format)
    R = kron(B, identity(A.shape[0],dtype=dtype), format=format)

    return (L+R).asformat(format) #since L + R is not always same format