[Scipy-svn] r3448 - in trunk/scipy/sparse: . tests

[scipy-svn at scipy.org]

Author: wnbell
Date: 2007-10-19 16:54:27 -0500 (Fri, 19 Oct 2007)
New Revision: 3448

Modified:
   trunk/scipy/sparse/sparse.py
   trunk/scipy/sparse/tests/test_sparse.py
Log:
added sparse.spkron(a,b) - sparse kronecker product 
resolves ticket #42


Modified: trunk/scipy/sparse/sparse.py
===================================================================
--- trunk/scipy/sparse/sparse.py	2007-10-19 18:05:30 UTC (rev 3447)
+++ trunk/scipy/sparse/sparse.py	2007-10-19 21:54:27 UTC (rev 3448)
@@ -7,7 +7,7 @@
 
 __all__ = ['spmatrix','csc_matrix','csr_matrix','coo_matrix',
             'lil_matrix','dok_matrix',
-            'spdiags','speye','spidentity','extract_diagonal',
+            'spdiags','speye','spidentity','spkron','extract_diagonal',
             'isspmatrix','issparse','isspmatrix_csc','isspmatrix_csr',
             'isspmatrix_lil','isspmatrix_dok', 'isspmatrix_coo',
             'lil_eye', 'lil_diags' ]
@@ -2854,6 +2854,65 @@
     diags = ones((1, n), dtype = dtype)
     return spdiags(diags, k, n, m)
 
+def spkron(a,b):
+    """kronecker product of sparse matrices a and b 
+    in COOrdinate format.
+    
+    *Parameters*:
+        a,b : sparse matrices 
+            E.g. csr_matrix, csc_matrix, coo_matrix, etc.
+
+    *Returns*:
+        coo_matrix 
+            kronecker product in COOrdinate format
+
+    *Example*:
+    -------
+
+    >>> a = csr_matrix(array([[0,2],[5,0]]))
+    >>> b = csr_matrix(array([[1,2],[3,4]]))
+    >>> spkron(a,b).todense()
+    matrix([[  0.,   0.,   2.,   4.],
+            [  0.,   0.,   6.,   8.],
+            [  5.,  10.,   0.,   0.],
+            [ 15.,  20.,   0.,   0.]])
+
+    """
+    if not isspmatrix(a) and isspmatrix(b):
+        raise ValueError,'expected sparse matrix'
+
+    a,b = a.tocoo(),b.tocoo()
+    output_shape = (a.shape[0]*b.shape[0],a.shape[1]*b.shape[1])  
+
+    if a.nnz == 0 or b.nnz == 0:
+        # kronecker product is the zero matrix
+        return coo_matrix(None, dims=output_shape)
+            
+
+    # expand entries of a into blocks
+    row  = a.row.repeat(b.nnz)
+    col  = a.col.repeat(b.nnz)
+    data = a.data.repeat(b.nnz)
+    
+    row *= b.shape[0]
+    col *= b.shape[1]
+
+    # increment block indices
+    row,col = row.reshape(-1,b.nnz),col.reshape(-1,b.nnz)
+    row += b.row
+    col += b.col
+    row,col = row.reshape(-1),col.reshape(-1)
+
+    # compute block entries
+    data = data.reshape(-1,b.nnz)
+    data *= b.data
+    data = data.reshape(-1)
+
+    return coo_matrix((data,(row,col)), dims=output_shape)
+
+
+
+
 def lil_eye((r,c), k=0, dtype=float):
     """Generate a lil_matrix of dimensions (r,c) with the k-th
     diagonal set to 1.

Modified: trunk/scipy/sparse/tests/test_sparse.py
===================================================================
--- trunk/scipy/sparse/tests/test_sparse.py	2007-10-19 18:05:30 UTC (rev 3447)
+++ trunk/scipy/sparse/tests/test_sparse.py	2007-10-19 21:54:27 UTC (rev 3448)
@@ -22,7 +22,7 @@
 from numpy.testing import *
 set_package_path()
 from scipy.sparse import csc_matrix, csr_matrix, dok_matrix, coo_matrix, \
-     spidentity, speye, extract_diagonal, lil_matrix, lil_eye, lil_diags
+     spidentity, speye, spkron, extract_diagonal, lil_matrix, lil_eye, lil_diags
 from scipy.linsolve import splu
 restore_path()
 
@@ -1072,9 +1072,32 @@
         b = array([[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype='d')
         assert_array_equal(a.toarray(), b)
 
-                   
+    def check_spkron(self):
+        from numpy import kron
 
-class TestCoo(NumpyTestCase):
+        cases = []
+
+        cases.append(array([[ 0]]))
+        cases.append(array([[-1]]))
+        cases.append(array([[ 4]]))
+        cases.append(array([[10]]))
+        cases.append(array([[0],[0]]))
+        cases.append(array([[0,0]]))
+        cases.append(array([[1,2],[3,4]]))
+        cases.append(array([[0,2],[5,0]]))
+        cases.append(array([[0,2,-6],[8,0,14]]))
+        cases.append(array([[5,4],[0,0],[6,0]]))
+        cases.append(array([[5,4,4],[1,0,0],[6,0,8]]))
+        cases.append(array([[0,1,0,2,0,5,8]]))
+
+        for a in cases:
+            for b in cases:
+                result = spkron(csr_matrix(a),csr_matrix(b)).todense()
+                expected = kron(a,b)
+        
+                assert_array_equal(result,expected)
+
+class TestCOO(NumpyTestCase):
     def check_constructor1(self):
         row  = numpy.array([2, 3, 1, 3, 0, 1, 3, 0, 2, 1, 2])
         col  = numpy.array([0, 1, 0, 0, 1, 1, 2, 2, 2, 2, 1])