MAG_postprocessing.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "c8f57984-6d7c-454c-8ef9-80ad3d9b1d68",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "import os\n",
    "import os.path as osp\n",
    "import glob\n",
    "from tqdm import tqdm\n",
    "from ogb.lsc import MAG240MEvaluator\n",
    "\n",
    "from root import ROOT\n",
    "\n",
    "train_set = 'train' # 'train' or 'train_val'\n",
    "data_dir = osp.join(ROOT, 'mag240m_kddcup2021')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "631d9151-b4bb-4025-be68-9e7fdf4cce1f",
   "metadata": {},
   "outputs": [],
   "source": [
    "__meta__ = torch.load(osp.join(data_dir, 'meta.pt'))\n",
    "num_papers = __meta__['paper']\n",
    "num_authors = __meta__['author']\n",
    "num_institutions = __meta__['institution']\n",
    "__split__ = torch.load(osp.join(data_dir, 'split_dict.pt'))\n",
    "train_idx = __split__['train']\n",
    "val_idx = __split__['valid']\n",
    "test_idx = __split__['test']\n",
    "num_classes = 153\n",
    "\n",
    "paper_labels = np.load(osp.join(data_dir, 'processed', 'paper', 'node_label.npy')).astype(np.int16)\n",
    "val_reverse_idx = np.zeros(num_papers, dtype=np.int32)\n",
    "val_reverse_idx[val_idx] = np.arange(val_idx.shape[0])\n",
    "val_labels = paper_labels[val_idx]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "4f081b7b-9bde-47b7-90ef-f682f6db15fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "def edge_index(dir, id1: str, id2: str,\n",
    "                   id3: str = None) -> np.ndarray:\n",
    "        src = id1\n",
    "        rel, dst = (id3, id2) if id3 is None else (id2, id3)\n",
    "        rel = self.__rels__[(src, dst)] if rel is None else rel\n",
    "        name = f'{src}___{rel}___{dst}'\n",
    "        path = osp.join(dir, 'processed', name, 'edge_index.npy')\n",
    "        return np.load(path)\n",
    "    \n",
    "author_writes_authors, author_writes_papers = edge_index(data_dir, 'author', 'writes', 'paper')\n",
    "author_writes_papers_argsort = np.argsort(author_writes_papers)\n",
    "papers_written_papers = author_writes_papers[author_writes_papers_argsort]\n",
    "papers_written_authors = author_writes_authors[author_writes_papers_argsort]\n",
    "papers, author_counts = np.unique(papers_written_papers, return_counts=True)\n",
    "paper_row_start = np.insert(np.cumsum(author_counts), 0, 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "c30a14e6-8ba2-46f8-85fb-acda0f61ce8f",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1251341/1251341 [00:09<00:00, 135651.84it/s]\n"
     ]
    }
   ],
   "source": [
    "path = osp.join(data_dir, f'author_{train_set}_label_probs.npy')\n",
    "if not osp.exists(path)\n",
    "    if train_set == 'train':\n",
    "        paper_idx = train_idx\n",
    "    else:\n",
    "        paper_idx = np.concatenate([train_idx, val_idx], axis=0)\n",
    "\n",
    "    author_label_counts = np.zeros((num_authors, num_classes), dtype=np.int16)\n",
    "    for paper in tqdm(paper_idx):\n",
    "        cur_paper_row_start = paper_row_start[paper]\n",
    "        paper_authors = papers_written_authors[cur_paper_row_start:cur_paper_row_start+author_counts[paper]]\n",
    "        author_label_counts[paper_authors, paper_labels[paper]] += 1\n",
    "\n",
    "    author_label_probs = author_label_counts.astype(np.float32) + 1e-10\n",
    "    author_label_probs = author_label_probs / author_label_probs.sum(axis=-1, keepdims=True)\n",
    "    np.save(path, author_label_probs)\n",
    "else:\n",
    "    author_label_probs = np.load(path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f65381fb-e351-4b5a-af61-f0162f97e3ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_eval(pred_prob):\n",
    "    preds = np.argmax(pred_prob, axis=-1)\n",
    "    res = {'y_pred': preds, 'y_true': val_labels}\n",
    "    eval = MAG240MEvaluator().eval(res)\n",
    "    return eval"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "43148a12-bbd7-42ad-bab1-7caeda5e4161",
   "metadata": {},
   "outputs": [],
   "source": [
    "results_folder = osp.join('results_rgnn', 'valid', 'rgat')\n",
    "model = '*'\n",
    "pred_probs_paths = osp.join(base_folder, model, '*', '*', 'pred_probs.npy')\n",
    "pred_probs_paths = glob.glob(preds_probs_paths)\n",
    "pred_probs = []\n",
    "for pred_probs_path in pred_probs_paths:\n",
    "    pred_prob = np.load(pred_probs_path)\n",
    "    pred_prob = np.exp(pred_prob)\n",
    "    pred_prob = pred_prob / pred_prob.sum(axis=-1, keepdims=True)\n",
    "    print(pred_probs_path, get_eval(pred_prob))\n",
    "    pred_probs.append(np.expand_dims(pred_prob, axis=-1))\n",
    "pred_probs = np.concatenate(pred_probs, axis=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92c1ee67-2798-4f40-8fbb-e7e990c0fdb5",
   "metadata": {},
   "outputs": [],
   "source": [
    "mean_pred_prob = pred_probs.mean(axis=-1)\n",
    "print(\"Mean logit aggregation: \", get_eval(mean_pred_prob))\n",
    "max_pred_prob = pred_probs.max(axis=-1)\n",
    "print(\"Max logit aggregation: \", get_eval(max_pred_prob))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3ccee7b-6544-43bb-a0ba-cb82d0480d1e",
   "metadata": {},
   "outputs": [],
   "source": [
    "post_process_paper_idx = val_idx\n",
    "pred_prob = max_pred_prob\n",
    "new_pred_probs = np.zeros_like(pred_prob)\n",
    "for paper in tqdm(post_process_paper_idx):\n",
    "    cur_paper_row_start = paper_row_start[paper]\n",
    "    paper_authors = papers_written_authors[cur_paper_row_start:cur_paper_row_start+author_counts[paper]]\n",
    "    cur_paper_author_label_probs = author_label_probs[paper_authors]\n",
    "    cur_paper_author_label_probs = cur_paper_author_label_probs.mean(axis=0, keepdims=1)\n",
    "    cur_paper_pred_probs = pred_prob[val_reverse_idx[paper]]\n",
    "    new_pred_probs[val_reverse_idx[paper]] = 0.4 * cur_paper_author_label_probs + 0.6 * cur_paper_pred_probs\n",
    "print(get_eval(new_pred_probs))\n",
    "\n"
   ]
  }
 ],
 "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}