Explore CoDEx.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you're running this notebook locally on your own machine, you can ignore the following cell. If you want to [run this notebook remotely on Google's Colab platform](https://colab.research.google.com/github/tsafavi/codex/blob/master/Explore%20CoDEx.ipynb), uncomment and run the lines in the following cell: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !git clone https://github.com/tsafavi/codex\n",
    "# %cd codex\n",
    "# !./extract.sh en  # unzip all KG text in English\n",
    "# !pip install -r requirements.txt\n",
    "# !pip install -e ."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you're running on Colab, you will need to restart the runtime after running the previous cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "from codex.codex import Codex"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sizes = Codex.SIZES\n",
    "codes = Codex.CODES"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Data exploration\n",
    "Load our data and do basic exploratory analysis."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for size in sizes:\n",
    "    codex = Codex(size=size)\n",
    "    \n",
    "    train, valid, test = [\n",
    "        codex.split(split) for split in (\"train\", \"valid\", \"test\")]\n",
    "    triples = codex.triples()\n",
    "    \n",
    "    print(codex.name())\n",
    "    \n",
    "    print(\n",
    "        \" \", len(codex.entities()), \"entities /\",\n",
    "        len(codex.relations()), \"relations\"\n",
    "    )\n",
    "        \n",
    "    print(\n",
    "        \" \", len(train), \"train /\",\n",
    "        len(valid), \"validation /\",\n",
    "        len(test), \"test\"\n",
    "    )\n",
    "    print(\" \", len(triples), \"total triples\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Inspect a single entity."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eid = \"Q51\"\n",
    "\n",
    "for code in codes:\n",
    "    codex = Codex(code=code)\n",
    "    print(codex.entity_label(eid))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "codex = Codex(code=\"en\")\n",
    "print(f\"From {codex.entity_wikipedia_url(eid)}:\")\n",
    "print(f\"  '{codex.entity_extract(eid)}'\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Explore entity types."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "codex = Codex(code=\"en\")\n",
    "types = codex.entity_types(eid)\n",
    "for etype in types:\n",
    "    print(codex.entity_label(eid), \"is of type\", codex.entity_type_label(etype))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "type_id = \"Q5\"\n",
    "print(f\"From {codex.entity_type_wikipedia_url(type_id)}:\")\n",
    "print(f\"  '{codex.entity_type_extract(type_id)}'\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Inspect a single relation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rid = \"P530\"\n",
    "print(codex.relation_label(rid) + \":\", codex.relation_description(rid))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Explore some of the negative triples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for size in (\"s\", \"m\"):\n",
    "    codex = Codex(size=size)\n",
    "    valid_neg, test_neg = (\n",
    "        codex.negative_split(\"valid\"),\n",
    "        codex.negative_split(\"test\")\n",
    "    )\n",
    "    print(\n",
    "        codex.name(), \"has\", \n",
    "        len(valid_neg), \"validation negatives and\",\n",
    "        len(test_neg), \"test negatives\"\n",
    "    )\n",
    "    \n",
    "    print(\"Samples:\")\n",
    "    sample = test_neg.sample(n=5)\n",
    "\n",
    "    rows = [(tuple(row)) for row in sample.values]\n",
    "    for h, r, t in rows:\n",
    "        print(\" \",\n",
    "              f\"({codex.entity_label(h)},\",\n",
    "              f\"{codex.relation_label(r)},\", \n",
    "              f\"{codex.entity_label(t)})\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Make sure none of the negative triples are actually positives."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_triples = {tuple(row) for row in pd.concat(\n",
    "    [Codex(size=size).triples() for size in sizes]).values}\n",
    "    \n",
    "for size in (\"s\", \"m\"):\n",
    "    codex = Codex(size=size)\n",
    "    valid_neg, test_neg = (\n",
    "        codex.negative_split(\"valid\"),\n",
    "        codex.negative_split(\"test\")\n",
    "    )\n",
    "    \n",
    "    for neg_triples in (valid_neg, test_neg):\n",
    "        neg_triples = {tuple(row) for row in neg_triples.values}\n",
    "        assert not len(neg_triples.intersection(all_triples))\n",
    "        \n",
    "print(\"Passed all checks successfully\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Plotting\n",
    "Plot various properties like degree distribution and top-k entities/relations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "\n",
    "import logging\n",
    "logging.getLogger(\"matplotlib.font_manager\").disabled = True\n",
    "\n",
    "# Styles from https://scentellegher.github.io/visualization/2018/10/10/beautiful-bar-plots-matplotlib.html\n",
    "\n",
    "plt.rcParams[\"font.family\"] = \"sans-serif\"\n",
    "plt.rcParams[\"font.sans-serif\"] = \"Helvetica\"\n",
    "\n",
    "plt.rcParams[\"axes.linewidth\"] = 0.8\n",
    "plt.rcParams[\"legend.frameon\"] = False\n",
    "\n",
    "dark_grey = \"#333F4B\"\n",
    "plt.rcParams[\"text.color\"] = dark_grey\n",
    "keys = [\n",
    "    \"axes.edgecolor\",\n",
    "    \"axes.labelcolor\",\n",
    "    \"xtick.color\",\n",
    "    \"ytick.color\"\n",
    "]\n",
    "for key in keys:\n",
    "    plt.rcParams[key] = dark_grey\n",
    "\n",
    "palette = sns.color_palette(\"RdBu\", 10)\n",
    "# sns.palplot(palette)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot the degree distribution of CoDEx datasets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def count_entities(triples):\n",
    "    \"\"\"Get the degree of all entities across train/valid/test\"\"\"\n",
    "    entities = pd.concat((triples[\"head\"], triples[\"tail\"]))\n",
    "    df = pd.DataFrame(\n",
    "        entities.value_counts().reset_index(),\n",
    "    )\n",
    "    df.columns = [\"entity\", \"count\"]\n",
    "    return df\n",
    "\n",
    "def count_relations(triples):\n",
    "    \"\"\"Get relation mention counts\"\"\"\n",
    "    df = pd.DataFrame(\n",
    "        triples[\"relation\"].value_counts().reset_index(),\n",
    "    )\n",
    "    df.columns = [\"relation\", \"count\"]\n",
    "    return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_degree_distribution(\n",
    "        triples, step=5, ax=None, \n",
    "        figsize=(5,3), **kwargs):\n",
    "    \"\"\"Plot the degree distribution of all entities\n",
    "    in a set of triples\"\"\"\n",
    "    counts = count_entities(triples)[\"count\"]\n",
    "    print(f\"Average degree: {counts.mean():.2f}\")\n",
    "    \n",
    "    x = np.arange(0, max(counts), step)\n",
    "    y = [len(counts[counts > val]) for val in x]\n",
    "    y = np.array(y, dtype=np.float32)\n",
    "\n",
    "    if not ax:\n",
    "        fig, ax = plt.subplots(figsize=figsize)\n",
    "    \n",
    "    ax.scatter(x, y, alpha=0.5, **kwargs)\n",
    "    ax.set_xscale(\"symlog\")\n",
    "    ax.set_yscale(\"symlog\")\n",
    "    \n",
    "    ax.set_xlabel(\"Entity degree\", fontsize=14)\n",
    "    ax.set_ylabel(\"Cumulative count\", fontsize=12)\n",
    "    \n",
    "    ax.tick_params(\n",
    "        which=\"both\",\n",
    "        labelsize=12\n",
    "    )\n",
    "    \n",
    "    ax.spines[\"top\"].set_color(\"none\")\n",
    "    ax.spines[\"right\"].set_color(\"none\")\n",
    "    ax.spines[\"left\"].set_smart_bounds(True)\n",
    "    ax.spines[\"bottom\"].set_smart_bounds(True)\n",
    "\n",
    "    return ax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for size in sizes:\n",
    "    codex = Codex(size=size)\n",
    "    triples = codex.triples()\n",
    "\n",
    "    ax = plot_degree_distribution(\n",
    "        triples, \n",
    "        color=palette[-1],\n",
    "        step=5,\n",
    "        figsize=(4, 2),\n",
    "        label=codex.name(),\n",
    "    )\n",
    "\n",
    "    legend = ax.legend(\n",
    "        frameon=False, \n",
    "        fontsize=14,\n",
    "        handletextpad=0.001,\n",
    "        labelspacing=0.2,\n",
    "        loc=3\n",
    "    )\n",
    "\n",
    "    for text in legend.get_texts():\n",
    "        text.set_color(dark_grey)\n",
    "\n",
    "    for lh in legend.legendHandles: \n",
    "        lh.set_alpha(0.7)\n",
    "\n",
    "    xlim = ax.get_xlim()\n",
    "    ax.set_xlim(3, max(xlim))\n",
    "\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot the top-k entities and relations in a CoDEx dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_top_k(count_df, k=30, \n",
    "               label_col=\"label\",\n",
    "               count_col=\"count\",\n",
    "               figsize=(5,8.5),\n",
    "               color=\"#007acc\",\n",
    "               ax=None,\n",
    "               **kwargs):\n",
    "    count_df = count_df.sort_values(count_col, ascending=False).head(k)\n",
    "    top_labels, top_counts = count_df[label_col], count_df[count_col]\n",
    "    \n",
    "    if ax is None:\n",
    "        fig, ax = plt.subplots(figsize=figsize)\n",
    "\n",
    "    x = np.arange(k)\n",
    "    y = top_counts[::-1]\n",
    "        \n",
    "    ax.hlines(\n",
    "        y=x,\n",
    "        xmin=0, \n",
    "        xmax=y, \n",
    "        color=color, \n",
    "        alpha=0.2, \n",
    "        linewidth=5,\n",
    "        **kwargs\n",
    "    )\n",
    "    \n",
    "    ax.plot(\n",
    "        y, x, \"o\", \n",
    "        markersize=7, \n",
    "        color=color, \n",
    "        alpha=0.6)\n",
    "    \n",
    "    ax.spines[\"top\"].set_color(\"none\")\n",
    "    ax.spines[\"right\"].set_color(\"none\")\n",
    "\n",
    "    ax.spines[\"left\"].set_bounds(0, k - 1)\n",
    "    ax.spines[\"bottom\"].set_bounds(0, max(ax.get_xticks()))\n",
    "    \n",
    "    ax.set_xlabel(\"Count\", fontsize=14)\n",
    "\n",
    "    ax.set_yticks(x)\n",
    "    ax.set_yticklabels(top_labels[::-1])\n",
    "    \n",
    "    return ax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "size = \"l\"\n",
    "codex = Codex(size=size)\n",
    "\n",
    "triples = codex.triples()\n",
    "count_df = count_entities(triples)\n",
    "\n",
    "count_df[\"label\"] = [\n",
    "    codex.entity_label(eid) for eid in count_df[\"entity\"]]\n",
    "\n",
    "k = 30\n",
    "\n",
    "ax = plot_top_k(\n",
    "    count_df, \n",
    "    k=k, \n",
    "    color=palette[-1],\n",
    "    linewidths=6\n",
    ")\n",
    "\n",
    "ax.set_xscale(\"linear\")\n",
    "ax.set_xlabel(\"Entity degree\", fontsize=14)\n",
    "ax.set_title(codex.name(), fontsize=16)\n",
    "ax.tick_params(\"x\", labelsize=12)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "count_df = count_relations(triples)\n",
    "count_df[\"label\"] = [\n",
    "    codex.relation_label(rid) for rid in count_df[\"relation\"]]\n",
    "\n",
    "k = 15\n",
    "\n",
    "ax = plot_top_k(\n",
    "    count_df, \n",
    "    k=k, \n",
    "    color=palette[-1],\n",
    "    linewidths=6,\n",
    "    figsize=(5, 4)\n",
    ")\n",
    "\n",
    "ax.set_xscale(\"linear\")\n",
    "ax.set_xlabel(\"Mention count\", fontsize=14)\n",
    "ax.set_title(codex.name(), fontsize=16)\n",
    "ax.tick_params(\"x\", labelsize=12)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Investigate symmetry."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def entity_pairs(triples, reverse=False):\n",
    "    \"\"\"Get all unique head/tail entity pairs in a set of triples\"\"\"\n",
    "    col1, col2 = \"head\", \"tail\"\n",
    "    if reverse:\n",
    "        col1, col2 = col2, col1\n",
    "    return set(zip(triples[col1], triples[col2]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "threshold = 0.5\n",
    "\n",
    "for size in sizes: \n",
    "    codex = Codex(size=size)\n",
    "    print(\"----------\" + codex.name() + \"----------\")\n",
    "    relations = codex.relations()\n",
    "    triples = codex.triples()\n",
    "    tot = 0\n",
    "    \n",
    "    for relation in relations:\n",
    "        triples_r = triples[triples[\"relation\"] == relation]\n",
    "        \n",
    "        forward, reverse = entity_pairs(triples_r), \\\n",
    "            entity_pairs(triples_r, reverse=True)\n",
    "        symmetry = len(forward.intersection(reverse)) / len(forward)\n",
    "        percent = len(triples_r) / len(triples)\n",
    "        \n",
    "        if symmetry > threshold:\n",
    "            print(\n",
    "                f\"  {codex.relation_label(relation)} ({relation})\", \n",
    "                f\"| {symmetry * 100:.2f}% symmetry\", \n",
    "                f\"| {percent * 100:.3f}% of {codex.name()}\"\n",
    "            )\n",
    "            \n",
    "            tot += (percent * 100)\n",
    "                \n",
    "    print(f\"Symmetric relations > {threshold:.1f} make up {tot:.2f}% of {codex.name()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Investigate compositionality."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def rule_relations(rule):\n",
    "    return [item for item in rule.split(\" \") if item.startswith(\"P\")]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for size in sizes:\n",
    "    codex = Codex(size=size)\n",
    "    name = codex.name()\n",
    "    print(\"----------\" + name + \"----------\")\n",
    "    \n",
    "    rule_df = pd.read_csv(\n",
    "        os.path.join(\"analysis/paths\", f\"codex-{size}\", \"paths.tsv\"), sep=\"\\t\"\n",
    "    )\n",
    "        \n",
    "    # Count number of paths in this core\n",
    "    # Get the average confidence across all rules\n",
    "    print(\n",
    "        len(rule_df), \"paths in\", name\n",
    "    )\n",
    "    \n",
    "    avg_conf = rule_df[\"PCA Confidence\"].mean()\n",
    "    print(f\"Average confidence: {avg_conf:.3f}\")\n",
    "    \n",
    "    frac_of_total = (\n",
    "        np.sum(np.multiply(rule_df[\"Positive Examples\"], rule_df[\"Length\"])) / \n",
    "        len(codex.triples()) * 100\n",
    "    )\n",
    "    print(f\"n pos / n triples: {frac_of_total:.2f}%\")\n",
    "        \n",
    "    # Print a few examples\n",
    "    print(\"Examples:\")\n",
    "    for length in (2, 3):\n",
    "        df = rule_df[rule_df[\"Length\"] == length]\n",
    "        sample = df.sample(n=1)\n",
    "        rule, conf = sample[\"Rule\"].iloc[0], sample[\"PCA Confidence\"].iloc[0]\n",
    "        relations = rule_relations(rule)\n",
    "        \n",
    "        print(\n",
    "            \" \",\n",
    "            \"/\".join([codex.relation_label(relation) for relation in relations[:-1]]), \n",
    "            \"->\", \n",
    "            codex.relation_label(relations[-1]),\n",
    "            f\"(confidence {conf:.3f})\"\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}