From 67cc48a6a302c9ac9b08f54b3fc2a955c685f89c Mon Sep 17 00:00:00 2001
From: Nick Johnson <24689722+ntjohnson1@users.noreply.github.com>
Date: Sat, 25 Feb 2023 15:11:00 -0500
Subject: [PATCH 1/7] PYTTB_UTILS: Fix and enforce pylint

---
 pyttb/pyttb_utils.py  | 150 ++++++++++++++++++++++--------------------
 tests/test_package.py |   1 +
 2 files changed, 78 insertions(+), 73 deletions(-)

diff --git a/pyttb/pyttb_utils.py b/pyttb/pyttb_utils.py
index c043f989..f3f7ddfe 100644
--- a/pyttb/pyttb_utils.py
+++ b/pyttb/pyttb_utils.py
@@ -1,19 +1,20 @@
 # Copyright 2022 National Technology & Engineering Solutions of Sandia,
 # LLC (NTESS). Under the terms of Contract DE-NA0003525 with NTESS, the
 # U.S. Government retains certain rights in this software.
-
+"""PYTTB shared utilities across tensor types"""
 from inspect import signature
 from typing import Optional, Tuple, overload
 
 import numpy as np
-import scipy.sparse as sparse
+from scipy import sparse
 
 import pyttb as ttb
 
 
 def tt_to_dense_matrix(tensorInstance, mode, transpose=False):
     """
-    Helper function to unwrap tensor into dense matrix, should replace the core need for tenmat
+    Helper function to unwrap tensor into dense matrix, should replace the core need
+    for tenmat
 
     Parameters
     ----------
@@ -46,7 +47,8 @@ def tt_to_dense_matrix(tensorInstance, mode, transpose=False):
 
 def tt_from_dense_matrix(matrix, shape, mode, idx):
     """
-    Helper function to wrap dense matrix into tensor. Inverse of :class:`pyttb.tt_to_dense_matrix`
+    Helper function to wrap dense matrix into tensor.
+    Inverse of :class:`pyttb.tt_to_dense_matrix`
 
     Parameters
     ----------
@@ -72,7 +74,8 @@ def tt_from_dense_matrix(matrix, shape, mode, idx):
 
 def tt_to_sparse_matrix(sptensorInstance, mode, transpose=False):
     """
-    Helper function to unwrap sptensor into sparse matrix, should replace the core need for sptenmat
+    Helper function to unwrap sptensor into sparse matrix, should replace the core need
+    for sptenmat
 
     Parameters
     ----------
@@ -92,13 +95,13 @@ def tt_to_sparse_matrix(sptensorInstance, mode, transpose=False):
     ).spmatrix()
     if transpose:
         return spmatrix.transpose()
-    else:
-        return spmatrix
+    return spmatrix
 
 
 def tt_from_sparse_matrix(spmatrix, shape, mode, idx):
     """
-    Helper function to wrap sparse matrix into sptensor. Inverse of :class:`pyttb.tt_to_sparse_matrix`
+    Helper function to wrap sparse matrix into sptensor.
+    Inverse of :class:`pyttb.tt_to_sparse_matrix`
 
     Parameters
     ----------
@@ -118,7 +121,8 @@ def tt_from_sparse_matrix(spmatrix, shape, mode, idx):
 
     # This expands the compressed dimension back to full size
     sptensorInstance = sptensorInstance.reshape(siz[old], idx)
