tensor: Add coverage for improved indexing capability

pyttb/tensor.py

@@ -1276,7 +1276,6 @@ def __setitem__(self, key, value):
  """
  # Figure out if we are doing a subtensor, a list of subscripts or a list of
  # linear indices
- # print(f"Key: {key} {type(key)}")
  access_type = "error"
  # TODO pull out this big decision tree into a function
  if isinstance(key, (float, int, np.generic, slice)):
@@ -1296,15 +1295,11 @@ def __setitem__(self, key, value):
  validSubtensor = [
  isinstance(keyElement, (int, slice, Iterable)) for keyElement in key
  ]
- # TODO probably need to confirm the Iterable is in fact numeric
  if np.all(validSubtensor):
  access_type = "subtensor"
  elif isinstance(key, Iterable):
  key = np.array(key)
- # Clean up copy paste
- if len(key.shape) > 1 and key.shape[1] >= self.ndims:
- access_type = "subscripts"
- elif len(key.shape) == 1 or key.shape[1] == 1:
+ if len(key.shape) == 1 or key.shape[1] == 1:
  access_type = "linear indices"
 
  # Case 1: Rectangular Subtensor
@@ -1351,6 +1346,12 @@ def _set_subtensor(self, key, value):
  else:
  sliceCheck.append(element.stop)
  elif isinstance(element, Iterable):
+ if any(
+ not isinstance(entry, (float, int, np.generic)) for entry in element
+ ):
+ raise ValueError(
+ f"Entries for setitem must be numeric but recieved, {element}"
+ )
  sliceCheck.append(max(element))
  else:
  sliceCheck.append(element)

tests/test_tensor.py

@@ -251,6 +251,13 @@ def test_tensor__setitem__(sample_tensor_2way):
  # Subtensor add dimension
  empty_tensor[0, 0, 0] = 2
 
+ # Subtensor with lists
+ some_tensor = ttb.tenones((3, 3))
+ some_tensor[[0, 1], [0, 1]] = 11
+ assert some_tensor[0, 0] == 11
+ assert some_tensor[1, 1] == 11
+ assert np.all(some_tensor[[0, 1], [0, 1]].data == 11)
+
  # Subscripts with constant
  tensorInstance[np.array([[1, 1]])] = 13.0
  dataGrowth[1, 1] = 13.0
@@ -293,6 +300,13 @@ def test_tensor__setitem__(sample_tensor_2way):
  dataGrowth[np.unravel_index([0, 3, 4], dataGrowth.shape, "F")] = 13
  assert (tensorInstance.data == dataGrowth).all()
 
+ # Linear index with multiple indicies
+ some_tensor = ttb.tenones((3, 3))
+ some_tensor[[0, 1]] = 2
+ assert some_tensor[0] == 2
+ assert some_tensor[1] == 2
+ assert np.array_equal(some_tensor[[0, 1]], [2, 2])
+
  # Test Empty Tensor Set Item, subtensor
  emptyTensor = ttb.tensor.from_data(np.array([]))
  emptyTensor[0, 0, 0] = 0
@@ -313,10 +327,17 @@ def test_tensor__setitem__(sample_tensor_2way):
  )
 
  # Attempting to set some other way
- # TODO either catch this error ourselves or specify more specific exception we expect here
- with pytest.raises(Exception) as excinfo:
+ with pytest.raises(ValueError) as excinfo:
  tensorInstance[0, "a", 5] = 13.0
- # assert "Invalid use of tensor setitem" in str(excinfo)
+ assert "must be numeric" in str(excinfo)
+
+ with pytest.raises(AssertionError) as excinfo:
+
+ class BadKey:
+ pass
+
+ tensorInstance[BadKey] = 13.0
+ assert "Invalid use of tensor setitem" in str(excinfo)
 
 
 @pytest.mark.indevelopment
@@ -346,6 +367,11 @@ def test_tensor__getitem__(sample_tensor_2way):
  tensorInstance[np.array([[0, 0], [1, 1]]), "extract"]
  == params["data"][([0, 0], [1, 1])]
  ).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], [1, 1]])] == params["data"][([0, 0], [1, 1])]
+ ).all()
 
  # Case 2b: Linear Indexing
  assert tensorInstance[np.array([0])] == params["data"][0, 0]