Nick/sum tensor

docs/source/sumtensor.rst

@@ -1,2 +1,8 @@
 pyttb.sumtensor
-=======================
+-----------------------
+
+.. autoclass:: pyttb.sumtensor
+    :members:
+    :special-members:
+    :exclude-members: __dict__, __weakref__, __slots__, __init__
+    :show-inheritance:

docs/source/tensor_classes.rst

@@ -6,6 +6,7 @@ Tensor Classes
 
     ktensor.rst
     sptensor.rst
+    sumtensor.rst
     tensor.rst
     ttensor.rst
     tenmat.rst

docs/source/tutorial/class_sumtensor.ipynb

@@ -0,0 +1,346 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "a2683b4e",
+   "metadata": {},
+   "source": [
+    "# Sum of Structured Tensors\n",
+    "```\n",
+    "Copyright 2022 National Technology & Engineering Solutions of Sandia,\n",
+    "LLC (NTESS). Under the terms of Contract DE-NA0003525 with NTESS, the\n",
+    "U.S. Government retains certain rights in this software.\n",
+    "```"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "6c96f0b4",
+   "metadata": {},
+   "source": [
+    "When certain operations are performed on a tensor which is formed as a sum of tensors, it can be beneficial to avoid explicitly forming the sum. For example, if a tensor is formed as a sum of a low rank tensor and a sparse tensor, the structure of the summands can make storage, decomposition and operations with other tensors significantly more efficient. A `sumtensor` exploits this structure. Here we explain the basics of defining and using sumtensors."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8c2ad211",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import pyttb as ttb\n",
+    "import numpy as np\n",
+    "import pickle"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a808d286",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def estimate_mem_usage(variable) -> int:\n",
+    "    \"\"\"\n",
+    "    Python variables contain references to memory.\n",
+    "    Quickly estimate memory usage of custom types.\n",
+    "    \"\"\"\n",
+    "    return len(pickle.dumps(variable))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d169e48b",
+   "metadata": {},
+   "source": [
+    "## Creating sumtensors\n",
+    "A sumtensor `T` can only be delared as a sum of same-shaped tensors T1, T2,...,TN. The summand tensors are stored internally, which define the \"parts\" of the `sumtensor`. The parts of a `sumtensor` can be (dense) tensors (`tensor`), sparse tensors (` sptensor`), Kruskal tensors (`ktensor`), or Tucker tensors (`ttensor`). An example of the use of the sumtensor constructor follows."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "75797e16",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "T1 = ttb.tenones((3, 3, 3))\n",
+    "T2 = ttb.sptensor(\n",
+    "    subs=np.array([[0, 0, 0], [1, 1, 1], [2, 2, 1], [1, 0, 0]]),\n",
+    "    vals=np.ones((4, 1)),\n",
+    "    shape=(3, 3, 3),\n",
+    ")\n",
+    "T = ttb.sumtensor([T1, T2])\n",
+    "print(T)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "0a17f6a8",
+   "metadata": {},
+   "source": [
+    "## A Magnitude Example\n",
+    "For large-scale problems, the `sumtensor` class may make the difference as to whether or not a tensor can be stored in memory. Consider the following example, where $\\mathcal{T}$ is of size $1000 x 1000 x 1000$, formed from the sum of a `ktensor` and an `sptensor`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1215e3d2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "np.random.seed(0)\n",
+    "X1 = np.random.rand(50, 3)\n",
+    "X2 = np.random.rand(50, 3)\n",
+    "X3 = np.random.rand(50, 3)\n",
+    "K = ttb.ktensor([X1, X2, X3], np.ones((3,)), copy=False)\n",
+    "S = ttb.sptenrand((50, 50, 50), 1e-100)\n",
+    "\n",
+    "ST = ttb.sumtensor([K, S])\n",
+    "TT = ST.full()\n",
+    "print(\n",
+    "    f\"Size of sumtensor: {estimate_mem_usage(ST)}\\n\"\n",
+    "    f\"Size of tensor: {estimate_mem_usage(TT)}\"\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "9650dccc",
+   "metadata": {},
+   "source": [
+    "## Further examples of the sumtensor constructor\n",
+    "We can declare an empty sumtensor, with no parts."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "36ea6eb5",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "P = ttb.sumtensor()\n",
+    "print(P)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "346374a7",
+   "metadata": {},
+   "source": [
+    "`sumtensor` also supports a copy constructor."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5f309064",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "S = P.copy()\n",
+    "print(S)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c28b7cc6",
+   "metadata": {},
+   "source": [
+    "## Ndims and shape for dimensions of a sumtensor\n",
+    "For a given `sumtensor`, `ndims` returns the number of modes and `shape` returns the shape in each dimension of the `sumtensor`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8e484181",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "print(f\"Ndims: {T.ndims}\\n\" f\"Shape: {T.shape}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ec82fc57",
+   "metadata": {},
+   "source": [
+    "## Use full to convert sumtensor to a dense tensor\n",
+    "The `full` method can convert all the parts of a `sumtensor` to a dense tensor. Note that for large tensors, this can use a large amount of memory to expand then sum the parts."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d247884d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "print(T.full())"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b1cfbffa",
+   "metadata": {},
+   "source": [
+    "## Use double to convert to a numpy array\n",
+    "The `double` method can convert the parts of a `sumtensor` to a dense numpy array. Similar warnings for memory usages as `full`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3be2ba03",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "print(T.double())"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "bd9f0550",
+   "metadata": {},
+   "source": [
+    "## Matricized Khatri-Rao product of a sumtensor\n",
+    "The `mttkrp` method computes the Khatri-Rao product of a matricized tensor and `sumtensor`. The required arguments are:\n",
+    "* A list of matrices (or a `ktensor`)\n",
+    "* A mode n\n",
+    "\n",
+    "The list of matrices must consist of m matrices, where m is the number of modes in the `sumtensor`. The number of columns in all matrices should be the same and the number of rows of matrix i should match the dimension of the `sumtensor` shape in mode i. For more details see the documentation of `tensor.mttkrp`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3dd198fa",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "matrices = [np.eye(3), np.ones((3, 3)), np.random.rand(3, 3)]\n",
+    "n = 1\n",
+    "T.mttkrp(matrices, n)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "dd35fc03",
+   "metadata": {},
+   "source": [
+    "## Innerproducts of sumtensors\n",
+    "The `innerprod` method computes the inner product of a `sumtensors` parts with other tensor types."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "dfa6d728",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "S = ttb.sptensor(\n",
+    "    subs=np.array([[0, 0, 0], [1, 1, 1], [2, 2, 1], [1, 0, 0]]),\n",
+    "    vals=np.ones((4, 1)),\n",
+    "    shape=(3, 3, 3),\n",
+    ")\n",
+    "T.innerprod(S)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "802902bb",
+   "metadata": {},
+   "source": [
+    "## Norm compatibility interface\n",
+    "The `norm` method just returns 0 and issues a warning. Norm cannot be distributed, but some algorithms access the norm for verbose details."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f945dea9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "T.norm()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "4647329a",
+   "metadata": {},
+   "source": [
+    "## Use CP-ALS with sumtensor\n",
+    "One of the primary motivations for defining the `sumtensor` class is for efficient decomposition. In particular, when trying to find a CP decomposition of a tensor using alternating least squares, the subproblems can be efficiently created and solved using mttkrp and innerprod. Both of these operations can be performed more efficiently by exploiting extra structure in the tensors which form the sum, so the performance of `cp_als` is also improved. Consider the following example, where a cp_als is run on a sumtensor."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8d980252",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "result, _, _ = ttb.cp_als(T, 2, maxiters=10)\n",
+    "print(result)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b93b33f9",
+   "metadata": {},
+   "source": [
+    "It follows that in cases where $\\mathcal{T}$ is too large for its full expansion to be stored in memory, we may still be able find a CP decomposition by exploiting the sumtensor structure.\n",
+    "\n",
+    "_Note_ that the fit returned by `cp_als` is not correct for `sumtensor`, because the norm operation is not supported."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3b3793eb",
+   "metadata": {},
+   "source": [
+    "## Addition with sumtensors\n",
+    "Sumtensors can be added to any other type of tensor. The result is a new `sumtensor` with the tensor appended to the parts of the original `sumtensor`. Note that the tensor is always appended, despite the order of the operation."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3514aa4b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Equivalent additions despite the order\n",
+    "print(f\"T+S:\\n{T+S}\\n\")\n",
+    "print(f\"S+T:\\n{S+T}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "bd401c2d",
+   "metadata": {},
+   "source": [
+    "## Accessing sumtensor parts\n",
+    "Subscripted reference can be used to access individual parts of the `sumtensor`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "47938e4f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "print(f\"Part 0:\\n{T.parts[0]}\\n\\n\" f\"Part 1:\\n{T.parts[1]}\")"
+   ]
+  }
+ ],
+ "metadata": {},
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

pyttb/cp_als.py

@@ -12,7 +12,7 @@
 
 
 def cp_als(  # noqa: PLR0912,PLR0913,PLR0915
-    input_tensor: Union[ttb.tensor, ttb.sptensor, ttb.ttensor],
+    input_tensor: Union[ttb.tensor, ttb.sptensor, ttb.ttensor, ttb.sumtensor],
     rank: int,
     stoptol: float = 1e-4,
     maxiters: int = 1000,
@@ -156,6 +156,9 @@ def cp_als(  # noqa: PLR0912,PLR0913,PLR0915
             )
         init = ttb.ktensor(factor_matrices)
     elif isinstance(init, str) and init.lower() == "nvecs":
+        assert not isinstance(
+            input_tensor, ttb.sumtensor
+        ), "Sumtensor doesn't support nvecs"
         factor_matrices = []
         for n in range(N):
             factor_matrices.append(input_tensor.nvecs(n, rank))