TTENSOR: Half-hearted support for sparse cores and factors (#177)

* We should revisit this if exercised often
* Tests are mostly smoke tests to make sure extreme sparsity doesn't break solvers

Closes #61 
Co-authored-by: Danny Dunlavy <[email protected]>

pyttb/tensor.py

@@ -14,6 +14,7 @@
 import numpy as np
 import scipy.sparse.linalg
 from numpy_groupies import aggregate as accumarray
+from scipy import sparse
 
 import pyttb as ttb
 from pyttb.pyttb_utils import tt_dimscheck, tt_ind2sub
@@ -1059,7 +1060,7 @@ def ttm(
                 Y = Y.ttm(matrix[vidx[k]], dims[k], transpose=transpose)
             return Y
 
-        if not isinstance(matrix, np.ndarray):
+        if not isinstance(matrix, (np.ndarray, sparse.spmatrix)):
             assert False, f"matrix must be of type numpy.ndarray but got:\n{matrix}"
 
         dims, _ = ttb.tt_dimscheck(self.ndims, dims=dims, exclude_dims=exclude_dims)

pyttb/ttensor.py

@@ -4,15 +4,18 @@
 
 import logging
 import textwrap
+from typing import Union
 
 import numpy as np
 import scipy
+import scipy.sparse as sparse
 
 from pyttb import ktensor
 from pyttb import pyttb_utils as ttb_utils
 from pyttb import sptenmat, sptensor, tenmat, tensor
+from pyttb.sptensor import tt_to_sparse_matrix
 
-ALT_CORE_ERROR = "TTensor doesn't support non-tensor cores yet"
+ALT_CORE_ERROR = "TTensor doesn't support non-tensor cores yet. Only tensor/sptensor."
 
 
 class ttensor:
@@ -30,7 +33,8 @@ def __init__(self):
         :class:`pyttb.ttensor`
         """
         # Empty constructor
-        self.core = tensor()
+        # TODO explore replacing with typing protocol
+        self.core: Union[tensor, sptensor] = tensor()
         self.u = []
 
     @classmethod
@@ -63,12 +67,12 @@ def from_data(cls, core, factors):
         >>> K0 = ttb.ttensor.from_data(core, factors)
         """
         ttensorInstance = ttensor()
-        if isinstance(core, tensor):
-            ttensorInstance.core = tensor.from_data(core.data, core.shape)
+        if isinstance(core, (tensor, sptensor)):
+            ttensorInstance.core = core.from_tensor_type(core)
             ttensorInstance.u = factors.copy()
         else:
             # TODO support any tensor type with supported ops
-            raise ValueError("TTENSOR doesn't yet support generic cores, only tensor")
+            raise ValueError(ALT_CORE_ERROR)
         ttensorInstance._validate_ttensor()
         return ttensorInstance
 
@@ -98,7 +102,7 @@ def _validate_ttensor(self):
         """
         # Confirm all factors are matrices
         for factor_idx, factor in enumerate(self.u):
-            if not isinstance(factor, np.ndarray):
+            if not isinstance(factor, (np.ndarray, sparse.coo_matrix)):
                 raise ValueError(
                     f"Factor matrices must be numpy arrays but factor {factor_idx} was {type(factor)}"
                 )
@@ -163,7 +167,7 @@ def full(self):
 
         # There is a small chance tensor could be sparse so ensure we cast that to dense.
         if not isinstance(recomposed_tensor, tensor):
-            raise ValueError(ALT_CORE_ERROR)
+            recomposed_tensor = tensor.from_tensor_type(recomposed_tensor)
         return recomposed_tensor
 
     def double(self):
@@ -554,19 +558,27 @@ def nvecs(self, n, r, flipsign=True):
         H = self.core.ttm(V)
 
         if isinstance(H, sptensor):
-            raise NotImplementedError(ALT_CORE_ERROR)
+            HnT = tt_to_sparse_matrix(H, n, True)
         else:
             HnT = tenmat.from_tensor_type(H.full(), cdims=np.array([n])).double()
 
         G = self.core
 
         if isinstance(G, sptensor):
-            raise NotImplementedError(ALT_CORE_ERROR)
+            GnT = tt_to_sparse_matrix(G, n, True)
         else:
             GnT = tenmat.from_tensor_type(G.full(), cdims=np.array([n])).double()
 
         # Compute Xn * Xn'
-        Y = HnT.transpose().dot(GnT.dot(self.u[n].transpose()))
+        # Big hack because if RHS is sparse wrong dot product is used
+        if sparse.issparse(self.u[n]):
+            XnT = sparse.coo_matrix.dot(GnT, self.u[n].transpose())
+        else:
+            XnT = GnT.dot(self.u[n].transpose())
+        if sparse.issparse(XnT):
+            Y = sparse.coo_matrix.dot(HnT.transpose(), XnT)
+        else:
+            Y = HnT.transpose().dot(XnT)
 
         # TODO: Lifted from tensor, consider common location
         if r < Y.shape[0] - 1:
@@ -577,6 +589,8 @@ def nvecs(self, n, r, flipsign=True):
             logging.debug(
                 "Greater than or equal to tensor.shape[n] - 1 eigenvectors requires cast to dense to solve"
             )
+            if sparse.issparse(Y):
+                Y = Y.toarray()
             w, v = scipy.linalg.eigh(Y)
             v = v[:, (-np.abs(w)).argsort()]
             v = v[:, :r]

tests/test_ttensor.py

@@ -4,7 +4,7 @@
 
 import numpy as np
 import pytest
-import scipy.sparse as sparse
+from scipy import sparse
 
 import pyttb as ttb
 
@@ -47,6 +47,15 @@ def test_ttensor_initialization_from_data(sample_ttensor):
     assert isinstance(ttensorInstance.core, ttb.tensor)
     assert all([isinstance(a_factor, np.ndarray) for a_factor in ttensorInstance.u])
 
+    # Sparse ttensor smoke test
+    sparse_u = [
+        sparse.coo_matrix(np.ones(factor.shape)) for factor in ttensorInstance.u
+    ]
+    sparse_ttensor = ttb.ttensor.from_data(
+        ttb.sptensor.from_tensor_type(ttensorInstance.core), sparse_u
+    )
+    assert isinstance(sparse_ttensor, ttb.ttensor)
+
     # Negative Tests
     non_array_factor = ttensorInstance.u + [1]
     with pytest.raises(ValueError):
@@ -66,12 +75,9 @@ def test_ttensor_initialization_from_data(sample_ttensor):
     wrong_shape_factor[0] = np.random.random((row + 1, col + 1))
     with pytest.raises(ValueError):
         ttb.ttensor.from_data(ttensorInstance.core, wrong_shape_factor)
-
-    # Enforce error until sptensor core/other cores supported
+    # Invalid core type
     with pytest.raises(ValueError):
-        ttb.ttensor.from_data(
-            ttb.sptensor.from_tensor_type(ttensorInstance.core), ttensorInstance.u
-        )
+        ttb.ttensor.from_data(ttensorInstance, ttensorInstance.u)
 
 
 @pytest.mark.indevelopment
@@ -91,18 +97,24 @@ def test_ttensor_full(sample_ttensor):
     # This sanity check only works for all 1's
     assert tensor.double() == np.prod(ttensorInstance.core.shape)
 
-    # Negative tests
+    # Sparse tests
+    sparse_core = ttb.sptensor.from_tensor_type(ttensorInstance.core)
+    sparse_u = [
+        sparse.coo_matrix(np.ones(factor.shape)) for factor in ttensorInstance.u
+    ]
+    # We could probably make these properties to avoid this edge case but expect to eventually cover these alternate
+    # cores
+    sparse_ttensor = ttb.ttensor.from_data(sparse_core, sparse_u)
+    assert sparse_ttensor.full().isequal(tensor)
+
+    # Empty sparse ttensor components
     sparse_core = ttb.sptensor()
     sparse_core.shape = ttensorInstance.core.shape
     sparse_u = [
         sparse.coo_matrix(np.zeros(factor.shape)) for factor in ttensorInstance.u
     ]
-    # We could probably make these properties to avoid this edge case but expect to eventually cover these alternate
-    # cores
-    ttensorInstance.core = sparse_core
-    ttensorInstance.u = sparse_u
-    with pytest.raises(ValueError):
-        ttensorInstance.full()
+    sparse_ttensor = ttb.ttensor.from_data(sparse_core, sparse_u)
+    assert sparse_ttensor.full().double().item() == 0
 
 
 @pytest.mark.indevelopment
@@ -357,35 +369,66 @@ def test_ttensor_reconstruct(random_ttensor):
 @pytest.mark.indevelopment
 def test_ttensor_nvecs(random_ttensor):
     ttensorInstance = random_ttensor
+
+    sparse_core = ttb.sptensor.from_tensor_type(ttensorInstance.core)
+    sparse_core_ttensor = ttb.ttensor.from_data(sparse_core, ttensorInstance.u)
+
+    sparse_u = [sparse.coo_matrix(factor) for factor in ttensorInstance.u]
+    sparse_factor_ttensor = ttb.ttensor.from_data(sparse_core, sparse_u)
+
+    # Smaller number of eig vals
     n = 0
     r = 2
     ttensor_eigvals = ttensorInstance.nvecs(n, r)
     full_eigvals = ttensorInstance.full().nvecs(n, r)
     assert np.allclose(ttensor_eigvals, full_eigvals)
 
+    sparse_core_ttensor_eigvals = sparse_core_ttensor.nvecs(n, r)
+    assert np.allclose(ttensor_eigvals, sparse_core_ttensor_eigvals)
+
+    sparse_factors_ttensor_eigvals = sparse_factor_ttensor.nvecs(n, r)
+    assert np.allclose(ttensor_eigvals, sparse_factors_ttensor_eigvals)
+
     # Test for eig vals larger than shape-1
     n = 1
     r = 2
     full_eigvals = ttensorInstance.full().nvecs(n, r)
     ttensor_eigvals = ttensorInstance.nvecs(n, r)
     assert np.allclose(ttensor_eigvals, full_eigvals)
 
-    # Negative Tests
-    sparse_core = ttb.sptensor()
-    sparse_core.shape = ttensorInstance.core.shape
-    ttensorInstance.core = sparse_core
+    sparse_core_ttensor_eigvals = sparse_core_ttensor.nvecs(n, r)
+    assert np.allclose(ttensor_eigvals, sparse_core_ttensor_eigvals)
+
+    sparse_factors_ttensor_eigvals = sparse_factor_ttensor.nvecs(n, r)
+    assert np.allclose(ttensor_eigvals, sparse_factors_ttensor_eigvals)
 
-    # Sparse core
-    with pytest.raises(NotImplementedError):
-        ttensorInstance.nvecs(0, 1)
 
-    # Sparse factors
+def test_ttensor_nvecs_all_zeros(random_ttensor):
+    """Perform nvecs calculation on all zeros tensor to exercise sparsity edge cases"""
+    ttensorInstance = random_ttensor
+    n = 0
+    r = 2
+
+    # Construct all zeros ttensors
+    dense_u = [np.zeros(factor.shape) for factor in ttensorInstance.u]
+    dense_core = ttb.tenzeros(ttensorInstance.core.shape)
+    dense_ttensor = ttb.ttensor.from_data(dense_core, dense_u)
+    dense_nvecs = dense_ttensor.nvecs(n, r)
+
     sparse_u = [
         sparse.coo_matrix(np.zeros(factor.shape)) for factor in ttensorInstance.u
     ]
-    ttensorInstance.u = sparse_u
-    with pytest.raises(NotImplementedError):
-        ttensorInstance.nvecs(0, 1)
+    sparse_core = ttb.sptensor()
+    sparse_core.shape = ttensorInstance.core.shape
+    sparse_ttensor = ttb.ttensor.from_data(sparse_core, sparse_u)
+    sparse_nvecs = sparse_ttensor.nvecs(n, r)
+    # Currently just a smoke test need to find a very sparse
+    # but well formed tensor for nvecs
+    assert isinstance(dense_nvecs, np.ndarray)
+    assert isinstance(sparse_nvecs, np.ndarray)
+
+    higher_value_sparse_nvecs = sparse_ttensor.nvecs(1, r)
+    assert isinstance(higher_value_sparse_nvecs, np.ndarray)
 
 
 @pytest.mark.indevelopment