-    # This puts the modes in the right order, reshape places modified modes after the unchanged ones
+    # This puts the modes in the right order, reshape places modified modes after the
+    # unchanged ones
     sptensorInstance = sptensorInstance.reshape(
         shape,
         np.concatenate((np.arange(1, mode + 1), [0], np.arange(mode + 1, len(shape)))),
@@ -206,8 +210,8 @@ def tt_dimscheck(
     # Save dimensions of dims
     P = len(dims)
 
-    # Reorder dims from smallest to largest
-    # (this matters in particular for the vector multiplicand case, where the order affects the result)
+    # Reorder dims from smallest to largest (this matters in particular for the vector
+    # multiplicand case, where the order affects the result)
     sidx = np.argsort(dims)
     sdims = dims[sidx]
     vidx = None
@@ -217,24 +221,25 @@ def tt_dimscheck(
         if M > N:
             assert False, "Cannot have more multiplicands than dimensions"
 
-        # Check that the number of multiplicands must either be full dimensional or equal to the specified dimensions
-        # (M==N) or M(==P) respectively
-        if M != N and M != P:
+        # Check that the number of multiplicands must either be full dimensional or
+        # equal to the specified dimensions (M==N) or M(==P) respectively
+        if M not in (N, P):
             assert False, "Invalid number of multiplicands"
 
         # Check sizes to determine how to index multiplicands
         if P == M:
-            # Case 1: Number of items in dims and number of multiplicands are equal; therfore, index in order of sdims
+            # Case 1: Number of items in dims and number of multiplicands are equal;
+            # therfore, index in order of sdims
             vidx = sidx
         else:
-            # Case 2: Number of multiplicands is equal to the number of dimensions of tensor;
-            # therefore, index multiplicands by dimensions in dims argument.
+            # Case 2: Number of multiplicands is equal to the number of dimensions of
+            # tensor; therefore, index multiplicands by dimensions in dims argument.
             vidx = sdims
 
     return sdims, vidx
 
 
-def tt_tenfun(function_handle, *inputs):
+def tt_tenfun(function_handle, *inputs):  # pylint:disable=too-many-branches
     """
     Apply a function to each element in a tensor
 
@@ -256,13 +261,13 @@ def tt_tenfun(function_handle, *inputs):
     assert callable(function_handle), "function_handle must be callable"
 
     # Convert inputs to tensors if they aren't already
-    for i in range(0, len(inputs)):
-        if isinstance(inputs[i], ttb.tensor) or isinstance(inputs[i], (float, int)):
+    for i, an_input in enumerate(inputs):
+        if isinstance(an_input, (ttb.tensor, float, int)):
             continue
-        elif isinstance(inputs[i], np.ndarray):
-            inputs[i] = ttb.tensor.from_data(inputs[i])
+        if isinstance(an_input, np.ndarray):
+            inputs[i] = ttb.tensor.from_data(an_input)
         elif isinstance(
-            inputs[i],
+            an_input,
             (
                 ttb.ktensor,
                 ttb.ttensor,
@@ -272,11 +277,12 @@ def tt_tenfun(function_handle, *inputs):
                 ttb.symktensor,
             ),
         ):
-            inputs[i] = ttb.tensor.from_tensor_type(inputs[i])
+            inputs[i] = ttb.tensor.from_tensor_type(an_input)
         else:
             assert False, "Invalid input to ten fun"
 
-    # It's ok if there are two input and one is a scalar; otherwise all inputs have to be the same size
+    # It's ok if there are two input and one is a scalar; otherwise all inputs have to
+    # be the same size
     if (
         (len(inputs) == 2)
         and isinstance(inputs[0], (float, int))
@@ -290,15 +296,15 @@ def tt_tenfun(function_handle, *inputs):
     ):
         sz = inputs[0].shape
     else:
-        for i in range(0, len(inputs)):
-            if isinstance(inputs[i], (float, int)):
-                assert False, "Argument {} is a scalar but expected a tensor".format(i)
+        for i, an_input in enumerate(inputs):
+            if isinstance(an_input, (float, int)):
+                assert False, f"Argument {i} is a scalar but expected a tensor"
             elif i == 0:
-                sz = inputs[i].shape
-            elif sz != inputs[i].shape:
+                sz = an_input.shape
+            elif sz != an_input.shape:
                 assert (
                     False
-                ), "Tensor {} is not the same size as the first tensor input".format(i)
+                ), f"Tensor {i} is not the same size as the first tensor input"
 
     # Number of inputs for function handle
     nfunin = len(signature(function_handle).parameters)
@@ -322,8 +328,8 @@ def tt_tenfun(function_handle, *inputs):
         X = np.reshape(X, (1, -1))
     else:
         X = np.zeros((len(inputs), np.prod(sz)))
-        for i in range(0, len(inputs)):
-            X[i, :] = np.reshape(inputs[i].data, (np.prod(sz)))
+        for i, an_input in enumerate(inputs):
+            X[i, :] = np.reshape(an_input.data, (np.prod(sz)))
     data = function_handle(X)
     data = np.reshape(data, sz)
     Z = ttb.tensor.from_data(data)
@@ -395,7 +401,7 @@ def tt_intersect_rows(MatrixA, MatrixB):
     return location[np.where(location >= 0)]
 
 
-def tt_irenumber(t, shape, number_range):
+def tt_irenumber(t, shape, number_range):  # pylint: disable=unused-argument
     """
     RENUMBER indices for sptensor subsasgn
 
@@ -409,25 +415,25 @@ def tt_irenumber(t, shape, number_range):
     -------
     newsubs: :class:`numpy.ndarray`
     """
-    # TODO shape is unused. Should it be used? I don't particularly understand what this is meant to be doing
+    # TODO shape is unused. Should it be used? I don't particularly understand what
+    #  this is meant to be doing
     nz = t.nnz
     if nz == 0:
         newsubs = np.array([])
         return newsubs
-    else:
-        newsubs = t.subs.astype(int)
-        for i in range(0, len(number_range)):
-            r = number_range[i]
-            if isinstance(r, slice):
-                newsubs[:, i] = (newsubs[:, i])[r]
-            elif isinstance(r, int):
-                # This appears to be inserting new keys as rows to our subs here
-                newsubs = np.insert(newsubs, obj=i, values=r, axis=1)
-            else:
-                if isinstance(r, list):
-                    r = np.array(r)
-                newsubs[:, i] = r[newsubs[:, i]]
-        return newsubs
+
+    newsubs = t.subs.astype(int)
+    for i, r in enumerate(number_range):
+        if isinstance(r, slice):
+            newsubs[:, i] = (newsubs[:, i])[r]
+        elif isinstance(r, int):
+            # This appears to be inserting new keys as rows to our subs here
+            newsubs = np.insert(newsubs, obj=i, values=r, axis=1)
+        else:
+            if isinstance(r, list):
+                r = np.array(r)
+            newsubs[:, i] = r[newsubs[:, i]]
+    return newsubs
 
 
 def tt_assignment_type(x, subs, rhs):
@@ -444,13 +450,12 @@ def tt_assignment_type(x, subs, rhs):
     -------
     objectType
     """
-    if type(x) == type(rhs):
+    if type(x) is type(rhs):
         return "subtensor"
     # If subscripts is a tuple that contains an nparray
-    elif isinstance(subs, tuple) and len(subs) >= 2:
+    if isinstance(subs, tuple) and len(subs) >= 2:
         return "subtensor"
-    else:
-        return "subscripts"
+    return "subscripts"
 
 
 def tt_renumber(subs, shape, number_range):
@@ -476,8 +481,8 @@ def tt_renumber(subs, shape, number_range):
     """
     newshape = np.array(shape)
     newsubs = subs
-    for i in range(0, len(shape)):
-        if not (number_range[i] == slice(None, None, None)):
+    for i in range(0, len(shape)):  # pylint: disable=consider-using-enumerate
+        if not number_range[i] == slice(None, None, None):
             if subs.size == 0:
                 if not isinstance(number_range[i], slice):
                     if isinstance(number_range[i], (int, float)):
@@ -529,12 +534,14 @@ def tt_renumberdim(idx, shape, number_range):
     return newidx, newshape
 
 
+# TODO make more efficient, decide if we want to support the multiple response
+#  matlab does
+# pylint: disable=line-too-long
+# https://stackoverflow.com/questions/22699756/python-version-of-ismember-with-rows-and-index
+# For thoughts on how to speed this up
 def tt_ismember_rows(search, source):
     """
     Find location of search rows in source array
-    https://stackoverflow.com/questions/22699756/python-version-of-ismember-with-rows-and-index
-    For thoughts on how to speed this up
-    #TODO make more efficient, decide if we want to support the multiple response matlab does
 
     Parameters
     ----------
@@ -551,10 +558,10 @@ def tt_ismember_rows(search, source):
     Examples
     --------
     >>> a = np.array([[4, 6], [1, 9], [2, 6]])
-    >>> b = np.array([[1, 7],[1, 8],[2, 6],[2, 1],[2, 4],[4, 6],[4, 7],[5, 9],[5, 2],[5, 1]])
+    >>> b = np.array([[2, 6],[2, 1],[2, 4],[4, 6],[4, 7],[5, 9],[5, 2],[5, 1]])
     >>> results = tt_ismember_rows(a,b)
     >>> print(results)
-    [ 5 -1  2]
+    [ 3 -1  0]
 
     """
     results = np.ones(shape=search.shape[0]) * -1
@@ -585,7 +592,7 @@ def tt_ind2sub(shape: Tuple[int, ...], idx: np.ndarray) -> np.ndarray:
     return np.array(np.unravel_index(idx, shape, order="F")).transpose()
 
 
-def tt_subsubsref(obj, s):
+def tt_subsubsref(obj, s):  # pylint: disable=unused-argument
     """
     Helper function for tensor toolbox subsref.
 
@@ -598,7 +605,8 @@ def tt_subsubsref(obj, s):
     -------
     Still uncertain to this functionality
     """
-    # TODO figure out when subsref yields key of length>1 for now ignore this logic and just return
+    # TODO figure out when subsref yields key of length>1 for now ignore this logic and
+    #  just return
     # if len(s) == 1:
     #    return obj
     # else:
@@ -608,7 +616,8 @@ def tt_subsubsref(obj, s):
 
 def tt_intvec2str(v):
     """
-    Print integer vector to a string with brackets. Numpy should already handle this so it is a placeholder stub
+    Print integer vector to a string with brackets. Numpy should already handle this so
+    it is a placeholder stub
 
     Parameters
     ----------
@@ -774,10 +783,7 @@ def isrow(v):
     -------
     bool
     """
-    if v.ndim == 2 and v.shape[0] == 1 and v.shape[1] >= 1:
-        return True
-    else:
-        return False
+    return v.ndim == 2 and v.shape[0] == 1 and v.shape[1] >= 1
 
 
 def isvector(a):
@@ -794,13 +800,11 @@ def isvector(a):
     -------
     bool
     """
-    if a.ndim == 1 or (a.ndim == 2 and (a.shape[0] == 1 or a.shape[1] == 1)):
-        return True
-    else:
-        return False
+    return a.ndim == 1 or (a.ndim == 2 and (a.shape[0] == 1 or a.shape[1] == 1))
 
 
-# TODO: this is a challenge, since it may need to apply to either Python built in types or numpy types
+# TODO: this is a challenge, since it may need to apply to either Python built in types
+#  or numpy types
 def islogical(a):
     """
     ISLOGICAL Checks if vector is a logical vector.
@@ -815,4 +819,4 @@ def islogical(a):
     -------
     bool
     """
-    return type(a) == bool
+    return isinstance(a, bool)
diff --git a/tests/test_package.py b/tests/test_package.py
index 245246e5..c963a9c9 100644
--- a/tests/test_package.py
+++ b/tests/test_package.py
@@ -23,6 +23,7 @@ def test_linting():
 
     enforced_files = [
         os.path.join(os.path.dirname(ttb.__file__), f"{ttb.tensor.__name__}.py"),
+        ttb.pyttb_utils.__file__,
     ]
     # TODO pylint fails to import pyttb in tests
     # add mypy check

From a1f24dffbc0d22e7bb15c32ee045701049d7e94c Mon Sep 17 00:00:00 2001
From: Nick Johnson <24689722+ntjohnson1@users.noreply.github.com>
Date: Sat, 25 Feb 2023 15:32:14 -0500
Subject: [PATCH 2/7] PYTTB_UTILS: Pull out utility only used internally in
 sptensor

---
 pyttb/pyttb_utils.py      | 60 -------------------------------------
 pyttb/sptensor.py         | 63 +++++++++++++++++++++++++++++++++++++--
 tests/test_pyttb_utils.py | 36 ----------------------
 tests/test_sptensor.py    | 37 +++++++++++++++++++++++
 4 files changed, 98 insertions(+), 98 deletions(-)

diff --git a/pyttb/pyttb_utils.py b/pyttb/pyttb_utils.py
index f3f7ddfe..40ed7147 100644
--- a/pyttb/pyttb_utils.py
+++ b/pyttb/pyttb_utils.py
@@ -6,7 +6,6 @@
 from typing import Optional, Tuple, overload
 
 import numpy as np
-from scipy import sparse
 
 import pyttb as ttb
 
@@ -72,65 +71,6 @@ def tt_from_dense_matrix(matrix, shape, mode, idx):
     return tensorInstance
 
 
-def tt_to_sparse_matrix(sptensorInstance, mode, transpose=False):
-    """
-    Helper function to unwrap sptensor into sparse matrix, should replace the core need
-    for sptenmat
-
-    Parameters
-    ----------
-    sptensorInstance: :class:`pyttb.sptensor`
-    mode: int
-        Mode around which to unwrap tensor
-    transpose: bool
-        Whether or not to tranpose unwrapped tensor
-
-    Returns
-    -------
-    spmatrix: :class:`Scipy.sparse.coo_matrix`
-    """
-    old = np.setdiff1d(np.arange(sptensorInstance.ndims), mode).astype(int)
-    spmatrix = sptensorInstance.reshape(
-        (np.prod(np.array(sptensorInstance.shape)[old]),), old
-    ).spmatrix()
-    if transpose:
-        return spmatrix.transpose()
-    return spmatrix
-
-
-def tt_from_sparse_matrix(spmatrix, shape, mode, idx):
-    """
-    Helper function to wrap sparse matrix into sptensor.
-    Inverse of :class:`pyttb.tt_to_sparse_matrix`
-
-    Parameters
-    ----------
-    spmatrix: :class:`Scipy.sparse.coo_matrix`
-    mode: int
-        Mode around which tensor was unwrapped
-    idx: int
-        in {0,1}, idx of mode in spmatrix, s.b. 0 for tranpose=True
-
-    Returns
-    -------
-    sptensorInstance: :class:`pyttb.sptensor`
-    """
-    siz = np.array(shape)
-    old = np.setdiff1d(np.arange(len(shape)), mode).astype(int)
-    sptensorInstance = ttb.sptensor.from_tensor_type(sparse.coo_matrix(spmatrix))
-
-    # This expands the compressed dimension back to full size
-    sptensorInstance = sptensorInstance.reshape(siz[old], idx)
-    # This puts the modes in the right order, reshape places modified modes after the
-    # unchanged ones
-    sptensorInstance = sptensorInstance.reshape(
-        shape,
-        np.concatenate((np.arange(1, mode + 1), [0], np.arange(mode + 1, len(shape)))),
-    )
-
-    return sptensorInstance
-
-
 def tt_union_rows(MatrixA, MatrixB):
     """
     Helper function to reproduce functionality of MATLABS intersect(a,b,'rows')
diff --git a/pyttb/sptensor.py b/pyttb/sptensor.py
index 84e2cc43..22b8573e 100644
--- a/pyttb/sptensor.py
+++ b/pyttb/sptensor.py
@@ -14,6 +14,65 @@
 from .pyttb_utils import *
 
 
+def tt_to_sparse_matrix(sptensorInstance, mode, transpose=False):
+    """
+    Helper function to unwrap sptensor into sparse matrix, should replace the core need
+    for sptenmat
+
+    Parameters
+    ----------
+    sptensorInstance: :class:`pyttb.sptensor`
+    mode: int
+        Mode around which to unwrap tensor
+    transpose: bool
+        Whether or not to tranpose unwrapped tensor
+
+    Returns
+    -------
+    spmatrix: :class:`Scipy.sparse.coo_matrix`
+    """
+    old = np.setdiff1d(np.arange(sptensorInstance.ndims), mode).astype(int)
+    spmatrix = sptensorInstance.reshape(
+        (np.prod(np.array(sptensorInstance.shape)[old]),), old
+    ).spmatrix()
+    if transpose:
+        return spmatrix.transpose()
+    return spmatrix
+
+
+def tt_from_sparse_matrix(spmatrix, shape, mode, idx):
+    """
+    Helper function to wrap sparse matrix into sptensor.
+    Inverse of :class:`pyttb.tt_to_sparse_matrix`
+
+    Parameters
+    ----------
+    spmatrix: :class:`Scipy.sparse.coo_matrix`
+    mode: int
+        Mode around which tensor was unwrapped
+    idx: int
+        in {0,1}, idx of mode in spmatrix, s.b. 0 for tranpose=True
+
+    Returns
+    -------
+    sptensorInstance: :class:`pyttb.sptensor`
+    """
+    siz = np.array(shape)
+    old = np.setdiff1d(np.arange(len(shape)), mode).astype(int)
+    sptensorInstance = ttb.sptensor.from_tensor_type(sparse.coo_matrix(spmatrix))
+
+    # This expands the compressed dimension back to full size
+    sptensorInstance = sptensorInstance.reshape(siz[old], idx)
+    # This puts the modes in the right order, reshape places modified modes after the
+    # unchanged ones
+    sptensorInstance = sptensorInstance.reshape(
+        shape,
+        np.concatenate((np.arange(1, mode + 1), [0], np.arange(mode + 1, len(shape)))),
+    )
+
+    return sptensorInstance
+
+
 class sptensor(object):
     """
     SPTENSOR Class for sparse tensors.
@@ -2433,7 +2492,7 @@ def ttm(self, matrices, dims=None, transpose=False):
         siz[dims] = matrices.shape[0]
 
         # Compute self[mode]'
-        Xnt = ttb.tt_to_sparse_matrix(self, dims, True)
+        Xnt = tt_to_sparse_matrix(self, dims, True)
 
         # Reshape puts the reshaped things after the unchanged modes, transpose then puts it in front
         idx = 0
@@ -2442,7 +2501,7 @@ def ttm(self, matrices, dims=None, transpose=False):
         Z = Xnt.dot(matrices.transpose())
 
         # Rearrange back into sparse tensor of correct shape
-        Ynt = ttb.tt_from_sparse_matrix(Z, siz, dims, idx)
+        Ynt = tt_from_sparse_matrix(Z, siz, dims, idx)
 
         if not isinstance(Z, np.ndarray) and Z.nnz <= 0.5 * np.prod(siz):
             return Ynt
diff --git a/tests/test_pyttb_utils.py b/tests/test_pyttb_utils.py
index 6ba9da4e..cd31eccf 100644
--- a/tests/test_pyttb_utils.py
+++ b/tests/test_pyttb_utils.py
@@ -51,42 +51,6 @@ def test_sptensor_from_dense_matrix():
         assert tensorCopy.isequal(Ynt)
 
 
-@pytest.mark.indevelopment
-def test_sptensor_to_sparse_matrix():
-    subs = np.array([[1, 1, 1], [1, 1, 3], [2, 2, 2], [3, 3, 3]])
-    vals = np.array([[0.5], [1.5], [2.5], [3.5]])
-    shape = (4, 4, 4)
-    mode0 = sparse.coo_matrix(([0.5, 1.5, 2.5, 3.5], ([5, 13, 10, 15], [1, 1, 2, 3])))
-    mode1 = sparse.coo_matrix(([0.5, 1.5, 2.5, 3.5], ([5, 13, 10, 15], [1, 1, 2, 3])))
-    mode2 = sparse.coo_matrix(([0.5, 1.5, 2.5, 3.5], ([5, 5, 10, 15], [1, 3, 2, 3])))
-    Ynt = [mode0, mode1, mode2]
-    sptensorInstance = ttb.sptensor().from_data(subs, vals, shape)
-
-    for mode in range(sptensorInstance.ndims):
-        Xnt = ttb.tt_to_sparse_matrix(sptensorInstance, mode, True)
-        assert (Xnt != Ynt[mode]).nnz == 0
-        assert Xnt.shape == Ynt[mode].shape
-
-
-@pytest.mark.indevelopment
-def test_sptensor_from_sparse_matrix():
-    subs = np.array([[1, 1, 1], [1, 1, 3], [2, 2, 2], [3, 3, 3]])
-    vals = np.array([[0.5], [1.5], [2.5], [3.5]])
-    shape = (4, 4, 4)
-    sptensorInstance = ttb.sptensor().from_data(subs, vals, shape)
-    for mode in range(sptensorInstance.ndims):
-        sptensorCopy = ttb.sptensor.from_tensor_type(sptensorInstance)
-        Xnt = ttb.tt_to_sparse_matrix(sptensorCopy, mode, True)
-        Ynt = ttb.tt_from_sparse_matrix(Xnt, sptensorCopy.shape, mode, 0)
-        assert sptensorCopy.isequal(Ynt)
-
-    for mode in range(sptensorInstance.ndims):
-        sptensorCopy = ttb.sptensor.from_tensor_type(sptensorInstance)
-        Xnt = ttb.tt_to_sparse_matrix(sptensorCopy, mode, False)
-        Ynt = ttb.tt_from_sparse_matrix(Xnt, sptensorCopy.shape, mode, 1)
-        assert sptensorCopy.isequal(Ynt)
-
-
 @pytest.mark.indevelopment
 def test_tt_union_rows():
     a = np.array([[4, 6], [1, 9], [2, 6], [2, 6], [99, 0]])
diff --git a/tests/test_sptensor.py b/tests/test_sptensor.py
index f73de267..c9a47b96 100644
--- a/tests/test_sptensor.py
+++ b/tests/test_sptensor.py
@@ -9,6 +9,7 @@
 import scipy.sparse as sparse
 
 import pyttb as ttb
+from pyttb.sptensor import tt_from_sparse_matrix, tt_to_sparse_matrix
 
 
 @pytest.fixture()
@@ -1708,3 +1709,39 @@ def test_sptensor_ttm(sample_sptensor):
         4,
         1,
     )
+
+
+@pytest.mark.indevelopment
+def test_sptensor_to_sparse_matrix():
+    subs = np.array([[1, 1, 1], [1, 1, 3], [2, 2, 2], [3, 3, 3]])
+    vals = np.array([[0.5], [1.5], [2.5], [3.5]])
+    shape = (4, 4, 4)
+    mode0 = sparse.coo_matrix(([0.5, 1.5, 2.5, 3.5], ([5, 13, 10, 15], [1, 1, 2, 3])))
+    mode1 = sparse.coo_matrix(([0.5, 1.5, 2.5, 3.5], ([5, 13, 10, 15], [1, 1, 2, 3])))
+    mode2 = sparse.coo_matrix(([0.5, 1.5, 2.5, 3.5], ([5, 5, 10, 15], [1, 3, 2, 3])))
+    Ynt = [mode0, mode1, mode2]
+    sptensorInstance = ttb.sptensor().from_data(subs, vals, shape)
+
+    for mode in range(sptensorInstance.ndims):
+        Xnt = tt_to_sparse_matrix(sptensorInstance, mode, True)
+        assert (Xnt != Ynt[mode]).nnz == 0
+        assert Xnt.shape == Ynt[mode].shape
+
+
+@pytest.mark.indevelopment
+def test_sptensor_from_sparse_matrix():
+    subs = np.array([[1, 1, 1], [1, 1, 3], [2, 2, 2], [3, 3, 3]])
+    vals = np.array([[0.5], [1.5], [2.5], [3.5]])
+    shape = (4, 4, 4)
+    sptensorInstance = ttb.sptensor().from_data(subs, vals, shape)
+    for mode in range(sptensorInstance.ndims):
+        sptensorCopy = ttb.sptensor.from_tensor_type(sptensorInstance)
+        Xnt = tt_to_sparse_matrix(sptensorCopy, mode, True)
+        Ynt = tt_from_sparse_matrix(Xnt, sptensorCopy.shape, mode, 0)
+        assert sptensorCopy.isequal(Ynt)
+
+    for mode in range(sptensorInstance.ndims):
+        sptensorCopy = ttb.sptensor.from_tensor_type(sptensorInstance)
+        Xnt = tt_to_sparse_matrix(sptensorCopy, mode, False)
+        Ynt = tt_from_sparse_matrix(Xnt, sptensorCopy.shape, mode, 1)
+        assert sptensorCopy.isequal(Ynt)

From 6f384c40c9a370ab1490314f8c19352a0a22a546 Mon Sep 17 00:00:00 2001
From: Nick Johnson <24689722+ntjohnson1@users.noreply.github.com>
Date: Sat, 25 Feb 2023 16:30:29 -0500
Subject: [PATCH 3/7] SPTENSOR: Fix and enforce pylint

---
 pyttb/sptensor.py      | 334 +++++++++++++++++++++--------------------
 tests/test_package.py  |   1 +
 tests/test_sptensor.py |   2 +-
 3 files changed, 177 insertions(+), 160 deletions(-)

diff --git a/pyttb/sptensor.py b/pyttb/sptensor.py
index 22b8573e..08af1965 100644
--- a/pyttb/sptensor.py
+++ b/pyttb/sptensor.py
@@ -1,17 +1,26 @@
 # Copyright 2022 National Technology & Engineering Solutions of Sandia,
 # LLC (NTESS). Under the terms of Contract DE-NA0003525 with NTESS, the
 # U.S. Government retains certain rights in this software.
-
+"""Sparse Tensor Implementation"""
 import warnings
 
 import numpy as np
-import scipy.sparse as sparse
 import scipy.sparse.linalg
 from numpy_groupies import aggregate as accumarray
+from scipy import sparse
 
 import pyttb as ttb
-
-from .pyttb_utils import *
+from pyttb.pyttb_utils import (
+    tt_assignment_type,
+    tt_dimscheck,
+    tt_ind2sub,
+    tt_intvec2str,
+    tt_sizecheck,
+    tt_sub2ind,
+    tt_subscheck,
+    tt_subsubsref,
+    tt_valscheck,
+)
 
 
 def tt_to_sparse_matrix(sptensorInstance, mode, transpose=False):
@@ -73,7 +82,7 @@ def tt_from_sparse_matrix(spmatrix, shape, mode, idx):
     return sptensorInstance
 
 
-class sptensor(object):
+class sptensor:
     """
     SPTENSOR Class for sparse tensors.
     """
@@ -143,8 +152,8 @@ def from_tensor_type(cls, source):
 
         Parameters
         ----------
-        source: :class:`pyttb.sptensor`, :class:`pyttb.tensor`, :class:`pyttb.sptenmat`,\
-            or :class:`pyttb.sptensor3`
+        source: :class:`pyttb.sptensor`, :class:`pyttb.tensor`, \
+            :class:`pyttb.sptenmat`, or :class:`pyttb.sptensor3`
 
         Returns
         -------
@@ -182,11 +191,13 @@ def from_tensor_type(cls, source):
     @classmethod
     def from_function(cls, function_handle, shape, nonzeros):
         """
-        Creates a sparse tensor of the specified shape with NZ nonzeros created from the specified function handle
+        Creates a sparse tensor of the specified shape with NZ nonzeros created from
+        the specified function handle
 
         Parameters
         ----------
-        function_handle: function that accepts 2 arguments and generates :class:`numpy.ndarray` of length nonzeros
+        function_handle: function that accepts 2 arguments and generates
+            :class:`numpy.ndarray` of length nonzeros
         shape: tuple
         nonzeros: int or float
 
@@ -198,9 +209,10 @@ def from_function(cls, function_handle, shape, nonzeros):
         assert callable(function_handle), "function_handle must be callable"
 
         if (nonzeros < 0) or (nonzeros >= np.prod(shape)):
-            assert (
-                False
-            ), "Requested number of non-zeros must be positive and less than the total size"
+            assert False, (
+                "Requested number of non-zeros must be positive "
+                "and less than the total size"
+            )
         elif nonzeros < 1:
             nonzeros = int(np.ceil(np.prod(shape) * nonzeros))
         else:
@@ -226,7 +238,8 @@ def from_function(cls, function_handle, shape, nonzeros):
     @classmethod
     def from_aggregator(cls, subs, vals, shape=None, function_handle="sum"):
         """
-        Construct an sptensor from fully defined SUB, VAL and shape matrices, after an aggregation is applied
+        Construct an sptensor from fully defined SUB, VAL and shape matrices,
+        after an aggregation is applied
 
         Parameters
         ----------
@@ -266,8 +279,8 @@ def from_aggregator(cls, subs, vals, shape=None, function_handle="sum"):
             assert False, "More subscripts than specified by shape"
 
         # Check for subscripts out of range
-        for j in range(len(shape)):
-            if subs.size > 0 and np.max(subs[:, j]) > shape[j]:
+        for j, dim in enumerate(shape):
+            if subs.size > 0 and np.max(subs[:, j]) > dim:
                 assert False, "Subscript exceeds sptensor shape"
 
         if subs.size == 0:
@@ -343,8 +356,7 @@ def collapse(self, dims=None, fun="sum"):
         if remdims.size == 0:
             if fun == "sum":
                 return sum(self.vals.transpose()[0])
-            else:
-                return fun(self.vals.transpose()[0])
+            return fun(self.vals.transpose()[0])
 
         # Calculate the size of the result
         newsize = np.array(self.shape)[remdims]
@@ -358,16 +370,14 @@ def collapse(self, dims=None, fun="sum"):
                     size=newsize[0],
                     func=fun,
                 )
-            else:
-                return np.zeros((newsize[0], 1))
+            return np.zeros((newsize[0], 1))
 
         # Create Result
         if self.subs.size > 0:
             return ttb.sptensor.from_aggregator(
                 self.subs[:, remdims], self.vals, tuple(newsize), fun
             )
-        else:
-            return ttb.sptensor.from_data(np.array([]), np.array([]), tuple(newsize))
+        return ttb.sptensor.from_data(np.array([]), np.array([]), tuple(newsize))
 
     def contract(self, i, j):
         """
@@ -417,8 +427,7 @@ def contract(self, i, j):
         if y.nnz > 0.5 * np.prod(y.shape):
             # Final result is a dense tensor
             return ttb.tensor.from_tensor_type(y)
-        else:
-            return y
+        return y
 
     def double(self):
         """
@@ -450,8 +459,7 @@ def elemfun(self, function):
         idx = np.where(vals > 0)[0]
         if idx.size == 0:
             return ttb.sptensor.from_data(np.array([]), np.array([]), self.shape)
-        else:
-            return ttb.sptensor.from_data(self.subs[idx, :], vals[idx], self.shape)
+        return ttb.sptensor.from_data(self.subs[idx, :], vals[idx], self.shape)
 
     def end(self, k=None):
         """
@@ -467,8 +475,7 @@ def end(self, k=None):
         """
         if k is not None:
             return self.shape[k] - 1
-        else:
-            return np.prod(self.shape) - 1
+        return np.prod(self.shape) - 1
 
     def extract(self, searchsubs):
         """
@@ -580,19 +587,18 @@ def innerprod(self, other):
                 valsSelf = self.extract(subsOther)
             return valsOther.transpose().dot(valsSelf)
 
-        elif isinstance(other, ttb.tensor):
+        if isinstance(other, ttb.tensor):
             if self.shape != other.shape:
                 assert False, "Sptensor and tensor must be same shape for innerproduct"
             [subsSelf, valsSelf] = self.find()
             valsOther = other[subsSelf, "extract"]
             return valsOther.transpose().dot(valsSelf)
 
-        elif isinstance(other, (ttb.ktensor, ttb.ttensor)):  # pragma: no cover
+        if isinstance(other, (ttb.ktensor, ttb.ttensor)):  # pragma: no cover
             # Reverse arguments to call ktensor/ttensor implementation
             return other.innerprod(self)
 
-        else:
-            assert False, "Inner product between sptensor and that class not supported"
+        assert False, f"Inner product between sptensor and {type(other)} not supported"
 
     def isequal(self, other):
         """
@@ -608,12 +614,11 @@ def isequal(self, other):
         """
         if self.shape != other.shape:
             return False
-        elif isinstance(other, ttb.sptensor):
+        if isinstance(other, ttb.sptensor):
             return (self - other).nnz == 0
-        elif isinstance(other, ttb.tensor):
+        if isinstance(other, ttb.tensor):
             return other.isequal(self)
-        else:
-            return False
+        return False
 
     def logical_and(self, B):
         """
@@ -633,17 +638,17 @@ def logical_and(self, B):
         # Case 2: Argument is a tensor of some sort
         if isinstance(B, sptensor):
             # Check that the shapes match
-            if not (self.shape == B.shape):
+            if not self.shape == B.shape:
                 assert False, "Must be tensors of the same shape"
 
-            def isLength2(x):
+            def is_length_2(x):
                 return len(x) == 2
 
             C = sptensor.from_aggregator(
                 np.vstack((self.subs, B.subs)),
                 np.vstack((self.vals, B.vals)),
                 self.shape,
-                isLength2,
+                is_length_2,
             )
 
             return C
@@ -664,7 +669,8 @@ def logical_not(self):
 
         Returns
         -------
-        :class:`pyttb.sptensor` Sparse tensor with all zero-values marked from original sparse tensor
+        :class:`pyttb.sptensor` Sparse tensor with all zero-values marked from original
+        sparse tensor
         """
         allsubs = self.allsubs()
         subsIdx = ttb.tt_setdiff_rows(allsubs, self.subs)
@@ -678,11 +684,12 @@ def logical_or(self, B):
 
         Returns
         -------
-        :class:'pyttb.sptensor` or :class:'pyttb.tensor` sptensor.logical_or(<int,float,tensor>) yields
-        tensor, sptensor.logical_or(sptensor) yields sptensor.
+        :class:'pyttb.sptensor` or :class:'pyttb.tensor`
+            sptensor.logical_or(<int,float,tensor>) yields tensor
+            sptensor.logical_or(sptensor) yields sptensor
         """
         # Case 1: Argument is a scalar or tensor
-        if isinstance(B, (float, int)) or isinstance(B, ttb.tensor):
+        if isinstance(B, (float, int, ttb.tensor)):
             return self.full().logical_or(B)
 
         # Case 2: Argument is an sptensor
@@ -691,14 +698,14 @@ def logical_or(self, B):
 
         if isinstance(B, ttb.sptensor):
 
-            def isLengthGE1(x):
+            def is_length_ge_1(x):
                 return len(x) >= 1
 
             return sptensor.from_aggregator(
                 np.vstack((self.subs, B.subs)),
                 np.ones((self.subs.shape[0] + B.subs.shape[0], 1)),
                 self.shape,
-                isLengthGE1,
+                is_length_ge_1,
             )
 
         assert False, "Sptensor Logical Or argument must be scalar or sptensor"
@@ -715,7 +722,7 @@ def logical_xor(self, other):
 
         """
         # Case 1: Argument is a scalar or dense tensor
-        if isinstance(other, (float, int)) or isinstance(other, ttb.tensor):
+        if isinstance(other, (float, int, ttb.tensor)):
             return self.full().logical_xor(other)
 
         # Case 2: Argument is an sptensor
@@ -854,8 +861,7 @@ def nnz(self):
         """
         if self.subs.size == 0:
             return 0
-        else:
-            return self.subs.shape[0]
+        return self.subs.shape[0]
 
     def norm(self):
         """
@@ -892,7 +898,8 @@ def nvecs(self, n, r, flipsign=True):
             _, v = scipy.sparse.linalg.eigs(y, r)
         else:
             warnings.warn(
-                "Greater than or equal to sptensor.shape[n] - 1 eigenvectors requires cast to dense to solve"
+                "Greater than or equal to sptensor.shape[n] - 1 eigenvectors requires"
+                " cast to dense to solve"
             )
             w, v = scipy.linalg.eig(y.toarray())
             v = v[(-np.abs(w)).argsort()]
@@ -973,18 +980,17 @@ def reshape(self, new_shape, old_modes=None):
                 np.array([]),
                 tuple(np.concatenate((keep_shape, new_shape))),
             )
+        if np.isscalar(old_shape):
+            old_shape = (old_shape,)
+            inds = ttb.tt_sub2ind(old_shape, self.subs[:, old_modes][:, None])
         else:
-            if np.isscalar(old_shape):
-                old_shape = (old_shape,)
-                inds = ttb.tt_sub2ind(old_shape, self.subs[:, old_modes][:, None])
-            else:
-                inds = ttb.tt_sub2ind(old_shape, self.subs[:, old_modes])
-            new_subs = ttb.tt_ind2sub(new_shape, inds)
-            return ttb.sptensor.from_data(
-                np.concatenate((self.subs[:, keep_modes], new_subs), axis=1),
-                self.vals,
-                tuple(np.concatenate((keep_shape, new_shape))),
-            )
+            inds = ttb.tt_sub2ind(old_shape, self.subs[:, old_modes])
+        new_subs = ttb.tt_ind2sub(new_shape, inds)
+        return ttb.sptensor.from_data(
+            np.concatenate((self.subs[:, keep_modes], new_subs), axis=1),
+            self.vals,
+            tuple(np.concatenate((keep_shape, new_shape))),
+        )
 
     def scale(self, factor, dims):
         """
@@ -1012,14 +1018,14 @@ def scale(self, factor, dims):
                 self.vals * factor[self.subs[:, dims[0]], "extract"][:, None],
                 self.shape,
             )
