{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tucker 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": [
    "Tucker format is a decomposition of a tensor $\\mathcal{X}$ as the product of a core tensor $\\mathcal{G}$ and matrices (e.g., $A$, $B$, $C$) in each dimension. In other words, a tensor $\\mathcal{X}$ is expressed as:\n",
    "\n",
    "$\n",
    "\\mathcal{X} = \\mathcal{G} \\times_1 A \\times_2 B \\times_3 C\n",
    "$\n",
    "\n",
    "In MATLAB notation: `X=ttm(G,{A,B,C})`"
   ]
  },
  {
   "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": [
    "# Upcoming ttensors will be generated with this same initialization.\n",
    "def generate_sample_ttensor() -> ttb.ttensor:\n",
    "    np.random.seed(0)\n",
    "    core = ttb.tensor(np.random.rand(3, 2, 1), shape=(3, 2, 1))  # The core tensor.\n",
    "    U = [\n",
    "        np.random.rand(5, 3),\n",
    "        np.random.rand(4, 2),\n",
    "        np.random.rand(3, 1),\n",
    "    ]  # The factor matrices.\n",
    "    X = ttb.ttensor(core, U)  # Create the ttensor.\n",
    "\n",
    "    return X\n",
    "\n",
    "\n",
    "X = generate_sample_ttensor()  # Create the ttensor.\n",
    "X"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating a `ttensor` with a `tensor` core\n",
    "Alternate core formats: `sptensor` or `tensor`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(0)\n",
    "sptensor_core = ttb.sptenrand([3, 2, 1], nonzeros=3)  # Create a 3 x 2 x 1 sptensor.\n",
    "U = [\n",
    "    np.random.rand(5, 3),\n",
    "    np.random.rand(4, 2),\n",
    "    np.random.rand(3, 1),\n",
    "]  # The factor matrices.\n",
    "Y = ttb.ttensor(sptensor_core, U)  # Core is a sptensor.\n",
    "Y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dense_tensor = ttb.tensor(np.random.rand(3, 2, 1), (3, 2, 1))  # Core is a tensor.\n",
    "U = [\n",
    "    np.random.rand(5, 3),\n",
    "    np.random.rand(4, 2),\n",
    "    np.random.rand(3, 1),\n",
    "]  # The factor matrices.\n",
    "Y = ttb.ttensor(dense_tensor, U)  # Create the ttensor.\n",
    "Y"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating a one-dimensional `ttensor`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(0)\n",
    "dense_tensor = ttb.tensor(2 * np.random.rand(2, 1), (2,))  # Core tensor.\n",
    "Z = ttb.ttensor(dense_tensor, [np.random.rand(4, 2)])  # One-dimensional ttensor.\n",
    "Z"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Constituent parts of a `ttensor`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X = generate_sample_ttensor()  # Create the ttensor.\n",
    "X.core  # Core tensor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X.factor_matrices  # List of matrices."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating a `ttensor` from its constituent parts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X = generate_sample_ttensor()  # Create the ttensor.\n",
    "Y = ttb.ttensor(X.core, X.factor_matrices)  # Recreate a ttensor from its parts.\n",
    "Y"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating an empty `ttensor`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X = ttb.ttensor()  # Empty ttensor.\n",
    "X"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Use `full` or `to_tensor` to convert a `ttensor` to a `tensor`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X = generate_sample_ttensor()  # Create a ttensor.\n",
    "X"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X.full()  # Converts to a tensor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X.to_tensor()  # Also converts to a tensor."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Use `reconstruct` to compute part of a full `tensor`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X = generate_sample_ttensor()  # Create a ttensor.\n",
    "X.reconstruct(1, 2)  # Extract first front slice."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Use `double` to convert a `ttensor` to a (multidimensional) array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X = generate_sample_ttensor()  # Create the ttensor.\n",
    "X.double()  # Converts to an array."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Use `ndims` and `shape` to get the shape of a `ttensor`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X = generate_sample_ttensor()  # Create the ttensor.\n",
    "X.ndims  # Number of dimensions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X.shape  # Row vector 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 `ttensor`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X = generate_sample_ttensor()  # Create the ttensor.\n",
    "X.core[0, 0, 0]  # Access an element of the core."
   ]
  },
  {
   "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 `ttensor`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X = generate_sample_ttensor()  # Create a ttensor.\n",
    "X.core = ttb.tenones(X.core.shape)  # Insert a new core.\n",
    "X"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X.core[1, 1, 0] = 7  # Change a single element.\n",
    "X"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X.factor_matrices[2][0:2, 0] = [1, 1]  # change slice of factor matrix\n",
    "X.factor_matrices[2]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using negative indexing for the last array index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X.core[-1]  # last element of core"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X.factor_matrices[-1]  # last factor matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X.factor_matrices[-1][-1]  # last element of last factor matrix"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Basic operations (uplus, uminus, mtimes, etc.) on a `ttensor`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X = generate_sample_ttensor()  # Create ttensor."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Addition"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "+X  # Calls uplus."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Subtraction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "-X  # Calls uminus."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Multiplication"
   ]
  },
  {
   "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 `ttensor`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X = generate_sample_ttensor()  # Create ttensor.\n",
    "X.permute(np.array([2, 1, 0]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Displaying a `ttensor`\n",
    "The `ttensor` displays by displaying the core and each of the component matrices."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X = generate_sample_ttensor()  # Create ttensor.\n",
    "print(X)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X  # In the python interface"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Partial Reconstruction of a Tucker Tensor\n",
    "### Benefits of Partial Reconstruction\n",
    "An advantage of Tucker decomposition is that the tensor can be partially reconstructed without ever forming the *full* tensor. The `reconstruct()` member function does this, resulting in significant time and memory savings, as we demonstrate below.\n",
    "### Create a random `tensor` for the data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "shape = (20, 30, 50)\n",
    "X = ttb.tenrand(shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Compute HOSVD\n",
    "We compute the Tucker decomposition using ST-HOSVD with target relative error `0.001`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "T = ttb.hosvd(X, tol=0.001)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Note: If the result is < 1.0 x, it will be unsurprising\n",
    "# since the random generation process below wasn't expected\n",
    "# to return a low-rank approximation\n",
    "print(\n",
    "    f\"Compression: {X.data.nbytes/(T.core.data.nbytes + np.sum([i.nbytes for i in T.factor_matrices]))} x\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Full reconstruction\n",
    "We can create a full reconstruction of the data using the `full` command. Not only is this expensive in computational time but also in memory. Now, let's see how long it takes to reconstruct the approximation to `X`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "Xf = T.full()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Partial reconstruction\n",
    "If we really only want part of the tensor, we can reconstruct just that part. Suppose we only want the `[:,15,:]` slice. The reconstruct function can do this much more efficiently with no loss in accuracy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "Xslice = T.reconstruct(modes=[1], samples=[15])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Compression: {Xf.data.nbytes/Xslice.data.nbytes} x\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Down-sampling\n",
    "Additionally, we may want to downsample high-dimensional data to something lower resolution. For example, here we downsample in modes 0 and 2 by a factor of 2 and see even further speed-up and memory savings. There is no loss of accuracy as compared to downsampling after constructing the full tensor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "S0 = np.kron(np.eye(int(shape[0] / 2)), 0.5 * np.ones((1, 2)))\n",
    "S2 = np.kron(np.eye(int(shape[2] / 2)), 0.5 * np.ones((1, 2)))\n",
    "S1 = np.array([15])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "Xds = T.reconstruct(modes=[0, 1, 2], samples=[S0, S1, S2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Compression: {Xf.data.nbytes/Xds.data.nbytes} x\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Compare visualizations\n",
    "We can compare the results of reconstruction. There is no degradation in doing only a partial reconstruction. Downsampling is obviously lower resolution, but the same result as first doing the full reconstruction and then downsampling."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "fig, axs = plt.subplots(2, 2, sharey=True)\n",
    "\n",
    "axs[0, 0].imshow(np.rot90(X[:, 15, :].double().squeeze()), vmin=1, vmax=3)\n",
    "axs[0, 1].imshow(np.rot90(Xf[:, 15, :].double().squeeze()), vmin=1, vmax=3)\n",
    "axs[1, 0].imshow(np.rot90(Xslice.double().squeeze()), vmin=1, vmax=3)\n",
    "axs[1, 1].imshow(np.rot90(Xds.double().squeeze()), vmin=1, vmax=3)\n",
    "\n",
    "axs[0, 0].set_aspect(aspect=\"equal\")\n",
    "axs[0, 1].set_aspect(aspect=\"equal\")\n",
    "axs[1, 0].set_aspect(aspect=\"equal\")\n",
    "axs[1, 1].set_aspect(aspect=\"equal\")\n",
    "\n",
    "axs[0, 0].set_axis_off()\n",
    "axs[0, 1].set_axis_off()\n",
    "axs[1, 0].set_axis_off()\n",
    "axs[1, 1].set_axis_off()\n",
    "\n",
    "axs[0, 0].set_title(\"Original Slice\")\n",
    "axs[0, 1].set_title(\"Full Reconstruction\")\n",
    "axs[1, 0].set_title(\"Partial Reconstruction\")\n",
    "axs[1, 1].set_title(\"Partial+Downsampled Reconstruction\")\n",
    "\n",
    "axs[1, 1].set_xlim = axs[1, 0].get_xlim()\n",
    "axs[1, 1].set_ylim = axs[1, 0].get_ylim()"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 1
}