Added a maintain_sparsity argument to SparseMatrixSimilarity.

gensim/similarities/docsim.py

@@ -562,13 +562,18 @@ class SparseMatrixSimilarity(interfaces.SimilarityABC):
     The matrix is internally stored as a `scipy.sparse.csr` matrix. Unless the entire
     matrix fits into main memory, use `Similarity` instead.
 
+    Takes an optional `maintain_sparsity` argument, setting this to True
+    causes `get_similarities` to return a sparse matrix instead of a
+    dense representation if possible.
+
     See also `Similarity` and `MatrixSimilarity` in this module.
     """
     def __init__(self, corpus, num_features=None, num_terms=None, num_docs=None, num_nnz=None,
-                 num_best=None, chunksize=500, dtype=numpy.float32):
+                 num_best=None, chunksize=500, dtype=numpy.float32, maintain_sparsity=False):
         self.num_best = num_best
         self.normalize = True
         self.chunksize = chunksize
+        self.maintain_sparsity = maintain_sparsity
 
         if corpus is not None:
             logger.info("creating sparse index")
@@ -633,6 +638,9 @@ def get_similarities(self, query):
         if result.shape[1] == 1 and not is_corpus:
             # for queries of one document, return a 1d array
             result = result.toarray().flatten()
+        elif self.maintain_sparsity:
+            # avoid converting to dense array if maintaining sparsity
+            result = result.T
         else:
             # otherwise, return a 2d matrix (#queries x #index)
             result = result.toarray().T

gensim/test/test_similarities.py

@@ -15,6 +15,7 @@
 import tempfile
 
 import numpy
+import scipy
 
 from gensim.corpora import mmcorpus, Dictionary
 from gensim import matutils, utils, similarities
@@ -262,6 +263,20 @@ class TestSparseMatrixSimilarity(unittest.TestCase, _TestSimilarityABC):
     def setUp(self):
         self.cls = similarities.SparseMatrixSimilarity
 
+    def testMaintainSparsity(self):
+        """Sparsity is correctly maintained when maintain_sparsity=True"""
+        num_features = len(dictionary)
+
+        index = self.cls(corpus, num_features=num_features)
+        dense_sims = index[corpus]
+
+        index = self.cls(corpus, num_features=num_features, maintain_sparsity=True)
+        sparse_sims = index[corpus]
+
+        self.assertFalse(scipy.sparse.issparse(dense_sims))
+        self.assertTrue(scipy.sparse.issparse(sparse_sims))
+        numpy.testing.assert_array_equal(dense_sims, sparse_sims.todense())
+
 
 class TestSimilarity(unittest.TestCase, _TestSimilarityABC):
     def setUp(self):