-        elif isinstance(factor, ttb.sptensor):
+        if isinstance(factor, ttb.sptensor):
             shapeArray = np.array(self.shape)
             if np.any(factor.shape != shapeArray[dims]):
                 assert False, "Size mismatch in scale"
             return ttb.sptensor.from_data(
                 self.subs, self.vals * factor.extract(self.subs[:, dims[0]]), self.shape
             )
-        elif isinstance(factor, np.ndarray):
+        if isinstance(factor, np.ndarray):
             shapeArray = np.array(self.shape)
             if factor.shape[0] != shapeArray[dims]:
                 assert False, "Size mismatch in scale"
@@ -1028,12 +1034,12 @@ def scale(self, factor, dims):
                 self.vals * factor[self.subs[:, dims[0]].transpose()[0]],
                 self.shape,
             )
-        else:
-            assert False, "Invalid scaling factor"
+        assert False, "Invalid scaling factor"
 
     def spmatrix(self):
         """
-        Converts a two-way sparse tensor to a sparse matrix in scipy.sparse.coo_matrix format
+        Converts a two-way sparse tensor to a sparse matrix in
+        scipy.sparse.coo_matrix format
 
         Returns
         -------
@@ -1044,10 +1050,9 @@ def spmatrix(self):
 
         if self.subs.size == 0:
             return sparse.coo_matrix(self.shape)
-        else:
-            return sparse.coo_matrix(
-                (self.vals.transpose()[0], self.subs.transpose()), self.shape
-            )
+        return sparse.coo_matrix(
+            (self.vals.transpose()[0], self.subs.transpose()), self.shape
+        )
 
     def squeeze(self):
         """
