sptensor: Add coverage for improved indexing capability

Author: ntjohnson1

pyttb/sptensor.py

@@ -1647,6 +1647,8 @@ def _set_subtensor(self, key, value):
                     newsz.append(self.shape[n])
                 else:
                     newsz.append(max([self.shape[n], key[n].stop]))
+            elif isinstance(key[n], Iterable):
+                newsz.append(max([self.shape[n], max(key[n]) + 1]))
             else:
                 newsz.append(max([self.shape[n], key[n] + 1]))
 

tests/test_sptensor.py

@@ -533,6 +533,18 @@ def test_sptensor_setitem_Case1(sample_sptensor):
     assert (sptensorInstance.vals == np.vstack((data["vals"], np.array([[7]])))).all()
     assert sptensorInstance.shape == data["shape"]
 
+    # Case I(b)ii: Set with scalar, iterable index, empty sptensor
+    someTensor = ttb.sptensor()
+    someTensor[[0, 1], 0] = 1
+    assert someTensor[0, 0] == 1
+    assert someTensor[1, 0] == 1
+    assert np.all(someTensor[[0, 1], 0].vals == 1)
+    # Case I(b)ii: Set with scalar, iterable index, non-empty sptensor
+    someTensor[[0, 1], 1] = 2
+    assert someTensor[0, 1] == 2
+    assert someTensor[1, 1] == 2
+    assert np.all(someTensor[[0, 1], 1].vals == 2)
+
     # Case I: Assign with non-scalar or sptensor
     sptensorInstanceLarger = ttb.sptensor.from_tensor_type(sptensorInstance)
     with pytest.raises(AssertionError) as excinfo: