{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {
    "vscode": {
     "languageId": "plaintext"
    }
   },
   "source": [
    "# Linear Algebra"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1",
   "metadata": {},
   "source": [
    "That linear algebra is fun, is a widely accepted fact. This notebooks will guide you through some of the linear algebra fun you can realize with Heat. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4 engines found\n"
     ]
    }
   ],
   "source": [
    "from ipyparallel import Client\n",
    "rc = Client(profile=\"default\")\n",
    "rc.ids\n",
    "\n",
    "if len(rc.ids) == 0:\n",
    "    print(\"No engines found\")\n",
    "else:\n",
    "    print(f\"{len(rc.ids)} engines found\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%px\n",
    "import heat as ht"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "5",
   "metadata": {},
   "source": [
    "## Matrix-Matrix Multiplication"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6",
   "metadata": {},
   "source": [
    "The most basic operation in linear algebra is matrix-matrix multiplication (\"matmul\"). Doing it by hand for a small matrix is not difficult and in fact not very spectacular. However, in the distributed setting (e.g., on 4 GPUs) even such a simple operation is not trivial any more: just imagine you work together with 3 other people and each of you only knows one fourth of the columns of a matrix $A$ and one fourth of the rows of a matrix $B$. Together, you have to compute the product $AB$ such that in the end each of you only has one fourth of the columns of $AB$..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": []
     },
     "metadata": {
      "after": null,
      "completed": null,
      "data": {},
      "engine_id": 3,
      "engine_uuid": "e9fac413-5ad332b3f3d1b43d6c63c9ee",
      "error": null,
      "execute_input": "split_A=0 \nsplit_B=1 \nM = 10000\nN = 10000\nK = 10000\nA = ht.random.randn(M, N, split=split_A, device=\"gpu\")\nB = ht.random.randn(N, K, split=split_B, device=\"gpu\")\nC = ht.matmul(A, B)\nC\n",
      "execute_result": {
       "data": {
        "text/plain": ""
       },
       "execution_count": 21,
       "metadata": {}
      },
      "follow": null,
      "msg_id": null,
      "outputs": [],
      "received": null,
      "started": null,
      "status": null,
      "stderr": "",
      "stdout": "",
      "submitted": "2025-11-21T14:40:53.538635Z"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": []
     },
     "metadata": {
      "after": null,
      "completed": null,
      "data": {},
      "engine_id": 2,
      "engine_uuid": "f85a0b27-08b851a725daa8d52a074bfd",
      "error": null,
      "execute_input": "split_A=0 \nsplit_B=1 \nM = 10000\nN = 10000\nK = 10000\nA = ht.random.randn(M, N, split=split_A, device=\"gpu\")\nB = ht.random.randn(N, K, split=split_B, device=\"gpu\")\nC = ht.matmul(A, B)\nC\n",
      "execute_result": {
       "data": {
        "text/plain": ""
       },
       "execution_count": 21,
       "metadata": {}
      },
      "follow": null,
      "msg_id": null,
      "outputs": [],
      "received": null,
      "started": null,
      "status": null,
      "stderr": "",
      "stdout": "",
      "submitted": "2025-11-21T14:40:53.538422Z"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": []
     },
     "metadata": {
      "after": null,
      "completed": null,
      "data": {},
      "engine_id": 1,
      "engine_uuid": "a1070193-d3a142f960197ef957f774ee",
      "error": null,
      "execute_input": "split_A=0 \nsplit_B=1 \nM = 10000\nN = 10000\nK = 10000\nA = ht.random.randn(M, N, split=split_A, device=\"gpu\")\nB = ht.random.randn(N, K, split=split_B, device=\"gpu\")\nC = ht.matmul(A, B)\nC\n",
      "execute_result": {
       "data": {
        "text/plain": ""
       },
       "execution_count": 21,
       "metadata": {}
      },
      "follow": null,
      "msg_id": null,
      "outputs": [],
      "received": null,
      "started": null,
      "status": null,
      "stderr": "",
      "stdout": "",
      "submitted": "2025-11-21T14:40:53.538350Z"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "\u001b[0;31mOut[0:21]: \u001b[0m\n",
       "DNDarray([[ 7.2963e+01, -1.4339e+02,  1.1302e+02,  ...,  2.1846e+01,  1.1797e+02,  2.3660e+01],\n",
       "          [-7.5554e+01,  3.2531e+01,  1.3144e+01,  ..., -1.0171e+02, -2.5141e+02,  1.4358e+02],\n",
       "          [ 3.4069e+01, -2.3046e+01,  1.9501e+02,  ...,  8.7619e+01, -7.5115e+01, -1.7974e+02],\n",
       "          ...,\n",
       "          [-8.0653e+00, -7.5908e+01,  6.8068e+01,  ..., -5.8282e+01, -2.5834e+01, -4.9030e+01],\n",
       "          [ 5.2255e+01,  2.8832e+01, -7.5600e+01,  ..., -2.5936e+01, -3.2145e-01, -4.5561e+01],\n",
       "          [ 9.0530e+01,  1.3209e-01,  7.0801e+01,  ...,  1.2266e+02,  1.5415e+01, -1.0515e+02]], dtype=ht.float32, device=gpu:0, split=0)"
      ]
     },
     "metadata": {
      "after": null,
      "completed": null,
      "data": {},
      "engine_id": 0,
      "engine_uuid": "2f5ee137-24f77a09dc6e8333cba67c27",
      "error": null,
      "execute_input": "split_A=0 \nsplit_B=1 \nM = 10000\nN = 10000\nK = 10000\nA = ht.random.randn(M, N, split=split_A, device=\"gpu\")\nB = ht.random.randn(N, K, split=split_B, device=\"gpu\")\nC = ht.matmul(A, B)\nC\n",
      "execute_result": {
       "data": {
        "text/plain": "DNDarray([[ 7.2963e+01, -1.4339e+02,  1.1302e+02,  ...,  2.1846e+01,  1.1797e+02,  2.3660e+01],\n          [-7.5554e+01,  3.2531e+01,  1.3144e+01,  ..., -1.0171e+02, -2.5141e+02,  1.4358e+02],\n          [ 3.4069e+01, -2.3046e+01,  1.9501e+02,  ...,  8.7619e+01, -7.5115e+01, -1.7974e+02],\n          ...,\n          [-8.0653e+00, -7.5908e+01,  6.8068e+01,  ..., -5.8282e+01, -2.5834e+01, -4.9030e+01],\n          [ 5.2255e+01,  2.8832e+01, -7.5600e+01,  ..., -2.5936e+01, -3.2145e-01, -4.5561e+01],\n          [ 9.0530e+01,  1.3209e-01,  7.0801e+01,  ...,  1.2266e+02,  1.5415e+01, -1.0515e+02]], dtype=ht.float32, device=gpu:0, split=0)"
       },
       "execution_count": 21,
       "metadata": {}
      },
      "follow": null,
      "msg_id": null,
      "outputs": [],
      "received": null,
      "started": null,
      "status": null,
      "stderr": "",
      "stdout": "",
      "submitted": "2025-11-21T14:40:53.537927Z"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "%%px\n",
    "split_A=0 \n",
    "split_B=1 \n",
    "M = 10000\n",
    "N = 10000\n",
    "K = 10000\n",
    "A = ht.random.randn(M, N, split=split_A, device=\"gpu\")\n",
    "B = ht.random.randn(N, K, split=split_B, device=\"gpu\")\n",
    "C = ht.matmul(A, B)\n",
    "C"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8",
   "metadata": {},
   "source": [
    "## QR Decomposition and Triangular Solve\n",
    "\n",
    "Given a matrix $A$, its QR decomposition is given by $A=QR$ where $Q$ is an orthogonal matrix (i.e. its columns are pairwise orthonormal) and $R$ is an upper triangular matrix. \n",
    "\n",
    "Further information: [QR on Wikipedia](https://en.wikipedia.org/wiki/QR_decomposition)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[stdout:3] (100000, 1000) (1000, 1000)\n",
       "\n",
       "\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "[stdout:2] (100000, 1000) (1000, 1000)\n",
       "\n",
       "\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "[stdout:1] (100000, 1000) (1000, 1000)\n",
       "\n",
       "\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "[stdout:0] (100000, 1000) (1000, 1000)\n",
       "DNDarray([[ 1.0000e+00, -2.7940e-09,  2.3283e-09,  ...,  2.6776e-09,  5.8208e-09,  7.1595e-09],\n",
       "          [-2.7940e-09,  1.0000e+00,  4.6566e-10,  ..., -1.3970e-09, -8.1491e-10,  6.7521e-09],\n",
       "          [ 2.3283e-09,  4.6566e-10,  1.0000e+00,  ..., -1.0477e-09,  5.8208e-10,  2.9104e-09],\n",
       "          ...,\n",
       "          [ 2.6776e-09, -1.3970e-09, -1.0477e-09,  ...,  1.0000e+00,  9.3132e-10, -1.3970e-09],\n",
       "          [ 5.8208e-09, -8.1491e-10,  5.8208e-10,  ...,  9.3132e-10,  1.0000e+00, -8.7311e-10],\n",
       "          [ 7.1595e-09,  6.7521e-09,  2.9104e-09,  ..., -1.3970e-09, -8.7311e-10,  1.0000e+00]], dtype=ht.float32, device=gpu:0, split=1)\n",
       "DNDarray([[ 3.1642e+02,  3.9908e-01, -1.7685e+00,  ...,  4.7194e-01,  1.5867e+00, -1.0417e+00],\n",
       "          [ 0.0000e+00, -3.1472e+02, -1.0726e-01,  ...,  3.7303e-01,  1.2852e+00, -2.5343e-01],\n",
       "          [ 0.0000e+00,  0.0000e+00, -3.1664e+02,  ..., -9.1426e-01,  7.1782e-01, -1.1222e+00],\n",
       "          ...,\n",
       "          [ 0.0000e+00,  0.0000e+00,  0.0000e+00,  ...,  3.1383e+02, -9.6070e-01,  3.5290e-01],\n",
       "          [ 0.0000e+00,  0.0000e+00,  0.0000e+00,  ...,  0.0000e+00,  3.1372e+02,  8.9383e-01],\n",
       "          [ 0.0000e+00,  0.0000e+00,  0.0000e+00,  ...,  0.0000e+00,  0.0000e+00,  3.1452e+02]], dtype=ht.float32, device=gpu:0, split=None)\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "%%px\n",
    "A = ht.random.randn(100000, 1000, split=0, device=\"gpu\")\n",
    "Q,R = ht.linalg.qr(A)\n",
    "print(Q.shape, R.shape)\n",
    "print(Q.T @ Q) \n",
    "print(R)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10",
   "metadata": {},
   "source": [
    "With a little bit of linear algebra fun, you find out that a linear least squares problem of type $\\min \\lVert Ax - b \\rVert_2$ boils down to computing the QR decomposition $A=QR$ and then solving for $Rx = Q^T b$. (Of course, we need to assume that if $A \\in \\mathbb{R}^{m \\times n}$ that $m \\geq n$ and $R$ is invertible...)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[stdout:3] \n",
       "\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "[stdout:0] DNDarray(315.8452, dtype=ht.float32, device=gpu:0, split=None)\n",
       "DNDarray(0.0107, dtype=ht.float32, device=gpu:0, split=None)\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "[stdout:1] \n",
       "\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "[stdout:2] \n",
       "\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "%%px\n",
    "b = ht.random.randn(100000,1,split=None, device=\"gpu\")\n",
    "Qtb = Q.T @ b\n",
    "x = ht.linalg.solve_triangular(R,Qtb.resplit_(None)) \n",
    "print(ht.norm(A @ x - b))\n",
    "print(ht.norm(A.T @ A @ x - A.T @ b))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12",
   "metadata": {},
   "source": [
    "If you want to solve a LASSO-regularized version of this linear regression problem, try out `heat.regression.Lasso`!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13",
   "metadata": {},
   "source": [
    "## Singular Value Decomposition\n",
    "\n",
    "Let $X \\in \\mathbb{R}^{m \\times n}$ be a matrix, e.g., given by a data set consisting of $m$ data points $\\in \\mathbb{R}^n$ stacked together. The so-called **singular value decomposition (SVD)** of $X$ is given by \n",
    "\n",
    "$$\n",
    "\tX = U \\Sigma V^T\n",
    "$$\n",
    "\n",
    "where $U \\in \\mathbb{R}^{m \\times r_X}$ and $V \\in \\mathbb{R}^{n \\times r_X}$ have orthonormal columns, $\\Sigma = \\text{diag}(\\sigma_1,...,\\sigma_{r_X}) \\in \\mathbb{R}^{r_X \\times r_X}$ is a diagonal matrix containing the so-called singular values $\\sigma_1 \\geq \\sigma_2 \\geq ... \\geq \\sigma_{r_X} > 0$, and $r_X \\leq \\min(m,n)$ denotes the rank of $X$ (i.e. the dimension of the subspace of $\\mathbb{R}^m$ spanned by the columns of $X$). Since $\\Sigma = U^T X V$ is diagonal, one can imagine this decomposition as finding orthogonal coordinate transformations under which $X$ looks \"linear\". Further information: [SVD on Wikipedia](https://en.wikipedia.org/wiki/Singular_value_decomposition)\n",
    "\n",
    "Computing the **full** SVD in a distributed environment can be quite **expensive**; nevertheless, we have implemented it: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a0ddc4f8b5494d4b8c33a1283c447bd5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "%px:   0%|          | 0/4 [00:00<?, ?tasks/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "[stdout:2]   \n",
       "\n",
       "\n",
       "\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "[stdout:3]   \n",
       "\n",
       "\n",
       "\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "[stdout:1]   \n",
       "\n",
       "\n",
       "\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "[stdout:0] DNDarray([[-2.1787e-02,  1.2095e-02, -5.7444e-02,  ...,  1.9402e-02, -1.6275e-02, -4.4408e-02],\n",
       "          [-2.1879e-02,  1.9144e-02,  1.5135e-03,  ..., -2.3827e-02,  1.2645e-02,  1.0633e-02],\n",
       "          [-2.2175e-02, -9.7101e-03,  4.5591e-02,  ..., -2.2914e-03, -6.4245e-03, -5.0843e-03],\n",
       "          ...,\n",
       "          [-2.2211e-02,  3.5838e-02, -2.5349e-02,  ..., -1.8349e-02, -4.0364e-02, -2.6403e-02],\n",
       "          [-2.2765e-02, -1.4892e-02,  7.2096e-03,  ...,  9.6319e-03, -1.0798e-03, -6.2438e-02],\n",
       "          [-2.2729e-02,  2.3975e-02,  4.2114e-02,  ..., -1.3341e-02,  1.3794e-02, -1.1054e-02]], dtype=ht.float32, device=gpu:0, split=1) DNDarray([1.0002e+03, 2.5776e+01, 2.5720e+01,  ..., 1.1860e-02, 6.3026e-03, 1.8313e-03], dtype=ht.float32, device=cpu:0, split=0) DNDarray([[-2.2504e-02,  8.9078e-03, -2.9508e-03,  ...,  2.4425e-03, -6.2131e-03, -3.3841e-02],\n",
       "          [-2.2063e-02,  4.0797e-02, -5.1930e-03,  ...,  6.2515e-02,  1.3404e-02, -2.4298e-02],\n",
       "          [-2.2753e-02,  6.7618e-03, -9.4777e-03,  ...,  1.1973e-02,  9.8253e-03, -1.4280e-03],\n",
       "          ...,\n",
       "          [-2.2317e-02, -1.6574e-02, -3.4864e-02,  ...,  9.1662e-04,  7.8013e-04,  1.0299e-02],\n",
       "          [-2.2796e-02,  1.0439e-02,  1.7705e-02,  ...,  4.6873e-02, -1.2069e-02,  1.8161e-03],\n",
       "          [-2.2331e-02, -1.8023e-02, -1.8411e-02,  ..., -4.9757e-04,  1.3030e-02,  1.9065e-03]], dtype=ht.float32, device=gpu:0, split=1)\n",
       "DNDarray([[ 1.0000e+00, -4.3772e-08, -1.5716e-08,  ..., -4.6566e-09,  1.2503e-07,  5.7183e-07],\n",
       "          [-4.3772e-08,  1.0000e+00,  3.8417e-08,  ...,  3.4831e-07, -1.3912e-07,  5.8045e-07],\n",
       "          [-1.5716e-08,  3.8417e-08,  1.0000e+00,  ..., -4.1444e-07,  1.1455e-07, -2.3399e-07],\n",
       "          ...,\n",
       "          [-4.6566e-09,  3.4831e-07, -4.1444e-07,  ...,  1.0000e+00, -7.2643e-08,  3.7253e-08],\n",
       "          [ 1.2503e-07, -1.3912e-07,  1.1455e-07,  ..., -7.2643e-08,  1.0000e+00, -4.0978e-07],\n",
       "          [ 5.7183e-07,  5.8045e-07, -2.3399e-07,  ...,  3.7253e-08, -4.0978e-07,  1.0000e+00]], dtype=ht.float32, device=gpu:0, split=0)\n",
       "DNDarray([[ 1.0000e+00, -3.9116e-08, -2.0489e-08,  ..., -1.0245e-08,  1.4249e-07,  5.9232e-07],\n",
       "          [-3.9116e-08,  1.0000e+00,  3.2247e-08,  ...,  3.3900e-07, -1.4063e-07,  6.0257e-07],\n",
       "          [-2.0489e-08,  3.2247e-08,  1.0000e+00,  ..., -4.1351e-07,  1.6205e-07, -2.5448e-07],\n",
       "          ...,\n",
       "          [-1.0245e-08,  3.3900e-07, -4.1351e-07,  ...,  1.0000e+00, -5.5414e-08,  2.2119e-09],\n",
       "          [ 1.4249e-07, -1.4063e-07,  1.6205e-07,  ..., -5.5414e-08,  1.0000e+00, -3.7765e-07],\n",
       "          [ 5.9232e-07,  6.0257e-07, -2.5448e-07,  ...,  2.2119e-09, -3.7765e-07,  1.0000e+00]], dtype=ht.float32, device=gpu:0, split=0)\n",
       "DNDarray(0.0001, dtype=ht.float32, device=gpu:0, split=None)\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "%%px\n",
    "A = ht.random.rand(2000,2000, split=1, device=\"gpu\")\n",
    "U, S, V = ht.linalg.svd(A)\n",
    "print(U, S, V)\n",
    "print(U.T @ U)\n",
    "print(V.T @ V)\n",
    "print(ht.norm(A - U @ ht.diag(S.gpu()) @  V.T)/ht.norm(A))  # tiny little bug: S is on cpu, even if A was on gpu... "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15",
   "metadata": {},
   "source": [
    "If the number of rows is much higher than the number of columns (we call such matrices \"tall-skinny\"), a more efficient implementation of SVD is available than in the general case:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "28d105dde39149bab7bf6d4d2b2af23f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "%px:   0%|          | 0/4 [00:00<?, ?tasks/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "[stdout:3]   \n",
       "\n",
       "\n",
       "\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "[stdout:2]   \n",
       "\n",
       "\n",
       "\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "[stdout:1]   \n",
       "\n",
       "\n",
       "\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "[stdout:0] DNDarray([[ 3.1664e-03,  7.7945e-04, -4.9356e-03,  ...,  4.3825e-04,  1.5350e-03,  2.6044e-03],\n",
       "          [ 3.1429e-03,  1.3692e-03, -5.2835e-03,  ...,  2.2554e-03,  6.3171e-03, -5.7571e-04],\n",
       "          [ 3.0861e-03,  1.8266e-04, -1.0699e-03,  ..., -2.9055e-03,  7.8701e-04, -1.5580e-03],\n",
       "          ...,\n",
       "          [ 3.0706e-03,  8.8718e-04, -3.2530e-03,  ..., -1.2671e-03,  4.5658e-03, -3.6211e-04],\n",
       "          [ 3.1243e-03, -4.6736e-03,  1.3189e-03,  ..., -1.3343e-03,  2.6499e-03,  4.7289e-03],\n",
       "          [ 3.1990e-03, -4.8147e-03,  5.9925e-04,  ..., -2.8875e-03,  5.8510e-03,  5.3508e-03]], dtype=ht.float32, device=gpu:0, split=0) DNDarray([7069.9585,  104.0775,  104.0356,  ...,   78.6280,   78.5792,   78.4621], dtype=ht.float32, device=gpu:0, split=None) DNDarray([[ 0.0224,  0.0002,  0.0186,  ..., -0.0204, -0.0025, -0.0104],\n",
       "          [ 0.0223,  0.0555, -0.0191,  ...,  0.0153,  0.0183, -0.0196],\n",
       "          [ 0.0224,  0.0138, -0.0458,  ..., -0.0074,  0.0018, -0.0049],\n",
       "          ...,\n",
       "          [ 0.0223,  0.0106, -0.0275,  ..., -0.0206,  0.0149,  0.0230],\n",
       "          [ 0.0224,  0.0003, -0.0055,  ...,  0.0358,  0.0066,  0.0073],\n",
       "          [ 0.0223,  0.0027, -0.0182,  ...,  0.0047,  0.0081,  0.0037]], dtype=ht.float32, device=gpu:0, split=None)\n",
       "DNDarray([[ 9.9933e-01,  2.9057e-07,  2.7590e-07,  ...,  3.1386e-07, -5.5830e-07, -1.9814e-07],\n",
       "          [ 2.9057e-07,  9.9946e-01, -7.2676e-06,  ..., -1.1833e-06, -1.0792e-07, -1.2032e-07],\n",
       "          [ 2.7590e-07, -7.2676e-06,  9.9964e-01,  ..., -2.7654e-07, -3.5056e-07,  4.5385e-07],\n",
       "          ...,\n",
       "          [ 3.1386e-07, -1.1833e-06, -2.7654e-07,  ...,  9.9985e-01, -1.4408e-06,  2.2130e-05],\n",
       "          [-5.5830e-07, -1.0792e-07, -3.5056e-07,  ..., -1.4408e-06,  9.9978e-01,  2.3541e-05],\n",
       "          [-1.9814e-07, -1.2032e-07,  4.5385e-07,  ...,  2.2130e-05,  2.3541e-05,  9.9989e-01]], dtype=ht.float32, device=gpu:0, split=1)\n",
       "DNDarray([[ 1.0003e+00, -5.4576e-07, -4.3120e-07,  ..., -2.8871e-07, -4.6566e-08, -8.5449e-08],\n",
       "          [-5.4576e-07,  1.0008e+00, -1.1735e-07,  ...,  2.9430e-07,  3.4925e-07,  9.1270e-08],\n",
       "          [-4.3120e-07, -1.1735e-07,  1.0008e+00,  ...,  2.1979e-07, -1.3690e-07, -1.3504e-07],\n",
       "          ...,\n",
       "          [-2.8871e-07,  2.9430e-07,  2.1979e-07,  ...,  1.0008e+00, -1.9521e-06, -3.4748e-06],\n",
       "          [-4.6566e-08,  3.4925e-07, -1.3690e-07,  ..., -1.9521e-06,  1.0008e+00, -5.1726e-06],\n",
       "          [-8.5449e-08,  9.1270e-08, -1.3504e-07,  ..., -3.4748e-06, -5.1726e-06,  1.0008e+00]], dtype=ht.float32, device=gpu:0, split=None)\n",
       "DNDarray(0.0004, dtype=ht.float32, device=gpu:0, split=None)\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "%%px \n",
    "X = ht.random.rand(100000,2000, split=0, device=\"gpu\")\n",
    "U, S, V = ht.linalg.svd(X) \n",
    "print(U, S, V)\n",
    "print(U.T @ U)\n",
    "print(V.T @ V)\n",
    "print(ht.norm(X - U @ ht.diag(S) @  V.T)/ht.norm(X))\n",
    "del X, U, S, V"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17",
   "metadata": {},
   "source": [
    "## Truncated SVD \n",
    "\n",
    "So, what to do if your matrix is very large but not tall-skinny? \n",
    "\n",
    "Let us start with a brief background on how SVD is actually used in applications... In data science, SVD is more often known as **principle component analysis (PCA)**, the columns of $U$ being called the principle components of $X$. In fact, in many applications **truncated SVD/PCA** suffices: to reduce $X$ to the \"essential\" information, one chooses a truncation rank $0 < r \\leq r_X$ and considers the truncated SVD/PCA given by \n",
    "\n",
    "$$\n",
    "X \\approx X_r := U_{[:,:r]} \\Sigma_{[:r,:r]} V_{[:,:r]}^T\n",
    "$$\n",
    "\n",
    "where we have used `numpy`-like notation for selecting only the first $r$ columns of $U$ and $V$, respectively. The rationale behind this is that if the first $r$ singular values of $X$ are much larger than the remaining ones, $X_r$ will still contain all \"essential\" information contained in $X$; in mathematical terms: \n",
    "\n",
    "$$\n",
    "\\lVert X_r - X \\rVert_{F}^2 = \\sum_{i=r+1}^{r_X} \\sigma_i^2, \n",
    "$$\n",
    "\n",
    "where $\\lVert \\cdot \\rVert_F$ denotes the Frobenius norm. Thus, truncated SVD/PCA may be used for, e.g.,  \n",
    "* filtering away non-essential information in order to get a \"feeling\" for the main characteristics of your data set, \n",
    "* to detect linear (or \"almost\" linear) dependencies in your data, \n",
    "* to generate features for further processing of your data. \n",
    "\n",
    "In the following we present two options available in Heat to compute an **approximate truncated SVD for a large matrix**. \n",
    "\n",
    "### Randomized SVD \n",
    "\n",
    "The most easy option is so-called randomized SVD which allows you to specify the rank $r$ you are targeting as second positional argument. Here, we use randomized SVD with a certain number of oversamples and one power iteration in order to compute an approximation to the leading 10 singular values and vectors of the same matrix $A$ as above:   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\u001b[0;31mOut[2:45]: \u001b[0m(, , )"
      ]
     },
     "metadata": {
      "after": null,
      "completed": null,
      "data": {},
      "engine_id": 2,
      "engine_uuid": "f85a0b27-08b851a725daa8d52a074bfd",
      "error": null,
      "execute_input": "Ur, Sr, Vr = ht.linalg.rsvd(A, 10, n_oversamples=10, power_iter=1)\nUr, Sr, Vr\n",
      "execute_result": {
       "data": {
        "text/plain": "(, , )"
       },
       "execution_count": 45,
       "metadata": {}
      },
      "follow": null,
      "msg_id": null,
      "outputs": [],
      "received": null,
      "started": null,
      "status": null,
      "stderr": "",
      "stdout": "",
      "submitted": "2025-11-21T14:53:30.542234Z"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "\u001b[0;31mOut[3:45]: \u001b[0m(, , )"
      ]
     },
     "metadata": {
      "after": null,
      "completed": null,
      "data": {},
      "engine_id": 3,
      "engine_uuid": "e9fac413-5ad332b3f3d1b43d6c63c9ee",
      "error": null,
      "execute_input": "Ur, Sr, Vr = ht.linalg.rsvd(A, 10, n_oversamples=10, power_iter=1)\nUr, Sr, Vr\n",
      "execute_result": {
       "data": {
        "text/plain": "(, , )"
       },
       "execution_count": 45,
       "metadata": {}
      },
      "follow": null,
      "msg_id": null,
      "outputs": [],
      "received": null,
      "started": null,
      "status": null,
      "stderr": "",
      "stdout": "",
      "submitted": "2025-11-21T14:53:30.542360Z"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "\u001b[0;31mOut[1:45]: \u001b[0m(, , )"
      ]
     },
     "metadata": {
      "after": null,
      "completed": null,
      "data": {},
      "engine_id": 1,
      "engine_uuid": "a1070193-d3a142f960197ef957f774ee",
      "error": null,
      "execute_input": "Ur, Sr, Vr = ht.linalg.rsvd(A, 10, n_oversamples=10, power_iter=1)\nUr, Sr, Vr\n",
      "execute_result": {
       "data": {
        "text/plain": "(, , )"
       },
       "execution_count": 45,
       "metadata": {}
      },
      "follow": null,
      "msg_id": null,
      "outputs": [],
      "received": null,
      "started": null,
      "status": null,
      "stderr": "",
      "stdout": "",
      "submitted": "2025-11-21T14:53:30.542105Z"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "\u001b[0;31mOut[0:45]: \u001b[0m\n",
       "(DNDarray([[-0.0218,  0.0066,  0.0240,  ..., -0.0276,  0.0048,  0.0121],\n",
       "          [-0.0219,  0.0288,  0.0158,  ..., -0.0010,  0.0068, -0.0420],\n",
       "          [-0.0222,  0.0445,  0.0010,  ...,  0.0100,  0.0118, -0.0079],\n",
       "          ...,\n",
       "          [-0.0222, -0.0460,  0.0116,  ...,  0.0155, -0.0024, -0.0150],\n",
       "          [-0.0228, -0.0173,  0.0391,  ...,  0.0006,  0.0284, -0.0470],\n",
       "          [-0.0227,  0.0193, -0.0302,  ...,  0.0241, -0.0231,  0.0173]], dtype=ht.float32, device=gpu:0, split=1),\n",
       " DNDarray([1000.2399,   22.8685,   22.4281,   22.2607,   22.1712,   22.0842,   22.0111,   21.7669,   21.7218,   21.5960], dtype=ht.float32, device=gpu:0, split=None),\n",
       " DNDarray([[-2.2502e-02, -1.9996e-02, -1.9707e-02,  ...,  2.1129e-02, -2.5087e-02, -4.2121e-03],\n",
       "          [-2.2063e-02, -2.5445e-02, -6.8166e-04,  ...,  1.9730e-02, -9.9687e-03,  6.8673e-02],\n",
       "          [-2.2750e-02,  1.0383e-02, -8.7034e-03,  ...,  7.8081e-05, -1.6337e-02,  3.9759e-03],\n",
       "          ...,\n",
       "          [-2.2317e-02, -1.6521e-03,  1.0914e-03,  ..., -1.4690e-04,  1.2948e-02, -3.8937e-02],\n",
       "          [-2.2794e-02,  1.7277e-02,  4.5935e-02,  ..., -2.4583e-02,  2.3801e-02,  2.4243e-02],\n",
       "          [-2.2332e-02,  7.4138e-03, -9.2279e-05,  ..., -1.8467e-02,  3.3541e-02, -9.2743e-03]], dtype=ht.float32, device=gpu:0, split=None))"
      ]
     },
     "metadata": {
      "after": null,
      "completed": null,
      "data": {},
      "engine_id": 0,
      "engine_uuid": "2f5ee137-24f77a09dc6e8333cba67c27",
      "error": null,
      "execute_input": "Ur, Sr, Vr = ht.linalg.rsvd(A, 10, n_oversamples=10, power_iter=1)\nUr, Sr, Vr\n",
      "execute_result": {
       "data": {
        "text/plain": "(DNDarray([[-0.0218,  0.0066,  0.0240,  ..., -0.0276,  0.0048,  0.0121],\n          [-0.0219,  0.0288,  0.0158,  ..., -0.0010,  0.0068, -0.0420],\n          [-0.0222,  0.0445,  0.0010,  ...,  0.0100,  0.0118, -0.0079],\n          ...,\n          [-0.0222, -0.0460,  0.0116,  ...,  0.0155, -0.0024, -0.0150],\n          [-0.0228, -0.0173,  0.0391,  ...,  0.0006,  0.0284, -0.0470],\n          [-0.0227,  0.0193, -0.0302,  ...,  0.0241, -0.0231,  0.0173]], dtype=ht.float32, device=gpu:0, split=1),\n DNDarray([1000.2399,   22.8685,   22.4281,   22.2607,   22.1712,   22.0842,   22.0111,   21.7669,   21.7218,   21.5960], dtype=ht.float32, device=gpu:0, split=None),\n DNDarray([[-2.2502e-02, -1.9996e-02, -1.9707e-02,  ...,  2.1129e-02, -2.5087e-02, -4.2121e-03],\n          [-2.2063e-02, -2.5445e-02, -6.8166e-04,  ...,  1.9730e-02, -9.9687e-03,  6.8673e-02],\n          [-2.2750e-02,  1.0383e-02, -8.7034e-03,  ...,  7.8081e-05, -1.6337e-02,  3.9759e-03],\n          ...,\n          [-2.2317e-02, -1.6521e-03,  1.0914e-03,  ..., -1.4690e-04,  1.2948e-02, -3.8937e-02],\n          [-2.2794e-02,  1.7277e-02,  4.5935e-02,  ..., -2.4583e-02,  2.3801e-02,  2.4243e-02],\n          [-2.2332e-02,  7.4138e-03, -9.2279e-05,  ..., -1.8467e-02,  3.3541e-02, -9.2743e-03]], dtype=ht.float32, device=gpu:0, split=None))"
       },
       "execution_count": 45,
       "metadata": {}
      },
      "follow": null,
      "msg_id": null,
      "outputs": [],
      "received": null,
      "started": null,
      "status": null,
      "stderr": "",
      "stdout": "",
      "submitted": "2025-11-21T14:53:30.541644Z"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "%%px \n",
    "Ur, Sr, Vr = ht.linalg.rsvd(A, 10, n_oversamples=10, power_iter=1)\n",
    "Ur, Sr, Vr"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19",
   "metadata": {},
   "source": [
    "### Hierarchical SVD \n",
    "\n",
    "Hierarchical SVD allows a bit more of freedom. Truncation takes place either w.r.t. a fixed truncation-rank (`heat.linalg.hsvd_rank`) or w.r.t. a desired accuracy (`heat.linalg.hsvd_rtol`). In the latter case it can be ensured that it holds for the \"reconstruction error\": \n",
    "\n",
    "$$\n",
    "\\frac{\\lVert X - U U^T X \\rVert_F}{\\lVert X \\rVert_F} \\overset{!}{\\leq} \\text{rtol},\n",
    "$$\n",
    "\n",
    "where $U$ denotes the approximate left-singular vectors of $X$ computed by `heat.linalg.hsvd_rtol`. \n",
    "\n",
    "Let's first compute the truncated SVD by setting the relative tolerance: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[stdout:0] hSVD level 0...\t processes  0\t\t1\t\t2\t\t3\n",
       "              current ranks: 500\t\t500\t\t500\t\t500\n",
       "hSVD level 1...\t processes  0\t\t2\n",
       "              current ranks: 1000\t\t1000\n",
       "hSVD level 2...\t processes  0\n",
       "relative residual: DNDarray(0.0282, dtype=ht.float32, device=gpu:0, split=None) rank:  1918\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "[stdout:3] relative residual:  rank:  1918\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "[stdout:2] relative residual:  rank:  1918\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "[stdout:1] relative residual:  rank:  1918\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "%%px\n",
    "# compute truncated SVD w.r.t. relative tolerance \n",
    "svd_with_reltol = ht.linalg.hsvd_rtol(A,rtol=1.0e-2,compute_sv=True,silent=False)\n",
    "print(\"relative residual:\", svd_with_reltol[3], \"rank: \", svd_with_reltol[0].shape[1])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "21",
   "metadata": {},
   "source": [
    "Alternatively, you can compute a truncated SVD with a fixed truncation rank:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[stdout:0] hSVD level 0...\t processes  0\t\t1\t\t2\t\t3\n",
       "              current ranks: 8\t\t8\t\t8\t\t8\n",
       "hSVD level 1...\t processes  0\n",
       "relative residual: DNDarray(0.5046, dtype=ht.float32, device=gpu:0, split=None) rank:  3\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "[stdout:1] relative residual:  rank:  3\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "[stdout:2] relative residual:  rank:  3\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "[stdout:3] relative residual:  rank:  3\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "%%px\n",
    "# compute truncated SVD w.r.t. a fixed truncation rank \n",
    "svd_with_rank = ht.linalg.hsvd_rank(A, maxrank=3,compute_sv=True,silent=False)\n",
    "print(\"relative residual:\", svd_with_rank[3], \"rank: \", svd_with_rank[0].shape[1])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23",
   "metadata": {},
   "source": [
    "Once we have computed the truncated SVD, we can use it to approximate the original data matrix `X` by the truncated matrix `X_r`. \n",
    "\n",
    "Check out the plot below to see how Heat's truncated SVD algorithm scales with the number of MPI processes and size of the dataset.\n",
    "\n",
    "<div>\n",
    "  <img src=https://github.com/helmholtz-analytics/heat/blob/main/doc/source/_static/images/hSVD_bench_rank5.png?raw=true title=\"hSVD rank 5\" width=\"30%\" style=\"float:left\"/>\n",
    "  <img src=https://github.com/helmholtz-analytics/heat/blob/main/doc/source/_static/images/hSVD_bench_rank50.png?raw=true title=\"hSVD rank 50\" width=\"30%\" style=\"float:center \"/>\n",
    "  <img src=https://github.com/helmholtz-analytics/heat/blob/main/doc/source/_static/images/hSVD_bench_rank500.png?raw=true title=\"HSVD rank 500\" width=\"30%\" style=\"float:center\"/>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "24",
   "metadata": {},
   "source": [
    "## Exercises \n",
    "\n",
    "1. Try out different split combinations `split_A=0,1,None`, `split_B=0,1,None` for matrix-matrix multiplication. What do you observe for the split of the outcome? Does every combination take the same computing time? (Whats the likely cause for this?)\n",
    "\n",
    "2. Compare the approximate singular values computed by randomized SVD / hierarchiacl SVD with the reference ones computed by full SVD. What do you observe for different choices of the parameters (number of oversamples, power iterations, truncation rank, truncation tolerance)? "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "25",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}