TENSOR: Indexing is complicated.

pyttb/sptensor.py

@@ -624,7 +624,7 @@ def innerprod(
             if self.shape != other.shape:
                 assert False, "Sptensor and tensor must be same shape for innerproduct"
             [subsSelf, valsSelf] = self.find()
-            valsOther = other[subsSelf.transpose(), "extract"]
+            valsOther = other[subsSelf, "extract"]
             return valsOther.transpose().dot(valsSelf)
 
         if isinstance(other, (ttb.ktensor, ttb.ttensor)):  # pragma: no cover
@@ -686,7 +686,7 @@ def logical_and(self, B: Union[float, sptensor, ttb.tensor]) -> sptensor:
 
         if isinstance(B, ttb.tensor):
             BB = sptensor.from_data(
-                self.subs, B[self.subs.transpose(), "extract"][:, None], self.shape
+                self.subs, B[self.subs, "extract"][:, None], self.shape
             )
             C = self.logical_and(BB)
             return C
@@ -1047,7 +1047,7 @@ def scale(self, factor: np.ndarray, dims: Union[float, np.ndarray]) -> sptensor:
                 assert False, "Size mismatch in scale"
             return ttb.sptensor.from_data(
                 self.subs,
-                self.vals * factor[self.subs[:, dims].transpose(), "extract"][:, None],
+                self.vals * factor[self.subs[:, dims], "extract"][:, None],
                 self.shape,
             )
         if isinstance(factor, ttb.sptensor):
@@ -1807,7 +1807,7 @@ def __eq__(self, other):
             ]
 
             # Find where their nonzeros intersect
-            othervals = other[self.subs.transpose(), "extract"]
+            othervals = other[self.subs, "extract"]
             znzsubs = self.subs[(othervals[:, None] == self.vals).transpose()[0], :]
 
             return sptensor.from_data(
@@ -1888,7 +1888,7 @@ def __ne__(self, other):
                 subs1 = np.empty((0, self.subs.shape[1]))
             # find entries where x is nonzero but not equal to y
             subs2 = self.subs[
-                self.vals.transpose()[0] != other[self.subs.transpose(), "extract"], :
+                self.vals.transpose()[0] != other[self.subs, "extract"], :
             ]
             if subs2.size == 0:
                 subs2 = np.empty((0, self.subs.shape[1]))
@@ -2003,7 +2003,7 @@ def __mul__(self, other):
             )
         if isinstance(other, ttb.tensor):
             csubs = self.subs
-            cvals = self.vals * other[csubs.transpose(), "extract"][:, None]
+            cvals = self.vals * other[csubs, "extract"][:, None]
             return ttb.sptensor.from_data(csubs, cvals, self.shape)
         if isinstance(other, ttb.ktensor):
             csubs = self.subs
@@ -2125,7 +2125,7 @@ def __le__(self, other):
 
             # self nonzero
             subs2 = self.subs[
-                self.vals.transpose()[0] <= other[self.subs.transpose(), "extract"], :
+                self.vals.transpose()[0] <= other[self.subs, "extract"], :
             ]
 
             # assemble
@@ -2213,9 +2213,7 @@ def __lt__(self, other):
             subs1 = subs1[ttb.tt_setdiff_rows(subs1, self.subs), :]
 
             # self nonzero
-            subs2 = self.subs[
-                self.vals.transpose()[0] < other[self.subs.transpose(), "extract"], :
-            ]
+            subs2 = self.subs[self.vals.transpose()[0] < other[self.subs, "extract"], :]
 
             # assemble
             subs = np.vstack((subs1, subs2))
@@ -2270,9 +2268,7 @@ def __ge__(self, other):
 
             # self nonzero
             subs2 = self.subs[
-                (
-                    self.vals >= other[self.subs.transpose(), "extract"][:, None]
-                ).transpose()[0],
+                (self.vals >= other[self.subs, "extract"][:, None]).transpose()[0],
                 :,
             ]
 
@@ -2331,9 +2327,7 @@ def __gt__(self, other):
 
             # self and other nonzero
             subs2 = self.subs[
-                (
-                    self.vals > other[self.subs.transpose(), "extract"][:, None]
-                ).transpose()[0],
+                (self.vals > other[self.subs, "extract"][:, None]).transpose()[0],
                 :,
             ]
 
@@ -2437,7 +2431,7 @@ def __truediv__(self, other):
 
         if isinstance(other, ttb.tensor):
             csubs = self.subs
