77 lines
1.9 KiB
Python
77 lines
1.9 KiB
Python
|
""" Functions that operate on sparse matrices
|
||
|
|
"""
|
||
|
|
|
||
|
|
__all__ = ['count_blocks','estimate_blocksize']
|
||
|
|
|
||
|
|
from ._csr import isspmatrix_csr, csr_matrix
|
||
|
|
from ._csc import isspmatrix_csc
|
||
|
|
from ._sparsetools import csr_count_blocks
|
||
|
|
|
||
|
|
|
||
|
|
def estimate_blocksize(A,efficiency=0.7):
|
||
|
|
"""Attempt to determine the blocksize of a sparse matrix
|
||
|
|
|
||
|
|
Returns a blocksize=(r,c) such that
|
||
|
|
- A.nnz / A.tobsr( (r,c) ).nnz > efficiency
|
||
|
|
"""
|
||
|
|
if not (isspmatrix_csr(A) or isspmatrix_csc(A)):
|
||
|
|
A = csr_matrix(A)
|
||
|
|
|
||
|
|
if A.nnz == 0:
|
||
|
|
return (1,1)
|
||
|
|
|
||
|
|
if not 0 < efficiency < 1.0:
|
||
|
|
raise ValueError('efficiency must satisfy 0.0 < efficiency < 1.0')
|
||
|
|
|
||
|
|
high_efficiency = (1.0 + efficiency) / 2.0
|
||
|
|
nnz = float(A.nnz)
|
||
|
|
M,N = A.shape
|
||
|
|
|
||
|
|
if M % 2 == 0 and N % 2 == 0:
|
||
|
|
e22 = nnz / (4 * count_blocks(A,(2,2)))
|
||
|
|
else:
|
||
|
|
e22 = 0.0
|
||
|
|
|
||
|
|
if M % 3 == 0 and N % 3 == 0:
|
||
|
|
e33 = nnz / (9 * count_blocks(A,(3,3)))
|
||
|
|
else:
|
||
|
|
e33 = 0.0
|
||
|
|
|
||
|
|
if e22 > high_efficiency and e33 > high_efficiency:
|
||
|
|
e66 = nnz / (36 * count_blocks(A,(6,6)))
|
||
|
|
if e66 > efficiency:
|
||
|
|
return (6,6)
|
||
|
|
else:
|
||
|
|
return (3,3)
|
||
|
|
else:
|
||
|
|
if M % 4 == 0 and N % 4 == 0:
|
||
|
|
e44 = nnz / (16 * count_blocks(A,(4,4)))
|
||
|
|
else:
|
||
|
|
e44 = 0.0
|
||
|
|
|
||
|
|
if e44 > efficiency:
|
||
|
|
return (4,4)
|
||
|
|
elif e33 > efficiency:
|
||
|
|
return (3,3)
|
||
|
|
elif e22 > efficiency:
|
||
|
|
return (2,2)
|
||
|
|
else:
|
||
|
|
return (1,1)
|
||
|
|
|
||
|
|
|
||
|
|
def count_blocks(A,blocksize):
|
||
|
|
"""For a given blocksize=(r,c) count the number of occupied
|
||
|
|
blocks in a sparse matrix A
|
||
|
|
"""
|
||
|
|
r,c = blocksize
|
||
|
|
if r < 1 or c < 1:
|
||
|
|
raise ValueError('r and c must be positive')
|
||
|
|
|
||
|
|
if isspmatrix_csr(A):
|
||
|
|
M,N = A.shape
|
||
|
|
return csr_count_blocks(M,N,r,c,A.indptr,A.indices)
|
||
|
|
elif isspmatrix_csc(A):
|
||
|
|
return count_blocks(A.T,(c,r))
|
||
|
|
else:
|
||
|
|
return count_blocks(csr_matrix(A),blocksize)
|