@@ -1062,16 +1067,13 @@ def squeeze(self):
         # No singleton dimensions
         if np.all(shapeArray > 1):
             return ttb.sptensor.from_tensor_type(self)
-        else:
-            idx = np.where(shapeArray > 1)[0]
-            if idx.size == 0:
-                return self.vals[0].copy()
-            else:
-                siz = tuple(shapeArray[idx])
-                if self.vals.size == 0:
-                    return ttb.sptensor.from_data(np.array([]), np.array([]), siz)
-                else:
-                    return ttb.sptensor.from_data(self.subs[:, idx], self.vals, siz)
+        idx = np.where(shapeArray > 1)[0]
+        if idx.size == 0:
+            return self.vals[0].copy()
+        siz = tuple(shapeArray[idx])
+        if self.vals.size == 0:
+            return ttb.sptensor.from_data(np.array([]), np.array([]), siz)
+        return ttb.sptensor.from_data(self.subs[:, idx], self.vals, siz)
 
     def subdims(self, region):
         """
@@ -1108,7 +1110,8 @@ def subdims(self, region):
         # Error check that range is valid
         # TODO I think only accepting numeric arrays fixes this
 
-        # TODO we use this empty check a lot, do we want a boolean we store in the class for this?
+        # TODO we use this empty check a lot, do we want a boolean we store in the
+        #  class for this?
         if self.subs.size == 0:
             loc = np.array([])
             return loc
@@ -1135,6 +1138,7 @@ def subdims(self, region):
                 loc = loc[tf]
         return loc
 
+    # pylint: disable=too-many-branches, too-many-locals
     def ttv(self, vector, dims=None):
         """
         Sparse tensor times vector
@@ -1154,7 +1158,8 @@ def ttv(self, vector, dims=None):
         elif isinstance(dims, (float, int)):
             dims = np.array([dims])
 
-        # Check that vector is a list of vectors, if not place single vector as element in list
+        # Check that vector is a list of vectors,
+        # if not place single vector as element in list
         if len(vector) > 0 and isinstance(vector[0], (int, float, np.int_, np.float_)):
             return self.ttv([vector], dims)
 
@@ -1173,8 +1178,8 @@ def ttv(self, vector, dims=None):
         if subs.size == 0:  # No non-zeros in tensor
             newsubs = np.array([], dtype=int)
         else:
-            for n in range(len(dims)):
-                idx = subs[:, dims[n]]  # extract indices for dimension n
+            for n, dims_n in enumerate(dims):
+                idx = subs[:, dims_n]  # extract indices for dimension n
                 w = vector[vidx[n]]  # extract the nth vector
                 bigw = w[idx][:, None]  # stretch out the vector
                 newvals = newvals * bigw
@@ -1198,8 +1203,7 @@ def ttv(self, vector, dims=None):
                 return ttb.sptensor.from_aggregator(
                     np.arange(0, newsiz)[:, None], c, tuple(newsiz)
                 )
-            else:
-                return ttb.tensor.from_data(c, tuple(newsiz))
+            return ttb.tensor.from_data(c, tuple(newsiz))
 
         # Case 2: Result is a multiway array
         c = ttb.sptensor.from_aggregator(newsubs, newvals, tuple(newsiz))
@@ -1210,6 +1214,7 @@ def ttv(self, vector, dims=None):
 
         return c
 
+    # pylint: disable=too-many-branches
     def __getitem__(self, item):
         """
         Subscripted reference for a sparse tensor.
@@ -1244,13 +1249,16 @@ def __getitem__(self, item):
 
         Examples
         --------
-        >>> X = sptensor.from_data(np.array([[3,3,3],[1,1,0],[1,2,1]]),np.array([3,5,1]),(4,4,4))
+        >>> subs = np.array([[3,3,3],[1,1,0],[1,2,1]])
+        >>> vals = np.array([3,5,1])
+        >>> shape = (4,4,4)
+        >>> X = sptensor.from_data(subs,vals,shape)
         >>> _ = X[0,1,0] #<-- returns zero
         >>> _ = X[3,3,3] #<-- returns 3
         >>> _ = X[2:3,:,:] #<-- returns 1 x 4 x 4 sptensor
         """
-        # This does not work like MATLAB TTB; you must call sptensor.extract to get this functionality
-        # X([1:6]','extract') %<-- extracts a vector of 6 elements
+        # This does not work like MATLAB TTB; you must call sptensor.extract to get
+        # this functionality: X([1:6]','extract') %<-- extracts a vector of 6 elements
 
         # TODO IndexError for value outside of indices
         # TODO Key error if item not in container
@@ -1275,12 +1283,12 @@ def __getitem__(self, item):
             rmdims = []  # dimensions to remove
 
             # Determine the new size and what dimensions to keep
-            for i in range(0, len(region)):
-                if isinstance(region[i], slice):
+            for i, a_region in enumerate(region):
+                if isinstance(a_region, slice):
                     newsiz.append(self.shape[i])
                     kpdims.append(i)
-                elif not isinstance(region[i], (int, float)):
-                    newsiz.append(np.prod(region[i]))
+                elif not isinstance(a_region, (int, float)):
+                    newsiz.append(np.prod(a_region))
                     kpdims.append(i)
                 else:
                     rmdims.append(i)
@@ -1341,6 +1349,7 @@ def __getitem__(self, item):
 
         return a
 
