stability_plot.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Analyze & plot from wandb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import itertools\n",
    "import wandb\n",
    "import yaml\n",
    "import pickle\n",
    "import tempfile\n",
    "from tqdm import tqdm\n",
    "import os\n",
    "\n",
    "from scipy import stats\n",
    "import numpy as np\n",
    "import itertools\n",
    "\n",
    "# set parameters you wish to plot: for example\n",
    "models = ['LeNet']\n",
    "hardness_methods = ['uniform'] # place in a list 'asymmetric', \"adjacent\", 'instance', 'domain_shift', 'ood_covariate', \"crop_shift\", \"zoom_shift\", \"far_ood\"]\n",
    "datasets = ['mnist']\n",
    "\n",
    "exp_combos = list(itertools.product(models, datasets, hardness_methods))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "failed = []\n",
    "sensitivity_dict = {method: [] for method in hardness_methods}\n",
    "corr_dict = {method: [] for method in hardness_methods}\n",
    "for exp_combo in tqdm(exp_combos):\n",
    "\n",
    "    try:\n",
    "        model = exp_combo[0]\n",
    "        dataset_name = exp_combo[1]\n",
    "        hardness = exp_combo[2]\n",
    "\n",
    "        title = f\"{hardness}_{model}_{dataset_name}\"\n",
    "\n",
    "        with open('./wandb.yaml') as file:\n",
    "            wandb_data = yaml.load(file, Loader=yaml.FullLoader)\n",
    "\n",
    "        os.environ[\"WANDB_API_KEY\"] = wandb_data['wandb_key'] \n",
    "        wandb_entity = str(wandb_data['wandb_entity'])\n",
    "\n",
    "        seed=0\n",
    "        runid=0\n",
    "        metainfo = f\"{dataset_name}_{hardness}_{0.1}_{seed}_{runid}\" \n",
    "       \n",
    "        # check file exists for one run\n",
    "        project_name = f'{hardness}_{dataset_name}_{model}'\n",
    "        api = wandb.Api()\n",
    "        artifact = api.artifact(f'{wandb_entity}/{project_name}/scores_dict_{metainfo}:latest', type='pickle')\n",
    "        temp_dir = tempfile.TemporaryDirectory()\n",
    "        artifact.download(root=temp_dir.name)\n",
    "        temp_dir.cleanup()\n",
    "\n",
    "\n",
    "        props = [0.1,0.2,0.3,0.4,0.5]\n",
    "        runids = [0,1,2]\n",
    "\n",
    "        overall_results = {}\n",
    "\n",
    "        for prop in props:\n",
    "\n",
    "            prop_list = []\n",
    "\n",
    "            for runid in runids:\n",
    "                \n",
    "\n",
    "                try:\n",
    "                    temp_dir = tempfile.TemporaryDirectory()\n",
    "                    seed=runid\n",
    "                    metainfo = f\"{dataset_name}_{hardness}_{prop}_{seed}_{runid}\" \n",
    "                    api = wandb.Api()\n",
    "                    artifact = api.artifact(f'{wandb_entity}/{project_name}/scores_dict_{metainfo}:latest', type='pickle')\n",
    "                    artifact_dir = artifact.download(root=temp_dir.name)\n",
    "\n",
    "                    with open(artifact_dir + f\"/scores_dict_{metainfo}.pkl\", \"rb\") as f:\n",
    "                        data = pickle.load(f)\n",
    "                    temp_dir.cleanup()\n",
    "\n",
    "                    prop_list.append(data)\n",
    "\n",
    "                except:\n",
    "                    print(f'FAILED {metainfo} - {model}')\n",
    "                    continue\n",
    "\n",
    "            overall_results[prop] = prop_list\n",
    "\n",
    "\n",
    "\n",
    "        combos = list(itertools.combinations(range(len(overall_results[0.1])), 2))\n",
    "\n",
    "        sensitivity = {}\n",
    "\n",
    "    \n",
    "        for prop in props:\n",
    "\n",
    "            corr_dict = {}\n",
    "\n",
    "            for model in list(overall_results[0.1][0]['scores'].keys()):\n",
    "                corr = []\n",
    "                for combo in combos:\n",
    "                    try:\n",
    "                        c1, c2 = combo[0], combo[1]\n",
    "                        rvs1 = overall_results[prop][c1]['scores'][model]\n",
    "                        rvs2 = overall_results[prop][c2]['scores'][model]\n",
    "                        corr.append(stats.spearmanr(rvs1, rvs2)[0])\n",
    "                    except:\n",
    "                        continue\n",
    "\n",
    "                corr_dict[model] = np.mean(corr)\n",
    "\n",
    "            sensitivity[prop] = corr_dict\n",
    "\n",
    "        sensitivity_dict[hardness].append(sensitivity)\n",
    "\n",
    "    except Exception as e:\n",
    "        import traceback\n",
    "        print(traceback.format_exc())\n",
    "        print(f'failed {title}')\n",
    "        failed.append(title)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_mean_dict(dict_list):\n",
    "    from collections import defaultdict\n",
    "\n",
    "    # Initialize dictionaries to store the total and count of each metric\n",
    "    totals = defaultdict(lambda: defaultdict(int))\n",
    "    counts = defaultdict(lambda: defaultdict(int))\n",
    "\n",
    "    # Iterate over the list of dictionaries\n",
    "    for d in dict_list:\n",
    "        # Iterate over each metric in the dictionary\n",
    "        for key, metrics in d.items():\n",
    "            for metric, value in metrics.items():\n",
    "                # Add the value to the total and increment the count\n",
    "                totals[key][metric] += value\n",
    "                counts[key][metric] += 1\n",
    "\n",
    "    # Compute the means\n",
    "    means = {key: {metric: total / counts[key][metric]\n",
    "                for metric, total in metrics.items()}\n",
    "            for key, metrics in totals.items()}\n",
    "    \n",
    "    return means"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "groups = {'allsh': \"Statistical\",\n",
    " 'aum': \"LB:Output\",\n",
    " 'cleanlab': \"Statistical\",\n",
    " 'conf_agree': \"Statistical\",\n",
    " 'dataiq':\"LB:Output\",\n",
    " 'datamaps': \"LB:Output\",\n",
    " 'detector': \"LB:Stats\",\n",
    " 'el2n': \"LB:Stats\",\n",
    " 'forgetting': \"LB:Other\",\n",
    " 'grand': \"LB:Grad\",\n",
    " 'loss': \"LB:Other\",\n",
    " \"protypicality\": \"Dist\",\n",
    " 'vog': \"LB:Grad\"}\n",
    "\n",
    "\n",
    "def get_sorted(dictionary, title, save='png'):\n",
    "        import matplotlib.pyplot as plt\n",
    "        import collections\n",
    "\n",
    "        # Convert the dictionary to a 2D NumPy array\n",
    "        data = np.array([[value for value in level.values()] for level in dictionary.values()])\n",
    "\n",
    "        # Create a list of keys and levels for labeling the heatmap\n",
    "        keys = list(dictionary.values())[0].keys()\n",
    "        keys = [key.split(\".\")[0] for key in keys]\n",
    "        levels = list(dictionary.keys())\n",
    "\n",
    "        # Create a list of tuples (group, method) for each key\n",
    "        grouped_keys = [(groups.get(key, \"Dist\"), key) for key in keys]\n",
    "\n",
    "        # Sort the keys and data by group\n",
    "        sorted_grouped_keys = sorted(grouped_keys, key=lambda x: (x[0], keys.index(x[1])))\n",
    "\n",
    "        # Extract the sorted keys and group labels\n",
    "        sorted_keys = [key[1] for key in sorted_grouped_keys]\n",
    "        group_labels = [key[0] for key in sorted_grouped_keys]\n",
    "\n",
    "        # Create an ordered dictionary that maps the original keys to their corresponding data columns\n",
    "        data_dict = collections.OrderedDict(zip(keys, data.T))\n",
    "\n",
    "        # Create a new sorted dictionary and a new sorted data array\n",
    "        sorted_data_dict = collections.OrderedDict(sorted(data_dict.items(), key=lambda x: sorted_grouped_keys.index((groups.get(x[0], \"Dist\"), x[0]))))\n",
    "\n",
    "        # Get the indices of the sorted keys\n",
    "        indices = [keys.index(key) for key in sorted_keys]\n",
    "\n",
    "        # Reorder the data array based on the sorted keys\n",
    "        sorted_data = data[:, indices]\n",
    "\n",
    "        return sorted_data, sorted_keys, levels, group_labels, sorted_data_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "import numpy as np\n",
    "\n",
    "def plot_bar_plot(sorted_data, sorted_keys, levels, group_labels, sorted_data_dict, title):\n",
    "    s_keys = []\n",
    "    for key in sorted_keys:\n",
    "        if key =='prototypicality':\n",
    "            s_keys.append('proto')\n",
    "        elif key =='forgetting':\n",
    "            s_keys.append('forget')\n",
    "        elif key=='conf_agree':\n",
    "            s_keys.append('conf_agr')\n",
    "        else:\n",
    "            s_keys.append(key)\n",
    "\n",
    "    sorted_keys = s_keys\n",
    "    # create pandas dataframe from sorted_data, where rows are levels and columns are sorted_keys\n",
    "    df = pd.DataFrame(sorted_data, index=levels, columns=sorted_keys)\n",
    "\n",
    "    # compute mean and standard deviation across levels\n",
    "    mean = df.mean(axis=0)\n",
    "    std = df.std(axis=0)\n",
    "\n",
    "    min_val = mean.min()\n",
    "\n",
    "    fs=18\n",
    "    # plot bar plot\n",
    "    plt.figure(figsize=(7, 3))\n",
    "    ax = sns.barplot(x=mean.index, y=mean.values, yerr=std.values)\n",
    "    ax.set_xticklabels(ax.get_xticklabels(), rotation=45, ha=\"right\", fontsize=fs-3)\n",
    "    \n",
    "    # set fontsize of yticks\n",
    "    ax.tick_params(axis='y', which='major', labelsize=fs)\n",
    "\n",
    "    #set ylim\n",
    "    ax.set_ylim(min_val-0.1, 1.0)\n",
    "    ax.set_ylabel(\"Spearman correlation\", fontsize=fs)\n",
    "\n",
    "    # add a vertical grid to the plot, but make it very light in color\n",
    "    # so we can use it for reading data values but not be distracting\n",
    "    ax.grid(True, which='major', axis='y', linestyle='--', color='grey', alpha=0.33)\n",
    "\n",
    "    # add extra xticks below based on groups\n",
    "    group_positions = {}\n",
    "    for i, key in enumerate(sorted_data_dict.keys()):\n",
    "        group = groups.get(key, \"Dist\")\n",
    "        if group not in group_positions:\n",
    "            group_positions[group] = [i]\n",
    "        else:\n",
    "            group_positions[group].append(i)\n",
    "\n",
    "    # Now, for each group, find the midpoint and add the label there\n",
    "    for group, positions in group_positions.items():\n",
    "        midpoint = np.mean(positions)\n",
    "        ax.text(midpoint, -0.425, group, ha='center', va='top', transform=ax.get_xaxis_transform(), fontsize=fs-5)\n",
    "\n",
    "    # if results_folder doesn't exist create it\n",
    "    if not os.path.exists('results_folder'):\n",
    "        os.makedirs('results_folder')\n",
    "    \n",
    "    if not os.path.exists('results_folder/model_compare'):\n",
    "        os.makedirs('results_folder/model_compare')\n",
    "\n",
    "    plt.savefig(f\"results_folder/model_compare/{title}.pdf\", bbox_inches='tight', pad_inches=0.1)\n",
    "\n",
    "    plt.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for hardness in hardness_methods:\n",
    "    try:\n",
    "        sensitivity_indiv = compute_mean_dict(sensitivity_dict[hardness])\n",
    "        title = f\"{hardness}\"\n",
    "        sorted_data, sorted_keys, levels, group_labels, sorted_data_dict = get_sorted(dictionary=sensitivity_indiv, title=title, save='png')\n",
    "        plot_bar_plot(sorted_data, sorted_keys, levels, group_labels, sorted_data_dict, title=title)\n",
    "        \n",
    "    except Exception as e:\n",
    "        print(f'Unable to get results: - {hardness} has not been logged to wandb... Continuing')\n",
    "        continue"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "hardness_env",
   "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.10.10"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "6a7e60037fa700cd8eaefe68718883a7c5484a19cbbdd784476a6a2fe63bc7d7"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}