-            cvals = self.vals / other[csubs.transpose(), "extract"][:, None]
+            cvals = self.vals / other[csubs, "extract"][:, None]
             return ttb.sptensor.from_data(csubs, cvals, self.shape)
         if isinstance(other, ttb.ktensor):
             # TODO consider removing epsilon and generating nans consistent with above

pyttb/tensor.py

@@ -1270,12 +1270,17 @@ def __setitem__(self, key, value):
         case.
 
         Examples
-        X = tensor(rand(3,4,2))
-        X(1:2,1:2,1) = ones(2,2) <-- replaces subtensor
-        X([1 1 1;1 1 2]) = [5;7] <-- replaces two elements
-        X([1;13]) = [5;7] <-- does the same thing
-        X(1,1,2:3) = 1 <-- grows tensor
-        X(1,1,4) = 1 %<- grows the size of the tensor
+        --------
+        >>> X = tensor.from_data(np.ones((3,4,2)))
+        >>> # replaces subtensor
+        >>> X[0:2,0:2,0] = np.ones((2,2))
+        >>> # replaces two elements
+        >>> X[np.array([[1, 1, 1], [1, 1, 2]])] = [5, 7]
+        >>> # replaces two elements with linear indices
+        >>> X[np.array([1, 13])] = [5, 7]
+        >>> # grows tensor to accept new element
+        >>> X[1,1,2:3] = 1
+        >>> X[1,1,4] = 1
         """
         # Figure out if we are doing a subtensor, a list of subscripts or a list of
         # linear indices
@@ -1421,7 +1426,7 @@ def _set_subscripts(self, key, value):
         elif key.shape[0] == 1:  # and len(key.shape) == 1:
             self.data[tuple(key[0, :])] = value
         else:
-            self.data[tuple(key)] = value
+            self.data[tuple(key.transpose())] = value
 
     def __getitem__(self, item):
         """
@@ -1434,7 +1439,7 @@ def __getitem__(self, item):
         scalar.
 
         Case 1b: Y = X(R1,R2,...,RN), where one or more Rn is a range and
-        the rest are indices, returns a sparse tensor.
+        the rest are indices, returns a tensor.
 
         Case 2a: V = X(S) or V = X(S,'extract'), where S is a p x n array
         of subscripts, returns a vector of p values.
@@ -1446,18 +1451,30 @@ def __getitem__(self, item):
         is particularly an issue if ndims(X)==1.
 
         Examples