+    # pylint:disable=too-many-statements, too-many-branches, too-many-locals
     def __setitem__(self, key, value):
         """
         Subscripted assignment for sparse tensor.
@@ -1396,37 +1405,38 @@ def __setitem__(self, key, value):
         objectType = tt_assignment_type(self, key, value)
 
         # Case 1: Replace a sub-tensor
-        if objectType == "subtensor":
+        if objectType == "subtensor":  # pylint:disable=too-many-nested-blocks
             # Case I(a): RHS is another sparse tensor
             if isinstance(value, ttb.sptensor):
-                # First, Resize the tensor and check the size match with the tensor that's being inserted.
+                # First, Resize the tensor and check the size match with the tensor
+                # that's being inserted.
                 m = 0
                 newsz = []
-                for n in range(0, len(key)):
-                    if isinstance(key[n], slice):
+                for n, key_n in enumerate(key):
+                    if isinstance(key_n, slice):
                         if self.ndims <= n:
-                            if key[n].stop is None:
+                            if key_n.stop is None:
                                 newsz.append(value.shape[m])
                             else:
-                                newsz.append(key[n].stop)
+                                newsz.append(key_n.stop)
                         else:
-                            if key[n].stop is None:
+                            if key_n.stop is None:
                                 newsz.append(max([self.shape[n], value.shape[m]]))
                             else:
-                                newsz.append(max([self.shape[n], key[n].stop]))
+                                newsz.append(max([self.shape[n], key_n.stop]))
                         m = m + 1
-                    elif isinstance(key[n], (float, int)):
+                    elif isinstance(key_n, (float, int)):
                         if self.ndims <= n:
-                            newsz.append(key[n] + 1)
+                            newsz.append(key_n + 1)
                         else:
-                            newsz.append(max([self.shape[n], key[n] + 1]))
+                            newsz.append(max([self.shape[n], key_n + 1]))
                     else:
-                        if len(key[n]) != value.shape[m]:
+                        if len(key_n) != value.shape[m]:
                             assert False, "RHS does not match range size"
                         if self.ndims <= n:
-                            newsz.append(max(key[n]) + 1)
+                            newsz.append(max(key_n) + 1)
                         else:
-                            newsz.append(max([self.shape[n], max(key[n]) + 1]))
+                            newsz.append(max([self.shape[n], max(key_n) + 1]))
                 self.shape = tuple(newsz)
 
                 # Expand subs array if there are new modes, i.e., if the order
@@ -1485,9 +1495,10 @@ def __setitem__(self, key, value):
             for n in range(self.ndims, len(key)):
                 if isinstance(key[n], slice):
                     if key[n].stop is None:
-                        assert (
-                            False
-                        ), "Must have well defined slice when expanding sptensor shape with setitem"
+                        assert False, (
+                            "Must have well defined slice when expanding sptensor "
+                            "shape with setitem"
+                        )
                     else:
                         newsz.append(key[n].stop)
                 elif isinstance(key[n], np.ndarray):
@@ -1496,7 +1507,7 @@ def __setitem__(self, key, value):
                     newsz.append(key[n] + 1)
             self.shape = tuple(newsz)
 
-            # Expand subs array if there are new modes, i.e.m if the order has increasesd
+            # Expand subs array if there are new modes, i.e. if the order has increased
             if self.subs.size > 0 and len(self.shape) > self.subs.shape[1]:
                 self.subs = np.append(
                     self.subs,
@@ -1606,15 +1617,18 @@ def __setitem__(self, key, value):
             if isinstance(newvals, (float, int)):
                 newvals = np.expand_dims([newvals], axis=1)
 
-            # Error check the rhs is a column vector. We don't bother to handle any other type with sparse tensors
+            # Error check the rhs is a column vector. We don't bother to handle any
+            # other type with sparse tensors
             tt_valscheck(newvals, nargout=False)
 
-            # Determine number of nonzeros being inserted. (This is determined by number of subscripts)
+            # Determine number of nonzeros being inserted.
+            # (This is determined by number of subscripts)
             newnnz = newsubs.shape[0]
 
             # Error check on size of newvals
             if newvals.size == 1:
-                # Special case where newvals is a single element to be assigned to multiple LHS. Fix to correct size
+                # Special case where newvals is a single element to be assigned
+                # to multiple LHS. Fix to correct size
                 newvals = newvals * np.ones((newnnz, 1))
 
             elif newvals.shape[0] != newnnz:
@@ -1644,7 +1658,8 @@ def __setitem__(self, key, value):
             # processing of Group B may change the locations of the
             # remaining elements.
 
-            # TF+1 for logical consideration because 0 is valid index and -1 is our null flag
+            # TF+1 for logical consideration because 0 is valid index
+            # and -1 is our null flag
             idxa = np.logical_and(tf + 1, newvals)[0]
             idxb = np.logical_and(tf + 1, np.logical_not(newvals))[0]
             idxc = np.logical_and(np.logical_not(tf + 1), newvals)[0]
@@ -1665,9 +1680,9 @@ def __setitem__(self, key, value):
 
             # Resize the tensor
             newshape = []
-            for n in range(0, len(self.shape)):
+            for n, dim in enumerate(self.shape):
                 smax = max(newsubs[:, n] + 1)
-                newshape.append(max(self.shape[n], smax))
+                newshape.append(max(dim, smax))
             self.shape = tuple(newshape)
 
             return
@@ -1688,13 +1703,12 @@ def __eq__(self, other):
         if isinstance(other, (float, int)):
             if other == 0:
                 return self.logical_not()
-            else:
-                idx = self.vals == other
-                return sptensor.from_data(
-                    self.subs[idx.transpose()[0]],
-                    True * np.ones((self.subs.shape[0], 1)).astype(bool),
-                    self.shape,
-                )
+            idx = self.vals == other
+            return sptensor.from_data(
+                self.subs[idx.transpose()[0]],
+                True * np.ones((self.subs.shape[0], 1)).astype(bool),
+                self.shape,
+            )
 
         # Case 2: other is a tensor type
         # Check sizes
@@ -1713,7 +1727,8 @@ def __eq__(self, other):
             zzerosubs = self.allsubs()[xzerosubs][zzerosubsIdx]
 
             # Find where their nonzeros intersect
-            # TODO consider if intersect rows should return 3 args so we don't have to call it twice
+            # TODO consider if intersect rows should return 3 args so we don't have to
+            #  call it twice
             nzsubsIdx = ttb.tt_intersect_rows(self.subs, other.subs)
             nzsubs = self.subs[nzsubsIdx]
             iother = ttb.tt_intersect_rows(other.subs, self.subs)
@@ -1733,7 +1748,7 @@ def __eq__(self, other):
         # Case 2b: other is a dense tensor
         if isinstance(other, ttb.tensor):
             # Find where their zeros interact
-            otherzerosubs, otherzerosubsflag = (other == 0).find()
+            otherzerosubs, _ = (other == 0).find()
             zzerosubs = otherzerosubs[
                 (self.extract(otherzerosubs) == 0).transpose()[0], :
             ]
@@ -1768,15 +1783,14 @@ def __ne__(self, other):
                 return ttb.sptensor.from_data(
                     self.subs, True * np.ones((self.subs.shape[0], 1)), self.shape
                 )
-            else:
-                subs1 = self.subs[self.vals.transpose()[0] != other, :]
-                subs2Idx = ttb.tt_setdiff_rows(self.allsubs(), self.subs)
-                subs2 = self.allsubs()[subs2Idx, :]
-                return ttb.sptensor.from_data(
-                    np.vstack((subs1, subs2)),
-                    True * np.ones((self.subs.shape[0], 1)).astype(bool),
-                    self.shape,
-                )
+            subs1 = self.subs[self.vals.transpose()[0] != other, :]
+            subs2Idx = ttb.tt_setdiff_rows(self.allsubs(), self.subs)
+            subs2 = self.allsubs()[subs2Idx, :]
+            return ttb.sptensor.from_data(
+                np.vstack((subs1, subs2)),
+                True * np.ones((self.subs.shape[0], 1)).astype(bool),
+                self.shape,
+            )
 
         # Case 2: Both x and y are tensors or some sort
         # Check that the sizes match
@@ -1852,7 +1866,7 @@ def __sub__(self, other):
         # a dense result, even if the scalar is zero.
 
         # Case 1: Second argument is a scalar or a dense tensor
-        if isinstance(other, (float, int)) or isinstance(other, ttb.tensor):
+        if isinstance(other, (float, int, ttb.tensor)):
             return self.full() - other
 
         # Case 2: Both are sparse tensors
@@ -1934,11 +1948,11 @@ def __mul__(self, other):
                 self.vals[idxSelf] * other.vals[idxOther],
                 self.shape,
             )
-        elif isinstance(other, ttb.tensor):
+        if isinstance(other, ttb.tensor):
             csubs = self.subs
             cvals = self.vals * other[csubs, "extract"][:, None]
             return ttb.sptensor.from_data(csubs, cvals, self.shape)
-        elif isinstance(other, ttb.ktensor):
+        if isinstance(other, ttb.ktensor):
             csubs = self.subs
             cvals = np.zeros(self.vals.shape)
             R = other.weights.size
@@ -1946,13 +1960,13 @@ def __mul__(self, other):
             for r in range(R):
                 tvals = other.weights[r] * self.vals
                 for n in range(N):
-                    # Note other[n][:, r] extracts 1-D instead of column vector, which necessitates [:, None]
+                    # Note other[n][:, r] extracts 1-D instead of column vector,
+                    # which necessitates [:, None]
                     v = other[n][:, r][:, None]
                     tvals = tvals * v[csubs[:, n]]
                 cvals += tvals
             return ttb.sptensor.from_data(csubs, cvals, self.shape)
-        else:
-            assert False, "Sptensor cannot be multiplied by that type of object"
+        assert False, "Sptensor cannot be multiplied by that type of object"
 
     def __rmul__(self, other):
         """
@@ -1968,9 +1982,9 @@ def __rmul__(self, other):
         """
         if isinstance(other, (float, int, np.number)):
             return self.__mul__(other)
-        else:
-            assert False, "This object cannot be multiplied by sptensor"
+        assert False, "This object cannot be multiplied by sptensor"
 
+    # pylint:disable=too-many-branches
     def __le__(self, other):
         """
         Less than or equal (<=) for sptensor
@@ -1983,7 +1997,8 @@ def __le__(self, other):
         -------
         :class:`pyttb.sptensor`
         """
-        # TODO le,lt,ge,gt have a lot of code duplication, look at generalizing them for future maintainabilty
+        # TODO le,lt,ge,gt have a lot of code duplication, look at generalizing them
+        #  for future maintainabilty
         # Case 1: One argument is a scalar
         if isinstance(other, (float, int)):
             subs1 = self.subs[(self.vals <= other).transpose()[0], :]
@@ -2069,6 +2084,7 @@ def __le__(self, other):
         # Otherwise
         assert False, "Cannot compare sptensor with that type"
 
+    # pylint:disable=too-many-branches
     def __lt__(self, other):
         """
         Less than (<) for sptensor
@@ -2270,6 +2286,7 @@ def __gt__(self, other):
         # Otherwise
         assert False, "Cannot compare sptensor with that type"
 
+    # pylint:disable=too-many-statements, too-many-branches, too-many-locals
     def __truediv__(self, other):
         """
         Division for sparse tensors (sptensor/other).
@@ -2287,7 +2304,8 @@ def __truediv__(self, other):
         if isinstance(other, (float, int)):
             # Inline mrdivide
             newsubs = self.subs
-            # We ignore the divide by zero errors because np.inf/np.nan is an appropriate representation
+            # We ignore the divide by zero errors because np.inf/np.nan is an
+            # appropriate representation
             with np.errstate(divide="ignore", invalid="ignore"):
                 newvals = self.vals / other
             if other == 0:
@@ -2356,11 +2374,11 @@ def __truediv__(self, other):
 
             return ttb.sptensor.from_data(newsubs, newvals, self.shape)
 
-        elif isinstance(other, ttb.tensor):
+        if isinstance(other, ttb.tensor):
             csubs = self.subs
             cvals = self.vals / other[csubs, "extract"][:, None]
             return ttb.sptensor.from_data(csubs, cvals, self.shape)
-        elif isinstance(other, ttb.ktensor):
+        if isinstance(other, ttb.ktensor):
             # TODO consider removing epsilon and generating nans consistent with above
             epsilon = np.finfo(float).eps
             subs = self.subs
@@ -2376,8 +2394,7 @@ def __truediv__(self, other):
             return ttb.sptensor.from_data(
                 self.subs, self.vals / np.maximum(epsilon, vals), self.shape
             )
-        else:
-            assert False, "Invalid arguments for sptensor division"
+        assert False, "Invalid arguments for sptensor division"
 
     def __rtruediv__(self, other):
         """
@@ -2394,8 +2411,7 @@ def __rtruediv__(self, other):
         # Scalar divided by a tensor -> result is dense
         if isinstance(other, (float, int)):
             return other / self.full()
-        else:
-            assert False, "Dividing that object by an sptensor is not supported"
+        assert False, "Dividing that object by an sptensor is not supported"
 
     def __repr__(self):  # pragma: no cover
         """
@@ -2414,17 +2430,17 @@ def __repr__(self):  # pragma: no cover
                 return s
             s += (" x ").join([str(int(d)) for d in self.shape])
             return s
-        else:
-            s = "Sparse tensor of shape "
-            s += (" x ").join([str(int(d)) for d in self.shape])
-            s += " with {} nonzeros \n".format(nz)
+
+        s = "Sparse tensor of shape "
+        s += (" x ").join([str(int(d)) for d in self.shape])
+        s += f" with {nz} nonzeros \n"
 
         # Stop insane printouts
         if nz > 10000:
             r = input("Are you sure you want to print all nonzeros? (Y/N)")
             if r.upper() != "Y":
                 return s
-        for i, j in enumerate(range(0, self.subs.shape[0])):
+        for i in range(0, self.subs.shape[0]):
             s += "\t"
             s += "["
             idx = self.subs[i, :]
@@ -2494,7 +2510,8 @@ def ttm(self, matrices, dims=None, transpose=False):
         # Compute self[mode]'
         Xnt = tt_to_sparse_matrix(self, dims, True)
 
-        # Reshape puts the reshaped things after the unchanged modes, transpose then puts it in front
+        # 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
@@ -2505,7 +2522,6 @@ def ttm(self, matrices, dims=None, transpose=False):
 
         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
-            # using spare matrix representation
-            return ttb.tensor.from_tensor_type(Ynt)
+        # TODO evaluate performance loss by casting into sptensor then tensor.
+        #  I assume minimal since we are already using spare matrix representation
+        return ttb.tensor.from_tensor_type(Ynt)
diff --git a/tests/test_package.py b/tests/test_package.py
index c963a9c9..8d8ca619 100644
--- a/tests/test_package.py
+++ b/tests/test_package.py
@@ -23,6 +23,7 @@ def test_linting():
 
     enforced_files = [
         os.path.join(os.path.dirname(ttb.__file__), f"{ttb.tensor.__name__}.py"),
+        os.path.join(os.path.dirname(ttb.__file__), f"{ttb.sptensor.__name__}.py"),
         ttb.pyttb_utils.__file__,
     ]
     # TODO pylint fails to import pyttb in tests
diff --git a/tests/test_sptensor.py b/tests/test_sptensor.py
index c9a47b96..16125371 100644
--- a/tests/test_sptensor.py
+++ b/tests/test_sptensor.py
@@ -1132,7 +1132,7 @@ def test_sptensor_innerprod(sample_sptensor):
     # Wrong type for innerprod
     with pytest.raises(AssertionError) as excinfo:
         sptensorInstance.innerprod(5)
-    assert "Inner product between sptensor and that class not supported" in str(excinfo)
+    assert f"Inner product between sptensor and {type(5)} not supported" in str(excinfo)
 
 
 @pytest.mark.indevelopment

From 587d317fb2216145e275f69d0d807cf2db32356a Mon Sep 17 00:00:00 2001
From: Nick Johnson <24689722+ntjohnson1@users.noreply.github.com>
Date: Sat, 25 Feb 2023 17:47:23 -0500
Subject: [PATCH 4/7] SPTENSOR: Initial pass a typing support

---
 pyttb/sptensor.py | 223 ++++++++++++++++++++++------------------------
 pyttb/tensor.py   |   6 +-
 2 files changed, 111 insertions(+), 118 deletions(-)

diff --git a/pyttb/sptensor.py b/pyttb/sptensor.py
index 08af1965..91936788 100644
--- a/pyttb/sptensor.py
+++ b/pyttb/sptensor.py
@@ -2,7 +2,10 @@
 # LLC (NTESS). Under the terms of Contract DE-NA0003525 with NTESS, the
 # U.S. Government retains certain rights in this software.
 """Sparse Tensor Implementation"""
+from __future__ import annotations
+
 import warnings
+from typing import Any, Callable, Optional, Tuple, Union, overload
 
 import numpy as np
 import scipy.sparse.linalg
@@ -23,22 +26,22 @@
 )
 
 
-def tt_to_sparse_matrix(sptensorInstance, mode, transpose=False):
+def tt_to_sparse_matrix(
+    sptensorInstance: sptensor, mode: int, transpose: bool = False
+) -> sparse.coo_matrix:
     """
     Helper function to unwrap sptensor into sparse matrix, should replace the core need
     for sptenmat
 
     Parameters
     ----------
