diff --git a/pyttb/sptensor.py b/pyttb/sptensor.py index 3b3fa810..3107726a 100644 --- a/pyttb/sptensor.py +++ b/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 diff --git a/pyttb/tensor.py b/pyttb/tensor.py index 72991171..9f5e675b 100644 --- a/pyttb/tensor.py +++ b/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 diff --git a/tests/test_tensor.py b/tests/test_tensor.py index d27463c1..1fa28bd9 100644 --- a/tests/test_tensor.py +++ b/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