abstraction_alignment_toy_example.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Abstraction Alignment on a Toy Example\n",
    "\n",
    "We apply abstraction alignment in a synthetic setting where we have enforced a hierarchical relationship in the model's behavior. First, we generating synthetic image data with inherent hierarchical structure. Then we train a simple neural network classifier on the data and analyze how well the model's behavior aligns with the known hierarchical structure."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import os\n",
    "from matplotlib import pyplot as plt\n",
    "import random\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from tqdm import tqdm\n",
    "from treelib import Tree\n",
    "from scipy import stats\n",
    "\n",
    "import metrics\n",
    "from abstraction_graph_toy_example import make_abstraction_graph, show_abstraction_graph, propagate\n",
    "from util.toy_example.data_generation import SyntheticDataGenerator, create_dataloaders"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generating the Synthetic Dataset\n",
    "We generate synthetic 4x4 RGB images with three classes. These classes are the result of a higher-order structure. We can think of these classes are image generation functions, where sometimes two generation functions share a higher-level parent function.\n",
    "\n",
    "```\n",
    "-Parent 0: Images with three corners sharing the same color\n",
    "--Class 0: Images with three corners sharing the same color\n",
    "-Parent 1: Images with two four-pixel same-color segments\n",
    "--Class 1: Images with two specific quadrants sharing the same color\n",
    "--Class 2: Images with alternating columns of the same color\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAq4AAAGqCAYAAADKhsaNAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjQuMywgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/MnkTPAAAACXBIWXMAAAsTAAALEwEAmpwYAAAZ80lEQVR4nO3de7Ctd13f8c+XXIiRkI6ccBmmkAEiCFQsopCKykhEiWAisVPlor3oJDXU2hAoAs3eBwQzgJlMR9IpFqpcSoe7BBqhabhEhTYUoUBUDCExRFID5HJCk/QkfPvHeg7sHvY+t5zz7PXb5/WaOXPOfp611vN7nv2std/57WetVHcHAACW3b02ewAAALAvhCsAAEMQrgAADEG4AgAwBOEKAMAQhCsAAEMQrgB8S1VdU1XnbvY4ANYjXIFDqqp+v6p6+rOzqv62qj5cVWdX1VH7+VhPmR5n26Ea7wbbPXHa7hP24ba9wZ+z5hgrwFZ25GYPADgsXJrkeUmOSHJCkp9Isj3J86rqqd39jc0c3CHwq0nev9uyWzZjIABbiRlXYA53dvcN3X19d3+6uy9I8pQkj0/yol03qqrnVtUVVbVjmpl9R1U9eFp3YpIPTze9cZrF/P1p3U9X1eVVdVNVfb2qPlhV37d2AFV1XlVdW1V3VtUNVfWmNeuqql5UVV+sqtur6rNV9dw1d//S9PcV03Y/spf9vXna37V/bp+29Yaq+nxVfdf09RFV9cdV9a3Qrarzq+ovp7FcU1Wvrqpj1qxfrarPVdUvT+tvq6r/WFVHV9WvVdV1VfW1qrqgqu615n7XTPd9y3SfG/Z2WUBVHV9Vr5++Hzuq6qNrZ56n9W+e1t9RVVdX1W/s5fgAHBDhCmyK7v5ckj9KcsaaxUcnWUnyuCTPSLItydumddetue1jkjwoyb+cvv7uJBcm+eEsgviWJBdX1dFJUlVnJDk3ya8lOWl67P+xZru/leSfJTk7yaOT/HaSf19VPzOt/+Hp75+etvusA93vJL+e5Kgkr52+fmmSRyT5p2tu843p6++bxvwL0+3WOjHJadO+nJHkHyb5wyQ/lORpSX4lyb9I8nO73e+cJH+exX80rCR5VVWtuz9VVUk+kOTB03b+fpKPJbmsqh403ey3kvy9af2jpnFfv7eDAHAgqrs3ewzAFjbNim7r7mess+78JL/e3cducN9HZRFZf7e7v1xVT8li1vWE7v7qHrb53UluTfLj3f3HVXVOkjOTPLa7d65z268meVp3X75m+YVJvre7T51me7+U5Ie6+5N72d9OckeSu3dbdXJ3f3a6zROS/GmS85P8ZpKf7e5L9vCYZyU5t7sfMX29muRfJ3lgd98yLXtnkh9P8uDu/r/Tso8k+Vx3P3/6+pokf9XdP7nmsf9Dkkd195PX3OZ3u/u1VfUTSd6XxfG+fc19Pp3kP3X3q6vqfUm+1t3/ZE/HBeBgcI0rsJkqybf+67mqds0C/kCS75nWJ8lDknx5wwepeniSVyR5YhbX0N5r+vOQ6SbvyGJ29ktV9cEsZnrf1913ZjHDekySP5qic5ejklxzgPv1wmkba/31rn909yer6pVJVpNctHu0VtXPJ/mNLGZi75PFtcFH7P54u6J18r+TfGFXtK5Zdv/d7vfxdb7eaAb5B5Mcm8WlGWuXH5Pk4dO//12Sd07fu/+a5OLu/ugGjwdwjwhXYDM9OsnVybdmPj+Yb7+R62+zuFTg8iwuIdiTi7P49fSZ0993Jbly1/26+7qqemSSpyY5JcnvJFmpqifm25dMPTNr4nKyMwfmhu6+aqOV06/gn5zFrOzDq6p6+vVXVT0pyX/O4s1r/yrJzUl+Nt++tGCjsfUGy3YP3v1xryzi90fXWXdrknT3JVX10CRPz+L4fqCq3mEGFjgUhCuwKarqsVlcM/pb06JHZRGqL+nuL0232X0mcNds4rdirKrul8W1oGd394enZY/Pbq9v3X1HFtdrfmC6ROGGJD+SxYzjnUke2t2XbTDc79juPXROFteY/liS/5LFtaj/dlr3I0mu7+5X7LrxFIYHy5PW+frPN7jtp5I8IMk3u/vqjR5wumzjzUneXFWXJHlbVZ01zWgDHDTCFZjDvavqgVnM4J2QxczcS5L8z3x7JvGvswjI51fV67KI0Vfs9jjXZjGL+DNVdXGS25PclMU1qr9aVddl8Uai12Qx65okqap/nMXr3X9PcluSf5TF7ORfdfeOqnptktdOM6Efy+LX80/KIthen8Xs7+1Jfmq6BvSO3X5Nv7u/M+3vWrd1921V9bgkr0zynO7+06r650neWFWXTW9Y+0KSB1fVc7KI6p9K8ot72Nb+elJV/WaSd2bxRrZfSvKcDW57aZI/SfKHVfWiJH+R5IFZ/AfHpd19eVW9PIvA/XwWx/hZSa4WrcCh4FMFgDmckuQrWcTpf8viV9/bk/zYrs9w7e4bk/xyktOz+DX/ShYzk9/S3ddPy1+Zxa+wf7e7v5lFiH5/ks8leV2Sf5NFBO9ycxafGnD5dJszkjxr18zudPvVLD554PNZXKt5RqaPweruu7L4NIBfSfI3Wbx7f09+b9rftX9ePH2k1VuzeGPTu6bHflsWEfnWqrp3d1+cRXhfmOR/JfnJJOftZXv744IsjtWfZTHbfV53v3O9G06XL5ya5LJpn/4yyduTPDKL45AsjvMrk3wmi8g9LovLLgAOOp8qAHCYWPuJAZs9FoADYcYVAIAhCFcAAIbgUgEAAIZgxhUAgCEIVwAAhiBcAQAYgnAFAGAIwhUAgCEIVwAAhiBcAQAYgnAFAGAIwhUAgCEIVwAAhiBcAQAYgnAFAGAIQ4drVa1W1Vs2exwAABx6Sx+uVfXsqvpkVd1WVV+pqkuq6smbNJYTq+rDVfV/quovquqUzRgHAMDhaKnDtarOSXJhklcleUCShyS5KMlpmzSktyX5syT3S/LSJO+sqhM2aSwAAIeVpQ3Xqjo+ycuTnN3d7+7ub3T3zu6+uLtfuMF93lFVN1TVLVX1sap6zJp1p1bVlVW1o6qur6pzp+Xbqur9VXVzVX29qi6vqu84LlX1vUken2Slu2/v7ncl+WySMw7F/gMA8P9b2nBNcnKSY5K8Zz/uc0mSk5LcP8mnkrx1zbo3JDmzu49L8tgkl03LX5Dky0lOyGJW9yVJep3HfkySq7t7x5pln5mWAwBwiC1zuN4vyVe7+659vUN3v7G7d3T3nUlWkzxumrlNkp1JHl1V9+3um7r7U2uWPyjJQ6cZ3cu7e71wvU+SW3ZbdkuS4/ZjnwAAOEDLHK5fS7Ktqo7clxtX1RFVdX5VfbGqbk1yzbRq2/T3GUlOTXJtVX20qk6elr8myVVJPlRVV1fVizfYxG1J7rvbsvsm2bHObQEAOMiWOVw/nuSOJKfv4+2fncWbtk5JcnySE6fllSTdfUV3n5bFZQTvTfL2afmO7n5Bdz8syTOTnFNVT13n8T+f5GFVtXaG9XHTcgAADrGlDdfuviXJeUleV1WnV9WxVXVUVT29ql69zl2OS3JnFjO1x2bxSQRJkqo6uqqeU1XHd/fOJLcmuXta94yqekRV1Zrld68zni8k+XSSlao6pqp+Lsn3J3nXQdxtAAA2sLThmiTdfUGSc5K8LMmNSa5L8vwsZkx396Yk1ya5PsmVST6x2/rnJblmuozgrCTPnZaflOTSLC4F+HiSi7r7IxsM6ReSPCHJTUnOT/Lz3X3jAewaAAD7qdZ/HxIAACyXpZ5xBQCAXYQrAABDEK4AAAxBuAIAMIQ9frh/bV/3f3160PVqzbGZJEnN9V607fO96a1XMt8B3AfOmwPXc34ne9at7V3Nc5Rnew1IsjrTxmZ9CThcz5vVrfdG5jlfQ5fqvNmK58zK8hzeg6U3+IloxhUAgCEIVwAAhiBcAQAYgnAFAGAIwhUAgCEIVwAAhiBcAQAYgnAFAGAIwhUAgCEIVwAAhiBcAQAYgnAFAGAIwhUAgCEIVwAAhiBcAQAYgnAFAGAIwhUAgCEIVwAAhiBcAQAYgnAFAGAIwhUAgCEIVwAAhiBcAQAYgnAFAGAIwhUAgCFUd2+8MrXxyoNodZ7NbFkrndrsMay1Fc+blZkO8fbMuE/Ldt7MtfMrM+729pm+nzPuU6eX67yZ6fVmtu9lMt/3c8Z96rleRPfBlnytmcmsP3c3+BllxhUAgCEIVwAAhiBcAQAYgnAFAGAIwhUAgCEIVwAAhiBcAQAYgnAFAGAIwhUAgCEIVwAAhiBcAQAYgnAFAGAIwhUAgCEIVwAAhiBcAQAYgnAFAGAIwhUAgCEIVwAAhiBcAQAYgnAFAGAIwhUAgCEIVwAAhiBcAQAYgnAFAGAIwhUAgCEIVwAAhnDknlZ2zTWM2TY0m+3pzR7Cplmtw3ff76mVWZ8LS/Z9Wplp37fPuN8z7dOsz7klO2224uvN6spM21mdZztJkpn26XA11/NgznNmo1PGjCsAAEMQrgAADEG4AgAwBOEKAMAQhCsAAEMQrgAADEG4AgAwBOEKAMAQhCsAAEMQrgAADEG4AgAwBOEKAMAQhCsAAEMQrgAADEG4AgAwBOEKAMAQhCsAAEMQrgAADEG4AgAwBOEKAMAQhCsAAEMQrgAADEG4AgAwBOEKAMAQhCsAAEOo7t545fZsvPJgWqlZNpMkqzXTLmW+fUr3jBvbu7nOm15dqt0+KGp1nvMzSXrWk3TvvN4cOK83M3De3DPLdN7UPAd4ztfzuc7PWb+LG5wzZlwBABiCcAUAYAjCFQCAIQhXAACGIFwBABiCcAUAYAjCFQCAIQhXAACGIFwBABiCcAUAYAjCFQCAIQhXAACGIFwBABiCcAUAYAjCFQCAIQhXAACGIFwBABiCcAUAYAjCFQCAIQhXAACGIFwBABiCcAUAYAjCFQCAIQhXAACGIFwBABiCcAUAYAjV3Zs9BgAA2CszrgAADEG4AgAwBOEKAMAQhCsAAEMQrgAADEG4AgAwBOEKAMAQhCsAAEMQrgAADEG4AgAwhKHDtapWq+otmz0OAAAOvaUP16p6dlV9sqpuq6qvVNUlVfXkTRrLK6rqs1V1V1WtbsYYAAAOV0sdrlV1TpILk7wqyQOSPCTJRUlO26QhXZXkRUk+sEnbBwA4bC1tuFbV8UlenuTs7n53d3+ju3d298Xd/cIN7vOOqrqhqm6pqo9V1WPWrDu1qq6sqh1VdX1VnTst31ZV76+qm6vq61V1eVWte1y6+w+6+5IkOw7BLgMAsAdLG65JTk5yTJL37Md9LklyUpL7J/lUkreuWfeGJGd293FJHpvksmn5C5J8OckJWczqviRJ36ORAwBw0B252QPYg/sl+Wp337Wvd+juN+7693QN6k1VdXx335JkZ5JHV9VnuvumJDdNN92Z5EFJHtrdVyW5/GDtAAAAB88yz7h+Lcm2qtqnuK6qI6rq/Kr6YlXdmuSaadW26e8zkpya5Nqq+mhVnTwtf00W165+qKqurqoXH7xdAADgYFnmcP14kjuSnL6Pt392Fm/aOiXJ8UlOnJZXknT3Fd19WhaXEbw3ydun5Tu6+wXd/bAkz0xyTlU99eDsAgAAB8vShuv06/3zkryuqk6vqmOr6qiqenpVvXqduxyX5M4sZmqPzeKTCJIkVXV0VT1numxgZ5Jbk9w9rXtGVT2iqmrN8rvXG9O0/WOyOG5HVtUxVXXEwdtrAAA2srThmiTdfUGSc5K8LMmNSa5L8vwsZkx396Yk1ya5PsmVST6x2/rnJblmuozgrCTPnZaflOTSJLdlMct7UXd/ZIMh/V6S25P8YpKXTv9+3v7vGQAA+6u6vYEeAIDlt9QzrgAAsItwBQBgCMIVAIAhCFcAAIawxw/3r9SWe+fW6ky7tLL4+Nh5dM+4sb2r7TP9L3NX5tvtuc6bOa30nCfpPqitd5C33h4lneV6vdmK582WtEQ/p2b7GTWj1dXNHsHBt9HPKDOuAAAMQbgCADAE4QoAwBCEKwAAQxCuAAAMQbgCADAE4QoAwBCEKwAAQxCuAAAMQbgCADAE4QoAwBCEKwAAQxCuAAAMQbgCADAE4QoAwBCEKwAAQxCuAAAMQbgCADAE4QoAwBCEKwAAQxCuAAAMQbgCADAE4QoAwBCEKwAAQxCuAAAM4cjNHsDcVlfn2lDPtKFkZbYt7ZterVm2U5nvGHPobd+S3895ngvA4W1lpteamvFlemWDnwlmXAEAGIJwBQBgCMIVAIAhCFcAAIYgXAEAGIJwBQBgCMIVAIAhCFcAAIYgXAEAGIJwBQBgCMIVAIAhCFcAAIYgXAEAGIJwBQBgCMIVAIAhCFcAAIYgXAEAGIJwBQBgCMIVAIAhCFcAAIYgXAEAGIJwBQBgCMIVAIAhCFcAAIYgXAEAGIJwBQBgCEfuaeVq9SyDWEnNsp0kqZ5nW6srs2wmSbKSeb5P+6pWZxrPynznTebap8PYXK8D25fs+cI9sxVfb+b62bs608/DJJ51h9hMp8xSMOMKAMAQhCsAAEMQrgAADEG4AgAwBOEKAMAQhCsAAEMQrgAADEG4AgAwBOEKAMAQhCsAAEMQrgAADEG4AgAwBOEKAMAQhCsAAEMQrgAADEG4AgAwBOEKAMAQhCsAAEMQrgAADEG4AgAwBOEKAMAQhCsAAEMQrgAADEG4AgAwBOEKAMAQqrs3Xrk9G68c1Upt9ggOuk4v105VbbnzZks+FTrOm0Ns6+3R8r3e1JY8ylvPMp03W7FtVlfn2c7KnD82ev1zxowrAABDEK4AAAxBuAIAMAThCgDAEIQrAABDEK4AAAxBuAIAMAThCgDAEIQrAABDEK4AAAxBuAIAMAThCgDAEIQrAABDEK4AAAxBuAIAMAThCgDAEIQrAABDEK4AAAxBuAIAMAThCgDAEIQrAABDEK4AAAxBuAIAMAThCgDAEIQrAABDEK4AAAyhunuzxwAAAHtlxhUAgCEIVwAAhiBcAQAYgnAFAGAIwhUAgCEIVwAAhiBcAQAYgnAFAGAIwhUAgCEIVwAAhjB0uFbValW9ZbPHAQDAobf04VpVz66qT1bVbVX1laq6pKqevAnjuH9Vva2q/qaqbqmqP6mqJ849DgCAw9VSh2tVnZPkwiSvSvKAJA9JclGS0zZhOPdJckWSH0zyPUn+IMkHquo+mzAWAIDDztKGa1Udn+TlSc7u7nd39ze6e2d3X9zdL9zgPu+oqhumGdGPVdVj1qw7taqurKodVXV9VZ07Ld9WVe+vqpur6utVdXlVfcdx6e6ru/uC7v5Kd9/d3a9PcnSSRx6aIwAAwFpLG65JTk5yTJL37Md9LklyUpL7J/lUkreuWfeGJGd293FJHpvksmn5C5J8OckJWczqviRJ721DVfUDWYTrVfsxPgAADtCRmz2APbhfkq929137eofufuOuf1fVapKbqur47r4lyc4kj66qz3T3TUlumm66M8mDkjy0u69KcvnetlNV903y5iTbp8cGAOAQW+YZ168l2VZV+xTXVXVEVZ1fVV+sqluTXDOt2jb9fUaSU5NcW1UfraqTp+WvyWLW9ENVdXVVvXgv2/muJBcn+UR3//b+7RIAAAdqmcP140nuSHL6Pt7+2Vm8aeuUJMcnOXFaXknS3Vd092lZXEbw3iRvn5bv6O4XdPfDkjwzyTlV9dT1NlBV957ue32SM/d3hwAAOHBLG67Tr+DPS/K6qjq9qo6tqqOq6ulV9ep17nJckjuzmKk9NotPIkiSVNXRVfWc6bKBnUluTXL3tO4ZVfWIqqo1y+/e/cGr6qgk70xye5Jf6u5vHtQdBgBgj5Y2XJOkuy9Ick6SlyW5Mcl1SZ6fxazn7t6U5NosZkOvTPKJ3dY/L8k102UEZyV57rT8pCSXJrkti1nei7r7I+s8/j9I8owkT0ty8/S5srdV1Y8e6P4BALDvqnuvb6AHAIBNt9QzrgAAsItwBQBgCMIVAIAhCFcAAIawxw/3r+17/1+fHhQrNctmkmS1tt6b0VY68x3AfbAVz5tsn2eXVldn2UyS5TtvUvM8Oec6PZN5v59z6ZXlOm+21zzf0Fm/l3O+ts2k00uzU5WtFwJzHd05Xz83+hllxhUAgCEIVwAAhiBcAQAYgnAFAGAIwhUAgCEIVwAAhiBcAQAYgnAFAGAIwhUAgCEIVwAAhiBcAQAYgnAFAGAIwhUAgCEIVwAAhiBcAQAYgnAFAGAIwhUAgCEIVwAAhiBcAQAYgnAFAGAIwhUAgCEIVwAAhiBcAQAYgnAFAGAIwhUAgCEcuce1KzXTMGD/rVbPt63VmbbT8z3nVjLf8YNDZSXzPGdWPV9YYttnOj/n+lmYJCsbLDfjCgDAEIQrAABDEK4AAAxBuAIAMAThCgDAEIQrAABDEK4AAAxBuAIAMAThCgDAEIQrAABDEK4AAAxBuAIAMAThCgDAEIQrAABDEK4AAAxBuAIAMAThCgDAEIQrAABDEK4AAAxBuAIAMAThCgDAEIQrAABDEK4AAAxBuAIAMAThCgDAEIQrAABDOHJPK7vmGUSt9jwbYktZyUwnaJLVOEe3CucNsNXM9bq2DK9pZlwBABiCcAUAYAjCFQCAIQhXAACGIFwBABiCcAUAYAjCFQCAIQhXAACGIFwBABiCcAUAYAjCFQCAIQhXAACGIFwBABiCcAUAYAjCFQCAIQhXAACGIFwBABiCcAUAYAjCFQCAIQhXAACGIFwBABiCcAUAYAjCFQCAIQhXAACGIFwBABjCkZs9gCTJSs23rdWeb1uwn3rGp0KW7KlQM41n1mPMIVde02E2q6szbmxl/cVmXAEAGIJwBQBgCMIVAIAhCFcAAIYgXAEAGIJwBQBgCMIVAIAhCFcAAIYgXAEAGIJwBQBgCMIVAIAhCFcAAIYgXAEAGIJwBQBgCMIVAIAhCFcAAIYgXAEAGIJwBQBgCMIVAIAhCFcAAIYgXAEAGIJwBQBgCMIVAIAhCFcAAIYgXAEAGIJwBQBgCNXdmz0GAADYKzOuAAAMQbgCADAE4QoAwBCEKwAAQxCuAAAMQbgCADCE/we/pXGno5MBjgAAAABJRU5ErkJggg==\n",
      "text/plain": [
       "<Figure size 720x432 with 15 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "generator = SyntheticDataGenerator(image_size=4)\n",
    "generator.show_examples(5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training batches: 141\n",
      "Test batches: 47\n",
      "Batch shapes:\n",
      "  Images: torch.Size([32, 48])\n",
      "  Labels: torch.Size([32])\n"
     ]
    }
   ],
   "source": [
    "train_loader, test_loader = create_dataloaders(\n",
    "    train_samples_per_class = 1500,\n",
    "    test_samples_per_class= 500,\n",
    "    batch_size = 32,\n",
    "    num_workers = 4)\n",
    "\n",
    "# Print dataset information\n",
    "print(f\"Training batches: {len(train_loader)}\")\n",
    "print(f\"Test batches: {len(test_loader)}\")\n",
    "    \n",
    "# Example of batch structure\n",
    "for images, labels in train_loader:\n",
    "    print(f\"Batch shapes:\")\n",
    "    print(f\"  Images: {images.shape}\")\n",
    "    print(f\"  Labels: {labels.shape}\")\n",
    "    break  # Just show first batch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training a Small Model\n",
    "We train a linear model with ReLU on the dataset to classify images into one of the three classes. Since there is some overlap between classes 1 and 2, the model should experience confusion between the two."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "num_classes = 3\n",
    "input_size = train_loader.dataset[0][0].shape[-1]\n",
    "hidden_dim = 20\n",
    "\n",
    "model = nn.Sequential(\n",
    "    nn.Linear(input_size, hidden_dim),\n",
    "    nn.ReLU(),\n",
    "    nn.Linear(hidden_dim, num_classes)\n",
    ").to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate_model(model, data_loader, device):\n",
    "    \"\"\"\n",
    "    Evaluate model accuracy on a dataset.\n",
    "    \n",
    "    Args:\n",
    "        model: Neural network model\n",
    "        data_loader: DataLoader containing the evaluation dataset\n",
    "        device: Device to run evaluation on\n",
    "        \n",
    "    Returns:\n",
    "        float: Accuracy score (0-1)\n",
    "    \"\"\"\n",
    "    model.eval()\n",
    "    correct = 0\n",
    "    total = 0\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        for images, labels in data_loader:\n",
    "            images = images.to(device).float()\n",
    "            labels = labels.to(device)\n",
    "            \n",
    "            outputs = model(images)\n",
    "            _, predicted = torch.max(outputs.data, 1)\n",
    "            \n",
    "            total += labels.size(0)\n",
    "            correct += (predicted == labels).sum().item()\n",
    "    \n",
    "    return correct / total"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_model(model, train_loader, test_loader, num_epochs=200, learning_rate=0.001):\n",
    "    \"\"\"\n",
    "    Train the neural network model.\n",
    "    \n",
    "    Args:\n",
    "        model: Neural network model\n",
    "        train_loader: DataLoader for training data\n",
    "        test_loader: DataLoader for test data\n",
    "        num_epochs: Number of training epochs\n",
    "        learning_rate: Learning rate for optimizer\n",
    "        \n",
    "    Returns:\n",
    "        Training losses and accuracies per epoch\n",
    "    \"\"\"\n",
    "    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "    \n",
    "    model = model.to(device)\n",
    "    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)\n",
    "    criterion = nn.CrossEntropyLoss()\n",
    "    \n",
    "    train_losses = []\n",
    "    train_accuracies = []\n",
    "    \n",
    "    pbar = tqdm(range(num_epochs), desc=\"Training\")\n",
    "    for epoch in pbar:\n",
    "        model.train()\n",
    "        epoch_losses = []\n",
    "        correct = 0\n",
    "        total = 0\n",
    "        \n",
    "        for images, labels in train_loader:\n",
    "            images = images.to(device).float()\n",
    "            labels = labels.to(device)\n",
    "            \n",
    "            # Forward pass\n",
    "            outputs = model(images)\n",
    "            loss = criterion(outputs, labels)\n",
    "            \n",
    "            # Backward pass and optimize\n",
    "            optimizer.zero_grad()\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            \n",
    "            # Track statistics\n",
    "            epoch_losses.append(loss.item())\n",
    "            _, predicted = torch.max(outputs.data, 1)\n",
    "            total += labels.size(0)\n",
    "            correct += (predicted == labels).sum().item()\n",
    "        \n",
    "        # Calculate epoch statistics\n",
    "        epoch_loss = sum(epoch_losses) / len(epoch_losses)\n",
    "        epoch_accuracy = correct / total\n",
    "        train_losses.append(epoch_loss)\n",
    "        train_accuracies.append(epoch_accuracy)\n",
    "        \n",
    "        # Update progress bar every epoch\n",
    "        if epoch % 10 == 0:\n",
    "            test_accuracy = evaluate_model(model, test_loader, device)\n",
    "        else:\n",
    "            test_accuracy = train_accuracies[-1]  # Use training accuracy when not computing test\n",
    "            \n",
    "        pbar.set_postfix({\n",
    "            'epoch': f'{epoch:3d}',\n",
    "            'loss': f'{epoch_loss:.4f}',\n",
    "            'train_acc': f'{epoch_accuracy:.4f}',\n",
    "            'test_acc': f'{test_accuracy:.4f}'\n",
    "        })\n",
    "    \n",
    "    return train_losses, train_accuracies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Training: 100%|███████████████████████████████████████| 200/200 [03:29<00:00,  1.05s/it, epoch=199, loss=0.1450, train_acc=0.9376, test_acc=0.9376]\n"
     ]
    }
   ],
   "source": [
    "loss, accuracy = train_model(model, train_loader, test_loader, num_epochs=200, learning_rate=0.001)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████| 47/47 [00:00<00:00, 135.75it/s]\n"
     ]
    }
   ],
   "source": [
    "# Compute model outputs on test instances\n",
    "labels = []\n",
    "outputs = []\n",
    "\n",
    "for i, (image, label) in enumerate(tqdm(test_loader)):\n",
    "    with torch.no_grad():\n",
    "        images = image.to(device).float()\n",
    "    labels.extend(label.numpy())\n",
    "\n",
    "    # Compute model inferences\n",
    "    output = model(images)\n",
    "    output = torch.nn.functional.softmax(output, dim=1).squeeze(0).detach().cpu().numpy()\n",
    "    outputs.append(output)\n",
    "    \n",
    "outputs = np.vstack(outputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define the Abstraction Graph\n",
    "We define the abstraction graph based on human similarity. Images in classes 1 and 2 are more similar to each other than images in class 0."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "root (None)\n",
      "├── parent0 (None)\n",
      "│   └── child0 (None)\n",
      "└── parent1 (None)\n",
      "    ├── child1 (None)\n",
      "    └── child2 (None)\n",
      "\n",
      "Abstraction graph with 6 nodes across 3 levels.\n"
     ]
    }
   ],
   "source": [
    "abstraction_graph = make_abstraction_graph()\n",
    "print(show_abstraction_graph(abstraction_graph))\n",
    "print(f'Abstraction graph with {abstraction_graph.size()} nodes across {abstraction_graph.depth() + 1} levels.')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Measure Abstraction Alignment\n",
    "We show that the model is abstractio aligned by measuring its abstraction match --- i.e., how much of its entropy in the model's predictions is reduced by moving up a level of abstraction."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, we measure abstraction alignment using the real abstraction graph. This abstraction graph represents the abstractions we expect the model to learn."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "aligned_abstraction_graphs = []\n",
    "for i in range(len(labels)):\n",
    "    abstraction_graph = propagate(outputs[i], make_abstraction_graph())\n",
    "    aligned_abstraction_graphs.append(abstraction_graph)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Abstraction match = 75.51%\n"
     ]
    }
   ],
   "source": [
    "abstraction_match = metrics.abstraction_match(aligned_abstraction_graphs, 1)\n",
    "print(f'Abstraction match = {abstraction_match:.2%}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then, we measure abstraction match using a misaligned abstraction graph. The misaligned abstraction graph does not align with the abstractions we expect the model to learn."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Abstraction match with misaligned graph = 13.98%\n"
     ]
    }
   ],
   "source": [
    "misaligned_abstraction_graphs = []\n",
    "for i in range(len(labels)):\n",
    "    abstraction_graph = propagate(outputs[i], make_abstraction_graph(misaligned=True))\n",
    "    misaligned_abstraction_graphs.append(abstraction_graph)\n",
    "abstraction_match = metrics.abstraction_match(misaligned_abstraction_graphs, 1)\n",
    "print(f'Abstraction match with misaligned graph = {abstraction_match:.2%}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAOcAAADnCAYAAADl9EEgAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjQuMywgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/MnkTPAAAACXBIWXMAAAsTAAALEwEAmpwYAAAC80lEQVR4nO3YQWrDMBBAUav43mlOPr2AE7qo0W94b2lvRojPgNbMHEDP1+4BgGvihChxQpQ4IUqcEHW+/buWp9x/5mNv7PmpBzuOeRzr6rvNCVHihChxQpQ4IUqcECVOiBInRIkTosQJUeKEKHFClDghSpwQJU6IEidEiROixAlR4oQocUKUOCFKnBAlTogSJ0SJE6LECVHihChxQpQ4IUqcECVOiBInRIkTosQJUeKEKHFClDghSpwQJU6IEidEiROixAlR4oQocUKUOCFKnBAlTogSJ0SJE6LECVHihChxQpQ4IUqcEHXuHoA/9pzdE9zjsXZPcKPrO7M5IUqcECVOiBInRIkTosQJUeKEKHFClDghSpwQJU6IEidEiROixAlR4oQocUKUOCFKnBAlTogSJ0SJE6LECVHihChxQpQ4IUqcECVOiBInRIkTosQJUeKEKHFClDghSpwQJU6IEidEiROixAlR4oQocUKUOCFKnBAlTogSJ0SJE6LECVHihChxQpQ4IUqcECVOiBInRIkTosQJUefuAeA3vtfsHuE+L45mc0KUOCFKnBAlTogSJ0SJE6LECVHihChxQpQ4IUqcECVOiBInRIkTosQJUeKEKHFClDghSpwQJU6IEidEiROixAlR4oQocUKUOCFKnBAlTogSJ0SJE6LECVHihChxQpQ4IUqcECVOiBInRIkTosQJUeKEKHFClDghSpwQJU6IEidEiROixAlR4oQocUKUOCFKnBAlTogSJ0SJE6LECVHihChxQpQ4IUqcECVOiBInRIkTosQJUeKEKHFClDghSpwQJU6IEidEiROixAlR4oQocUKUOCFKnBAlTogSJ0SJE6LECVHihChxQpQ4IUqcECVOiBInRIkTosQJUeKEKHFClDghSpwQJU6IEidEiROixAlR4oQocUKUOCFKnBAlTogSJ0SJE6LECVHihChxQpQ4IUqcECVOiFozs3sG4ILNCVHihChxQpQ4IUqcECVOiPoB8AQSSxpjXOMAAAAASUVORK5CYII=\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MODEL PREDICTIONS:  ['1 (0.64)', '2 (0.36)', '0 (0.00)']\n",
      "ALIGNED ABSTRACTION GRAPH:\n",
      "root (1.00)\n",
      "├── parent0 (0.00)\n",
      "│   └── child0 (0.00)\n",
      "└── parent1 (1.00)\n",
      "    ├── child1 (0.64)\n",
      "    └── child2 (0.36)\n",
      "\n",
      "MISALIGNED ABSTRACTION GRAPH:\n",
      "root (1.00)\n",
      "├── parent0 (0.64)\n",
      "│   ├── child0 (0.00)\n",
      "│   └── child1 (0.64)\n",
      "└── parent1 (0.36)\n",
      "    └── child2 (0.36)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# An example comparing aligned and misaligned abstraction alignment\n",
    "index = 510\n",
    "image = test_loader.dataset[index][0].detach().cpu().numpy()\n",
    "image = np.reshape(image, (4, 4, 3)) * 255\n",
    "\n",
    "plt.imshow(image.astype(int))\n",
    "plt.axis('off')\n",
    "plt.show()\n",
    "\n",
    "top_predictions = np.argsort(outputs[index])[-3:]\n",
    "print('MODEL PREDICTIONS: ', [f'{i} ({outputs[index][i]:.2f})' for i in top_predictions][::-1])\n",
    "\n",
    "print('ALIGNED ABSTRACTION GRAPH:')\n",
    "print(show_abstraction_graph(aligned_abstraction_graphs[index]))\n",
    "\n",
    "print('MISALIGNED ABSTRACTION GRAPH:')\n",
    "print(show_abstraction_graph(misaligned_abstraction_graphs[index]))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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.7.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}