-    sptensorInstance: :class:`pyttb.sptensor`
-    mode: int
-        Mode around which to unwrap tensor
-    transpose: bool
-        Whether or not to tranpose unwrapped tensor
+    sptensorInstance: sparse tensor to unwrap
+    mode: Mode around which to unwrap tensor
+    transpose: Whether or not to tranpose unwrapped tensor
 
     Returns
     -------
-    spmatrix: :class:`Scipy.sparse.coo_matrix`
+    spmatrix: unwrapped tensor
     """
     old = np.setdiff1d(np.arange(sptensorInstance.ndims), mode).astype(int)
     spmatrix = sptensorInstance.reshape(
@@ -49,7 +52,9 @@ def tt_to_sparse_matrix(sptensorInstance, mode, transpose=False):
     return spmatrix
 
 
-def tt_from_sparse_matrix(spmatrix, shape, mode, idx):
+def tt_from_sparse_matrix(
+    spmatrix: sparse.coo_matrix, shape: Any, mode: int, idx: int
+) -> sptensor:
     """
     Helper function to wrap sparse matrix into sptensor.
     Inverse of :class:`pyttb.tt_to_sparse_matrix`
@@ -76,7 +81,7 @@ def tt_from_sparse_matrix(spmatrix, shape, mode, idx):
     # unchanged ones
     sptensorInstance = sptensorInstance.reshape(
         shape,
-        np.concatenate((np.arange(1, mode + 1), [0], np.arange(mode + 1, len(shape)))),
+        np.concatenate([np.arange(1, mode + 1), [0], np.arange(mode + 1, len(shape))]),
     )
 
     return sptensorInstance
@@ -109,19 +114,17 @@ def __init__(self):
         #    return
 
     @classmethod
-    def from_data(cls, subs, vals, shape):
+    def from_data(
+        cls, subs: np.ndarray, vals: np.ndarray, shape: Tuple[int, ...]
+    ) -> sptensor:
         """
         Construct an sptensor from fully defined SUB, VAL and SIZE matrices.
 
         Parameters
         ----------
-        subs: :class:`numpy.ndarray`
-        vals: :class:`numpy.ndarray`
-        shape: tuple
-
-        Returns
-        -------
-        :class:`pyttb.sptensor`
+        subs: location of non-zero entries
+        vals: values for non-zero entries
+        shape: shape of sparse tensor
 
         Examples
         --------
@@ -145,19 +148,15 @@ def from_data(cls, subs, vals, shape):
         return sptensorInstance
 
     @classmethod
-    def from_tensor_type(cls, source):
+    def from_tensor_type(
+        cls, source: Union[sptensor, ttb.tensor, sparse.coo_matrix]
+    ) -> sptensor:
         """
-        Contruct an :class:`pyttb.sptensor` from :class:`pyttb.sptensor`,
-        :class:`pyttb.tensor`, :class:`pyttb.sptenmat`, :class:`pyttb.sptensor3`
+        Contruct an :class:`pyttb.sptensor` from compatible tensor types
 
         Parameters
         ----------
-        source: :class:`pyttb.sptensor`, :class:`pyttb.tensor`, \
-            :class:`pyttb.sptenmat`, or :class:`pyttb.sptensor3`
-
-        Returns
-        -------
-        :class:`pyttb.sptensor`
+        source: Source tensor to create sptensor from
         """
         # Copy Constructor
         if isinstance(source, sptensor):
@@ -189,7 +188,12 @@ def from_tensor_type(cls, source):
         assert False, "Invalid Tensor Type To initialize Sptensor"
 
     @classmethod
-    def from_function(cls, function_handle, shape, nonzeros):
+    def from_function(
+        cls,
+        function_handle: Callable[[Tuple[float, float]], np.ndarray],
+        shape: Tuple[int, ...],
+        nonzeros: float,
+    ) -> sptensor:
         """
         Creates a sparse tensor of the specified shape with NZ nonzeros created from
         the specified function handle
@@ -217,6 +221,7 @@ def from_function(cls, function_handle, shape, nonzeros):
             nonzeros = int(np.ceil(np.prod(shape) * nonzeros))
         else:
             nonzeros = int(np.floor(nonzeros))
+        nonzeros = int(nonzeros)
 
         # Keep iterating until we find enough unique non-zeros or we give up
         subs = np.array([])
@@ -236,7 +241,7 @@ def from_function(cls, function_handle, shape, nonzeros):
         return cls().from_data(subs, vals, shape)
 
     @classmethod
-    def from_aggregator(cls, subs, vals, shape=None, function_handle="sum"):
+    def from_aggregator(cls, subs, vals, shape=None, function_handle="sum") -> sptensor:
         """
         Construct an sptensor from fully defined SUB, VAL and shape matrices,
         after an aggregation is applied
@@ -307,12 +312,13 @@ def from_aggregator(cls, subs, vals, shape=None, function_handle="sum"):
         return cls().from_data(newsubs, newvals, shape)
 
     # TODO decide if property
-    def allsubs(self):
+    def allsubs(self) -> np.ndarray:
         """
         Generate all possible subscripts for sparse tensor
+
         Returns
         -------
-        s: :class:`numpy.ndarray` all possible subscripts for sptensor
+        s: All possible subscripts for sptensor
         """
 
         # Generate all possible indices
@@ -429,13 +435,9 @@ def contract(self, i, j):
             return ttb.tensor.from_tensor_type(y)
         return y
 
-    def double(self):
+    def double(self) -> np.ndarray:
         """
         Convert sptensor to dense multidimensional array
-
-        Returns
-        -------
-        :class:`numpy.ndarray`
         """
         a = np.zeros(self.shape)
         if self.nnz > 0:
@@ -461,33 +463,25 @@ def elemfun(self, function):
             return ttb.sptensor.from_data(np.array([]), np.array([]), self.shape)
         return ttb.sptensor.from_data(self.subs[idx, :], vals[idx], self.shape)
 
-    def end(self, k=None):
+    def end(self, k: Optional[int] = None) -> int:
         """
         Last index of indexing expression for sparse tensor
 
         Parameters
         ----------
         k: int Dimension for subscript indexing
-
-        Returns
-        -------
-        int:
         """
         if k is not None:
             return self.shape[k] - 1
         return np.prod(self.shape) - 1
 
-    def extract(self, searchsubs):
+    def extract(self, searchsubs: np.ndarray) -> np.ndarray:
         """
         Extract value for a sptensor.
 
         Parameters
         ----------
-        searchsubs: :class:`numpy.ndarray` subscripts to find in sptensor
-
-        Returns
-        -------
-        :class:`numpy.ndarray`
+        searchsubs: subscripts to find in sptensor
 
         See Also
         --------
@@ -521,22 +515,20 @@ def extract(self, searchsubs):
         a[nzsubs] = self.vals[loc[nzsubs]]
         return a
 
-    def find(self):
+    def find(self) -> Tuple[np.ndarray, np.ndarray]:
         """
         FIND Find subscripts of nonzero elements in a sparse tensor.
 
         Returns
         -------
-        subs: :class:`numpy.ndarray`
-        vals: :class:`numpy.ndarray`
+        subs: Subscripts of nonzero elements
+        vals: Values at corresponding subscripts
         """
         return self.subs, self.vals
 
-    def full(self):
+    def full(self) -> ttb.tensor:
         """
         FULL Convert a sparse tensor to a (dense) tensor.
-
-        :return: tensor
         """
         # Handle the completely empty (no shape) case
         if len(self.shape) == 0:
@@ -555,18 +547,15 @@ def full(self):
         B[idx.astype(int)] = self.vals.transpose()[0]
         return B
 
-    def innerprod(self, other):
+    def innerprod(
+        self, other: Union[sptensor, ttb.tensor, ttb.ktensor, ttb.ttensor]
+    ) -> float:
         """
         Efficient inner product with a sparse tensor
 
         Parameters
         ----------
-        other: :class:`pyttb.tensor`, :class:`pyttb.sptensor`, :class:`pyttb.ktensor`,
-        :class:`pyttb.ttensor`
-
-        Returns
-        -------
-        float
+        other: Other tensor to take innerproduct with
         """
         # If all entries are zero innerproduct must be 0
         if self.nnz == 0:
@@ -600,17 +589,13 @@ def innerprod(self, other):
 
         assert False, f"Inner product between sptensor and {type(other)} not supported"
 
-    def isequal(self, other):
+    def isequal(self, other: Union[sptensor, ttb.tensor]) -> bool:
         """
         Exact equality for sptensors
 
         Parameters
         ----------
-        other: :class:`pyttb.tensor`, :class:`pyttb.sptensor`
-
-        Returns
-        -------
-        bool: True if sptensors are identical, false otherwise
+        other: Other tensor to compare against
         """
         if self.shape != other.shape:
             return False
@@ -620,12 +605,17 @@ def isequal(self, other):
             return other.isequal(self)
         return False
 
-    def logical_and(self, B):
+    def logical_and(self, B: Union[float, sptensor, ttb.tensor]) -> sptensor:
         """
         Logical and with self and another object
 
-        :param B: Scalar, tensor, or sptensor
-        :return: Indicator tensor
+        Parameters
+        ----------
+        B: Other value to compare with
+
+        Returns
+        ----------
+        Indicator tensor
         """
         # Case 1: One argument is a scalar
         if isinstance(B, (int, float)):
@@ -663,13 +653,13 @@ def is_length_2(x):
         # Otherwise
         assert False, "The arguments must be two sptensors or an sptensor and a scalar."
 
-    def logical_not(self):
+    def logical_not(self) -> sptensor:
         """
         Logical NOT for sptensors
 
         Returns
         -------
-        :class:`pyttb.sptensor` Sparse tensor with all zero-values marked from original
+        Sparse tensor with all zero-values marked from original
         sparse tensor
         """
         allsubs = self.allsubs()
@@ -678,15 +668,23 @@ def logical_not(self):
         trueVector = np.ones(shape=(subs.shape[0], 1), dtype=bool)
         return sptensor.from_data(subs, trueVector, self.shape)
 
-    def logical_or(self, B):
+    @overload
+    def logical_or(self, B: Union[float, ttb.tensor]) -> ttb.tensor:
+        ...  # pragma: no cover see coveragepy/issues/970
+
+    @overload
+    def logical_or(self, B: sptensor) -> sptensor:
+        ...  # pragma: no cover see coveragepy/issues/970
+
+    def logical_or(
+        self, B: Union[float, ttb.tensor, sptensor]
+    ) -> Union[ttb.tensor, sptensor]:
         """
-        Logical OR for sptensors
+        Logical OR for sptensor and another value
 
         Returns
         -------
-        :class:'pyttb.sptensor` or :class:'pyttb.tensor`
-            sptensor.logical_or(<int,float,tensor>) yields tensor
-            sptensor.logical_or(sptensor) yields sptensor
+        Indicator tensor
         """
         # Case 1: Argument is a scalar or tensor
         if isinstance(B, (float, int, ttb.tensor)):
