SPTENSOR:

* Fix argument mismatch for ttm (modes s.b. dims)
* Fix ttm for rectangular matrices
* Make error message consitent with tensor
TENSOR:
* Fix error message

Author: ntjohnson1

pyttb/sptensor.py

@@ -2173,15 +2173,14 @@ def __repr__(self):  # pragma: no cover
 
     __str__ = __repr__
 
-    def ttm(self, matrices, mode, dims=None, transpose=False):
+    def ttm(self, matrices, dims=None, transpose=False):
         """
         Sparse tensor times matrix.
 
         Parameters
         ----------
         matrices: A matrix or list of matrices
-        mode:
-        dims:
+        dims: :class:`Numpy.ndarray`, int
         transpose: Transpose matrices to be multiplied
 
         Returns
@@ -2190,10 +2189,15 @@ def ttm(self, matrices, mode, dims=None, transpose=False):
         """
         if dims is None:
             dims = np.arange(self.ndims)
+        elif isinstance(dims, list):
+            dims = np.array(dims)
+        elif np.isscalar(dims) or isinstance(dims, list):
+            dims = np.array([dims])
+
         # Handle list of matrices
         if isinstance(matrices, list):
             # Check dimensions are valid
-            [dims, vidx] = tt_dimscheck(mode, self.ndims, len(matrices))
+            [dims, vidx] = tt_dimscheck(dims, self.ndims, len(matrices))
             # Calculate individual products
             Y = self.ttm(matrices[vidx[0]], dims[0], transpose=transpose)
             for i in range(1, dims.size):
@@ -2208,33 +2212,34 @@ def ttm(self, matrices, mode, dims=None, transpose=False):
         if transpose:
             matrices = matrices.transpose()
 
-        # Check mode
-        if not np.isscalar(mode) or mode < 0 or mode > self.ndims-1:
-            assert False, "Mode must be in [0, ndims)"
+        # Ensure this is the terminal single dimension case
+        if not (dims.size == 1 and np.isin(dims, np.arange(self.ndims))):
+            assert False, "dims must contain values in [0,self.dims)"
+        dims = dims[0]
 
         # Compute the product
 
         # Check that sizes match
-        if self.shape[mode] != matrices.shape[1]:
+        if self.shape[dims] != matrices.shape[1]:
             assert False, "Matrix shape doesn't match tensor shape"
 
         # Compute the new size
         siz = np.array(self.shape)
-        siz[mode] = matrices.shape[0]
+        siz[dims] = matrices.shape[0]
 
         # Compute self[mode]'
-        Xnt = ttb.tt_to_sparse_matrix(self, mode, True)
+        Xnt = ttb.tt_to_sparse_matrix(self, dims, True)
 
         # Reshape puts the reshaped things after the unchanged modes, transpose then puts it in front
         idx = 0
 
         # Convert to sparse matrix and do multiplication; generally result is sparse
         Z = Xnt.dot(matrices.transpose())
 
-        # Rearrange back into sparse tensor of original shape
-        Ynt = ttb.tt_from_sparse_matrix(Z, self.shape, mode, idx)
+        # Rearrange back into sparse tensor of correct shape
+        Ynt = ttb.tt_from_sparse_matrix(Z, siz, dims, idx)
 
-        if Z.nnz <= 0.5 * np.prod(siz):
+        if not isinstance(Z, np.ndarray) and Z.nnz <= 0.5 * np.prod(siz):
             return Ynt
         else:
             # TODO evaluate performance loss by casting into sptensor then tensor. I assume minimal since we are already

pyttb/tensor.py

@@ -921,7 +921,7 @@ def ttm(self, matrix, dims=None, transpose=False):
             assert False, "matrix must be of type numpy.ndarray"
 
         if not (dims.size == 1 and np.isin(dims, np.arange(self.ndims))):
-            assert False, "dims must contain values in [0,self.dims]"
+            assert False, "dims must contain values in [0,self.dims)"
 
         # old version (ver=0)
         shape = np.array(self.shape)

tests/test_sptensor.py

@@ -1365,25 +1365,25 @@ def test_sptensor_ttm(sample_sptensor):
     result[:, 3, 3] = 3.5
     result = ttb.tensor.from_data(result)
     result = ttb.sptensor.from_tensor_type(result)
