GAN(gan=rgan,data=best(xing,stress),dis=deep,share=4,embed=12).ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "wZtYFEhHoBGp"
   },
   "source": [
    "# Import libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "id": "KJPfzpxLQI5O"
   },
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "from deepgd import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "z4eBEyB1oaJe"
   },
   "source": [
    "# Config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "id": "Q_-MLFhpNvxF"
   },
   "outputs": [],
   "source": [
    "cuda_idx = 0\n",
    "canonicalization = Canonicalization(normalize=Standardization(norm_ord=1),\n",
    "                                    scale=False)\n",
    "config = StaticConfig({\n",
    "    \"name\": 'GAN(gan=rgan,data=best(xing,stress),dis=deep,share=4,embed=12)',\n",
    "    \"uid\": None,\n",
    "    \"link\": None,\n",
    "    \"generator\": {\n",
    "        \"params\": {\n",
    "            \"num_blocks\": 9,\n",
    "            \"normalize\": canonicalization\n",
    "        },\n",
    "        \"pretrained\": {\n",
    "            \"name\": None,\n",
    "            \"epoch\": -1,\n",
    "        },\n",
    "        \"optim\": torch.optim.AdamW,\n",
    "        \"lr\" : {\n",
    "            \"initial\": 1e-3,\n",
    "            \"decay\": 0.99,\n",
    "        },\n",
    "    },\n",
    "    \"discriminator\": {\n",
    "        \"params\": {\n",
    "            \"conv\": [2, 16, 16, 16],\n",
    "            \"dense\": [2],\n",
    "            \"shared_depth\": 4,\n",
    "            \"enet_depth\": 12,\n",
    "            \"enet_width\": 64,\n",
    "            \"aggr\": \"add\",\n",
    "            \"normalize\": canonicalization\n",
    "        },\n",
    "        \"pretrained\": {\n",
    "            \"name\": None,\n",
    "            \"epoch\": -1,\n",
    "        },\n",
    "        \"optim\": torch.optim.AdamW,\n",
    "        \"lr\" : {\n",
    "            \"initial\": 1e-3,\n",
    "            \"decay\": 0.99,\n",
    "        },\n",
    "        \"noise\": {\n",
    "            \"std\": 0,\n",
    "            \"decay\": 0.95,\n",
    "        },\n",
    "        \"repeat\": 1,\n",
    "        \"complete\": True,\n",
    "        \"adaptive\": True\n",
    "    },\n",
    "    \"alternate\": \"epoch\",\n",
    "    \"batchsize\": 24,\n",
    "    \"epoch\": {\n",
    "        \"start\": -1,\n",
    "        \"end\": None,\n",
    "    },\n",
    "    \"log_interval\": 1,\n",
    "    \"test\": {\n",
    "        \"name\": \"test\",\n",
    "        \"epoch\": -1,\n",
    "    },\n",
    "    \"gan_flavor\": \"rgan\",\n",
    "    \"gp_weight\": 0,\n",
    "})\n",
    "data_config = StaticConfig({\n",
    "    \"sparse\": False,\n",
    "    \"pivot\": None,\n",
    "    \"init\": \"pmds\",\n",
    "    \"edge\": {\n",
    "        \"index\": \"full_edge_index\",\n",
    "        \"attr\": \"full_edge_attr\",\n",
    "    },\n",
    "})\n",
    "loss_fns = {\n",
    "    Stress(): 1\n",
    "}\n",
    "ctrler_params = {\n",
    "    \"tau\": 0.95,\n",
    "    \"beta\": 1,\n",
    "    \"exploit_rate\": 0.5,\n",
    "    \"warmup\": 2,\n",
    "}\n",
    "paths = StaticConfig({\n",
    "    \"root\": \"artifacts\",\n",
    "    \"checkpoints\": lambda: f\"{paths.root}/checkpoints/{config.name}\",\n",
    "    \"gen_pretrain\": lambda: f\"{paths.root}/checkpoints/{config.generator.pretrained.name}\",\n",
    "    \"dis_pretrain\": lambda: f\"{paths.root}/checkpoints/{config.discriminator.pretrained.name}\",\n",
    "    \"tensorboard\": lambda: f\"{paths.root}/tensorboards/{config.name}\",\n",
    "    \"visualization\": lambda: f\"{paths.root}/visualizations/{config.name}_{config.test.name}\",\n",
    "    \"log\": lambda: f\"{paths.root}/logs/{config.name}.log\",\n",
    "    \"metrics\": lambda suffix: f\"{paths.root}/metrics/{config.name}_{suffix}.pickle\",\n",
    "})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "id": "mdMEAbH10Qaq"
   },
   "outputs": [],
   "source": [
    "if \" \" in config.name:\n",
    "    raise Exception(\"Space is not allowed in model name.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "WynW4ZAdBhep"
   },
   "source": [
    "# Prepare"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "DR6-vYtr_i_P"
   },
   "source": [
    "## Get log command"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "9-QpinlslcTO",
    "outputId": "48ae3716-212c-4d46-9efb-b7b4070f73de"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cd /users/PAS0027/osu10203/deepgd && tail -n1000 -f 'artifacts/logs/GAN(gan=rgan,data=best(xing,stress),dis=deep,share=4,embed=12,gp=0,lr=1e-3).log'\n"
     ]
    }
   ],
   "source": [
    "print(f\"cd {os.getcwd()} && tail -n1000 -f '{paths.log()}'\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensorboard dev upload --logdir 'artifacts/tensorboards/GAN(gan=rgan,data=best(xing,stress),dis=deep,share=4,embed=12,gp=0,lr=1e-3)'\n"
     ]
    }
   ],
   "source": [
    "print(f\"tensorboard dev upload --logdir '{paths.tensorboard()}'\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "innqkwvH_ydD"
   },
   "source": [
    "## Set globals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "id": "Qjq7JlSZlkQR"
   },
   "outputs": [],
   "source": [
    "if cuda_idx is not None and torch.cuda.is_available():\n",
    "    device = f'cuda:{cuda_idx}'\n",
    "    pynvml.nvmlInit()\n",
    "    cuda = pynvml.nvmlDeviceGetHandleByIndex(cuda_idx)\n",
    "else:\n",
    "    device = 'cpu'\n",
    "    cuda =  None\n",
    "np.set_printoptions(precision=2)\n",
    "warnings.filterwarnings(\"ignore\", category=RuntimeWarning)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "0da286l_ApEL"
   },
   "source": [
    "## Load data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "RFuSHiZJMU4x",
    "outputId": "54065609-aaea-4536-d79c-73bc73b3f576"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Load from 'cache/G_list.pickle'\n",
      "Load from 'cache/generate_data_list(list,sparse=False,pivot_mode=None,init_mode=pmds,edge_index=full_edge_index,edge_attr=full_edge_attr,pmds_list=ndarray,gviz_list=ndarray,noisy_layout=True,device=cpu).pickle'\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/users/PAS0027/osu10203/.conda/envs/deepgd/lib/python3.9/site-packages/torch_geometric/deprecation.py:13: UserWarning: 'data.DataLoader' is deprecated, use 'loader.DataLoader' instead\n",
      "  warnings.warn(out)\n"
     ]
    }
   ],
   "source": [
    "G_list = load_G_list(data_path='data/rome', index_file='data_index.txt', cache='G_list', cache_prefix='cache/')\n",
    "data_list = generate_data_list(G_list, \n",
    "                               sparse=data_config.sparse, \n",
    "                               pivot_mode=data_config.pivot,\n",
    "                               init_mode=data_config.init,\n",
    "                               edge_index=data_config.edge.index,\n",
    "                               edge_attr=data_config.edge.attr,\n",
    "                               pmds_list=np.load('layouts/rome/pmds.npy', allow_pickle=True),\n",
    "                               gviz_list=np.load('layouts/rome/gviz.npy', allow_pickle=True),\n",
    "                               noisy_layout=True,\n",
    "                               device='cpu', \n",
    "                               cache=True,\n",
    "                               cache_prefix='cache/')\n",
    "train_loader = LazyDeviceMappingDataLoader(data_list[:10000], batch_size=config.batchsize, shuffle=True, device=device)\n",
    "val_loader = LazyDeviceMappingDataLoader(data_list[11000:], batch_size=config.batchsize, shuffle=False, device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "id": "djd1pLt2OxZq"
   },
   "outputs": [],
   "source": [
    "def draw_layout(G, method, draw=True):\n",
    "    if method == 'fa2':\n",
    "        layout = get_fa2_layout(G)\n",
    "    else:\n",
    "        try:\n",
    "            fn = getattr(nx.drawing.layout, f'{method}_layout')\n",
    "            layout = fn(G)\n",
    "        except:\n",
    "            layout = nx.drawing.nx_agraph.graphviz_layout(G, prog=method)\n",
    "    if draw:\n",
    "        nx.draw(G, pos=layout)\n",
    "    return layout"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "id": "2vBBw77lR7Em"
   },
   "outputs": [],
   "source": [
    "methods = ['neato', 'dot', 'fdp', 'sfdp', 'twopi', 'circo', 'shell', 'spring', 'circular', 'spectral', 'kamada_kawai', 'fa2', 'pmds']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "id": "dx6VVA9XIc2O"
   },
   "outputs": [],
   "source": [
    "from functools import lru_cache\n",
    "\n",
    "@lru_cache(maxsize=None)\n",
    "def load_pos(method):\n",
    "    return np.load(f'layouts/rome/{method}.npy', allow_pickle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000,
     "referenced_widgets": [
      "05a37aeeb20049dfa2f97287b3990556",
      "a75afae3adfd4810a0b0991a062535ee",
      "a24ae105ed124595992f7e2e5f909247",
      "01b4cabaf8004dd4b9269b4e6889c498",
      "f3c7c84f2b314bd8b97a7358deaf4fa1",
      "35e4a83e692b4a7bafb46fdf192b2586",
      "62847b5e2678459f845159eef6d4feab",
      "143e81f00313436dacfa246a57108944",
      "e05833f11803400b95ff22246cdc8f2b",
      "53aeede91e074378b5d7f626165655ee",
      "79a43768c1794a10a788a49e5883aa2c"
     ]
    },
    "id": "MSmP-jem9I0X",
    "outputId": "c698f0bf-e96e-4de9-aeda-0b046fc0a03f",
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# best_list = []\n",
    "# best_layout_list = []\n",
    "# for idx, (G, data) in enumerate(zip(tqdm(G_list), data_list)):\n",
    "#     xing, stress, layout = {}, {}, {}\n",
    "#     for m in methods:\n",
    "#         batch = Batch.from_data_list([data])\n",
    "#         pos = load_pos(m)\n",
    "#         p = CanonicalizationByStress()(torch.tensor(pos[idx]).float(), batch)\n",
    "#         x = Xing()(p, batch).item()\n",
    "#         s = Stress()(p, batch).item()\n",
    "#         xing[m] = x\n",
    "#         stress[m] = s\n",
    "#         layout[m] = p.numpy()\n",
    "#         # plt.figure()\n",
    "#         # graph_vis(G, pos[idx])\n",
    "#         # plt.title(f'{m} stress={s:.2f} xing={x}')\n",
    "#     best, *_ = sorted(methods, key=lambda m: (xing[m], stress[m]))\n",
    "#     best_list.append(best)\n",
    "#     best_layout_list.append(layout[best])\n",
    "#     print(f'{best}, xing={xing[best]}, stress={stress[best]:.2f}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Mx2bZNugdTqa",
    "outputId": "cec05a67-18e8-492f-aaa1-057274a8341a"
   },
   "outputs": [],
   "source": [
    "# pickle.dump(best_list, open('layouts/rome/best[xing,stress].pkl', 'wb'))\n",
    "# np.save('layouts/rome/best[xing,stress].npy', best_layout_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 49,
     "referenced_widgets": [
      "7e0ae84dd0f3458d893c81b029fe1729",
      "18bdc3bf80fb434abde68de5b72a8870",
      "2e3552ded9594635a51832387ae754a0",
      "54f4aa5080c6498080a63be2cdeefee5",
      "3937586d0879451da5e2138ca3bfd0a0",
      "a9df6e84ce8c499eb2b2cf984995d791",
      "7af6b718dd3d4e01b264377a2fa938e2",
      "a774c86e64004bd6951645f1c24abad4",
      "c86575117cd041f2837762f37d1f387b",
      "6e9907cefb1241688fbd248b8281ed9f",
      "9b8e5405d8264944acab57e81f2f5ad6"
     ]
    },
    "id": "A1-w1fspC53p",
    "outputId": "09c2d5c3-1eef-41eb-dcee-581af1294804"
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "78d59dfad66f4e97805de67af4dfef84",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/11531 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "best_layout_list = np.load('layouts/rome/best[xing,stress].npy', allow_pickle=True)\n",
    "for data, layout in zip(tqdm(data_list), best_layout_list):\n",
    "    data.gt_pos = torch.tensor(layout)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "id": "y6EGh4U7exL_"
   },
   "outputs": [],
   "source": [
    "train_loader = LazyDeviceMappingDataLoader(data_list[:10000], batch_size=config.batchsize, shuffle=True, device=device)\n",
    "val_loader = LazyDeviceMappingDataLoader(data_list[11000:], batch_size=config.batchsize, shuffle=False, device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 103,
     "referenced_widgets": [
      "38a33cf69aff48f4aabd4902a7ce48b1",
      "57e277bcb44c4e2fa2a12105d75153d1",
      "20ce017c46bd4a7c9b420270ac579405",
      "13f2df4fe0ce4af88b6e3bd1637065b0",
      "2867faa145ca4860a198d4d0bf830f89",
      "02838cf669174ed9afa41a1a8e680765",
      "7e9e3c031667455f8baca163d3a213e5",
      "4000135e7dfa46529fa6244fccc5cb57",
      "e523e9faac5b4ec38643c05ac5ba99f7",
      "e8e3ee9741e4480dae5003ccc19a7891",
      "9a569045221b43acb3ba2fbf116b2c74"
     ]
    },
    "id": "YEfRRUK0NiNm",
    "outputId": "93f53d05-7667-4cda-a38e-c3a953124f79"
   },
   "outputs": [],
   "source": [
    "# for m in methods:\n",
    "#     layouts = []\n",
    "#     for G in tqdm(G_list):\n",
    "#         layout = draw_layout(G, method=m, draw=False)\n",
    "#         layouts.append(np.array(list(layout.values())))\n",
    "#     np.save(f'layouts/rome/{m}.npy', layouts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 564
    },
    "id": "upK7AHYqPmC5",
    "outputId": "29a88e61-f55f-409c-9194-12c6d5762528"
   },
   "outputs": [],
   "source": [
    "# from collections import Counter\n",
    "\n",
    "# letter_counts = Counter(best_list)\n",
    "# df = pd.DataFrame.from_dict(letter_counts, orient='index')\n",
    "# df.plot(kind='bar', figsize=[12, 8])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "f-ODE2k8BFV6"
   },
   "source": [
    "## Create folders"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "id": "3ykUMINYBRog"
   },
   "outputs": [],
   "source": [
    "mkdirs(paths.checkpoints(), paths.tensorboard(), paths.visualization(), f\"{paths.root}/logs\", f\"{paths.root}/metrics\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "BFq9k5nzBIh5"
   },
   "source": [
    "## Load checkpoints"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "id": "W5NbsKkobIfA"
   },
   "outputs": [],
   "source": [
    "class EdgeFeatureDiscriminator(nn.Module):\n",
    "    def __init__(self, \n",
    "                 conv, \n",
    "                 dense,\n",
    "                 shared_depth,\n",
    "                 enet_depth,\n",
    "                 enet_width,\n",
    "                 aggr='add', \n",
    "                 root_weight=True,\n",
    "                 normalize=None):\n",
    "        super().__init__()\n",
    "        self.enet = nn.Sequential(*[\n",
    "            DenseLayer(in_dim=in_d,\n",
    "                       out_dim=out_d,\n",
    "                       skip=nonlin,\n",
    "                       bn=nonlin,\n",
    "                       act=nonlin,\n",
    "                       dp=None)\n",
    "            for in_d, out_d, nonlin \n",
    "            in zip([self._get_feature_dim()] + [enet_width] * (shared_depth-1),\n",
    "                   [enet_width] * shared_depth,\n",
    "                   [True] * (shared_depth-1) + [False])     \n",
    "        ])\n",
    "        self.blocks = nn.ModuleList([\n",
    "            GNNLayer(nfeat_dims=(in_d, out_d),\n",
    "                     efeat_dim=enet_width,\n",
    "                     edge_net=EdgeNet(nfeat_dims=(in_d, out_d), \n",
    "                                      efeat_dim=enet_width, \n",
    "                                      depth=enet_depth, \n",
    "                                      width=enet_width),\n",
    "                     aggr=aggr,\n",
    "                     dense=False,\n",
    "                     skip=nonlin,\n",
    "                     bn=nonlin,\n",
    "                     act=nonlin,\n",
    "                     root_weight=root_weight) \n",
    "            for in_d, out_d, nonlin \n",
    "            in zip(conv[:-1], conv[1:], [True] * (len(conv)-2) + [False])     \n",
    "        ])\n",
    "        self.pool = (gnn.global_mean_pool if aggr == 'mean' \n",
    "                     else gnn.global_add_pool if aggr == 'add' \n",
    "                     else None)\n",
    "        self.dense = nn.Identity() if not dense else nn.Sequential(*[\n",
    "            DenseLayer(in_dim=in_d,\n",
    "                       out_dim=out_d,\n",
    "                       skip=nonlin,\n",
    "                       bn=nonlin,\n",
    "                       act=nonlin,\n",
    "                       dp=None)\n",
    "            for in_d, out_d, nonlin \n",
    "            in zip(conv[-1:] + dense[:-1], dense, [True] * (len(dense)-1) + [False])     \n",
    "        ])\n",
    "        self.normalize = normalize or IdentityTransformation()\n",
    "    \n",
    "    def forward(self, batch):\n",
    "        x = torch.ones_like(batch.pos)\n",
    "        e = self.enet(self._get_features(self._get_edge_info(batch, layout='pos')))\n",
    "        for block in self.blocks:\n",
    "            x = block(x, e, batch)\n",
    "        x = self.pool(x, batch.batch)\n",
    "        x = self.dense(x)\n",
    "        return x\n",
    "    \n",
    "    def _get_edge_info(self, batch, layout='gt_pos'):\n",
    "        pos = self.normalize(batch[layout].float(), batch)\n",
    "        src, dst = get_edges(pos, batch)\n",
    "        v, u = l2_normalize(dst - src, return_norm=True)\n",
    "        d = batch.edge_attr[:, :1]\n",
    "        return {\n",
    "            \"src\": src,\n",
    "            \"dst\": dst,\n",
    "            \"v\": v,\n",
    "            \"u\": u,\n",
    "            \"d\": d,\n",
    "        }\n",
    "\n",
    "    def _get_features(self, edges):\n",
    "        return torch.cat([edges['src'], edges['dst'], edges['d']], dim=1)\n",
    "        \n",
    "    def _get_feature_dim(self):\n",
    "        return self._get_features({\n",
    "            \"src\": torch.zeros(1, 2),\n",
    "            \"dst\": torch.zeros(1, 2),\n",
    "            \"v\": torch.zeros(1, 2),\n",
    "            \"u\": torch.zeros(1, 1),\n",
    "            \"d\": torch.zeros(1, 1),\n",
    "        }).shape[-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "id": "96ZActxlL0fJ"
   },
   "outputs": [],
   "source": [
    "class StressDiscriminator(nn.Module):\n",
    "    def __init__(self, normalize=CanonicalizationByStress(), **kwargs):\n",
    "        super().__init__()\n",
    "        self.dummy = nn.Parameter(torch.zeros(1))\n",
    "        self.normalize = normalize\n",
    "        self.stress = Stress(reduce=None)\n",
    "\n",
    "    def forward(self, batch):\n",
    "        return -self.stress(self.normalize(batch.pos, batch), batch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "id": "IJ6xCYBUuhcw"
   },
   "outputs": [],
   "source": [
    "def get_ckpt_epoch(folder, epoch):\n",
    "    if not os.path.isdir(folder):\n",
    "        os.mkdir(folder)\n",
    "    if epoch >= 0:\n",
    "        return epoch\n",
    "    ckpt_files = os.listdir(folder)\n",
    "    last_epoch = 0\n",
    "    if ckpt_files:\n",
    "        last_epoch = sorted(list(map(lambda x: int(re.search('(?<=epoch_)(\\d+)(?=\\.)', x).group(1)), ckpt_files)))[-1]\n",
    "    return last_epoch + epoch + 1\n",
    "\n",
    "def start_epoch():\n",
    "    return get_ckpt_epoch(paths.checkpoints(), config.epoch.start)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "JacDYXka9U01",
    "outputId": "ef4e31ab-2144-497c-8e4f-4d3740ce2d52"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading from artifacts/checkpoints/GAN(gan=rgan,data=best(xing,stress),dis=deep,share=4,embed=12,gp=0,lr=1e-3)/gen_epoch_899.pt...\n",
      "Loading from artifacts/checkpoints/GAN(gan=rgan,data=best(xing,stress),dis=deep,share=4,embed=12,gp=0,lr=1e-3)/gen_optim_epoch_899.pt...\n"
     ]
    }
   ],
   "source": [
    "generator = Generator(**config.generator.params[...]).to(device)\n",
    "generator_optimizer = config.generator.optim(generator.parameters(), lr=config.generator.lr.initial * config.generator.lr.decay ** start_epoch())\n",
    "generator_scheduler = torch.optim.lr_scheduler.ExponentialLR(generator_optimizer, gamma=config.generator.lr.decay)\n",
    "if start_epoch() != 0:\n",
    "    gen_ckpt_epoch = start_epoch()\n",
    "elif config.generator.pretrained.name is not None and config.generator.pretrained.epoch != 0:\n",
    "    gen_pretrained_epoch = get_ckpt_epoch(paths.gen_pretrain(), config.generator.pretrained.epoch)\n",
    "    gen_ckpt_epoch = gen_pretrained_epoch \n",
    "else:\n",
    "    gen_ckpt_epoch = None\n",
    "if gen_ckpt_epoch is not None:\n",
    "    # Load generator\n",
    "    gen_ckpt_file = f\"{paths.checkpoints()}/gen_epoch_{gen_ckpt_epoch}.pt\"\n",
    "    print(f\"Loading from {gen_ckpt_file}...\")\n",
    "    generator.load_state_dict(torch.load(gen_ckpt_file, map_location=torch.device(device)))\n",
    "    # Load generator optimizer\n",
    "    gen_optim_ckpt_file = f\"{paths.checkpoints()}/gen_optim_epoch_{gen_ckpt_epoch}.pt\"\n",
    "    print(f\"Loading from {gen_optim_ckpt_file}...\")\n",
    "    generator_optimizer.load_state_dict(torch.load(gen_optim_ckpt_file, map_location=torch.device(device)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "oWQhD7-88VXe",
    "outputId": "3ba01600-b546-4e45-a6bd-c7a65af1ae9e"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading from artifacts/checkpoints/GAN(gan=rgan,data=best(xing,stress),dis=deep,share=4,embed=12,gp=0,lr=1e-3)/dis_epoch_899.pt...\n",
      "Loading from artifacts/checkpoints/GAN(gan=rgan,data=best(xing,stress),dis=deep,share=4,embed=12,gp=0,lr=1e-3)/dis_optim_epoch_899.pt...\n"
     ]
    }
   ],
   "source": [
    "discriminator = EdgeFeatureDiscriminator(**config.discriminator.params[...]).to(device)\n",
    "discriminator_optimizer = config.discriminator.optim(discriminator.parameters(), lr=config.discriminator.lr.initial * config.discriminator.lr.decay ** start_epoch())\n",
    "discriminator_scheduler = torch.optim.lr_scheduler.ExponentialLR(discriminator_optimizer, gamma=config.discriminator.lr.decay)\n",
    "if start_epoch() != 0:\n",
    "    dis_ckpt_epoch = start_epoch()\n",
    "elif config.discriminator.pretrained.name is not None and config.discriminator.pretrained.epoch != 0:\n",
    "    dis_pretrained_epoch = get_ckpt_epoch(paths.dis_pretrain(), config.discriminator.pretrained.epoch)\n",
    "    dis_ckpt_epoch = dis_pretrained_epoch # f\"{paths.dis_pretrain()}/dis_epoch_{dis_pretrained_epoch}.pt\"\n",
    "else:\n",
    "    dis_ckpt_epoch = None\n",
    "if dis_ckpt_epoch is not None:\n",
    "    # Load discriminator\n",
    "    dis_ckpt_file = f\"{paths.checkpoints()}/dis_epoch_{dis_ckpt_epoch}.pt\"\n",
    "    print(f\"Loading from {dis_ckpt_file}...\")\n",
    "    discriminator.load_state_dict(torch.load(dis_ckpt_file, map_location=torch.device(device)))\n",
    "    # Load discriminator optimizer\n",
    "    dis_optim_ckpt_file = f\"{paths.checkpoints()}/dis_optim_epoch_{dis_ckpt_epoch}.pt\"\n",
    "    print(f\"Loading from {dis_optim_ckpt_file}...\")\n",
    "    discriminator_optimizer.load_state_dict(torch.load(dis_optim_ckpt_file, map_location=torch.device(device)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "N39dDHraedM6"
   },
   "source": [
    "# Train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "id": "BdIJSDSJedM-"
   },
   "outputs": [],
   "source": [
    "stress_criterion = StressDiscriminator().to(device)\n",
    "val_criterion = Stress(reduce=None)\n",
    "xing_criterion = Xing(reduce=None)\n",
    "dis_convert = DiscriminatorDataConverter(complete_graph=config.discriminator.complete, normalize=config.discriminator.params.normalize)\n",
    "tensorboard = SummaryWriter(log_dir=paths.tensorboard())\n",
    "epoch = start_epoch() + 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "id": "OHRWj1h0edM-"
   },
   "outputs": [],
   "source": [
    "def gradient_penalty(interpolated, discriminator, weight=10):\n",
    "    interpolated.pos.requires_grad_()\n",
    "    prob_interpolated = discriminator(interpolated)\n",
    "    gradients = autograd.grad(outputs=prob_interpolated, \n",
    "                              inputs=interpolated.pos,\n",
    "                              grad_outputs=torch.ones_like(prob_interpolated),\n",
    "                              create_graph=True, \n",
    "                              retain_graph=True, \n",
    "                              allow_unused=True)[0]\n",
    "    gradients_norm = torch.sqrt(gnn.global_add_pool(gradients.square().sum(dim=1), batch.batch) + 1e-8)\n",
    "    return weight * ((gradients_norm - 1) ** 2).mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "id": "o5wJJN9AedM_"
   },
   "outputs": [],
   "source": [
    "def get_gp_loss(batch, fake_pos, weight):\n",
    "    if weight > 0:\n",
    "        interp = dis_convert(batch, fake_pos, random.random())\n",
    "        return gradient_penalty(interp, discriminator, weight).mean()\n",
    "    return 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "id": "QwuA6sRoedM_"
   },
   "outputs": [],
   "source": [
    "def get_sgan_loss(batch, fake_pos, mode='discriminator'):\n",
    "    real = dis_convert(batch)\n",
    "    fake = dis_convert(batch, fake_pos)\n",
    "    pred = discriminator(merge_batch(real, fake)).view(2, -1).T\n",
    "    if mode == 'discriminator':\n",
    "        label = torch.zeros(pred.shape[0]).long()\n",
    "    elif mode == 'generator':\n",
    "        label = torch.ones(pred.shape[0]).long()\n",
    "    else:\n",
    "        raise Exception\n",
    "    return nn.CrossEntropyLoss()(pred, label)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "id": "GuAP1YchedNA"
   },
   "outputs": [],
   "source": [
    "def get_rgan_loss(batch, fake_pos, mode='discriminator'):\n",
    "    real = dis_convert(batch)\n",
    "    fake = dis_convert(batch, fake_pos)\n",
    "    pred = discriminator(merge_batch(real, fake)).view(2, -1).T\n",
    "    real_pred, fake_pred = pred[:,0], pred[:,1]\n",
    "    if mode == 'discriminator':\n",
    "        losses = - F.logsigmoid(real_pred - fake_pred)\n",
    "    elif mode == 'generator':\n",
    "        losses = - F.logsigmoid(fake_pred - real_pred)\n",
    "    else:\n",
    "        raise Exception\n",
    "    return losses.mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "id": "KFllQ7BredNA"
   },
   "outputs": [],
   "source": [
    "def get_wgan_loss(batch, fake_pos, mode='discriminator'):\n",
    "    real = dis_convert(batch)\n",
    "    fake = dis_convert(batch, fake_pos)\n",
    "    pred = discriminator(merge_batch(real, fake)).view(2, -1).T\n",
    "    real_pred, fake_pred = pred[:,0], pred[:,1]\n",
    "    if mode == 'discriminator':\n",
    "        losses = fake_pred - real_pred \n",
    "    elif mode == 'generator':\n",
    "        losses = real_pred - fake_pred\n",
    "    else:\n",
    "        raise Exception\n",
    "    return losses.mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "id": "r5Kg_tYEedNA"
   },
   "outputs": [],
   "source": [
    "def get_ragan_loss(batch, fake_pos, mode='discriminator'):\n",
    "    real = dis_convert(batch)\n",
    "    fake = dis_convert(batch, fake_pos)\n",
    "    pred = discriminator(merge_batch(real, fake)).view(2, -1).T\n",
    "    real_pred, fake_pred = pred[:,0], pred[:,1]\n",
    "    if mode == 'discriminator':\n",
    "        losses = - F.logsigmoid(real_pred - fake_pred.mean()) - F.logsigmoid(real_pred.mean() - fake_pred)\n",
    "    elif mode == 'generator':\n",
    "        losses = - F.logsigmoid(fake_pred - real_pred.mean()) - F.logsigmoid(fake_pred.mean() - real_pred)\n",
    "    else:\n",
    "        raise Exception\n",
    "    return losses.mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_dgdv2_loss(batch, fake_pos, mode='discriminator'):\n",
    "    fake = dis_convert(batch, fake_pos)\n",
    "    pred = discriminator(fake)\n",
    "    if mode == 'discriminator':\n",
    "        gt = get_gt(batch, fake_pos)\n",
    "        loss = criterion(pred, gt)\n",
    "    elif mode == 'generator':\n",
    "        losses = pred.sum(dim=0)\n",
    "#         losses = torch.tensor(config.importance).to(device) * losses #/ losses.detach()\n",
    "        loss = losses.sum()\n",
    "    else:\n",
    "        raise Exception\n",
    "    return -loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "id": "3IE_84f-edNB"
   },
   "outputs": [],
   "source": [
    "def get_gan_loss(batch, fake_pos, mode='discriminator'):\n",
    "    return {\"sgan\": get_sgan_loss,\n",
    "            \"wgan\": get_wgan_loss,\n",
    "            \"rgan\": get_rgan_loss,\n",
    "            \"ragan\": get_ragan_loss,\n",
    "            \"dgdv2\": get_dgdv2_loss}[config.gan_flavor](batch, fake_pos, mode)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {
    "colab": {
     "background_save": true,
     "base_uri": "https://localhost:8080/",
     "height": 80
    },
    "id": "2NjekXgkedNB",
    "outputId": "da85b2ec-c0de-459c-c0dd-7ba3519419fa",
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d351ef6670d74438bfb8e85ab92c98fa",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Tab(children=(VBox(children=(Wrapper(), Hud())), HBox(children=(Output(),), layout=Layout(height='500px', over…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m/tmp/slurmtmp.5660320/ipykernel_39383/801679965.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m    133\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0mconfig\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0malternate\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;34m'epoch'\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    134\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mbatch\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mtrain_loader\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 135\u001b[0;31m             \u001b[0mtrain_gen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mbatch\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mepoch\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    136\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    137\u001b[0m     \u001b[0mdiscriminator_scheduler\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstep\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/tmp/slurmtmp.5660320/ipykernel_39383/801679965.py\u001b[0m in \u001b[0;36mtrain_gen\u001b[0;34m(batch, epoch)\u001b[0m\n\u001b[1;32m     25\u001b[0m     \u001b[0mgenerator\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mzero_grad\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     26\u001b[0m     \u001b[0mdiscriminator\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mzero_grad\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 27\u001b[0;31m     \u001b[0mgenerator_output\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mgenerator\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mbatch\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     28\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0mconfig\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdiscriminator\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnoise\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstd\u001b[0m \u001b[0;34m>\u001b[0m \u001b[0;36m0\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     29\u001b[0m         \u001b[0mgenerator_output\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mgenerator_output\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrandn_like\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mgenerator_output\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m*\u001b[0m \u001b[0mconfig\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdiscriminator\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnoise\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstd\u001b[0m \u001b[0;34m*\u001b[0m \u001b[0mconfig\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdiscriminator\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnoise\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdecay\u001b[0m \u001b[0;34m**\u001b[0m \u001b[0mepoch\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.conda/envs/deepgd/lib/python3.9/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1049\u001b[0m         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks\n\u001b[1;32m   1050\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1051\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1052\u001b[0m         \u001b[0;31m# Do not call functions when jit is used\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1053\u001b[0m         \u001b[0mfull_backward_hooks\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnon_full_backward_hooks\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/deepgd/deepgd/modules.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, data, weights, output_hidden, numpy)\u001b[0m\n\u001b[1;32m    631\u001b[0m             \u001b[0mvout\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mv\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdetach\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcpu\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnumpy\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mnumpy\u001b[0m \u001b[0;32melse\u001b[0m \u001b[0mv\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    632\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnormalize\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 633\u001b[0;31m                 \u001b[0mvout\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnormalize\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mvout\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdata\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    634\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    635\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mhidden\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0moutput_hidden\u001b[0m \u001b[0;32melse\u001b[0m \u001b[0mvout\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.conda/envs/deepgd/lib/python3.9/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1049\u001b[0m         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks\n\u001b[1;32m   1050\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1051\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1052\u001b[0m         \u001b[0;31m# Do not call functions when jit is used\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1053\u001b[0m         \u001b[0mfull_backward_hooks\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnon_full_backward_hooks\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/deepgd/deepgd/transform.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, pos, data)\u001b[0m\n\u001b[1;32m    151\u001b[0m         \u001b[0mbatch\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmake_batch\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdata\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    152\u001b[0m         \u001b[0mpos\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtranslate\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpos\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbatch\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 153\u001b[0;31m         \u001b[0mpos\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrotate\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpos\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbatch\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    154\u001b[0m         \u001b[0mpos\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnormalize\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpos\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbatch\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    155\u001b[0m         \u001b[0mpos\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mscale\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpos\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbatch\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.conda/envs/deepgd/lib/python3.9/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1049\u001b[0m         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks\n\u001b[1;32m   1050\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1051\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1052\u001b[0m         \u001b[0;31m# Do not call functions when jit is used\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1053\u001b[0m         \u001b[0mfull_backward_hooks\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnon_full_backward_hooks\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/deepgd/deepgd/transform.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, pos, data)\u001b[0m\n\u001b[1;32m     51\u001b[0m         \u001b[0mXXT\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0meinsum\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'ni,nj->nij'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpos\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpos\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     52\u001b[0m         \u001b[0mcov\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch_scatter\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mscatter\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mXXT\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbatch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbatch\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdim\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mreduce\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m'mean'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 53\u001b[0;31m         \u001b[0mcomponents\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mlinalg\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0meigh\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcov\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0meigenvectors\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     54\u001b[0m         rotated_pos = torch.einsum('ij,njk,nk->ni', \n\u001b[1;32m     55\u001b[0m                                    \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbase_rotation\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpos\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "def train_dis(batch, epoch):\n",
    "    generator.requires_grad_(False)\n",
    "    discriminator.zero_grad()\n",
    "    generator_output = generator(batch)\n",
    "    if config.discriminator.noise.std > 0:\n",
    "        generator_output = generator_output + torch.randn_like(generator_output) * config.discriminator.noise.std * config.discriminator.noise.decay ** epoch\n",
    "    discriminator_loss = get_gan_loss(batch, generator_output, mode='discriminator')\n",
    "    \n",
    "    # train discriminator\n",
    "    discriminator_loss.backward()\n",
    "    discriminator_optimizer.step()\n",
    "\n",
    "    # gradient penalty\n",
    "    if config.gp_weight > 0:\n",
    "        discriminator.zero_grad()\n",
    "        gp_loss = get_gp_loss(batch, generator_output, config.gp_weight)\n",
    "        gp_loss.backward()\n",
    "        discriminator_optimizer.step()\n",
    "\n",
    "    hud['dis_loss'] = format(discriminator_loss.item(), '.2e')\n",
    "    pbar().update()\n",
    "\n",
    "def train_gen(batch, epoch):\n",
    "    generator.requires_grad_(True)\n",
    "    generator.zero_grad()\n",
    "    discriminator.zero_grad()\n",
    "    generator_output = generator(batch)\n",
    "    if config.discriminator.noise.std > 0:\n",
    "        generator_output = generator_output + torch.randn_like(generator_output) * config.discriminator.noise.std * config.discriminator.noise.decay ** epoch\n",
    "    generator_loss = get_gan_loss(batch, generator_output, mode='generator') \n",
    "    \n",
    "    #train generator\n",
    "    generator_loss.backward()\n",
    "    generator_optimizer.step()\n",
    "\n",
    "    with torch.no_grad():\n",
    "        dis_batch = dis_convert(batch, generator_output)\n",
    "        stress = stress_criterion(dis_batch).mean()\n",
    "        critic = discriminator(dis_batch).mean()\n",
    "    hud.append({'gen_loss': format(generator_loss.item(), '.2e'),\n",
    "                'stress': format(stress.item(), '.2e'),\n",
    "                'critic': format(critic.item(), '.2e')})\n",
    "    pbar().update()\n",
    "\n",
    "def cuda_memsafe_map(fn, *iterables, summary=False):\n",
    "    total, failed = 0, 0\n",
    "    iterator = zip(*iterables)\n",
    "    items = None\n",
    "    while True:\n",
    "        try:\n",
    "            items = next(iterator)\n",
    "            yield fn(*items)\n",
    "        except StopIteration:\n",
    "            if summary:\n",
    "                print(f'Iteration finished. {failed} out of {total} failed!')\n",
    "            break\n",
    "        except RuntimeError:\n",
    "            print('CUDA memory overflow! Skip batch...')\n",
    "            del items\n",
    "            failed += 1\n",
    "        torch.cuda.empty_cache()\n",
    "        total += 1\n",
    "    \n",
    "def validate(model, data_loader, criterion=val_criterion):\n",
    "    def val_one_batch(batch):\n",
    "        batch = preprocess_batch(model, batch)\n",
    "        pred = CanonicalizationByStress()(model(batch), batch)\n",
    "        gt = CanonicalizationByStress()(batch.gt_pos, batch)\n",
    "        loss = criterion(pred, batch)\n",
    "        gt_loss = criterion(gt, batch)\n",
    "        spc = (loss - gt_loss) / torch.maximum(torch.maximum(loss, gt_loss), torch.ones_like(loss)*1e-5)\n",
    "        return loss.mean().item(), spc.mean().item()\n",
    "    loss_all, spc_all = zip(*cuda_memsafe_map(val_one_batch, data_loader))\n",
    "    return np.mean(loss_all), np.mean(spc_all)\n",
    "\n",
    "def log(msg):\n",
    "    msg = f\"[{epoch:03}] {msg}\"\n",
    "    print(msg, file=open(paths.log(), \"a\"))\n",
    "    with log_out: \n",
    "        print(msg)\n",
    "\n",
    "print(f\"{'='*10} {config.link} {'='*10}\", file=open(paths.log(), \"a\"))\n",
    "hud = Hud()\n",
    "pbar = Wrapper(tqdm, total=len(train_loader)*2, smoothing=0)\n",
    "plot_out = Output()\n",
    "log_out = Output()\n",
    "tabs = {\"status\": VBox([pbar, hud]), \n",
    "        \"plot\": HBox([plot_out], layout=Layout(height='500px', overflow_y='auto')),\n",
    "        \"log\": HBox([log_out], layout=Layout(height='500px', overflow_y='auto'))}\n",
    "tab_bar = Tab(children=list(tabs.values()))\n",
    "[tab_bar.set_title(i, name) for i, name in enumerate(tabs)]\n",
    "display(tab_bar)\n",
    "while True:\n",
    "    if epoch % config.log_interval == 0:\n",
    "        torch.save(generator.state_dict(), f\"{paths.checkpoints()}/gen_epoch_{epoch}.pt\")\n",
    "        torch.save(generator_optimizer.state_dict(), f\"{paths.checkpoints()}/gen_optim_epoch_{epoch}.pt\")\n",
    "        torch.save(discriminator.state_dict(), f\"{paths.checkpoints()}/dis_epoch_{epoch}.pt\")\n",
    "        torch.save(discriminator_optimizer.state_dict(), f\"{paths.checkpoints()}/dis_optim_epoch_{epoch}.pt\")\n",
    "        generator.eval()\n",
    "        with torch.no_grad():\n",
    "            val_stress, val_stress_spc = validate(model=generator, data_loader=val_loader)\n",
    "            val_xing, val_xing_spc = validate(model=generator, data_loader=val_loader, criterion=xing_criterion)\n",
    "            with plot_out:\n",
    "                fig = plt.figure()\n",
    "                graph_vis(G_list[11100], generator(make_batch(data_list[11100]).to(device)).cpu())\n",
    "                plt.show()\n",
    "        # tensorboard.add_scalars('loss', {'train': train_loss, \n",
    "        #                                  'validation': val_loss}, epoch)\n",
    "        # for i, fn in enumerate(loss_fns):\n",
    "        #     tensorboard.add_scalars(type(fn).__name__, {'train': train_loss_comp[i].item(), \n",
    "        #                                           'validation': val_loss_comp[i].item()}, epoch)\n",
    "        hud.append({\n",
    "            'val_stress': format(val_stress, '.2f'),\n",
    "            'val_stress_spc': format(val_stress_spc, '.2%'),\n",
    "            'val_xing': format(val_xing, '.2f'),\n",
    "            'val_xing_spc': format(val_xing_spc, '.2%'),\n",
    "        })\n",
    "        log(f\"stress={hud.data['val_stress']}({hud.data['val_stress_spc']}) xing={hud.data['val_xing']}({hud.data['val_xing_spc']})\")\n",
    "        \n",
    "    # handle.update(tab_bar)\n",
    "    pbar().reset()\n",
    "    pbar().set_description(desc=f\"[epoch {epoch}/{config.epoch.end}]\")\n",
    "    hud(title=f\"epoch {epoch}\")\n",
    "    generator.train()\n",
    "    discriminator.train()\n",
    "    # proper: proper layout\n",
    "    for _ in range(config.discriminator.repeat):\n",
    "        for batch in train_loader:\n",
    "            train_dis(batch, epoch)\n",
    "            if config.alternate == 'iteration':\n",
    "                train_gen(batch, epoch)\n",
    "\n",
    "    if config.alternate == 'epoch':\n",
    "        for batch in train_loader:\n",
    "            train_gen(batch, epoch)\n",
    "\n",
    "    discriminator_scheduler.step()\n",
    "    generator_scheduler.step()\n",
    "\n",
    "    if epoch == config.epoch.end:\n",
    "        break\n",
    "    epoch += 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "0qYK_X_hrq9o"
   },
   "source": [
    "# Test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "XEvmAtNjWhcY",
    "outputId": "9f4c86a4-a4b0-48f1-cb58-ab3679173184"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading from artifacts/checkpoints/GAN(gan=rgan,data=best(xing,stress),dis=deep,share=4,embed=12,gp=0,lr=1e-3)/gen_epoch_900.pt...\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test_epoch = -1\n",
    "\n",
    "test_generator = Generator(**config.generator.params[...]).to(device)\n",
    "test_ckpt_epoch = get_ckpt_epoch(paths.checkpoints(), test_epoch)\n",
    "test_ckpt_file = f\"{paths.checkpoints()}/gen_epoch_{test_ckpt_epoch}.pt\"\n",
    "print(f\"Loading from {test_ckpt_file}...\")\n",
    "test_generator.load_state_dict(torch.load(test_ckpt_file, map_location=torch.device(device)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {
    "colab": {
     "background_save": true,
     "base_uri": "https://localhost:8080/",
     "height": 1000,
     "referenced_widgets": [
      "9b7d18992dbf4a84931f916fbaaa37d8",
      "2c664d7e1e6d474fb87094fc0da270d3",
      "d375f5f82a4645d7b29c541596dc1c7f",
      "ecfa311bb06b4f6fa504598263ff69d3",
      "d1859990877149eab2b414b4e2a0e0c7",
      "d2be6733734844bca3e1db9b49eb1c1e",
      "4d98072bdf674792a13d77e85dd285d1",
      "f7ed92c383d741a0a1cef2536ff2db13",
      "2ba110f835c64ec4b82974996aed10af",
      "09397ceac9014af0aebca848d4d96d48",
      "fd0709f48d6c4527bc00b5f407f3c1b9"
     ]
    },
    "id": "ZubmSTcYrq9o",
    "outputId": "fc2f65a2-2085-43ac-b110-cbfe742318a8",
    "scrolled": true,
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "dc869004ce174446b89169b728760059",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "rotate = RotateByPrincipalComponents()\n",
    "def test_callback(*, idx, pred, metrics):\n",
    "    # graph_vis(G_list[idx], pred, file_name=f\"{paths.visualization()}/{idx}_{metrics['stress']:.2f}_{metrics['resolution_score']:.2f}.png\")\n",
    "    pred = rotate(torch.tensor(pred), data_list[idx])\n",
    "    graph_vis(G_list[idx], pred)\n",
    "    plt.title(f\"[pred] idx: {idx}, stress: {metrics['stress']:.2f}({metrics['stress_spc']:.2%}), xing: {metrics['xing']:.2f}({metrics['xing_spc']:.2%})\")\n",
    "    plt.show()\n",
    "    gt_pos = rotate(data_list[idx].gt_pos, data_list[idx])\n",
    "    graph_vis(G_list[idx], gt_pos, node_color='orange')\n",
    "    plt.title(f\"[gt] idx: {idx}, stress: {metrics['gt_stress']:.2f}, xing: {metrics['gt_xing']:.2f}\")\n",
    "    plt.show()\n",
    "    \n",
    "test_metrics = test(model=test_generator, \n",
    "                    criteria_list=[], \n",
    "                    dataset=data_list, \n",
    "                    idx_range=range(10000, 11000), \n",
    "#                     callback=test_callback,\n",
    "                    gt_pos=None)\n",
    "pickle.dump(test_metrics, open(paths.metrics(\"test\"), \"wb\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 318
    },
    "id": "t2rFXIr_PsrJ",
    "outputId": "6a4205b1-4630-4de4-ec72-bdba3d3ea114"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "stress: tensor(329.5504)\n",
      "stress_spc: tensor(0.0501, dtype=torch.float64)\n",
      "xing: tensor(28.5690)\n",
      "xing_spc: tensor(0.2044, dtype=torch.float64)\n",
      "l1_angle: tensor(81.3751)\n",
      "l1_angle_spc: tensor(0.0572, dtype=torch.float64)\n",
      "edge: tensor(0.1551)\n",
      "edge_spc: tensor(0.1297, dtype=torch.float64)\n",
      "ring: tensor(257.5183)\n",
      "ring_spc: tensor(0.0073, dtype=torch.float64)\n",
      "tsne: tensor(0.2073)\n",
      "tsne_spc: tensor(0.0921, dtype=torch.float64)\n",
      "reso_score: tensor(0.5994)\n",
      "min_angle: tensor(5.7359)\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<style type=\"text/css\">\n",
       "</style>\n",
       "<table id=\"T_4f3a0_\">\n",
       "  <thead>\n",
       "    <tr>\n",
       "      <th class=\"blank level0\" >&nbsp;</th>\n",
       "      <th class=\"col_heading level0 col0\" >stress</th>\n",
       "      <th class=\"col_heading level0 col1\" >stress_spc</th>\n",
       "      <th class=\"col_heading level0 col2\" >xing</th>\n",
       "      <th class=\"col_heading level0 col3\" >xing_spc</th>\n",
       "      <th class=\"col_heading level0 col4\" >l1_angle</th>\n",
       "      <th class=\"col_heading level0 col5\" >l1_angle_spc</th>\n",
       "      <th class=\"col_heading level0 col6\" >edge</th>\n",
       "      <th class=\"col_heading level0 col7\" >edge_spc</th>\n",
       "      <th class=\"col_heading level0 col8\" >ring</th>\n",
       "      <th class=\"col_heading level0 col9\" >ring_spc</th>\n",
       "      <th class=\"col_heading level0 col10\" >tsne</th>\n",
       "      <th class=\"col_heading level0 col11\" >tsne_spc</th>\n",
       "      <th class=\"col_heading level0 col12\" >reso_score</th>\n",
       "      <th class=\"col_heading level0 col13\" >min_angle</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th id=\"T_4f3a0_level0_row0\" class=\"row_heading level0 row0\" >0</th>\n",
       "      <td id=\"T_4f3a0_row0_col0\" class=\"data row0 col0\" >329.55</td>\n",
       "      <td id=\"T_4f3a0_row0_col1\" class=\"data row0 col1\" >5.01%</td>\n",
       "      <td id=\"T_4f3a0_row0_col2\" class=\"data row0 col2\" >28.57</td>\n",
       "      <td id=\"T_4f3a0_row0_col3\" class=\"data row0 col3\" >20.44%</td>\n",
       "      <td id=\"T_4f3a0_row0_col4\" class=\"data row0 col4\" >81.38</td>\n",
       "      <td id=\"T_4f3a0_row0_col5\" class=\"data row0 col5\" >5.72%</td>\n",
       "      <td id=\"T_4f3a0_row0_col6\" class=\"data row0 col6\" >0.16</td>\n",
       "      <td id=\"T_4f3a0_row0_col7\" class=\"data row0 col7\" >12.97%</td>\n",
       "      <td id=\"T_4f3a0_row0_col8\" class=\"data row0 col8\" >257.52</td>\n",
       "      <td id=\"T_4f3a0_row0_col9\" class=\"data row0 col9\" >0.73%</td>\n",
       "      <td id=\"T_4f3a0_row0_col10\" class=\"data row0 col10\" >0.21</td>\n",
       "      <td id=\"T_4f3a0_row0_col11\" class=\"data row0 col11\" >9.21%</td>\n",
       "      <td id=\"T_4f3a0_row0_col12\" class=\"data row0 col12\" >0.60</td>\n",
       "      <td id=\"T_4f3a0_row0_col13\" class=\"data row0 col13\" >5.74</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n"
      ],
      "text/plain": [
       "<pandas.io.formats.style.Styler at 0x2b98e529e8e0>"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "metrics = test_metrics\n",
    "print('stress:', metrics['stress'].mean())\n",
    "print('stress_spc:', metrics['stress_spc'].mean())\n",
    "print('xing:', metrics['xing'].mean())\n",
    "print('xing_spc:', metrics['xing_spc'].mean())\n",
    "print('l1_angle:', metrics['l1_angle'].mean())\n",
    "print('l1_angle_spc:', metrics['l1_angle_spc'].mean())\n",
    "print('edge:', metrics['edge'].mean())\n",
    "print('edge_spc:', metrics['edge_spc'].mean())\n",
    "print('ring:', metrics['ring'].mean())\n",
    "print('ring_spc:', metrics['ring_spc'].mean())\n",
    "print('tsne:', metrics['tsne'].mean())\n",
    "print('tsne_spc:', metrics['tsne_spc'].mean())\n",
    "print('reso_score:', metrics['resolution_score'].mean())\n",
    "print('min_angle:', metrics['min_angle'].mean())\n",
    "columns = [\n",
    "    'stress',\n",
    "    'stress_spc',\n",
    "    'xing',\n",
    "    'xing_spc',\n",
    "    'l1_angle',\n",
    "    'l1_angle_spc',\n",
    "    'edge',\n",
    "    'edge_spc',\n",
    "    'ring',\n",
    "    'ring_spc',\n",
    "    'tsne',\n",
    "    'tsne_spc',\n",
    "    'reso_score',\n",
    "    'min_angle'\n",
    "]\n",
    "df = pd.DataFrame(map(lambda m: metrics[m].mean().item(), list(metrics.keys())[:-1])).set_axis(columns).T\n",
    "df.style.format({c: \"{:.2f}\" for c in columns if 'spc' not in c} | {c: \"{:.2%}\" for c in columns if 'spc' in c})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {
    "id": "5e7vPr0M3mmG"
   },
   "outputs": [],
   "source": [
    "metrics = {}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "methods = ['fa2']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 433,
     "referenced_widgets": [
      "a0cf6e046add4e639e52db69656e0db9",
      "8e237731bb1241ac858581b25389c012",
      "0891547293ad4211a138a4d92e01bb0b",
      "34912c8e581e47a99a1c4829c38cb75c",
      "5aa034c99c94424e835a2ee954b829ed",
      "248f6393eff74c3eb8740ca6d698045b",
      "686687a29c7f4ec98c4d76ca2062b14b",
      "996a84d1c23b47e4a2af252c6b60f55d",
      "3507f883bb954412b21c1c5641763ed7",
      "9582a849d91e457cadd514c95d6e1b5b",
      "52bfe674c2c547e09b7400a32a96f703",
      "93d58b3a12394fe1aa6f0431fc7cb0db",
      "af1d7855c9d646b0bcc7d95c4cb147a6",
      "f820e2b898cd461da04a4f651d27273d",
      "c0a112cfaa3d4d0aa9530372de84d8fb",
      "fd03985495694fbf9feb3f31ef8d18ce",
      "e02a2b36aa4c4d6fb5d2ac6693c1a927",
      "06f388ff5e3a48828d465096fe69977b",
      "c0052778c9c940e88935593e9f94b83b",
      "99a5e77a192d4a989c4534237eb16730",
      "77c6da092a6647fc9e54f7e276d75616",
      "8c13c0e8b5934117a662a78296bcb1b6",
      "cdcc3bd4235247b48c468c0f6b7682cc",
      "a2c6fa6fef4549cc853fa93b4b18cf95",
      "15cc02a756314083801b5f6798a1948c",
      "24d1c67fe52f452ab0c767e3c1a818aa",
      "c117db25c6fd484a9d0c7f210ccb1ade",
      "b090b60b61084cf692b3066dfac7f0a8",
      "5eab9848059c491b9c6cbac88683f025",
      "e87ffbaa1f204b52b9226337a506b040",
      "cf52198a9e4b46c7b49537a79919c09e",
      "51501227bc064ccbab97c22a679a0750",
      "d1b794a997dc457888a90bf44faee9b4",
      "21a755b50e5047bea797eccf24b35a0c",
      "16e78c3823ec4faf9f9cd9095b218ff5",
      "4ccd332c6b734acd8eff236eaf568dfa",
      "bd60fd13748347e48629df78fa455765",
      "e9187e4ba80b4fafa0334dae2e39aded",
      "4c39ec7c92b4402e8eedcb27beabaab1",
      "59a8194269f943bba571f710ad856973",
      "8e41bc25b9f348b1b85b381547cd60d1",
      "34d0f1529e64437fb25e639b41096072",
      "7e832c7543df43668eb72c0aa6356349",
      "f94162c9c5fe4c2581e3b80beaeb820b",
      "0284946fccf34607b25e1d6d81d4f72a",
      "8e8d4243a3624ba58d760bdf601d5126",
      "b1f0932594fe44f4be4beb8880172871",
      "820bb554d139437b8f6a6997f9002261",
      "459e375d1fe4497cb4197ef88e303c6d",
      "5726814966da46bbba942a114de4b319",
      "448beea5445f4642809801da46aa7e7c",
      "3746c0f0daf04f51805b4c95db20063a",
      "2ff74d1c705348438a32da37910c3300",
      "52f03147b5ab4cb390f9802b7a17a508",
      "5428511d652a47348203e8f3cd2e076d",
      "7b58569bee2e44c5a61a81f9a32c4c01",
      "c5e05d74b9404333a8155304e3497ba2",
      "3ad4e058bde44d87abab780368d49118",
      "8e437049b5c74ff1beddad93e30b19aa",
      "0c61377f4cc34b27ae34bdd6d976d868",
      "9ce35f9a9a23436ea77604513b5c4659",
      "dafbc43ad79c4feda093aefe490fc4b3",
      "497160b6bf80469a89680a03a468c8bb",
      "3265725ecb914981b6f07f68b8a582df",
      "1f0ac20d99d04ffb9ee75613526d3de6",
      "dee0e5716560483b8e2b107566439d94",
      "ca99841a1abd4fbd94927c49aaa711cb",
      "163a44ff1be04583ae7bd62726c55204",
      "f3b727045a0b4b4fabd147ec9c843e14",
      "7c4b4fe859cf48ab99b7e2780a24a539",
      "d9700cffab66480cb81e3722faab4e8b",
      "2c49acc1e0a64aa4b8ec461b0596848e",
      "691817157f43412f9c36002167606b75",
      "d0ae688d71734b56bb348de54e841798",
      "07a364ebcdaa491dbcc2304e4a8dae96",
      "538b277307b1449eaa79a2a91b4741d3",
      "f7b7dd4cbaa5468687b988c5e416a154",
      "877be862a4f045ec8b221097626becb5",
      "1bf37012df0844338ac689d4f6cc8eab",
      "33fe06641be544d3aab2c25cb979416f",
      "68de0d9261fa4e6999967b6f196d149c",
      "7017c5ba1ffc467495aca2689e516f7d",
      "e0997460cc21428695a0c278c6691baf",
      "522443bef4904996b6361ea265c4c56f",
      "dfe044b67917466f971ba84a11956d5e",
      "70b2d0ba1c1d4102841fb963480de1c7",
      "7b158daacd754a40bc0a48fc9ea6fea2",
      "8904f61ff30741e488eacf0e49440f66",
      "685dee6f1e1f4c2983d2d0b8eb90ed00",
      "005155a4d62d48cbbaf3d69614b08d81",
      "d2779dfa620041e4b32a7731926a5982",
      "53048bf27cb24080b4cdf20379302223",
      "0b5bff7794a14a20964da70b87158c2c",
      "ec664e55ea16435590393ae1617cccea",
      "f10d62ef464a45a7935610d869c282e3",
      "5e7c09bc35db4a3a8b9847c6d8478883",
      "987ea331b617446facc81b08bab1f8c4",
      "446c5a687b0d43358ab036d43ce281c3",
      "d89e216f62ae4a9999da1a91dd5e48b9",
      "09d03e7e66f54cafacefbfeda5fd6e5c",
      "d58c680562a44bbd9252e7ec6eebc06c",
      "71ea2594fdd942acac1cc61055438544",
      "fecda4df301b437f9fbfb61a7b593dc0",
      "bcf6741d926248cb92acf627d6557d2b",
      "5bd9868b9bda4f68b500a43916e2d148",
      "4bc5e02035de44ad827c4f881fd75afc",
      "7d80033a955e4feb947ca9cea931d90b",
      "e975c796d80d498992a50f87a2b0105c",
      "028b2ff105904a19a9a3e8ebf66967bb",
      "10dbf2cbdb1943eda6a0acaf5da3fe85",
      "14d59c7acc5e4e07a289db44406a4851",
      "d9447df93d714907bc9c768eb4910352",
      "d4a9ace081bf488680a7e43e0eaed71d",
      "022a60cde19546bfa6ecfb4947615ad9",
      "10c01207d9ba4fe197dc80717fceecc0",
      "f4995936570d4ba1a65e9a66e410b505",
      "b4bf66303e61491385cd616076a11c04",
      "a24ab2bfe7a44263bf60936d6e53ec32",
      "59e2e8ecc9f84d57af7bf44f5bba5df7",
      "be2eb92aab8045a99fd19195c3546b60",
      "e6715db2b048434796e910952cb5eb62",
      "8de2335007ba4a6b9935fa0c514c9f6a",
      "493a4bbe8057437a8c8acf4d03f8dcb0",
      "20b03955acc74e238695717124705520",
      "a89f9507496b4cf8bdd63bdb9bff79c4",
      "9752a1f6c8b6488d9628b8e0c023f10f",
      "61a0bebc687341a7bb3352febc8d4986",
      "07e0ff9c82a444788c60edee70688ca7",
      "4427901f35e54cc48e08527bb402c22a",
      "bd9dafb825954091ba975d71ab8c89ea",
      "7ec0f07bfe51484da38c18dc66d77804",
      "03943afebc4d4b668b54c7a0fd014a37",
      "a7a6b38a76424d738146b6a5786cf2c8",
      "1d1aa4729784470584c0001c1a296008",
      "e0c68ac9b5b442eeb77e6d9ea3c1da42",
      "c858ef9fd9324ab0ab885d85e3c8582f",
      "627bc2899c994a0c96c5d43e89dc77ff",
      "0e2c2f21fa8041038151170992ecb848",
      "2d9bcbf7629a4671bf8755e790276eb7",
      "77a4edf6ab1848a6a6df20c1c07358c6",
      "acf2212d6bb44c729b63281a98377095",
      "0c7a72818da3454190bab7f30a29e249",
      "d865af915b8b4aaa8cec8d40f1f8351d"
     ]
    },
    "id": "s1UnzfID4Ai6",
    "outputId": "3e204e76-68fd-4111-f527-91f6e3b7a3ca"
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b94642c22fd840efb5825e8694bcbd16",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "for m in methods:\n",
    "    metrics[m] = test(model=test_generator, \n",
    "                      criteria_list=[], \n",
    "                      dataset=data_list, \n",
    "                      idx_range=range(10000, 11000), \n",
    "                      callback=None,\n",
    "                      gt_pos=load_pos(m))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {
    "id": "dQ7hdCtk5u6K"
   },
   "outputs": [],
   "source": [
    "mean_metrics = {key : list(map(lambda m: metrics[key][m].mean().item(), list(metrics[key].keys())[:-1])) for key in metrics}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 313
    },
    "id": "OGyaCmLXTuvA",
    "outputId": "6c42d9f3-9869-4504-bc9c-d0eedaecfd90"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<style type=\"text/css\">\n",
       "</style>\n",
       "<table id=\"T_b5d9b_\">\n",
       "  <thead>\n",
       "    <tr>\n",
       "      <th class=\"blank level0\" >&nbsp;</th>\n",
       "      <th class=\"col_heading level0 col0\" >stress</th>\n",
       "      <th class=\"col_heading level0 col1\" >stress_spc</th>\n",
       "      <th class=\"col_heading level0 col2\" >xing</th>\n",
       "      <th class=\"col_heading level0 col3\" >xing_spc</th>\n",
       "      <th class=\"col_heading level0 col4\" >l1_angle</th>\n",
       "      <th class=\"col_heading level0 col5\" >l1_angle_spc</th>\n",
       "      <th class=\"col_heading level0 col6\" >edge</th>\n",
       "      <th class=\"col_heading level0 col7\" >edge_spc</th>\n",
       "      <th class=\"col_heading level0 col8\" >ring</th>\n",
       "      <th class=\"col_heading level0 col9\" >ring_spc</th>\n",
       "      <th class=\"col_heading level0 col10\" >tsne</th>\n",
       "      <th class=\"col_heading level0 col11\" >tsne_spc</th>\n",
       "      <th class=\"col_heading level0 col12\" >reso_score</th>\n",
       "      <th class=\"col_heading level0 col13\" >min_angle</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th id=\"T_b5d9b_level0_row0\" class=\"row_heading level0 row0\" >fa2</th>\n",
       "      <td id=\"T_b5d9b_row0_col0\" class=\"data row0 col0\" >329.55</td>\n",
       "      <td id=\"T_b5d9b_row0_col1\" class=\"data row0 col1\" >-10.89%</td>\n",
       "      <td id=\"T_b5d9b_row0_col2\" class=\"data row0 col2\" >28.58</td>\n",
       "      <td id=\"T_b5d9b_row0_col3\" class=\"data row0 col3\" >8.25%</td>\n",
       "      <td id=\"T_b5d9b_row0_col4\" class=\"data row0 col4\" >81.36</td>\n",
       "      <td id=\"T_b5d9b_row0_col5\" class=\"data row0 col5\" >8.89%</td>\n",
       "      <td id=\"T_b5d9b_row0_col6\" class=\"data row0 col6\" >0.16</td>\n",
       "      <td id=\"T_b5d9b_row0_col7\" class=\"data row0 col7\" >-17.31%</td>\n",
       "      <td id=\"T_b5d9b_row0_col8\" class=\"data row0 col8\" >257.52</td>\n",
       "      <td id=\"T_b5d9b_row0_col9\" class=\"data row0 col9\" >-1.15%</td>\n",
       "      <td id=\"T_b5d9b_row0_col10\" class=\"data row0 col10\" >0.21</td>\n",
       "      <td id=\"T_b5d9b_row0_col11\" class=\"data row0 col11\" >5.86%</td>\n",
       "      <td id=\"T_b5d9b_row0_col12\" class=\"data row0 col12\" >0.60</td>\n",
       "      <td id=\"T_b5d9b_row0_col13\" class=\"data row0 col13\" >5.81</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n"
      ],
      "text/plain": [
       "<pandas.io.formats.style.Styler at 0x2b98e58c8a60>"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "columns = [\n",
    "    'stress',\n",
    "    'stress_spc',\n",
    "    'xing',\n",
    "    'xing_spc',\n",
    "    'l1_angle',\n",
    "    'l1_angle_spc',\n",
    "    'edge',\n",
    "    'edge_spc',\n",
    "    'ring',\n",
    "    'ring_spc',\n",
    "    'tsne',\n",
    "    'tsne_spc',\n",
    "    'reso_score',\n",
    "    'min_angle'\n",
    "]\n",
    "df = pd.DataFrame(mean_metrics).set_axis(columns).T\n",
    "df.style.format({c: \"{:.2f}\" for c in columns if 'spc' not in c} | {c: \"{:.2%}\" for c in columns if 'spc' in c})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "L5S1HWnHOWIl"
   },
   "source": [
    "# Large Graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "_pefPcV1OZdD"
   },
   "outputs": [],
   "source": [
    "scalability = pd.read_csv(f\"/__artifacts__/data/scalability.csv\", index_col=\"index\")\n",
    "scalability"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "oMhWAsdvPnTX"
   },
   "outputs": [],
   "source": [
    "rescale = CanonicalizationByStress()\n",
    "stressfn = Stress()\n",
    "rotate = RotateByPrincipalComponents()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "LR83OcAaPID5"
   },
   "outputs": [],
   "source": [
    "stress_list = []\n",
    "spc_list = []\n",
    "pmds_list = np.load(\"layouts/new_large_graph/pmds.npy\", allow_pickle=True)\n",
    "gviz_list = np.load(\"layouts/new_large_graph/gviz.npy\", allow_pickle=True)\n",
    "with torch.no_grad():\n",
    "    for idx, col in tqdm(scalability.iterrows(), total=len(scalability)):\n",
    "        # if idx not in [406, 516]: continue\n",
    "        torch.cuda.empty_cache()\n",
    "        G = load_mtx(col['file'])\n",
    "        G.remove_edges_from(nx.selfloop_edges(G))\n",
    "        data = generate_data_list(G, \n",
    "                                sparse=data_config.sparse, \n",
    "                                pivot_mode=data_config.pivot,\n",
    "                                init_mode=data_config.init,\n",
    "                                edge_index=data_config.edge.index,\n",
    "                                edge_attr=data_config.edge.attr,\n",
    "                                pmds_list=pmds_list[idx],\n",
    "                                gviz_list=gviz_list[idx],\n",
    "                                device=device)\n",
    "        batch = Batch.from_data_list([data]).to(device)\n",
    "        # generator.train()\n",
    "        # generator(batch)\n",
    "        generator.eval()\n",
    "        pred = generator(batch)\n",
    "        pos = rotate(rescale(pred, batch), batch)\n",
    "        gt = rotate(rescale(batch.gt_pos, batch), batch)\n",
    "        stress = stressfn(pos, batch).item()\n",
    "        gt_stress = stressfn(gt, batch).item()\n",
    "        spc = (stress - gt_stress) / np.maximum(stress, gt_stress)\n",
    "        stress_list.append(stress)\n",
    "        spc_list.append(spc)\n",
    "\n",
    "        np.save(f\"/__artifacts__/gan_result/data/scalability_{idx}.npy\", pos.cpu().numpy())\n",
    "        graph_attr = dict(node_size=1, \n",
    "                        with_labels=False, \n",
    "                        labels=dict(zip(list(G.nodes), map(lambda n: n if type(n) is int else n[1:], list(G.nodes)))),\n",
    "                        font_color=\"white\", \n",
    "                        font_weight=\"bold\",\n",
    "                        font_size=12,\n",
    "                        width=0.1)\n",
    "\n",
    "        # gt_pos = pickle.load(open(f\"/__artifacts__/data/scalability_{idx}_gt.pkl\", \"rb\"))\n",
    "\n",
    "        plt.figure(figsize=[12, 9])\n",
    "        nx.draw(G, pos=gt.cpu().numpy(), node_color='orange', **graph_attr)\n",
    "        plt.title(f\"neato: large_{idx}\")\n",
    "        plt.axis(\"equal\")\n",
    "        plt.savefig(f\"/__artifacts__/gan_result/output/{idx}_{col['name']}_{col['n']}_{spc}_nx.png\", dpi=300)\n",
    "        plt.show()\n",
    "\n",
    "        plt.figure(figsize=[12, 9])\n",
    "        graph_vis(G, pos.cpu().numpy(), **graph_attr)\n",
    "        plt.title(f\"dgd: large_{idx}, spc={spc:.2%}\")\n",
    "        plt.axis(\"equal\")\n",
    "        plt.savefig(f\"/__artifacts__/gan_result/output/{idx}_{col['name']}_{col['n']}_{spc}_dgd.png\", dpi=300)\n",
    "        plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "gRXTewkYeuq1"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [
    "w_p3OD0woiSM",
    "DR6-vYtr_i_P",
    "innqkwvH_ydD",
    "f-ODE2k8BFV6",
    "N39dDHraedM6"
   ],
   "machine_shape": "hm",
   "name": "GAN(gan=rgan,data=best(xing,stress),dis=deep,share=16,embed=0,gp=0).ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "DeepGD",
   "language": "python",
   "name": "deepgd"
  },
  "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.9.7"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}