Nick/pyttb utils improvement (#353)

* PYTTB_UTILS: Finish adding types without refactor

* PYTTB_UTILS: Fix typing on tt_tenfun

* PYTTB_UTILS: Doc cleanup
* Remove unneeded method
* Update examples/docs most places
* Mark remaining improvements

* PYTTB_UTILS: Doc cleanup
* Update doc strings for remaining utils

* Make tt_tenfun a tensor method to more closely align to elemfun and MATLAB

* RUFF: Enable Pydocstyle
* Mostly adding periods, newlines, and imperitive phrasing

* RUFF: forgot to run format after doc updates

* Fix merge errors and capture stricter requirements on main.

* Fix type hint for more flexible tensor collapse.

* TTENSOR: Fix missing newline and indentation.

Author: ntjohnson1

conftest.py

@@ -1,3 +1,4 @@
+"""Pyttb pytest configuration."""
 # Copyright 2024 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.
@@ -11,12 +12,12 @@
 
 
 @pytest.fixture(autouse=True)
-def add_packages(doctest_namespace):
+def add_packages(doctest_namespace):  # noqa: D103
     doctest_namespace["np"] = numpy
     doctest_namespace["ttb"] = pyttb
 
 
-def pytest_addoption(parser):
+def pytest_addoption(parser):  # noqa: D103
     parser.addoption(
         "--packaging",
         action="store_true",
@@ -26,6 +27,6 @@ def pytest_addoption(parser):
     )
 
 
-def pytest_configure(config):
+def pytest_configure(config):  # noqa: D103
     if not config.option.packaging:
         config.option.markexpr = "not packaging"

docs/source/matlab/common.rst

@@ -56,7 +56,7 @@ Methods
 +-----------------+----------------------+------------------------------------------------------------------------+
 | ``subsref``     | ``__getitem__``      | ``X[index]``                                                           |
 +-----------------+----------------------+------------------------------------------------------------------------+
-| ``tenfun``      | ``tt_tenfun``        | e.g., ``pyttb.tt_tenfun(lambda x: x + 1, A)``                          |
+| ``tenfun``      | ``tenfun``           | ``X.tenfun(lambda x: x + 1)``                                          |
 +-----------------+----------------------+------------------------------------------------------------------------+
 | ``times``       | ``__mul__``          | ``X * Y``                                                              |
 +-----------------+----------------------+------------------------------------------------------------------------+

docs/source/tutorial/algorithm_gcp_opt.ipynb

@@ -158,7 +158,6 @@
     "import sys\n",
     "import pyttb as ttb\n",
     "import numpy as np\n",
-    "from pyttb.pyttb_utils import tt_tenfun\n",
     "\n",
     "from pyttb.gcp.fg_setup import function_type, setup\n",
     "from pyttb.gcp.handles import Objectives\n",

docs/source/tutorial/class_tensor.ipynb

@@ -27,8 +27,7 @@
    "source": [
     "import pyttb as ttb\n",
     "import numpy as np\n",
-    "import sys\n",
-    "from pyttb.pyttb_utils import tt_tenfun"
+    "import sys"
    ]
   },
   {
@@ -847,8 +846,8 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## Using `tt_tenfun` for elementwise operations on one or more `tensor`s\n",
-    "The function `tt_tenfun` applies a specified function to a number of `tensor`s. This can be used for any function that is not predefined for `tensor`s."
+    "## Using `tenfun` for elementwise operations on one or more `tensor`s\n",
+    "The method `tenfun` applies a specified function to a number of `tensor`s. This can be used for any function that is not predefined for `tensor`s."
    ]
   },
   {
@@ -859,7 +858,7 @@
    "source": [
     "np.random.seed(0)\n",
     "A = ttb.tensor(np.floor(3 * np.random.rand(2, 2, 3)))  # Generate some data.\n",
-    "tt_tenfun(lambda x: x + 1, A)  # Increment every element of A by one."
+    "A.tenfun(lambda x: x + 1)  # Increment every element of A by one."
    ]
   },
   {
@@ -873,7 +872,7 @@
     "    return np.maximum(a, b)\n",
     "\n",
     "\n",
-    "tt_tenfun(max_elements, A, B)  # Max of A and B, elementwise."
+    "A.tenfun(max_elements, B)  # Max of A and B, elementwise."
    ]
   },
   {
@@ -891,7 +890,7 @@
     "    return np.floor(np.mean(X, axis=0))\n",
     "\n",
     "\n",
-    "tt_tenfun(elementwise_mean, A, B, C)  # Elementwise means for A, B, and C."
+    "A.tenfun(elementwise_mean, B, C)  # Elementwise means for A, B, and C."
    ]
   },
   {

pyproject.toml

@@ -67,7 +67,7 @@ requires = ["setuptools>=61.0", "numpy", "numpy_groupies", "scipy", "wheel"]
 build-backend = "setuptools.build_meta"
 
 [tool.ruff.lint]
-select = ["E", "F", "PL", "W", "I", "N", "NPY", "RUF", "B"]
+select = ["E", "F", "PL", "W", "I", "N", "NPY", "RUF", "B", "D"]
 ignore = [
     # Ignored in conversion to ruff since not previously enforced
     "PLR2004",
@@ -84,15 +84,18 @@ ignore = [
     # There is ongoing discussion about logging/warning etc
     "B028",
 ]
+[tool.ruff.lint.pydocstyle]
+convention = "numpy"
+
 [tool.ruff.lint.per-file-ignores]
 # See  see https://github.com/astral-sh/ruff/issues/3172 for details on this becoming simpler
 
 # Everything but I, F (to catch import mess and potential logic errors)
-"tests/**.py" = ["E", "PL", "W", "N", "NPY", "RUF", "B"]
+"tests/**.py" = ["E", "PL", "W", "N", "NPY", "RUF", "B", "D"]
 # Ignore everything for now
-"docs/**.py" = ["E", "F", "PL", "W", "I", "N", "NPY", "RUF", "B"]
-"docs/**.ipynb" = ["E", "F", "PL", "W", "I", "N", "NPY", "RUF", "B"]
-"profiling/**.ipynb" = ["E", "F", "PL", "W", "I", "N", "NPY", "RUF", "B"]
+"docs/**.py" = ["E", "F", "PL", "W", "I", "N", "NPY", "RUF", "B", "D"]
+"docs/**.ipynb" = ["E", "F", "PL", "W", "I", "N", "NPY", "RUF", "B", "D"]
+"profiling/**.ipynb" = ["E", "F", "PL", "W", "I", "N", "NPY", "RUF", "B", "D"]
 
 [tool.ruff.format]
 docstring-code-format = true

pyttb/__init__.py

@@ -1,4 +1,4 @@
-"""pyttb: Python Tensor Toolbox"""
+"""pyttb: Python Tensor Toolbox."""
 
 # Copyright 2024 National Technology & Engineering Solutions of Sandia,
 # LLC (NTESS). Under the terms of Contract DE-NA0003525 with NTESS, the
@@ -32,7 +32,7 @@
 
 
 def ignore_warnings(ignore=True):
-    """Helper to disable warnings"""
+    """Disable warnings."""
     if ignore:
         warnings.simplefilter("ignore")
     else:

pyttb/cp_als.py

@@ -1,4 +1,4 @@
-"""CP Decomposition via Alternating Least Squares"""
+"""CP Decomposition via Alternating Least Squares."""
 
 # Copyright 2024 National Technology & Engineering Solutions of Sandia,
 # LLC (NTESS). Under the terms of Contract DE-NA0003525 with NTESS, the
@@ -25,8 +25,7 @@ def cp_als(  # noqa: PLR0912,PLR0913,PLR0915
     printitn: int = 1,
     fixsigns: bool = True,
 ) -> Tuple[ttb.ktensor, ttb.ktensor, Dict]:
-    """
-    Compute CP decomposition with alternating least squares
+    """Compute CP decomposition with alternating least squares.
 
     Parameters
     ----------
@@ -129,7 +128,6 @@ def cp_als(  # noqa: PLR0912,PLR0913,PLR0915
      Iter 1: f = ... f-delta = ...
      Final f = ...
     """
-
     # Extract number of dimensions and norm of tensor
     N = input_tensor.ndims
     normX = input_tensor.norm()

pyttb/cp_apr.py

@@ -1,4 +1,4 @@
-"""Non-negative CP decomposition with alternating Poisson regression"""
+"""Non-negative CP decomposition with alternating Poisson regression."""
 
 # Copyright 2024 National Technology & Engineering Solutions of Sandia,
 # LLC (NTESS). Under the terms of Contract DE-NA0003525 with NTESS, the
@@ -230,9 +230,6 @@ def tt_cp_apr_mu(  # noqa: PLR0912,PLR0913,PLR0915
     kappatol:
         MU ALGORITHM PARAMETER: Tolerance on complementary slackness
 
-    Returns
-    -------
-
     Notes
     -----
     REFERENCE: E. C. Chi and T. G. Kolda. On Tensors, Sparsity, and
@@ -405,9 +402,9 @@ def tt_cp_apr_pdnr(  # noqa: PLR0912,PLR0913,PLR0915
     precompinds: bool,
     inexact: bool,
 ) -> Tuple[ttb.ktensor, Dict]:
-    """
-    Compute nonnegative CP with alternating Poisson regression
-    computes an estimate of the best rank-R
+    """Compute nonnegative CP with alternating Poisson regression.
+
+    Computes an estimate of the best rank-R
     CP model of a tensor X using an alternating Poisson regression.
     The algorithm solves "row subproblems" in each alternating subproblem,
     using a Hessian of size R^2.
@@ -1272,9 +1269,7 @@ def get_search_dir_pdnr(  # noqa: PLR0913
     mu: float,
     epsActSet: float,
 ) -> Tuple[np.ndarray, np.ndarray]:
-    """
-    Compute the search direction for PDNR using a two-metric projection with
-    damped Hessian
+    """Compute the search direction using a two-metric projection with damped Hessian.
 
     Parameters
     ----------
@@ -1372,8 +1367,7 @@ def tt_linesearch_prowsubprob(  # noqa: PLR0913
     phi_row: np.ndarray,
     display_warning: bool,
 ) -> Tuple[np.ndarray, float, float, float, int]:
-    """
-    Perform a line search on a row subproblem
+    """Perform a line search on a row subproblem.
 
     Parameters
     ----------
@@ -1488,9 +1482,9 @@ def tt_linesearch_prowsubprob(  # noqa: PLR0913
 def get_hessian(
     upsilon: np.ndarray, Pi: np.ndarray, free_indices: np.ndarray
 ) -> np.ndarray:
-    """
-    Return the Hessian for one PDNR row subproblem of Model[n], for just the rows and
-    columns corresponding to the free variables
+    """Return the Hessian for one PDNR row subproblem of Model[n].
+
+    Only for just the rows and columns corresponding to the free variables.
 
     Parameters
     ----------
@@ -1505,7 +1499,6 @@ def get_hessian(
         Sub-block of full Hessian identified by free-indices
 
     """
-
     num_free = len(free_indices)
     H = np.zeros((num_free, num_free))
     for i in range(num_free):
@@ -1523,8 +1516,7 @@ def tt_loglikelihood_row(
     model_row: np.ndarray,
     Pi: np.ndarray,
 ) -> float:
-    """
-    Compute log-likelihood of one row subproblem
+    """Compute log-likelihood of one row subproblem.
 
     Parameters
     ----------
@@ -1618,7 +1610,6 @@ def get_search_dir_pqnr(  # noqa: PLR0913
     URL: http://arxiv.org/abs/1304.4964. Submitted for publication.
 
     """
-
     lbfgsSize = delta_model.shape[1]
 
     # Determine active and free variables.
@@ -1679,8 +1670,7 @@ def calc_grad(
     data_row: np.ndarray,
     model_row: np.ndarray,
 ) -> Tuple[np.ndarray, np.ndarray]:
-    """
-    Compute the gradient for a PQNR row subproblem
+    """Compute the gradient for a PQNR row subproblem.
 
     Parameters
     ----------
@@ -1710,6 +1700,7 @@ def calc_grad(
     return grad_row, phi_row
 
 
+# TODO verify what pi is
 # Mu helper functions
 def calculate_pi(
     Data: Union[ttb.sptensor, ttb.tensor],
@@ -1718,9 +1709,7 @@ def calculate_pi(
     factorIndex: int,
     ndims: int,
 ) -> np.ndarray:
-    """
-    Helper function to calculate Pi matrix
-    # TODO verify what pi is
+    """Calculate Pi matrix.
 
     Parameters
     ----------
@@ -1758,7 +1747,7 @@ def calculate_phi(  # noqa: PLR0913
     Pi: np.ndarray,
     epsilon: float,
 ) -> np.ndarray:
-    """
+    """Calcualte Phi.
 
     Parameters
     ----------
@@ -1769,9 +1758,6 @@ def calculate_phi(  # noqa: PLR0913
     Pi:
     epsilon:
 
-    Returns
-    -------
-
     """
     if isinstance(Data, ttb.sptensor):
         Phi = -np.ones((Data.shape[factorIndex], rank))
@@ -1846,8 +1832,7 @@ def tt_loglikelihood(
 
 
 def vectorize_for_mu(matrix: np.ndarray) -> np.ndarray:
-    """
-    Helper Function to unravel matrix into vector
+    """Unravel matrix into vector.
 
     Parameters
     ----------

pyttb/export_data.py

@@ -1,4 +1,4 @@
-"""Utilities for saving tensor data"""
+"""Utilities for saving tensor data."""
 
 # Copyright 2024 National Technology & Engineering Solutions of Sandia,
 # LLC (NTESS). Under the terms of Contract DE-NA0003525 with NTESS, the
@@ -20,9 +20,7 @@ def export_data(
     fmt_data: Optional[str] = None,
     fmt_weights: Optional[str] = None,
 ):
-    """
-    Export tensor-related data to a file.
-    """
+    """Export tensor-related data to a file."""
     if not isinstance(data, (ttb.tensor, ttb.sptensor, ttb.ktensor, np.ndarray)):
         assert False, f"Invalid data type for export: {type(data)}"
 
@@ -58,36 +56,36 @@ def export_data(
 
 
 def export_size(fp: TextIO, shape: Shape):
-    """Export the size of something to a file"""
+    """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
 
 
 def export_rank(fp: TextIO, data: ttb.ktensor):
-    """Export the rank of a ktensor to a file"""
+    """Export the rank of a ktensor to a file."""
     print(f"{len(data.weights)}", file=fp)  # ktensor rank on one line
 
 
 def export_weights(fp: TextIO, data: ttb.ktensor, fmt_weights: Optional[str]):
-    """Export KTensor weights"""
+    """Export KTensor weights."""
     if not fmt_weights:
         fmt_weights = "%.16e"
     data.weights.tofile(fp, sep=" ", format=fmt_weights)
     print(file=fp)
 
 
 def export_array(fp: TextIO, data: np.ndarray, fmt_data: Optional[str]):
-    """Export dense data"""
+    """Export dense data."""
     if not fmt_data:
         fmt_data = "%.16e"
     data.tofile(fp, sep="\n", format=fmt_data)
     print(file=fp)
 
 
 def export_factor(fp: TextIO, data: np.ndarray, fmt_data: Optional[str]):
-    """Export KTensor factor"""
+    """Export KTensor factor."""
     if not fmt_data:
         fmt_data = "%.16e"
     for i in range(data.shape[0]):
@@ -97,15 +95,15 @@ def export_factor(fp: TextIO, data: np.ndarray, fmt_data: Optional[str]):
 
 
 def export_sparse_size(fp: TextIO, A: ttb.sptensor):
-    """Export the size of something to a file"""
+    """Export the size of something to a file."""
     print(f"{len(A.shape)}", file=fp)  # # of dimensions on one line
     shape_str = " ".join([str(d) for d in A.shape])
     print(f"{shape_str}", file=fp)  # size of each dimensions on the next line
     print(f"{A.nnz}", file=fp)  # number of nonzeros
 
 
 def export_sparse_array(fp: TextIO, A: ttb.sptensor, fmt_data: Optional[str]):
-    """Export sparse array data in coordinate format"""
+    """Export sparse array data in coordinate format."""
     if not fmt_data:
         fmt_data = "%.16e"
     # TODO: looping through all values may take a long time, can this be more efficient?

pyttb/gcp/__init__.py

@@ -0,0 +1 @@
+"""Generalized CP Decomposition Support Code."""