Cache reshape and transpose

This stores the last few objects returned from reshape and transpose
calls.  This allows efficiencies from in-place operations like
`sum_duplicates` and `sort_indices` to persist in interative workflows.

Modern NumPy programmers are accustomed to operations like .transpose()
being cheap and aren't accustomed to having to pay sorting costs after
many computations.  These assumptions are no longer true in sparse by
default.  However, by caching recent transpose and reshape objects we
can reuse their inplace modifications.  This greatly accelerates common
machine learning workloads.

Author: mrocklin

sparse/core.py

@@ -154,6 +154,7 @@ def __init__(self, coords, data=None, shape=None, has_duplicates=True,
         assert not self.shape or len(data) == self.coords.shape[1]
         self.has_duplicates = has_duplicates
         self.sorted = sorted
+        self._cache = defaultdict(lambda: deque(maxlen=3))
 
     @classmethod
     def from_numpy(cls, x):
@@ -209,6 +210,9 @@ def __getitem__(self, index):
             index = (index,)
         index = tuple(ind + self.shape[i] if isinstance(ind, numbers.Integral) and ind < 0 else ind
                       for i, ind in enumerate(index))
+        if (all(ind == slice(None) or ind == slice(0, d)
+                for ind, d in zip(index, self.shape))):
+            return self
         mask = np.ones(self.nnz, dtype=bool)
         for i, ind in enumerate([i for i in index if i is not None]):
             if ind == slice(None, None):
@@ -325,19 +329,15 @@ def transpose(self, axes=None):
         if axes == tuple(range(self.ndim)):
             return self
 
+        for ax, value in self._cache['transpose']:
+            if ax == axes:
+                return value
+
         shape = tuple(self.shape[ax] for ax in axes)
         result = COO(self.coords[axes, :], self.data, shape,
                      has_duplicates=self.has_duplicates)
 
-        if axes == (1, 0):
-            try:
-                result._csc = self._csr.T
-            except AttributeError:
-                pass
-            try:
-                result._csr = self._csc.T
-            except AttributeError:
-                pass
+        self._cache['transpose'].append((axes, result))
         return result
 
     @property
@@ -374,8 +374,14 @@ def reshape(self, shape):
             extra = int(np.prod(self.shape) /
                         np.prod([d for d in shape if d != -1]))
             shape = tuple([d if d != -1 else extra for d in shape])
+
         if self.shape == shape:
             return self
+
+        for sh, value in self._cache['reshape']:
+            if sh == shape:
+                return value
+
         # TODO: this np.prod(self.shape) enforces a 2**64 limit to array size
         linear_loc = self.linear_loc()
 
@@ -385,9 +391,12 @@ def reshape(self, shape):
             coords[-(i + 1), :] = (linear_loc // strides) % d
             strides *= d
 
-        return COO(coords, self.data, shape,
-                   has_duplicates=self.has_duplicates,
-                   sorted=self.sorted)
+        result = COO(coords, self.data, shape,
+                     has_duplicates=self.has_duplicates,
+                     sorted=self.sorted)
+
+        self._cache['reshape'].append((shape, result))
+        return result
 
     def to_scipy_sparse(self):
         assert self.ndim == 2

sparse/tests/test_core.py

@@ -418,14 +418,6 @@ def test_cache_csr():
     assert s.tocsr() is s.tocsr()
     assert s.tocsc() is s.tocsc()
 
-    st = s.T
-
-    assert_eq(st._csr, st)
-    assert_eq(st._csc, st)
-
-    assert isinstance(st.tocsr(), scipy.sparse.csr_matrix)
-    assert isinstance(st.tocsc(), scipy.sparse.csc_matrix)
-
 
 def test_empty_shape():
     x = COO([], [1.0])
@@ -469,3 +461,16 @@ def test_add_many_sparse_arrays():
     x = COO({(1, 1): 1})
     y = sum([x] * 100)
     assert y.nnz < np.prod(y.shape)
+
+
+def test_caching():
+    x = COO({(10, 10, 10): 1})
+
+    assert x[:].reshape((100, 10)).transpose().tocsr() is x[:].reshape((100, 10)).transpose().tocsr()
+
+    x = COO({(1, 1, 1, 1, 1, 1, 1, 2): 1})
+
+    for i in range(x.ndim):
+        x.reshape((1,) * i + (2,) + (1,) * (x.ndim - i - 1))
+
+    assert len(x._cache['reshape']) < 5