Nick/shape support

docs/source/tutorial/class_sptensor.ipynb

@@ -271,7 +271,7 @@
    "outputs": [],
    "source": [
     "indices = np.array([[0, 0, 0], [1, 0, 0]])\n",
-    "values = np.ones((2,))\n",
+    "values = np.ones((2, 1))\n",
     "\n",
     "Y = ttb.sptensor.from_aggregator(indices, values)  # Create a sparse tensor.\n",
     "Y"
@@ -300,7 +300,7 @@
    "outputs": [],
    "source": [
     "indices = np.array([[0, 0, 0], [2, 2, 2]])\n",
-    "values = np.array([1, 3])\n",
+    "values = np.array([[1], [3]])\n",
     "\n",
     "Y = ttb.sptensor.from_aggregator(indices, values)  # Create a sparse tensor.\n",
     "Y"
@@ -555,7 +555,7 @@
    "outputs": [],
    "source": [
     "indices = np.array([[0], [2], [4]])\n",
-    "values = np.array([1, 1, 1])\n",
+    "values = np.array([[1], [1], [1]])\n",
     "X = ttb.sptensor.from_aggregator(indices, values)\n",
     "X"
    ]

pyttb/cp_als.py

@@ -6,20 +6,21 @@
 
 from __future__ import annotations
 
-from typing import Dict, List, Literal, Optional, Tuple, Union
+from typing import Dict, Literal, Optional, Tuple, Union
 
 import numpy as np
 
 import pyttb as ttb
+from pyttb.pyttb_utils import OneDArray, parse_one_d
 
 
 def cp_als(  # noqa: PLR0912,PLR0913,PLR0915
     input_tensor: Union[ttb.tensor, ttb.sptensor, ttb.ttensor, ttb.sumtensor],
     rank: int,
     stoptol: float = 1e-4,
     maxiters: int = 1000,
-    dimorder: Optional[List[int]] = None,
-    optdims: Optional[List[int]] = None,
+    dimorder: Optional[OneDArray] = None,
+    optdims: Optional[OneDArray] = None,
     init: Union[Literal["random"], Literal["nvecs"], ttb.ktensor] = "random",
     printitn: int = 1,
     fixsigns: bool = True,
@@ -109,8 +110,8 @@ def cp_als(  # noqa: PLR0912,PLR0913,PLR0915
     [[0.1467... 0.0923...]
      [0.1862... 0.3455...]]
     >>> print(output["params"]) # doctest: +NORMALIZE_WHITESPACE
-    {'stoptol': 0.0001, 'maxiters': 1000, 'dimorder': [0, 1],\
-     'optdims': [0, 1], 'printitn': 1, 'fixsigns': True}
+    {'stoptol': 0.0001, 'maxiters': 1000, 'dimorder': array([0, 1]),\
+     'optdims': array([0, 1]), 'printitn': 1, 'fixsigns': True}
 
     Example using "nvecs" initialization:
 
@@ -135,15 +136,17 @@ def cp_als(  # noqa: PLR0912,PLR0913,PLR0915
 
     # Set up dimorder if not specified
     if dimorder is None:
-        dimorder = list(range(N))
-    elif not isinstance(dimorder, list):
-        assert False, "Dimorder must be a list"
-    elif tuple(range(N)) != tuple(sorted(dimorder)):
+        dimorder = np.arange(N)
+    else:
+        dimorder = parse_one_d(dimorder)
+    if tuple(range(N)) != tuple(sorted(dimorder)):
         assert False, "Dimorder must be a list or permutation of range(tensor.ndims)"
 
     # Set up optdims if not specified
     if optdims is None:
-        optdims = list(range(N))
+        optdims = np.arange(N)
+    else:
+        optdims = parse_one_d(optdims)
 
     # Error checking
     assert rank > 0, "Number of components requested must be positive"

pyttb/export_data.py

@@ -6,11 +6,12 @@
 
 from __future__ import annotations
 
-from typing import Optional, TextIO, Tuple, Union
+from typing import Optional, TextIO, Union
 
 import numpy as np
 
 import pyttb as ttb
+from pyttb.pyttb_utils import Shape, parse_shape
 
 
 def export_data(
@@ -56,8 +57,9 @@ def export_data(
             export_array(fp, data, fmt_data)
 
 
-def export_size(fp: TextIO, shape: Tuple[int, ...]):
+def export_size(fp: TextIO, shape: Shape):
     """Export the size of something to a file"""
+    shape = parse_shape(shape)
     print(f"{len(shape)}", file=fp)  # # of dimensions on one line
     shape_str = " ".join([str(d) for d in shape])
     print(f"{shape_str}", file=fp)  # size of each dimensions on the next line

pyttb/gcp_opt.py

@@ -8,7 +8,8 @@
 
 import logging
 import time
-from typing import Dict, List, Literal, Optional, Tuple, Union
+from math import prod
+from typing import Dict, Literal, Optional, Sequence, Tuple, Union
 
 import numpy as np
 
@@ -24,7 +25,7 @@ def gcp_opt(  # noqa:  PLR0912,PLR0913
     rank: int,
     objective: Union[Objectives, Tuple[function_type, function_type, float]],
     optimizer: Union[StochasticSolver, LBFGSB],
-    init: Union[Literal["random"], ttb.ktensor, List[np.ndarray]] = "random",
+    init: Union[Literal["random"], ttb.ktensor, Sequence[np.ndarray]] = "random",
     mask: Optional[Union[ttb.tensor, np.ndarray]] = None,
     sampler: Optional[GCPSampler] = None,
     printitn: int = 1,
@@ -74,7 +75,7 @@ def gcp_opt(  # noqa:  PLR0912,PLR0913
     if not isinstance(data, (ttb.tensor, ttb.sptensor)):
         raise ValueError("Input data must be tensor or sptensor.")
 
-    tensor_size = int(np.prod(data.shape))
+    tensor_size = prod(data.shape)
 
     if isinstance(data, ttb.tensor) and isinstance(mask, ttb.tensor):
         data *= mask
@@ -134,7 +135,7 @@ def gcp_opt(  # noqa:  PLR0912,PLR0913
 def _get_initial_guess(
     data: Union[ttb.tensor, ttb.sptensor],
     rank: int,
-    init: Union[Literal["random"], ttb.ktensor, List[np.ndarray]],
+    init: Union[Literal["random"], ttb.ktensor, Sequence[np.ndarray]],
 ) -> ttb.ktensor:
     """Get initial guess for gcp_opt
 
@@ -143,7 +144,7 @@ def _get_initial_guess(
         Normalized ktensor.
     """
     # TODO might be nice to merge with ALS/other CP methods
-    if isinstance(init, list):
+    if isinstance(init, Sequence) and not isinstance(init, str):
         return ttb.ktensor(init).normalize("all")
     if isinstance(init, ttb.ktensor):
         init.normalize("all")

pyttb/hosvd.py

@@ -7,21 +7,22 @@
 from __future__ import annotations
 
 import warnings
-from typing import List, Optional
+from typing import Optional
 
 import numpy as np
 import scipy
 
 import pyttb as ttb
+from pyttb.pyttb_utils import OneDArray, parse_one_d
 
 
 def hosvd(  # noqa: PLR0912,PLR0913,PLR0915
     input_tensor: ttb.tensor,
     tol: float,
     verbosity: float = 1,
-    dimorder: Optional[List[int]] = None,
+    dimorder: Optional[OneDArray] = None,
     sequential: bool = True,
-    ranks: Optional[List[int]] = None,
+    ranks: Optional[OneDArray] = None,
 ) -> ttb.ttensor:
     """Compute sequentially-truncated higher-order SVD (Tucker).
 
@@ -57,21 +58,22 @@ def hosvd(  # noqa: PLR0912,PLR0913,PLR0915
     # In tucker als this is N
     d = input_tensor.ndims
 
-    if ranks is not None:
-        if len(ranks) != d:
-            raise ValueError(
-                f"Ranks must be a list of length tensor ndims. Ndims: {d} but got "
-                f"ranks: {ranks}."
-            )
+    if ranks is None:
+        ranks = np.zeros((d,), dtype=int)
     else:
-        ranks = [0] * d
+        ranks = parse_one_d(ranks)
+
+    if len(ranks) != d:
+        raise ValueError(
+            "Ranks must be a sequence of length tensor ndims."
+            f" Ndims: {d} but got ranks: {ranks}."
+        )
 
     # Set up dimorder if not specified (this is copy past from tucker_als
-    if not dimorder:
-        dimorder = list(range(d))
+    if dimorder is None:
+        dimorder = np.arange(d)
     else:
-        if not isinstance(dimorder, list):
-            raise ValueError("Dimorder must be a list")
+        dimorder = parse_one_d(dimorder)
         if tuple(range(d)) != tuple(sorted(dimorder)):
             raise ValueError(
                 "Dimorder must be a list or permutation of range(tensor.ndims)"