@@ -710,16 +708,27 @@ def is_length_ge_1(x):
 
         assert False, "Sptensor Logical Or argument must be scalar or sptensor"
 
-    def logical_xor(self, other):
+    @overload
+    def logical_xor(self, other: Union[float, ttb.tensor]) -> ttb.tensor:
+        ...  # pragma: no cover see coveragepy/issues/970
+
+    @overload
+    def logical_xor(self, other: sptensor) -> sptensor:
+        ...  # pragma: no cover see coveragepy/issues/970
+
+    def logical_xor(
+        self, other: Union[float, ttb.tensor, sptensor]
+    ) -> Union[ttb.tensor, sptensor]:
         """
         Logical XOR for sptensors
 
         Parameters
         ----------
+        other: Other value to xor against
 
         Returns
         -------
-
+        Indicator tensor
         """
         # Case 1: Argument is a scalar or dense tensor
         if isinstance(other, (float, int, ttb.tensor)):
@@ -741,17 +750,17 @@ def length1(x):
 
         assert False, "The argument must be an sptensor, tensor or scalar"
 
-    def mask(self, W):
+    def mask(self, W: sptensor) -> np.ndarray:
         """
         Extract values as specified by a mask tensor
 
         Parameters
         ----------
-        W: :class:`pyttb.sptensor`
+        W: Mask tensor
 
         Returns
         -------
-        :class:`Numpy.ndarray`
+        Extracted values
         """
         # Error check
         if len(W.shape) != len(self.shape) or np.any(
@@ -839,36 +848,24 @@ def mttkrp(self, U, n):
         return V
 
     @property
-    def ndims(self):
+    def ndims(self) -> int:
         """
         NDIMS Number of dimensions of a sparse tensor.
-
-        Returns
-        -------
-        int
-            Number of dimensions of Sptensor
         """
         return len(self.shape)
 
     @property
-    def nnz(self):
+    def nnz(self) -> int:
         """
         Number of nonzeros in sparse tensor
-
-        Returns
-        -------
-        nnz: int
         """
         if self.subs.size == 0:
             return 0
         return self.subs.shape[0]
 
-    def norm(self):
+    def norm(self) -> np.floating:
         """
-        Compute the norm of a sparse tensor.
-        Returns
-        -------
-        norm: float, Frobenius norm of Tensor
+        Compute the Frobenius norm of a sparse tensor.
         """
         return np.linalg.norm(self.vals)
 
@@ -912,7 +909,7 @@ def nvecs(self, n, r, flipsign=True):
                     v[:, i] *= -1
         return v
 
-    def ones(self):
+    def ones(self) -> sptensor:
         """
         Replace nonzero elements of sparse tensor with ones
         """
@@ -920,17 +917,13 @@ def ones(self):
         oneVals.fill(1)
         return ttb.sptensor.from_data(self.subs, oneVals, self.shape)
 
-    def permute(self, order):
+    def permute(self, order: np.ndarray) -> sptensor:
         """
         Rearrange the dimensions of a sparse tensor
 
         Parameters
         ----------
-        order: :class:`Numpy.ndarray`
-
-        Returns
-        -------
-        :class:`pyttb.sptensor`
+        order: Updated order of dimensions
         """
         # Error check
         if self.ndims != order.size or np.any(
@@ -947,7 +940,11 @@ def permute(self, order):
             self.subs, self.vals, tuple(np.array(self.shape)[order])
         )
 
-    def reshape(self, new_shape, old_modes=None):
+    def reshape(
+        self,
+        new_shape: Tuple[int, ...],
+        old_modes: Optional[Union[np.ndarray, int]] = None,
+    ) -> sptensor:
         """
         Reshape specified modes of sparse tensor
 
@@ -992,7 +989,7 @@ def reshape(self, new_shape, old_modes=None):
             tuple(np.concatenate((keep_shape, new_shape))),
         )
 
-    def scale(self, factor, dims):
+    def scale(self, factor: np.ndarray, dims: Union[float, np.ndarray]) -> sptensor:
         """
         Scale along specified dimensions for sparse tensors
 
@@ -1007,7 +1004,7 @@ def scale(self, factor, dims):
         """
         if isinstance(dims, (float, int)):
             dims = np.array([dims])
-        dims = ttb.tt_dimscheck(dims, self.ndims)
+        dims, _ = ttb.tt_dimscheck(dims, self.ndims)
 
         if isinstance(factor, ttb.tensor):
             shapeArray = np.array(self.shape)
@@ -1015,7 +1012,7 @@ def scale(self, factor, dims):
                 assert False, "Size mismatch in scale"
             return ttb.sptensor.from_data(
                 self.subs,
-                self.vals * factor[self.subs[:, dims[0]], "extract"][:, None],
+                self.vals * factor[self.subs[:, dims], "extract"][:, None],
                 self.shape,
             )
         if isinstance(factor, ttb.sptensor):
@@ -1023,7 +1020,7 @@ def scale(self, factor, dims):
             if np.any(factor.shape != shapeArray[dims]):
                 assert False, "Size mismatch in scale"
             return ttb.sptensor.from_data(
-                self.subs, self.vals * factor.extract(self.subs[:, dims[0]]), self.shape
+                self.subs, self.vals * factor.extract(self.subs[:, dims]), self.shape
             )
         if isinstance(factor, np.ndarray):
             shapeArray = np.array(self.shape)
@@ -1031,19 +1028,15 @@ def scale(self, factor, dims):
                 assert False, "Size mismatch in scale"
             return ttb.sptensor.from_data(
                 self.subs,
-                self.vals * factor[self.subs[:, dims[0]].transpose()[0]],
+                self.vals * factor[self.subs[:, dims].transpose()[0]],
                 self.shape,
             )
         assert False, "Invalid scaling factor"
 
-    def spmatrix(self):
+    def spmatrix(self) -> sparse.coo_matrix:
         """
         Converts a two-way sparse tensor to a sparse matrix in
         scipy.sparse.coo_matrix format
-
-        Returns
-        -------
-        :class:`scipy.sparse.coo_matrix`
         """
         if self.ndims != 2:
             assert False, "Sparse tensor must be two dimensional"
@@ -1054,7 +1047,7 @@ def spmatrix(self):
             (self.vals.transpose()[0], self.subs.transpose()), self.shape
         )
 
-    def squeeze(self):
+    def squeeze(self) -> Union[sptensor, float]:
         """
         Remove singleton dimensions from a sparse tensor
 
diff --git a/pyttb/tensor.py b/pyttb/tensor.py
index 59b772b7..61948dd3 100644
--- a/pyttb/tensor.py
+++ b/pyttb/tensor.py
@@ -413,7 +413,7 @@ def innerprod(self, other: Union[tensor, ttb.sptensor, ttb.ktensor]) -> float:
             return other.innerprod(self)
         assert False, "Inner product between tensor and that class is not supported"
 
-    def isequal(self, other: Union[tensor, ttb.sptensor]) -> Union[bool, np.bool_]:
+    def isequal(self, other: Union[tensor, ttb.sptensor]) -> bool:
         """
         Exact equality for tensors
 
@@ -429,9 +429,9 @@ def isequal(self, other: Union[tensor, ttb.sptensor]) -> Union[bool, np.bool_]:
         False
         """
         if isinstance(other, ttb.tensor):
-            return np.all(self.data == other.data)
+            return bool(np.all(self.data == other.data))
         if isinstance(other, ttb.sptensor):
-            return np.all(self.data == other.full().data)
+            return bool(np.all(self.data == other.full().data))
         return False
 
     # TODO: We should probably always return details and let caller drop them

From f5de9d83f90a144710bcc3421a6d247aa349f0c2 Mon Sep 17 00:00:00 2001
From: Nick Johnson <24689722+ntjohnson1@users.noreply.github.com>
Date: Sat, 18 Mar 2023 17:15:36 -0400
Subject: [PATCH 5/7] SPTENSOR: Complete initial typing coverage

---
 pyttb/sptensor.py | 176 +++++++++++++++++++++++++++++-----------------
 pyttb/tensor.py   |  15 ++--
 2 files changed, 117 insertions(+), 74 deletions(-)

diff --git a/pyttb/sptensor.py b/pyttb/sptensor.py
index 91936788..791a2b09 100644
--- a/pyttb/sptensor.py
+++ b/pyttb/sptensor.py
@@ -5,7 +5,8 @@
 from __future__ import annotations
 
 import warnings
-from typing import Any, Callable, Optional, Tuple, Union, overload
+from collections.abc import Sequence
+from typing import Any, Callable, List, Optional, Tuple, Union, cast, overload
 
 import numpy as np
 import scipy.sparse.linalg
@@ -157,6 +158,10 @@ def from_tensor_type(
         Parameters
         ----------
         source: Source tensor to create sptensor from
+
+        Returns
+        -------
+        Generated Sparse Tensor
         """
         # Copy Constructor
         if isinstance(source, sptensor):
@@ -207,7 +212,7 @@ def from_function(
 
         Returns
         -------
-        :class:`pyttb.sptensor`
+        Generated Sparse Tensor
         """
         # Random Tensor
         assert callable(function_handle), "function_handle must be callable"
@@ -241,21 +246,28 @@ def from_function(
         return cls().from_data(subs, vals, shape)
 
     @classmethod
-    def from_aggregator(cls, subs, vals, shape=None, function_handle="sum") -> sptensor:
+    def from_aggregator(
+        cls,
+        subs: np.ndarray,
+        vals: np.ndarray,
+        shape: Optional[Tuple[int, ...]] = None,
+        function_handle: Union[str, Callable[[Any], Union[float, np.ndarray]]] = "sum",
+    ) -> sptensor:
         """
         Construct an sptensor from fully defined SUB, VAL and shape matrices,
         after an aggregation is applied
 
         Parameters
         ----------
-        subs: :class:`numpy.ndarray`
-        vals: :class:`numpy.ndarray`
-        shape: tuple
-        function_handle: callable
+        subs: location of non-zero entries
+        vals: values for non-zero entries
+        shape: shape of sparse tensor
+        function_handle: Aggregation function, or name of supported
+            aggregation function from numpy_groupies
 
         Returns
         -------
-        :class:`pyttb.sptensor`
+        Generated Sparse Tensor
 
         Examples
         --------
@@ -339,18 +351,33 @@ def allsubs(self) -> np.ndarray:
 
         return s.astype(int)
 
-    def collapse(self, dims=None, fun="sum"):
+    def collapse(
+        self,
+        dims: Optional[np.ndarray] = None,
+        fun: Callable[[np.ndarray], Union[float, np.ndarray]] = np.sum,
+    ) -> Union[float, np.ndarray, sptensor]:
         """
         Collapse sparse tensor along specified dimensions.
 
         Parameters
         ----------
-        dims:
-        fun: callable
+        dims: Dimensions to collapse
+        fun: Method used to collapse dimensions
 
         Returns
         -------
+        Collapsed value
 
+        Example
+        -------
+        >>> subs = np.array([[1, 2], [1, 3]])
+        >>> vals = np.array([[1], [1]])
+        >>> shape = np.array([4, 4])
+        >>> X = ttb.sptensor.from_data(subs, vals, shape)
+        >>> X.collapse()
+        2
+        >>> X.collapse(np.arange(X.ndims), sum)
+        2
         """
         if dims is None:
             dims = np.arange(0, self.ndims)
@@ -360,8 +387,6 @@ def collapse(self, dims=None, fun="sum"):
 
         # Check for the case where we accumulate over *all* dimensions
         if remdims.size == 0:
-            if fun == "sum":
-                return sum(self.vals.transpose()[0])
             return fun(self.vals.transpose()[0])
 
         # Calculate the size of the result
@@ -385,19 +410,25 @@ def collapse(self, dims=None, fun="sum"):
             )
         return ttb.sptensor.from_data(np.array([]), np.array([]), tuple(newsize))
 
-    def contract(self, i, j):
+    def contract(self, i: int, j: int) -> Union[np.ndarray, sptensor, ttb.tensor]:
         """
         Contract tensor along two dimensions (array trace).
 
         Parameters
         ----------
-        i: int
-        j: int
+        i: First dimension
+        j: Second dimension
 
         Returns
         -------
+        Contracted sptensor, converted to tensor if sufficiently dense
 
-
+        Example
+        -------
+        >>> X = ttb.tensor.from_data(np.ones((2,2)))
+        >>> Y = sptensor.from_tensor_type(X)
+        >>> Y.contract(0, 1)
+        2.0
         """
         if self.shape[i] != self.shape[j]:
             assert False, "Must contract along equally sized dimensions"
@@ -444,20 +475,28 @@ def double(self) -> np.ndarray:
             a[tuple(self.subs.transpose())] = self.vals.transpose()[0]
         return a
 
-    def elemfun(self, function):
+    def elemfun(self, function_handle: Callable[[np.ndarray], np.ndarray]) -> sptensor:
         """
         Manipulate the non-zero elements of a sparse tensor
 
         Parameters
         ----------
-        function: callable
+        function_handle: Function that updates all values.
 
         Returns
         -------
-        :class:`Tensortoolbox.sptensor`
+        Updated sptensor
+
+        Example
+        -------
+        >>> X = ttb.tensor.from_data(np.ones((2,2)))
+        >>> Y = sptensor.from_tensor_type(X)
+        >>> Z = Y.elemfun(lambda values: values*2)
+        >>> Z.isequal(Y*2)
+        True
         """
 
-        vals = function(self.vals)
+        vals = function_handle(self.vals)
         idx = np.where(vals > 0)[0]
         if idx.size == 0:
             return ttb.sptensor.from_data(np.array([]), np.array([]), self.shape)
@@ -780,18 +819,18 @@ def mask(self, W: sptensor) -> np.ndarray:
         vals[idx] = self.vals[idx]
         return vals
 
-    def mttkrp(self, U, n):
+    def mttkrp(self, U: Union[ttb.ktensor, List[np.ndarray]], n: int) -> np.ndarray:
         """
         Matricized tensor times Khatri-Rao product for sparse tensor.
 
         Parameters
         ----------
-        U: array of matrices or ktensor
-        n: multiplies by all modes except n
+        U: Matrices to create the Khatri-Rao product
+        n: Mode to matricize sptensor in
 
         Returns
         -------
-        :class:`numpy.ndarray`
+        Matrix product
 
         Examples
         --------
@@ -869,27 +908,28 @@ def norm(self) -> np.floating:
         """
         return np.linalg.norm(self.vals)
 
-    def nvecs(self, n, r, flipsign=True):
+    def nvecs(self, n: int, r: int, flipsign: bool = True) -> np.ndarray:
         """
         Compute the leading mode-n vectors for a sparse tensor.
 
         Parameters
         ----------
-        n: mode for tensor matricization
-        r: number of eigenvalues
-        flipsign: Make each column's largest element positive if true
-
-        Returns
-        -------
-
+        n: Mode to unfold
+        r: Number of eigenvectors to compute
+        flipsign: Make each eigenvector's largest element positive
         """
         old = np.setdiff1d(np.arange(self.ndims), n).astype(int)
-        tnt = (
-            self.reshape((np.prod(np.array(self.shape)[old]), 1), old)
+        # tnt calculation is a workaround for missing sptenmat
+        mutatable_sptensor = (
+            sptensor.from_tensor_type(self)
+            .reshape((np.prod(np.array(self.shape)[old]), 1), old)
             .squeeze()
-            .spmatrix()
-            .transpose()
         )
+        if isinstance(mutatable_sptensor, (int, float, np.generic)):
+            raise ValueError(
+                "Cannot call nvecs on sptensor with only singleton dimensions"
+            )
+        tnt = mutatable_sptensor.spmatrix().transpose()
         y = tnt.transpose().dot(tnt)
         if r < y.shape[0] - 1:
             _, v = scipy.sparse.linalg.eigs(y, r)
@@ -1068,7 +1108,7 @@ def squeeze(self) -> Union[sptensor, float]:
             return ttb.sptensor.from_data(np.array([]), np.array([]), siz)
         return ttb.sptensor.from_data(self.subs[:, idx], self.vals, siz)
 
-    def subdims(self, region):
+    def subdims(self, region: Sequence[Union[int, np.ndarray, slice]]) -> np.ndarray:
         """
         SUBDIMS Compute the locations of subscripts within a subdimension.
 
@@ -1116,34 +1156,39 @@ def subdims(self, region):
         loc = np.arange(0, len(self.subs))
 
         for i in range(0, self.ndims):
-            if not isinstance(region[i], slice):
-                # Find subscripts that match in dimension i
-                tf = np.isin(self.subs[loc, i], region[i])
-
-                # Pare down the list of indices
-                loc = loc[tf]
-            else:
+            # TODO: Consider cleaner typing coercion
+            # Find subscripts that match in dimension i
+            if isinstance(region[i], (int, np.generic)):
+                tf = np.isin(self.subs[loc, i], cast(int, region[i]))
+            elif isinstance(region[i], (np.ndarray, list)):
+                tf = np.isin(self.subs[loc, i], cast(np.ndarray, region[i]))
+            elif isinstance(region[i], slice):
                 sliceRegion = range(0, self.shape[i])[region[i]]
-
                 tf = np.isin(self.subs[loc, i], sliceRegion)
+            else:
+                raise ValueError(
+                    f"Unexpected type in region sequence. "
+                    f"At index: {i} got {region[i]} with type {type(region[i])}"
+                )
+
+            # Pare down the list of indices
+            loc = loc[tf]
 
-                # Pare down the list of indices
-                loc = loc[tf]
         return loc
 
     # pylint: disable=too-many-branches, too-many-locals
-    def ttv(self, vector, dims=None):
+    def ttv(
+        self,
+        vector: Union[np.ndarray, List[np.ndarray]],
+        dims: Optional[Union[int, np.ndarray]] = None,
+    ) -> Union[sptensor, ttb.tensor]:
         """
         Sparse tensor times vector
 
         Parameters
         ----------
-        vector
-        dims
-
-        Returns
-        -------
-
+        vector: Vector(s) to multiply against
+        dims: Dimensions to multiply with vector(s)
         """
 
         if dims is None:
@@ -1154,7 +1199,7 @@ def ttv(self, vector, dims=None):
         # Check that vector is a list of vectors,
         # if not place single vector as element in list
         if len(vector) > 0 and isinstance(vector[0], (int, float, np.int_, np.float_)):
-            return self.ttv([vector], dims)
+            return self.ttv(np.array([vector]), dims)
 
         # Get sorted dims and index for multiplicands
         dims, vidx = ttb.tt_dimscheck(dims, self.ndims, len(vector))
@@ -2446,7 +2491,12 @@ def __repr__(self):  # pragma: no cover
 
     __str__ = __repr__
 
-    def ttm(self, matrices, dims=None, transpose=False):
+    def ttm(
+        self,
+        matrices: Union[np.ndarray, List[np.ndarray]],
+        dims: Optional[Union[float, np.ndarray]] = None,
+        transpose: bool = False,
+    ):
         """
         Sparse tensor times matrix.
 
@@ -2464,7 +2514,7 @@ def ttm(self, matrices, dims=None, transpose=False):
             dims = np.arange(self.ndims)
         elif isinstance(dims, list):
             dims = np.array(dims)
-        elif np.isscalar(dims) or isinstance(dims, list):
+        elif isinstance(dims, (float, int, np.generic)):
             dims = np.array([dims])
 
         # Handle list of matrices
@@ -2488,20 +2538,20 @@ def ttm(self, matrices, dims=None, transpose=False):
         # 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]
+        final_dim: int = dims[0]
 
         # Compute the product
 
         # Check that sizes match
-        if self.shape[dims] != matrices.shape[1]:
+        if self.shape[final_dim] != matrices.shape[1]:
             assert False, "Matrix shape doesn't match tensor shape"
 
         # Compute the new size
         siz = np.array(self.shape)
-        siz[dims] = matrices.shape[0]
+        siz[final_dim] = matrices.shape[0]
 
         # Compute self[mode]'
-        Xnt = tt_to_sparse_matrix(self, dims, True)
+        Xnt = tt_to_sparse_matrix(self, final_dim, True)
 
         # Reshape puts the reshaped things after the unchanged modes, transpose then
         # puts it in front
@@ -2511,7 +2561,7 @@ def ttm(self, matrices, dims=None, transpose=False):
         Z = Xnt.dot(matrices.transpose())
 
         # Rearrange back into sparse tensor of correct shape
-        Ynt = tt_from_sparse_matrix(Z, siz, dims, idx)
+        Ynt = tt_from_sparse_matrix(Z, siz, final_dim, idx)
 
         if not isinstance(Z, np.ndarray) and Z.nnz <= 0.5 * np.prod(siz):
             return Ynt
diff --git a/pyttb/tensor.py b/pyttb/tensor.py
index 61948dd3..deb686d0 100644
--- a/pyttb/tensor.py
+++ b/pyttb/tensor.py
@@ -7,7 +7,7 @@
 import warnings
 from itertools import permutations
 from math import factorial
-from typing import Any, Callable, List, Literal, Optional, Tuple, Union
+from typing import Any, Callable, List, Optional, Tuple, Union
 
 import numpy as np
 import scipy.sparse.linalg
@@ -179,9 +179,7 @@ def from_function(
     def collapse(
         self,
         dims: Optional[np.ndarray] = None,
-        fun: Union[
-            Literal["sum"], Callable[[np.ndarray], Union[float, np.ndarray]]
-        ] = "sum",
+        fun: Callable[[np.ndarray], Union[float, np.ndarray]] = np.sum,
     ) -> Union[float, np.ndarray, tensor]:
         """
         Collapse tensor along specified dimensions.
@@ -200,7 +198,7 @@ def collapse(
         >>> X = ttb.tensor.from_data(np.ones((2,2)))
         >>> X.collapse()
         4.0
-        >>> X.collapse(np.arange(X.ndims), lambda values: sum(values))
+        >>> X.collapse(np.arange(X.ndims), sum)
         4.0
         """
         if self.data.size == 0:
@@ -217,8 +215,6 @@ def collapse(
 
         # Check for the case where we accumulate over *all* dimensions
         if remdims.size == 0:
-            if fun == "sum":
-                return sum(self.data.flatten("F"))
             return fun(self.data.flatten("F"))
 
         ## Calculate the shape of the result
@@ -230,10 +226,7 @@ def collapse(
         ## Apply the collapse function
         B = np.zeros((A.shape[0], 1))
         for i in range(0, A.shape[0]):
-            if fun == "sum":
-                B[i] = np.sum(A[i, :])
-            else:
-                B[i] = fun(A[i, :])
+            B[i] = fun(A[i, :])
 
         ## Form and return the final result
         return ttb.tensor.from_data(B, newshape)

From 3410d15e767ad8b85ceae3709c27dd19655f094a Mon Sep 17 00:00:00 2001
From: Nick Johnson <24689722+ntjohnson1@users.noreply.github.com>
Date: Sat, 18 Mar 2023 17:35:34 -0400
Subject: [PATCH 6/7] SPTENSOR: Fix test coverage from typing changes.

---
 tests/test_sptensor.py | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/tests/test_sptensor.py b/tests/test_sptensor.py
index 16125371..4d70ebc8 100644
--- a/tests/test_sptensor.py
+++ b/tests/test_sptensor.py
@@ -231,6 +231,9 @@ def test_sptensor_subdims(sample_sptensor):
         sptensorInstance.subdims([[1], [1, 3]])
     assert "Number of subdimensions must equal number of dimensions" in str(excinfo)
 
+    with pytest.raises(ValueError):
+        sptensorInstance.subdims(("bad", "region", "types"))
+
 
 @pytest.mark.indevelopment
 def test_sptensor_ndims(sample_sptensor):
@@ -1633,6 +1636,13 @@ def test_sptensor_nvecs(sample_sptensor):
         in str(record[0].message)
     )
 
+    # Negative test, check for only singleton dims
+    with pytest.raises(ValueError):
+        single_val_sptensor = ttb.sptensor.from_data(
+            np.array([[0, 0]]), np.array([1]), shape=(1, 1)
+        )
+        single_val_sptensor.nvecs(0, 0)
+
 
 @pytest.mark.indevelopment
 def test_sptensor_ttm(sample_sptensor):

From 7b75dd185d07dee8926afb405cbef5deee126edb Mon Sep 17 00:00:00 2001
From: Nick Johnson <24689722+ntjohnson1@users.noreply.github.com>
Date: Sat, 18 Mar 2023 17:58:19 -0400
Subject: [PATCH 7/7] PYLINT: Update test to lint files in parallel to improve
 dev experience.

---
 tests/test_package.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/tests/test_package.py b/tests/test_package.py
index 8d8ca619..a5aa95e0 100644
--- a/tests/test_package.py
+++ b/tests/test_package.py
@@ -30,8 +30,9 @@ def test_linting():
     # add mypy check
     root_dir = os.path.dirname(os.path.dirname(__file__))
     toml_file = os.path.join(root_dir, "pyproject.toml")
-    for a_file in enforced_files:
-        subprocess.run(f"pylint {a_file} --rcfile {toml_file}", check=True)
+    subprocess.run(
+        f"pylint {' '.join(enforced_files)} --rcfile {toml_file} -j0", check=True
+    )
 
 
 @pytest.mark.packaging