{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "831e4e1d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt \n",
    "\n",
    "plt.rcParams.update({\n",
    "        \"text.usetex\": True,\n",
    "        \"font.family\": \"serif\",\n",
    "        \"font.serif\": [\"Computer Modern Roman\"],\n",
    "        \"text.latex.preamble\": r\"\\usepackage{amsmath,amssymb,amsfonts,mathrsfs,bm}\",\n",
    "        \"axes.labelsize\": 26,\n",
    "        \"font.size\": 20,\n",
    "        \"legend.fontsize\": 20,\n",
    "        \"xtick.labelsize\": 20,\n",
    "        \"ytick.labelsize\": 20,\n",
    "        \"axes.titlesize\": 22\n",
    "    })"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "dab0b956",
   "metadata": {},
   "outputs": [],
   "source": [
    "data = {\n",
    "    \"Original\": {\n",
    "        3: 34,\n",
    "        4: 28,\n",
    "        5: 22,\n",
    "        6: 26,\n",
    "        7: 14,\n",
    "        8: 18,\n",
    "        9: 22,\n",
    "        10: 18,\n",
    "    },\n",
    "    \"Non-Curriculum\": {\n",
    "        3: 88,\n",
    "        4: 78,\n",
    "        5: 70,\n",
    "        6: 72,\n",
    "        7: 50,\n",
    "        8: 54,\n",
    "        9: 40,\n",
    "        10: 46,\n",
    "    },\n",
    "    \"Curriculum\": {\n",
    "        3: 100,\n",
    "        4: 92,\n",
    "        5: 86,\n",
    "        6: 90,\n",
    "        7: 70,\n",
    "        8: 70,\n",
    "        9: 56,\n",
    "        10: 64,\n",
    "    }\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "859a7f84",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 51388 characters]",
      "text/plain": [
       "<Figure size 1200x600 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "\n",
    "cats   = list(data.keys())         # y-axis\n",
    "grades = list(range(3, 11))        # x-axis\n",
    "\n",
    "mat = np.array([[data[cat].get(g, 0) for g in grades] for cat in cats])\n",
    "\n",
    "fig, ax = plt.subplots(\n",
    "    figsize=(max(12, len(grades) * 0.8), max(4, len(cats) * 2.0))\n",
    ")\n",
    "\n",
    "im = ax.imshow(mat, cmap=\"YlOrRd\", vmin=0, vmax=100, aspect=\"auto\")\n",
    "\n",
    "# numbers inside cells\n",
    "for i in range(mat.shape[0]):\n",
    "    for j in range(mat.shape[1]):\n",
    "        text_color = \"white\" if mat[i, j] >= 50 else \"black\"\n",
    "        ax.text(j, i, f\"{mat[i, j]}\", ha=\"center\", va=\"center\", fontsize=25, color=text_color)\n",
    "\n",
    "ax.set_xticks(range(len(grades)))\n",
    "ax.set_xticklabels(grades)\n",
    "ax.set_yticks(range(len(cats)))\n",
    "ax.set_yticklabels(cats)\n",
    "ax.set_xlabel(r\"\\textbf{Word Length}\")\n",
    "\n",
    "cbar = fig.colorbar(im, ax=ax)\n",
    "cbar.ax.set_ylabel(r\"\\textbf{Score (\\%)}\", rotation=-90, va=\"bottom\", fontsize=26)\n",
    "\n",
    "plt.tight_layout()\n",
    "\n",
    "# save as pdf with high resolution\n",
    "plt.savefig(\"curriculum_heatmap.pdf\", dpi=300, bbox_inches=\"tight\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11b374b2",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "reasoning_gym",
   "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.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
