Support for general slicing

sparsestack/StackedSparseArray.py

@@ -96,6 +96,11 @@ def __getitem__(self, key):
             if len(r) == 0:
                 return np.array([0])
             return d
+        elif isinstance(row, slice) and isinstance(col, slice):
+            if None not in (row.start, row.stop, col.start, col.stop):
+                clone = StackedSparseArray(row.stop - row.start, col.stop - col.start)
+                clone.add_sparse_data(r, c, d, name=self.score_names[0])
+                return clone
         return r, c, d
 
     def _getitem_method(self, row, col, name):
@@ -115,17 +120,21 @@ def _getitem_method(self, row, col, name):
             return self.row[idx], self.col[idx], self._slicing_data(name, idx)
         # matrix[:, :, "score_1"]
         if isinstance(row, slice) and isinstance(col, slice):
-            self._is_implemented_slice(row)
-            self._is_implemented_slice(col)
-            return self.row, self.col, self._slicing_data(name)
+            # TODO: self.row, col are sorted, use that knowledge to avoid masking entire array, damnit.
+            rmask = np.ma.masked_inside(self.row, row.start, row.stop).mask
+            cmask = np.ma.masked_inside(self.col, col.start, col.stop).mask
+            mask = rmask & cmask
+            idx = np.where(mask)
+            return self.row[mask], self.col[mask], self._slicing_data(name=name, idx=idx)
         if (row is None and col is None) and isinstance(name, str):
             return self.row, self.col, self._slicing_data(name)
         raise IndexError(_slicing_not_implemented_msg)
 
     def _is_implemented_slice(self, input_slice):
+        pass
         # Currently slices like matrix[2:4, :] or not implemented
-        if not (input_slice.start is None and input_slice.stop is None and input_slice.step is None):
-            raise IndexError(_slicing_not_implemented_msg)
+        # if not (input_slice.start is None and input_slice.stop is None and input_slice.step is None):
+        #     raise IndexError(_slicing_not_implemented_msg)
 
     def _slicing_data(self, name, idx=None):
         if isinstance(name, slice) and len(self.score_names) == 1:
@@ -147,6 +156,8 @@ def validate_index(index, shape):
                     raise IndexError(f"Index ({index}) out of range")
                 if index < 0:
                     index += shape
+            elif isinstance(index, slice):
+                index = slice(*index.indices(shape))
             elif not isinstance(index, slice):
                 index = self._asindices(index, shape)
             return index

tests/test_stacked_sparse_array.py

@@ -237,6 +237,10 @@ def test_sparsestack_slicing():
     matrix = StackedSparseArray(12, 10)
     matrix.add_dense_matrix(arr, "test_score")
 
+    # matrix[0,0,'test_score'] = 42
+    m = matrix[:1,:1, 'test_score']
+    assert m[0, 0] == 0
+
     # Test slicing
     assert matrix[0, 0] == 0
     assert matrix[2, 2] == 22
@@ -253,16 +257,34 @@ def test_sparsestack_slicing():
     assert np.all(r == 2)
     assert np.all(v2 == np.arange(20, 30))
     assert np.all(r2 == 2)
-    r, c, v = matrix["test_score"]
-    r2, c2, v2 = matrix[:, :]
-    r3, c3, v3 = matrix[:, :, 0]
-    r4, c4, v4 = matrix[:, :, :]
-    assert len(c) == len(c2) == len(c3) == len(c4) == 119
-    assert len(r) == len(r2) == len(r3) == len(r4) == 119
-    assert np.all(v == np.arange(1, 120))
-    assert np.all(v2 == np.arange(1, 120))
-    assert np.all(v3 == np.arange(1, 120))
-    assert np.all(v4 == np.arange(1, 120))
+    # r, c, v = matrix["test_score"]
+
+    m2 = matrix[:, :]
+    m3 = matrix[:, :, 0]
+    m4 = matrix[:, :, :]
+
+    assert matrix.shape == m2.shape == m3.shape == m4.shape == (12, 10, 1)
+
+    for m in [m2,m3,m4]:
+        assert np.allclose(m.row, matrix.row)
+        assert np.allclose(m.col, matrix.col)
+        assert np.allclose(m.data['test_score'], matrix.data['test_score'])
+
+    # assert len(r) == len(r2) == len(r3) == len(r4) == 119
+    # assert np.all(v == np.arange(1, 120))
+    # assert np.all(v2 == np.arange(1, 120))
+    # assert np.all(v3 == np.arange(1, 120))
+    # assert np.all(v4 == np.arange(1, 120))
+
+    # r2, c2, v2 = matrix[:, :]
+    # r3, c3, v3 = matrix[:, :, 0]
+    # r4, c4, v4 = matrix[:, :, :]
+    # assert len(c) == len(c2) == len(c3) == len(c4) == 119
+    # assert len(r) == len(r2) == len(r3) == len(r4) == 119
+    # # assert np.all(v == np.arange(1, 120))
+    # # assert np.all(v2 == np.arange(1, 120))
+    # # assert np.all(v3 == np.arange(1, 120))
+    # # assert np.all(v4 == np.arange(1, 120))
 
 
 def test_sparsestack_slicing_mostly_empty_array():
@@ -280,12 +302,12 @@ def test_sparsestack_slicing_mostly_empty_array():
 
 
 @pytest.mark.parametrize("slicing_option", [
-    "matrix[0, 1:3, 'scores1']",
-    "matrix[:2, :, 0]",
-    "matrix[:2, 1, 0]",
-    "matrix[:, 1:, 0]",
-    "matrix[1, 1:, 0]",
-    "matrix[1, 1, :1]",
+    # "matrix[0, 1:3, 'scores1']",
+    # "matrix[:2, :, 0]",
+    # "matrix[:2, 1, 0]",
+    # "matrix[:, 1:, 0]",
+    # "matrix[1, 1:, 0]",
+    # "matrix[1, 1, :1]",
     "matrix[None]",
 ])
 def test_sparsestack_slicing_exceptions(dense_array_sparse, slicing_option):