{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# COALS, all three steps\n",
    "\n",
    "The correlation matrix, reproducing the lecture tables on the woodchuck corpus.\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/coals.py`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## What you need\n",
    "\n",
    "Nothing. This example uses only the Python standard library."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Nothing to install: this example uses only the Python standard library.\n",
    "import sys; print(sys.version.split()[0])"
   ]
  },
  {
   "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": [
    "\"\"\"COALS, traced as Chapter 7 traces it.\n",
    "\n",
    "Counts become correlations, negatives are discarded, positives are square\n",
    "rooted. Run this and you get the same three tables the lectures show for the\n",
    "woodchuck corpus.\n",
    "\n",
    "    python3 coals.py              # all three steps\n",
    "    python3 coals.py --step 2     # just the correlation matrix\n",
    "    python3 coals.py --cell a if  # the arithmetic for one cell, shown long\n",
    "\n",
    "The point of the correlation step is that it factors out frequency. Watch the\n",
    "word \"a\", which is the most common word here, stop dominating.\n",
    "\n",
    "Install:  nothing, the Python standard library is enough\n",
    "\"\"\"\n",
    "\n",
    "import argparse\n",
    "import math\n",
    "\n",
    "WORDS = [\"a\", \"as\", \"chuck\", \"could\", \"how\", \"if\", \"much\", \"wood\",\n",
    "         \"woodchuck\", \"would\", \",\", \".\", \"?\"]\n",
    "\n",
    "# The step-1 count matrix from the lectures, built with a ramped window of 4\n",
    "# over: \"How much wood would a woodchuck chuck, if a woodchuck could chuck\n",
    "# wood? As much wood as a woodchuck would, if a woodchuck could chuck wood.\"\n",
    "COUNTS = [\n",
    "    [0,  5,  9,  6,  1, 10,  4,  8, 18,  9, 10,  0,  0],   # a\n",
    "    [5,  4,  2,  1,  0,  0,  7, 10,  3,  2,  1,  0,  5],   # as\n",
    "    [9,  2,  0,  8,  0,  0,  1,  9, 11,  2,  4,  3,  3],   # chuck\n",
    "    [6,  1,  8,  0,  0,  4,  0,  6,  8,  0,  2,  2,  2],   # could\n",
    "    [1,  0,  0,  0,  0,  4,  3,  0,  0,  2,  0,  0,  0],   # how\n",
    "    [10, 0,  5,  4,  0,  0,  4,  3, 10,  3,  0,  0,  0],   # if\n",
    "    [4,  7,  1,  0,  4,  0,  0, 10,  2,  3,  0,  0,  3],   # much\n",
    "    [8, 10,  9,  6,  3,  0, 10,  2,  8,  5,  0,  4,  6],   # wood\n",
    "    [18, 3, 11,  0,  0, 10,  2,  8,  8, 10,  1,  1,  1],   # woodchuck\n",
    "    [9,  2,  2,  0,  2,  3,  5,  8,  0,  5,  0,  0,  0],   # would\n",
    "    [10, 1,  4,  2,  0,  0,  0, 10,  5,  0,  0,  0,  0],   # ,\n",
    "    [0,  0,  3,  2,  0,  0,  4,  1,  1,  0,  0,  0,  0],   # .\n",
    "    [0,  5,  3,  2,  0,  0,  3,  6,  1,  0,  0,  0,  0],   # ?\n",
    "]\n",
    "\n",
    "\n",
    "def margins(M):\n",
    "    \"\"\"Row sums, column sums, and the grand total T.\"\"\"\n",
    "    rows = [sum(r) for r in M]\n",
    "    cols = [sum(M[i][j] for i in range(len(M))) for j in range(len(M[0]))]\n",
    "    return rows, cols, sum(rows)\n",
    "\n",
    "\n",
    "def correlation(w_ab, row_a, col_b, T):\n",
    "    \"\"\"Equation 7.7.\n",
    "\n",
    "    Numerator: what we observed (scaled by T) minus what independence predicts\n",
    "    (the product of the margins). Same comparison PMI makes.\n",
    "\n",
    "    Denominator: how much each margin could vary. This is what puts every cell\n",
    "    on the same [-1, 1] scale regardless of how common the words are, and it\n",
    "    is exactly what HAL lacked.\n",
    "    \"\"\"\n",
    "    num = T * w_ab - row_a * col_b\n",
    "    den = math.sqrt(row_a * (T - row_a) * col_b * (T - col_b))\n",
    "    return num / den if den else 0.0\n",
    "\n",
    "\n",
    "def coals_value(r):\n",
    "    \"\"\"Equation 7.8. Discard negatives, damp the survivors.\"\"\"\n",
    "    return math.sqrt(r) if r > 0 else 0.0\n",
    "\n",
    "\n",
    "def table(M, title, fmt=\"{:>7.3f}\"):\n",
    "    print(f\"\\n{title}\\n\")\n",
    "    head = \"  \" + \" \" * 10 + \"\".join(f\"{w[:6]:>7}\" for w in WORDS)\n",
    "    print(head)\n",
    "    for i, w in enumerate(WORDS):\n",
    "        print(f\"  {w:<10}\" + \"\".join(fmt.format(v) for v in M[i]))\n",
    "\n",
    "\n",
    "def one_cell(a, b):\n",
    "    \"\"\"Show the arithmetic for a single cell, the long way.\"\"\"\n",
    "    rows, cols, T = margins(COUNTS)\n",
    "    i, j = WORDS.index(a), WORDS.index(b)\n",
    "    w, ra, cb = COUNTS[i][j], rows[i], cols[j]\n",
    "    num = T * w - ra * cb\n",
    "    den = math.sqrt(ra * (T - ra) * cb * (T - cb))\n",
    "    r = num / den\n",
    "    print(f\"\\nCELL ({a}, {b})\\n\")\n",
    "    print(f\"  count w[{a}][{b}]        = {w}\")\n",
    "    print(f\"  row sum for '{a}'         = {ra}\")\n",
    "    print(f\"  column sum for '{b}'      = {cb}\")\n",
    "    print(f\"  grand total T            = {T}\")\n",
    "    print(f\"\\n  numerator   = T*w - row*col = {T}*{w} - {ra}*{cb}\"\n",
    "          f\" = {T*w} - {ra*cb} = {num}\")\n",
    "    print(f\"  denominator = sqrt({ra}*{T-ra} * {cb}*{T-cb}) = {den:.0f}\")\n",
    "    print(f\"\\n  r = {num} / {den:.0f} = {r:+.3f}\")\n",
    "    print(f\"  after clipping and square root: {coals_value(r):.3f}\")\n",
    "    if w == 0:\n",
    "        print(f\"\\n  Note the count is zero and r is negative. '{a}' and '{b}'\")\n",
    "        print(\"  do not merely fail to co-occur, they avoid each other.\")\n",
    "\n",
    "\n",
    "def main():\n",
    "    ap = argparse.ArgumentParser(description=__doc__)\n",
    "    ap.add_argument(\"--step\", type=int, choices=[1, 2, 3], default=None)\n",
    "    ap.add_argument(\"--cell\", nargs=2, metavar=(\"A\", \"B\"))\n",
    "    a = ap.parse_args()\n",
    "\n",
    "    if a.cell:\n",
    "        one_cell(*a.cell)\n",
    "        print()\n",
    "        return\n",
    "\n",
    "    rows, cols, T = margins(COUNTS)\n",
    "    R = [[correlation(COUNTS[i][j], rows[i], cols[j], T)\n",
    "          for j in range(len(WORDS))] for i in range(len(WORDS))]\n",
    "    C = [[coals_value(v) for v in row] for row in R]\n",
    "\n",
    "    if a.step in (None, 1):\n",
    "        table(COUNTS, \"STEP 1: raw counts, ramped window of 4\", \"{:>7d}\")\n",
    "        print(f\"\\n  grand total T = {T}\")\n",
    "        print(f\"  row sums    {dict(zip(WORDS, rows))}\")\n",
    "    if a.step in (None, 2):\n",
    "        table(R, \"STEP 2: counts converted to correlations\")\n",
    "        print(\"\\n  Every value now sits in [-1, 1], whatever the word's\"\n",
    "              \" frequency.\")\n",
    "        print(\"  'a' is the most frequent word here and no longer dominates.\")\n",
    "    if a.step in (None, 3):\n",
    "        table(C, \"STEP 3: negatives set to 0, positives square rooted\")\n",
    "        neg = sum(1 for row in R for v in row if v < 0)\n",
    "        print(f\"\\n  {neg} of {len(WORDS)**2} cells were negative and are now 0,\"\n",
    "              f\" which keeps the matrix sparse.\")\n",
    "        print(\"  The square root pulls in the large values, so a few strong\")\n",
    "        print(\"  pairings cannot dominate a vector.\")\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 = [\"coals.py\", *args]\n",
    "    main()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `python3 coals.py`\n",
    "\n",
    "All three steps, in order."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "run()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `python3 coals.py --cell a a`\n",
    "\n",
    "One cell the long way. The word never sits beside itself, so the count is zero and the correlation goes negative. The statistic noticed an avoidance, not just an absence."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "run(\"--cell\", \"a\", \"a\")"
   ]
  },
  {
   "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": "coals.ipynb",
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}