{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# SVD, and what truncation costs\n",
    "\n",
    "The factorisation, the truncation, and the dense vectors that come out the other side.\n",
    "\n",
    "From chapter 7, [Count Vectors, PPMI, and SVD](https://nlp.jcrlabz.com/book/countvectors/), of the course notes.\n",
    "\n",
    "Source: `book/code/worked_examples/svd.py`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## What you need\n",
    "\n",
    "`pip install numpy`. The next cell does it. Colab usually has numpy already, but nothing here assumes that."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "%pip install -q numpy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The script\n",
    "\n",
    "Everything the example defines, in one cell. Run it and the definitions\n",
    "are live in the notebook."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "\"\"\"The SVD itself: factorise, truncate, and see what the truncation cost.\n",
    "\n",
    "`svd_rank.py` answers \"how big should K be\". This one answers \"what actually\n",
    "happens when I cut at K\". They are the two halves of Chapter 7's second repair.\n",
    "\n",
    "    python3 svd.py            # factorise, truncate at every K, show the error\n",
    "    python3 svd.py --k 2      # dense word vectors at K=2, and similarities\n",
    "    python3 svd.py --ppmi     # run PPMI over the counts first, as LSA does\n",
    "\n",
    "Watch for the moment the truncation stops losing information and starts\n",
    "removing noise.\n",
    "\n",
    "Install:  pip install numpy\n",
    "\"\"\"\n",
    "\n",
    "import argparse\n",
    "\n",
    "try:\n",
    "    import numpy as np\n",
    "except ImportError:\n",
    "    raise SystemExit(\"this one needs numpy:  pip install numpy\")\n",
    "\n",
    "WORDS = [\"dog\", \"cat\", \"car\", \"truck\", \"pet\", \"vet\", \"drive\", \"fuel\"]\n",
    "CONTEXTS = [\"pet\", \"feed\", \"fur\", \"vet\", \"drive\", \"fuel\", \"road\", \"tyre\"]\n",
    "\n",
    "M = np.array([\n",
    "    [4, 3, 3, 2, 0, 0, 0, 0],   # dog\n",
    "    [4, 3, 4, 2, 0, 0, 0, 0],   # cat\n",
    "    [0, 0, 0, 0, 4, 3, 3, 2],   # car\n",
    "    [0, 0, 0, 0, 4, 3, 4, 2],   # truck\n",
    "    [3, 2, 2, 3, 0, 0, 0, 0],   # pet\n",
    "    [2, 1, 1, 4, 0, 0, 0, 0],   # vet\n",
    "    [0, 0, 0, 0, 3, 2, 3, 2],   # drive\n",
    "    [0, 0, 0, 0, 2, 3, 2, 3],   # fuel\n",
    "], dtype=float)\n",
    "\n",
    "\n",
    "def ppmi(counts):\n",
    "    \"\"\"Positive pointwise mutual information, Chapter 7 Equations 7.2 and 7.3.\"\"\"\n",
    "    total = counts.sum()\n",
    "    p_wc = counts / total\n",
    "    p_w = p_wc.sum(axis=1, keepdims=True)\n",
    "    p_c = p_wc.sum(axis=0, keepdims=True)\n",
    "    with np.errstate(divide=\"ignore\", invalid=\"ignore\"):\n",
    "        pmi = np.log2(p_wc / (p_w * p_c))\n",
    "    return np.nan_to_num(np.maximum(pmi, 0.0), neginf=0.0, posinf=0.0)\n",
    "\n",
    "\n",
    "def truncate(U, S, Vt, k):\n",
    "    \"\"\"Rebuild the matrix from only the top k singular values. Equation 7.6.\"\"\"\n",
    "    return U[:, :k] @ np.diag(S[:k]) @ Vt[:k, :]\n",
    "\n",
    "\n",
    "def cosine(u, v):\n",
    "    n = np.linalg.norm(u) * np.linalg.norm(v)\n",
    "    return float(u @ v / n) if n else 0.0\n",
    "\n",
    "\n",
    "def main():\n",
    "    ap = argparse.ArgumentParser(description=__doc__)\n",
    "    ap.add_argument(\"--k\", type=int, default=None,\n",
    "                    help=\"show dense word vectors at this K\")\n",
    "    ap.add_argument(\"--ppmi\", action=\"store_true\",\n",
    "                    help=\"reweight with PPMI before factorising, as LSA does\")\n",
    "    a = ap.parse_args()\n",
    "\n",
    "    A = ppmi(M) if a.ppmi else M\n",
    "    label = \"PPMI-weighted counts\" if a.ppmi else \"raw counts\"\n",
    "    U, S, Vt = np.linalg.svd(A)\n",
    "\n",
    "    print(f\"\\nMATRIX  {A.shape[0]} words x {A.shape[1]} contexts, {label}\")\n",
    "    print(f\"rank    {np.linalg.matrix_rank(A)}\"\n",
    "          f\"   (nonzero singular values: {(S > 1e-10).sum()})\")\n",
    "\n",
    "    print(f\"\\nFACTORS   A = U . diag(sigma) . V^T\")\n",
    "    print(f\"  U  {U.shape}   sigma  {S.shape}   V^T  {Vt.shape}\")\n",
    "    print(f\"  sigma = \" + \"  \".join(f\"{s:.3f}\" for s in S))\n",
    "\n",
    "    print(f\"\\nTRUNCATION  what each K costs\\n\")\n",
    "    print(f\"  {'K':>2}{'kept share':>12}{'error':>10}{'stored':>9}\"\n",
    "          f\"{'vs full':>9}\")\n",
    "    print(\"  \" + \"-\" * 42)\n",
    "    full_store = A.size\n",
    "    for k in range(1, len(S) + 1):\n",
    "        Ak = truncate(U, S, Vt, k)\n",
    "        err = np.linalg.norm(A - Ak) / np.linalg.norm(A)\n",
    "        share = S[:k].sum() / S.sum()\n",
    "        store = k * (A.shape[0] + A.shape[1] + 1)\n",
    "        print(f\"  {k:>2}{share:>11.1%}{err:>10.3f}{store:>9}\"\n",
    "              f\"{store / full_store:>8.1f}x\")\n",
    "\n",
    "    print(\"\\n  'error' is the relative Frobenius distance from the original.\")\n",
    "    print(\"  It falls fast to K=2, then only crawls. Where it flattens, the\")\n",
    "    print(\"  extra dimensions were describing noise, not structure.\")\n",
    "    print(\"\\n  'stored' counts the numbers you must keep: k*(rows+cols+1).\")\n",
    "    print(\"  Note it passes 1.0x at K=4. Truncation only saves space while K\")\n",
    "    print(\"  stays small, which on a real vocabulary it always does.\")\n",
    "\n",
    "    if a.k:\n",
    "        k = a.k\n",
    "        dense = U[:, :k] @ np.diag(S[:k])\n",
    "        print(f\"\\nDENSE WORD VECTORS at K={k}\"\n",
    "              f\"   ({A.shape[1]} dimensions down to {k})\\n\")\n",
    "        for w, v in zip(WORDS, dense):\n",
    "            print(f\"  {w:<8}\" + \"  \".join(f\"{x:+7.3f}\" for x in v))\n",
    "\n",
    "        print(f\"\\nCOSINE SIMILARITY in the compressed space\\n\")\n",
    "        pairs = [(\"dog\", \"cat\"), (\"car\", \"truck\"), (\"dog\", \"car\"),\n",
    "                 (\"pet\", \"vet\"), (\"drive\", \"fuel\")]\n",
    "        for x, y in pairs:\n",
    "            i, j = WORDS.index(x), WORDS.index(y)\n",
    "            print(f\"  {x:<7} {y:<7} {cosine(dense[i], dense[j]):+.3f}\")\n",
    "        print(\"\\n  Words from the same group score 1. Words from different\")\n",
    "        print(\"  groups score 0. Two dimensions were enough, because the data\")\n",
    "        print(\"  only ever had two, which is what svd_rank.py detects.\")\n",
    "    print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Running it\n",
    "\n",
    "The script is a command line tool, so each run below sets the arguments\n",
    "and calls `main()`. The heading shows the equivalent shell command."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "import sys\n",
    "\n",
    "def run(*args):\n",
    "    sys.argv = [\"svd.py\", *args]\n",
    "    main()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `python3 svd.py`\n",
    "\n",
    "The three factors and the reconstruction."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "run()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `python3 svd.py --k 2`\n",
    "\n",
    "What the chosen K buys. dog and cat score 1.000, dog and car score 0.000."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "run(\"--k\", \"2\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `python3 svd.py --ppmi`\n",
    "\n",
    "Reweight before factorising, which is what LSA does."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "run(\"--ppmi\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "Read the chapter this comes from: [Count Vectors, PPMI, and SVD](https://nlp.jcrlabz.com/book/countvectors/)."
   ]
  }
 ],
 "metadata": {
  "colab": {
   "name": "svd.ipynb",
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}