{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Kruskal 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",
   "metadata": {},
   "source": [
    "Kruskal format is a decomposition of a tensor $\\mathcal{X}$ as the sum of the outer products as the columns of matrices. For example, we might write:\n",
    "\n",
    "$\n",
    "\\mathcal{X} = \\sum_{r} a_r \\circ b_r \\circ c_r\n",
    "$\n",
    "\n",
    "where a subscript denotes column index and a circle denotes outer product. In other words, the tensor $\\mathcal{X}$ is built from the columns of the matrices $A$, $B$, and $C$. It's often helpful to explicitly specify a weight for each outer product, which we do here:\n",
    "\n",
    "$\n",
    "\\mathcal{X} = \\sum_{r} \\lambda_r \\, a_r \\circ b_r \\circ c_r\n",
    "$\n",
    "The `ktensor` class stores the components of the tensor $\\mathcal{X}$ and can perform many operations, e.g., `ttm`, without explicitly forming the tensor $\\mathcal{X}$.\n",
    "## Kruskal tensor format via `ktensor`\n",
    "Kruskal format stores a tensor as a sum of rank-1 outer products. For example, consider a tensor of the following form:\n",
    "\n",
    "$\n",
    "\\mathcal{X} = a_1 \\circ b_1 \\circ c_1 + a_2 \\circ b_2 \\circ c_2\n",
    "$\n",
    "\n",
    "This can be stored in Kruskal form as follows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pyttb as ttb\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(0)\n",
    "A = np.random.rand(4, 2)  # First column is a_1, second is a_2.\n",
    "B = np.random.rand(3, 2)  # Likewise for B.\n",
    "C = np.random.rand(2, 2)  # Likewise for C.\n",
    "X = ttb.ktensor([A, B, C])  # Create the ktensor.\n",
    "X"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Y = ttb.ktensor(\n",
    "    [np.random.rand(4, 1), np.random.rand(2, 1), np.random.rand(3, 1)]\n",
    ")  # Another ktensor.\n",
    "Y"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For Kruskal format, there can be any number of matrices, but every matrix must have the same number of columns. The number of rows can vary."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Specifying weights in a `ktensor`\n",
    "Weights for each rank-1 tensor can be specified by passing in a column vector. For example: \n",
    "\n",
    "$\n",
    "\\mathcal{X} = \\lambda_1 \\, a_1 \\circ b_1 \\circ c_1 + \\lambda_2 \\, a_2 \\circ b_2 \\circ c_2\n",
    "$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Upcoming ktensors will be generated with this same initialization.\n",
    "def generate_sample_ktensor() -> ttb.ktensor:\n",
    "    np.random.seed(0)\n",
    "    A = np.random.rand(4, 2)  # Create some data.\n",
    "    B = np.random.rand(3, 2)\n",
    "    C = np.random.rand(2, 2)\n",
    "    weights = np.array([5.0, 0.25])\n",
    "\n",
    "    return ttb.ktensor([A, B, C], weights)  # Create the ktensor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X = generate_sample_ktensor()\n",
    "X"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating a one-dimensional `ktensor`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(0)\n",
    "Y = ttb.ktensor([np.random.rand(4, 5)])  # A one-dimensional ktensor.\n",
    "Y"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Constituent parts of a `ktensor`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X = generate_sample_ktensor()\n",
    "X.weights  # Weights or multipliers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X.factor_matrices  # Cell array of matrices."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating a `ktensor` from its constituent parts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X = generate_sample_ktensor()\n",
    "Y = ttb.ktensor(X.factor_matrices, X.weights)  # Recreate X.\n",
    "Y"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating an empty `ktensor`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Z = ttb.ktensor()  # Empty ktensor.\n",
    "Z"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Use `full` or `to_tensor` to convert a `ktensor` to a `tensor`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X = generate_sample_ktensor()\n",
    "X.full()  # Converts to a tensor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X.to_tensor()  # Same as above."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Use `double` to convert a `ktensor` to a multidimensional array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X = generate_sample_ktensor()\n",
    "X.double()  # Converts to an array."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Use `tendiag` or `sptendiag` to convert a `ktensor` to a `ttensor`\n",
    "A `ktensor` can be regarded as a `ttensor` with a diagonal core."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X = generate_sample_ktensor()\n",
    "R = len(X.weights)  # Number of factors in X.\n",
    "core = ttb.tendiag(X.weights, ((R,) * X.ndims))  # Create a diagonal core.\n",
    "Y = ttb.ttensor(core, X.factor_matrices)  # Assemble the ttensor\n",
    "Y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(X.full() - Y.full()).norm()  # They are the same."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "core = ttb.sptendiag(X.weights, ((R,) * X.ndims))  # Sparse diagonal core.\n",
    "Y = ttb.ttensor(core, X.factor_matrices)  # Assemble the ttensor\n",
    "Y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(X.full() - Y.full()).norm()  # They are the same."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Use `ndims` and `shape` for the dimensions of a `ktensor`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X = generate_sample_ktensor()\n",
    "X.ndims  # Number of dimensions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X.shape  # Tuple of the shapes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X.shape[1]  # Shape of the 2nd mode."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Subscripted reference for a `ktensor`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X = generate_sample_ktensor()\n",
    "X.weights[1]  # Weight of the 2nd factor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X.factor_matrices[1]  # Extract a matrix."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Subscripted assignment for a `ktensor`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X = generate_sample_ktensor()\n",
    "X.weights = np.ones(X.weights.shape)  # Insert new multipliers.\n",
    "X"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X.weights[0] = 7  # Change a single element of weights."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X.factor_matrices[2][:, [0]] = np.ones((2, 1))  # Change the matrix for mode 3.\n",
    "X"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using negative indexing for the last array index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X = generate_sample_ktensor()\n",
    "X.factor_matrices[0][-1:, :]\n",
    "X.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X = generate_sample_ktensor()\n",
    "X.factor_matrices[0][0][\n",
    "    1 : (np.prod(X.shape) - 1)\n",
    "].item()  # Calculates factor_matrix[0][1]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Adding and subtracting `ktensor`s\n",
    "Adding two ktensors is the same as concatenating the matrices."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(0)\n",
    "X = ttb.ktensor(\n",
    "    [np.random.rand(4, 2), np.random.rand(2, 2), np.random.rand(3, 2)]\n",
    ")  # Data.\n",
    "Y = ttb.ktensor(\n",
    "    [np.random.rand(4, 2), np.random.rand(2, 2), np.random.rand(3, 2)]\n",
    ")  # More data.\n",
    "X, Y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Z = X + Y  # Concatenates the factor matrices.\n",
    "Z"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Z = X - Y  # Concatenates as with plus, but changes the weights.\n",
    "Z"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(Z.full() - (X.full() - Y.full())).norm()  # Should be zero."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Basic operations with a `ktensor`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "+X  # Calls uplus."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "-X  # Calls uminus."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "5 * X  # Calls mtimes."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Use `permute` to reorder the modes of a `ktensor`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(0)\n",
    "X = ttb.ktensor(\n",
    "    [np.random.rand(4, 2), np.random.rand(2, 2), np.random.rand(3, 2)]\n",
    ")  # Data.\n",
    "X.permute(np.array((1, 2, 0)))  # Reorders modes of X."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Use `arrange` to normalize the factors of a `ktensor`\n",
    "The function `arrange` normalizes the columns of the factors and then arranges the rank-one pieces in decreasing order of shape."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(0)\n",
    "X = ttb.ktensor(\n",
    "    [np.random.rand(3, 2), np.random.rand(4, 2), np.random.rand(2, 2)]\n",
    ")  # Unit weights.\n",
    "X"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X.arrange()  # Normalized and rearranged.\n",
    "X"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Use `fixsigns` for sign indeterminacies in a `ktensor`\n",
    "The largest magnitude entry for each factor is changed to be positive provided that we can flip the signs of pairs of vectors in that rank-1 component."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(0)\n",
    "X = ttb.ktensor(\n",
    "    [np.random.rand(4, 2), np.random.rand(2, 2), np.random.rand(3, 2)]\n",
    ")  # Data.\n",
    "Y = X\n",
    "Y.factor_matrices[0][:, 0] = -Y.factor_matrices[0][\n",
    "    :, 0\n",
    "]  # switch the sign on a pair of columns\n",
    "Y.factor_matrices[1][:, 0] = -Y.factor_matrices[1][:, 0]\n",
    "Y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Y.fixsigns()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Use `ktensor` to store the 'skinny' SVD of a matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(0)\n",
    "A = np.random.rand(4, 3)\n",
    "A"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "[U, S, Vh] = np.linalg.svd(A, full_matrices=False)  # Compute the SVD.\n",
    "# Numpy Expects U*S*Vh where pyttb expects U*S*V'\n",
    "X = ttb.ktensor([U, Vh.transpose()], S)  # Store the SVD as a ktensor.\n",
    "X"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"U*S*Vh:\\n{U@np.diag(S)@Vh}\")\n",
    "print(\n",
    "    f\"\\nX.factor_matrices[0]@np.diag(X.weights)@(X.factor_matrices[1].transpose()):\\n\\\n",
    "{X.factor_matrices[0]@np.diag(X.weights)@(X.factor_matrices[1].transpose())}\"\n",
    ")\n",
    "print(f\"\\nX.full():\\n{X.full()}\")  # Reassemble the original matrix."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Displaying a `ktensor`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X = generate_sample_ktensor()\n",
    "X"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(X)"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 1
}