_spfuncs.py 1.9 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576
  1. """ Functions that operate on sparse matrices
  2. """
  3. __all__ = ['count_blocks','estimate_blocksize']
  4. from ._csr import isspmatrix_csr, csr_matrix
  5. from ._csc import isspmatrix_csc
  6. from ._sparsetools import csr_count_blocks
  7. def estimate_blocksize(A,efficiency=0.7):
  8. """Attempt to determine the blocksize of a sparse matrix
  9. Returns a blocksize=(r,c) such that
  10. - A.nnz / A.tobsr( (r,c) ).nnz > efficiency
  11. """
  12. if not (isspmatrix_csr(A) or isspmatrix_csc(A)):
  13. A = csr_matrix(A)
  14. if A.nnz == 0:
  15. return (1,1)
  16. if not 0 < efficiency < 1.0:
  17. raise ValueError('efficiency must satisfy 0.0 < efficiency < 1.0')
  18. high_efficiency = (1.0 + efficiency) / 2.0
  19. nnz = float(A.nnz)
  20. M,N = A.shape
  21. if M % 2 == 0 and N % 2 == 0:
  22. e22 = nnz / (4 * count_blocks(A,(2,2)))
  23. else:
  24. e22 = 0.0
  25. if M % 3 == 0 and N % 3 == 0:
  26. e33 = nnz / (9 * count_blocks(A,(3,3)))
  27. else:
  28. e33 = 0.0
  29. if e22 > high_efficiency and e33 > high_efficiency:
  30. e66 = nnz / (36 * count_blocks(A,(6,6)))
  31. if e66 > efficiency:
  32. return (6,6)
  33. else:
  34. return (3,3)
  35. else:
  36. if M % 4 == 0 and N % 4 == 0:
  37. e44 = nnz / (16 * count_blocks(A,(4,4)))
  38. else:
  39. e44 = 0.0
  40. if e44 > efficiency:
  41. return (4,4)
  42. elif e33 > efficiency:
  43. return (3,3)
  44. elif e22 > efficiency:
  45. return (2,2)
  46. else:
  47. return (1,1)
  48. def count_blocks(A,blocksize):
  49. """For a given blocksize=(r,c) count the number of occupied
  50. blocks in a sparse matrix A
  51. """
  52. r,c = blocksize
  53. if r < 1 or c < 1:
  54. raise ValueError('r and c must be positive')
  55. if isspmatrix_csr(A):
  56. M,N = A.shape
  57. return csr_count_blocks(M,N,r,c,A.indptr,A.indices)
  58. elif isspmatrix_csc(A):
  59. return count_blocks(A.T,(c,r))
  60. else:
  61. return count_blocks(csr_matrix(A),blocksize)