-        X = tensor(rand(3,4,2,1),[3 4 2 1]);
-        X.data <-- returns multidimensional array
-        X.size <-- returns size
-        X(1,1,1,1) <-- produces a scalar
-        X(1,1,1,:) <-- produces a tensor of order 1 and size 1
-        X(:,1,1,:) <-- produces a tensor of size 3 x 1
-        X(1:2,[2 4],1,:) <-- produces a tensor of size 2 x 2 x 1
-        X(1:2,[2 4],1,1) <-- produces a tensor of size 2 x 2
-        X([1,1,1,1;3,4,2,1]) <-- returns a vector of length 2
-        X = tensor(rand(10,1),10);
-        X([1:6]') <-- extracts a subtensor
-        X([1:6]','extract') <-- extracts a vector of 6 elements
+        --------
+        >>> X = tensor.from_data(np.ones((3,4,2,1)))
+        >>> X[0,0,0,0] # produces a scalar
+        1.0
+        >>> # produces a tensor of order 1 and size 1
+        >>> X[1,1,1,:] # doctest: +NORMALIZE_WHITESPACE
+        tensor of shape 1
+        data[:] =
+        [1.]
+        >>> # produces a tensor of size 2 x 2 x 1
+        >>> X[0:2,[2, 3],1,:] # doctest: +NORMALIZE_WHITESPACE
+        tensor of shape 2 x 2 x 1
+        data[0, :, :] =
+        [[1.]
+         [1.]]
+        data[1, :, :] =
+        [[1.]
+         [1.]]
+        >>> # returns a vector of length 2
+        >>> # Equivalent to selecting [0,0,0,0] and [1,1,1,0] separately
+        >>> X[np.array([[0, 0, 0, 0], [1, 1, 1, 0]])]
+        array([1., 1.])
+        >>> X[[0,1,2]] # extracts the first three linearized indices
+        array([1., 1., 1.])
 
         Parameters
         ----------
@@ -1518,31 +1535,42 @@ def __getitem__(self, item):
             return a
 
         # *** CASE 2a: Subscript indexing ***
-        if isinstance(item, np.ndarray) and len(item) > 1:
-            # Extract array of subscripts
-            subs = np.array(item)
-            a = np.squeeze(self.data[tuple(subs)])
-            # TODO if is row make column?
-            return ttb.tt_subsubsref(a, subs)
-        if (
-            len(item) > 1
+        is_subscript = (
+            isinstance(item, np.ndarray)
+            and len(item.shape) == 2
+            and item.shape[-1] == self.ndims
+        )
+        is_extract_subscript = (
+            len(item) == 2
             and isinstance(item[0], np.ndarray)
             and isinstance(item[-1], str)
             and item[-1] == "extract"
-        ):
-            # TODO dry this up
-            subs = np.array(item[0])
-            a = np.squeeze(self.data[tuple(subs)])
+        )
+        if is_subscript or is_extract_subscript:
+            if is_extract_subscript:
+                item = item[0]
+            # Extract array of subscripts
+            subs = np.array(item)
+            a = np.squeeze(self.data[tuple(subs.transpose())])
             # TODO if is row make column?
             return ttb.tt_subsubsref(a, subs)
 
         # Case 2b: Linear Indexing
         if isinstance(item, tuple) and len(item) >= 2 and not isinstance(item[-1], str):
             assert False, "Linear indexing requires single input array"
-        idx = np.array(item)
-        a = np.squeeze(self.data[tuple(ttb.tt_ind2sub(self.shape, idx).transpose())])
-        # Todo if row make column?
-        return ttb.tt_subsubsref(a, idx)
+
+        if (isinstance(item, np.ndarray) and len(item.shape) == 1) or (
+            isinstance(item, list)
+            and all(isinstance(element, (int)) for element in item)
+        ):
+            idx = np.array(item)
+            a = np.squeeze(
+                self.data[tuple(ttb.tt_ind2sub(self.shape, idx).transpose())]
+            )
+            # Todo if row make column?
+            return ttb.tt_subsubsref(a, idx)
+
+        assert False, "Invalid use of tensor getitem"
 
     def __eq__(self, other):
         """
@@ -1978,7 +2006,7 @@ def tendiag(elements: np.ndarray, shape: Optional[Tuple[int, ...]] = None) -> te
     else:
         constructed_shape = tuple(max(N, dim) for dim in shape)
     X = tenzeros(constructed_shape)
-    subs = np.tile(np.arange(0, N).transpose(), (len(constructed_shape), 1))
+    subs = np.tile(np.arange(0, N)[:, None], (len(constructed_shape),))
     X[subs] = elements
     return X
 

tests/test_tensor.py

@@ -364,14 +364,19 @@ def test_tensor__getitem__(sample_tensor_2way):
     # Case 2a:
     assert tensorInstance[np.array([0, 0]), "extract"] == params["data"][0, 0]
     assert (
-        tensorInstance[np.array([[0, 0], [1, 1]]), "extract"]
-        == params["data"][([0, 0], [1, 1])]
+        tensorInstance[np.array([[0, 2], [1, 1], [1, 2]]), "extract"]
+        == params["data"][([0, 1, 1], [2, 1, 2])]
     ).all()
     # Case 2a: Extract doesn't seem to be needed
-    assert tensorInstance[np.array([0, 0])] == params["data"][0, 0]
+    assert tensorInstance[np.array([[0, 0]])] == params["data"][0, 0]
     assert (
-        tensorInstance[np.array([[0, 0], [1, 1]])] == params["data"][([0, 0], [1, 1])]
+        tensorInstance[np.array([[0, 2], [1, 1], [1, 2]])]
+        == params["data"][([0, 1, 1], [2, 1, 2])]
     ).all()
+    with pytest.raises(AssertionError) as excinfo:
+        # Must use numpy arrays for subscripts
+        tensorInstance[[[0, 2], [1, 1], [1, 2]]]
+    assert "Invalid use of tensor getitem" in str(excinfo)
 
     # Case 2b: Linear Indexing
     assert tensorInstance[np.array([0])] == params["data"][0, 0]
@@ -380,6 +385,9 @@ def test_tensor__getitem__(sample_tensor_2way):
     with pytest.raises(AssertionError) as excinfo:
         tensorInstance[np.array([0]), np.array([0]), np.array([0])]
     assert "Linear indexing requires single input array" in str(excinfo)
+    assert np.array_equal(
+        tensorInstance[[0, 1, 2]], tensorInstance[np.array([0, 1, 2])]
+    )
 
 
 @pytest.mark.indevelopment