-    assert sptensorInstance.ttm(sparse.coo_matrix(np.ones((4, 4))), mode=0).isequal(result)
-    assert sptensorInstance.ttm(sparse.coo_matrix(np.ones((4, 4))), mode=0, transpose=True).isequal(result)
+    assert sptensorInstance.ttm(sparse.coo_matrix(np.ones((4, 4))), dims=0).isequal(result)
+    assert sptensorInstance.ttm(sparse.coo_matrix(np.ones((4, 4))), dims=0, transpose=True).isequal(result)
 
     # This is a multiway multiplication yielding a sparse tensor, yielding a dense tensor relies on tensor.ttm
     matrix = sparse.coo_matrix(np.eye(4))
     list_of_matrices = [matrix, matrix, matrix]
-    assert sptensorInstance.ttm(list_of_matrices, mode=np.array([0, 1, 2])).isequal(sptensorInstance)
+    assert sptensorInstance.ttm(list_of_matrices, dims=np.array([0, 1, 2])).isequal(sptensorInstance)
 
     with pytest.raises(AssertionError) as excinfo:
-        sptensorInstance.ttm(sparse.coo_matrix(np.ones((5, 5))), mode=0)
+        sptensorInstance.ttm(sparse.coo_matrix(np.ones((5, 5))), dims=0)
     assert "Matrix shape doesn't match tensor shape" in str(excinfo)
 
     with pytest.raises(AssertionError) as excinfo:
-        sptensorInstance.ttm(np.array([1, 2, 3, 4]), mode=0)
+        sptensorInstance.ttm(np.array([1, 2, 3, 4]), dims=0)
     assert "Sptensor.ttm: second argument must be a matrix" in str(excinfo)
 
     with pytest.raises(AssertionError) as excinfo:
-        sptensorInstance.ttm(sparse.coo_matrix(np.ones((4, 4))), mode=4)
-    assert "Mode must be in [0, ndims)" in str(excinfo)
+        sptensorInstance.ttm(sparse.coo_matrix(np.ones((4, 4))), dims=4)
+    assert "dims must contain values in [0,self.dims)" in str(excinfo)
 
     sptensorInstance[0, :, :] = 1
     sptensorInstance[3, :, :] = 1
@@ -1397,17 +1397,20 @@ def test_sptensor_ttm(sample_sptensor):
     # TODO: Ensure mode mappings are consistent between matlab and numpy
     # MATLAB is opposite orientation so the mapping from matlab to numpy is
     # {3:0, 2:2, 1:1}
-    assert (sptensorInstance.ttm(sparse.coo_matrix(np.ones((4, 4))), mode=1).isequal(ttb.tensor.from_data(result)))
+    assert (sptensorInstance.ttm(sparse.coo_matrix(np.ones((4, 4))), dims=1).isequal(ttb.tensor.from_data(result)))
 
     result = 2*np.ones((4, 4, 4))
     result[:, 1, 1] = 2.5
     result[:, 1, 3] = 3.5
     result[:, 2, 2] = 4.5
-    assert (sptensorInstance.ttm(sparse.coo_matrix(np.ones((4, 4))), mode=0).isequal(ttb.tensor.from_data(result)))
+    assert (sptensorInstance.ttm(sparse.coo_matrix(np.ones((4, 4))), dims=0).isequal(ttb.tensor.from_data(result)))
 
     result = np.zeros((4, 4, 4))
     result[0, :, :] = 4.0
     result[3, :, :] = 4.0
     result[1, 1, :] = 2
     result[2, 2, :] = 2.5
-    assert (sptensorInstance.ttm(sparse.coo_matrix(np.ones((4, 4))), mode=2).isequal(ttb.tensor.from_data(result)))
+    assert (sptensorInstance.ttm(sparse.coo_matrix(np.ones((4, 4))), dims=2).isequal(ttb.tensor.from_data(result)))
+
+    # Confirm reshape for non-square matrix
+    assert sptensorInstance.ttm(sparse.coo_matrix(np.ones((1, 4))), dims=2).shape == (4,4,1)

tests/test_tensor.py

@@ -1057,7 +1057,7 @@ def test_tensor_ttm(sample_tensor_2way, sample_tensor_3way, sample_tensor_4way):
     # 3-way, dims must be in range [0,self.ndims]
     with pytest.raises(AssertionError) as excinfo:
         tensorInstance3.ttm(M2, tensorInstance3.ndims + 1)
-    assert "dims must contain values in [0,self.dims]" in str(excinfo)
+    assert "dims must contain values in [0,self.dims)" in str(excinfo)
 
 @pytest.mark.indevelopment
 def test_tensor_ttt(sample_tensor_2way, sample_tensor_3way, sample_tensor_4way):