Support single-integer slicing of DelayedArrays, SparseNdarrays.

This allows access with 'X[i]' to extract subarrays along the first dimension,
which is handy for generic functions like biocutils.show_as_cell().

Author: LTLA

src/delayedarray/DelayedArray.py

@@ -793,29 +793,31 @@ def __rand__(self, other) -> "DelayedArray":
         return _wrap_isometric_with_args(self, other, operation="logical_and", right=False)
 
     # Subsetting.
-    def __getitem__(self, subset: Tuple[Union[slice, Sequence], ...]) -> Union["DelayedArray", ndarray]:
+    def __getitem__(self, subset) -> Union["DelayedArray", ndarray]:
         """Take a subset of this ``DelayedArray``. This follows the same logic as NumPy slicing and will generate a
         :py:class:`~delayedarray.Subset.Subset` object when the subset operation preserves the dimensionality of the
         seed, i.e., ``args`` is defined using the :py:meth:`~numpy.ix_` function.
 
         Args:
             subset:
-                A :py:class:`tuple` of length equal to the dimensionality of
-                this ``DelayedArray``.  We attempt to support most types of
-                NumPy slicing; however, only subsets that preserve
-                dimensionality will generate a delayed subset operation.
+                A :py:class:`tuple` of length equal to the dimensionality of this ``DelayedArray``, or a single integer specifying an index on the first dimension.
+                We attempt to support most types of NumPy slicing; however, only subsets that preserve dimensionality will generate a delayed subset operation.
 
         Returns:
-            If the dimensionality is preserved by ``subset``, a
-            ``DelayedArray`` containing a delayed subset operation is returned.
-            Otherwise, a :py:class:`~numpy.ndarray` is returned containing the
-            realized subset.
+            If the dimensionality is preserved by ``subset``, a ``DelayedArray`` containing a delayed subset operation is returned.
+            Otherwise, a :py:class:`~numpy.ndarray` is returned containing the realized subset.
         """
+        if not isinstance(subset, Tuple):
+            replacement = [slice(None)] * len(self.shape)
+            replacement[0] = subset
+            subset = (*replacement,)
+
         cleaned = _getitem_subset_preserves_dimensions(self.shape, subset)
         if cleaned is not None:
             sout = Subset(self._seed, cleaned)
             sout = _simplify_subset(sout)
             return DelayedArray(sout)
+
         return _getitem_subset_discards_dimensions(self._seed, subset, extract_dense_array)
 
 

src/delayedarray/SparseNdarray.py

@@ -748,24 +748,28 @@ def __rand__(self, other) -> Union["SparseNdarray", numpy.ndarray]:
 
 
     # Subsetting.
-    def __getitem__(self, subset: Tuple[Union[slice, Sequence], ...]) -> Union["SparseNdarray", numpy.ndarray]:
+    def __getitem__(self, subset) -> Union["SparseNdarray", numpy.ndarray]:
         """Take a subset of this ``SparseNdarray``. This follows the same logic as NumPy slicing and will generate a
         :py:class:`~delayedarray.Subset.Subset` object when the subset operation preserves the dimensionality of the
         seed, i.e., ``args`` is defined using the :py:meth:`~numpy.ix_` function.
 
         Args:
-            args:
-                A :py:class:`tuple` of length equal to the dimensionality of this ``SparseNdarray``.
-                Any NumPy slicing is supported but only subsets that preserve dimensionality will generate a
-                delayed subset operation.
+            subset:
+                A :py:class:`tuple` of length equal to the dimensionality of this ``SparseNdarray``, or a single integer specfying an index on the first dimension. 
+                We attempt to support most types of NumPy slicing; however, only subsets that preserve dimensionality will generate a ``SparseNdarray``.
 
         Raises:
             ValueError: If ``args`` contain more dimensions than the shape of the array.
 
         Returns:
-            If the dimensionality is preserved by ``args``, a ``SparseNdarray`` containing a delayed subset operation is
-            returned. Otherwise, a :py:class:`~numpy.ndarray` is returned containing the realized subset.
+            If the dimensionality is preserved by ``subset``, a ``SparseNdarray`` containing the specified subset is returned.
+            Otherwise, a :py:class:`~numpy.ndarray` is returned containing the realized subset.
         """
+        if not isinstance(subset, Tuple):
+            replacement = [slice(None)] * len(self.shape)
+            replacement[0] = subset
+            subset = (*replacement,)
+
         cleaned = _getitem_subset_preserves_dimensions(self.shape, subset)
         if cleaned is not None:
             # No need to sanitize here, as the extractors can take unsorted subsets.

src/delayedarray/_subset.py

@@ -62,7 +62,7 @@ def _sanitize_subset(subset: Sequence):
     return san, remap
 
 
-def _getitem_subset_preserves_dimensions(shape: Tuple[int, ...], args: Tuple[Union[slice, Sequence], ...]):
+def _getitem_subset_preserves_dimensions(shape: Tuple[int, ...], args: Tuple): 
     ndim = len(shape)
     if not isinstance(args, tuple):
         args = [args] + [slice(None)] * (ndim - 1)
@@ -71,6 +71,11 @@ def _getitem_subset_preserves_dimensions(shape: Tuple[int, ...], args: Tuple[Uni
     elif len(args) > ndim:
         raise ValueError("more indices in 'args' than there are dimensions in 'seed'")
 
+    # Checking if there are any integers here.
+    for d, idx in enumerate(args):
+        if isinstance(idx, int) or isinstance(idx, integer):
+            return None
+
     # Checking if we're preserving the shape via a cross index.
     cross_index = True
     for d, idx in enumerate(args):
@@ -117,7 +122,7 @@ def _getitem_subset_preserves_dimensions(shape: Tuple[int, ...], args: Tuple[Uni
     return None
 
 
-def _getitem_subset_discards_dimensions(x, args: Tuple[Union[slice, Sequence], ...], injected_extract_dense_array: Callable):
+def _getitem_subset_discards_dimensions(x, args: Tuple, injected_extract_dense_array: Callable):
     failed = False
     sanitized = []
     remapping = []
@@ -129,7 +134,6 @@ def _getitem_subset_discards_dimensions(x, args: Tuple[Union[slice, Sequence], .
         if isinstance(idx, ndarray):
             if len(idx.shape) != 1: 
                 raise NotImplementedError("high-dimensional index arrays are not supported yet")
-
         elif isinstance(idx, slice):
             idx = range(*idx.indices(shape[d]))
         elif not isinstance(idx, Sequence):

tests/test_SparseNdarray.py

@@ -435,6 +435,11 @@ def test_SparseNdarray_subset_collapse(mask_rate):
     assert isinstance(first, numpy.ndarray)
     assert_identical_ndarrays(first, ref[:,1])
 
+    stuff = y[10]
+    assert_identical_ndarrays(stuff, ref[10])
+    stuff = y[numpy.int32(19)]
+    assert_identical_ndarrays(stuff, ref[numpy.int32(19)])
+
 
 #######################################################
 #######################################################

tests/test_Subset.py

@@ -147,6 +147,11 @@ def test_Subset_collapse(mask_rate):
     stuff = x[0, :, 2]
     assert_identical_ndarrays(stuff, y[0, :, 2])
 
+    stuff = x[10]
+    assert_identical_ndarrays(stuff, y[10])
+    stuff = x[numpy.int32(20)]
+    assert_identical_ndarrays(stuff, y[numpy.int32(20)])
+
 #    # Trying vectorized index.
 #    stuff = x[[1,2,3],[4,5,6],[7,8,9]]
 #    assert stuff.shape == (3,)