GraphIsomorphism.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#for some stupid reason, plt crashes unless I do this\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "\n",
    "plt.imshow(np.zeros((2,3)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "import math\n",
    "import random\n",
    "import matplotlib.pyplot as plt\n",
    "from collections import namedtuple, deque\n",
    "from itertools import count\n",
    "from PIL import Image\n",
    "import pickle\n",
    "import gzip\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torch.nn.functional as F\n",
    "import torchvision.transforms as T\n",
    " \n",
    "\n",
    "from tqdm import tqdm\n",
    "\n",
    "from torch_geometric.data import Data, Batch, DataLoader\n",
    "from torch_geometric.utils import*\n",
    "from torch_geometric.nn import GCNConv, CGConv, GINConv, Sequential, global_mean_pool, global_add_pool\n",
    "from scipy.sparse.csgraph import csgraph_from_dense, csgraph_to_dense\n",
    "from scipy.sparse import*\n",
    "\n",
    "import json\n",
    "import pdb\n",
    "import scipy.sparse\n",
    "import seaborn as sns\n",
    "\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "import copy\n",
    "import numpy as np\n",
    "import networkx as nx\n",
    "import time\n",
    "import os\n",
    "\n",
    "import AllClassicalAlgos\n",
    "from AllClassicalAlgos import *\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "warnings.filterwarnings(\"ignore\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.__version__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.cuda.is_available()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.cuda.device_count()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(torch.version.cuda)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Helper Funcs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch_geometric.utils import to_scipy_sparse_matrix, from_scipy_sparse_matrix\n",
    "from scipy.sparse import csr_matrix\n",
    "\n",
    "def pyg_to_csr(pyg_data):\n",
    "    edge_index = pyg_data.edge_index.numpy()\n",
    "    num_nodes = pyg_data.num_nodes\n",
    "\n",
    "    # Combine forward and reverse edges\n",
    "    edges = np.hstack([edge_index, np.flip(edge_index, axis=0)])\n",
    "\n",
    "    # Create a binary adjacency matrix in COO format, then convert to CSR\n",
    "    adj_coo = coo_matrix((np.ones(edges.shape[1]), edges), shape=(num_nodes, num_nodes))\n",
    "    adj_csr = adj_coo.tocsr()\n",
    "\n",
    "    return adj_csr\n",
    "\n",
    "def csr_to_pyg(csr_matrix, features, label):\n",
    "    # Converts a Scipy CSR matrix to a PyG Data object.\n",
    "    edge_index, _ = from_scipy_sparse_matrix(csr_matrix)\n",
    "    return Data(x=features, edge_index=edge_index, y=label)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def remove_redundant_edges_csr(adj_csr):\n",
    "    # Only keep the upper triangular part to remove redundant edges\n",
    "    adj_single = triu(adj_csr, format='csr')\n",
    "    return adj_single"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#   new version to handle edge cases of isolated nodes\n",
    "'''does not work'''\n",
    "def data_to_networkx(data):\n",
    "    '''this might not work'''\n",
    "    G = nx.Graph()\n",
    "    num_nodes = data.x.shape[0]\n",
    "    G.add_nodes_from(range(num_nodes))  # This line adds all nodes to the graph\n",
    "    edge_index = data.edge_index.numpy()\n",
    "    G.add_edges_from(edge_index.T)\n",
    "    return G"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load GNN-RE Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "re_adj_full = np.load(os.path.join('GNN-RE-DS','adj_full.npz')) \n",
    "re_adj_train = np.load(os.path.join('GNN-RE-DS','adj_train.npz')) \n",
    "re_cell = np.load(os.path.join('GNN-RE-DS','cell.npy')) \n",
    "re_feats = np.load(os.path.join('GNN-RE-DS','feats.npy')) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def gnn_re_load_data(prefix='GNN-RE-DS', normalize=True):\n",
    "    \"\"\"\n",
    "    FROM GNN-RE REPO:\n",
    "    Load the various data files residing in the `prefix` directory.\n",
    "    Files to be loaded:\n",
    "        adj_full.npz        sparse matrix in CSR format, stored as scipy.sparse.csr_matrix\n",
    "                            The shape is N by N. Non-zeros in the matrix correspond to all\n",
    "                            the edges in the full graph. It doesn't matter if the two nodes\n",
    "                            connected by an edge are training, validation or test nodes.\n",
    "                            For unweighted graph, the non-zeros are all 1.\n",
    "        adj_train.npz       sparse matrix in CSR format, stored as a scipy.sparse.csr_matrix\n",
    "                            The shape is also N by N. However, non-zeros in the matrix only\n",
    "                            correspond to edges connecting two training nodes. The graph\n",
    "                            sampler only picks nodes/edges from this adj_train, not adj_full.\n",
    "                            Therefore, neither the attribute information nor the structural\n",
    "                            information are revealed during training. Also, note that only\n",
    "                            a x N rows and cols of adj_train contains non-zeros. For\n",
    "                            unweighted graph, the non-zeros are all 1.\n",
    "        role.json           a dict of three keys. Key 'tr' corresponds to the list of all\n",
    "                              'tr':     list of all training node indices\n",
    "                              'va':     list of all validation node indices\n",
    "                              'te':     list of all test node indices\n",
    "                            Note that in the raw data, nodes may have string-type ID. You\n",
    "                            need to re-assign numerical ID (0 to N-1) to the nodes, so that\n",
    "                            you can index into the matrices of adj, features and class labels.\n",
    "        class_map.json      a dict of length N. Each key is a node index, and each value is\n",
    "                            either a length C binary list (for multi-class classification)\n",
    "                            or an integer scalar (0 to C-1, for single-class classification).\n",
    "        feats.npz           a numpy array of shape N by F. Row i corresponds to the attribute\n",
    "                            vector of node i.\n",
    "\n",
    "    Inputs:\n",
    "        prefix              string, directory containing the above graph related files\n",
    "        normalize           bool, whether or not to normalize the node features\n",
    "\n",
    "    Outputs:\n",
    "        adj_full            scipy sparse CSR (shape N x N, |E| non-zeros), the adj matrix of\n",
    "                            the full graph, with N being total num of train + val + test nodes.\n",
    "        adj_train           scipy sparse CSR (shape N x N, |E'| non-zeros), the adj matrix of\n",
    "                            the training graph. While the shape is the same as adj_full, the\n",
    "                            rows/cols corresponding to val/test nodes in adj_train are all-zero.\n",
    "        feats               np array (shape N x f), the node feature matrix, with f being the\n",
    "                            length of each node feature vector.\n",
    "        class_map           dict, where key is the node ID and value is the classes this node\n",
    "                            belongs to.\n",
    "        role                dict, where keys are: 'tr' for train, 'va' for validation and 'te'\n",
    "                            for test nodes. The value is the list of IDs of nodes belonging to\n",
    "                            the train/val/test sets.\n",
    "    \"\"\"\n",
    "    adj_full = scipy.sparse.load_npz('./{}/adj_full.npz'.format(prefix)).astype(np.bool)\n",
    "    adj_train = scipy.sparse.load_npz('./{}/adj_train.npz'.format(prefix)).astype(np.bool)\n",
    "    role = json.load(open('./{}/role.json'.format(prefix)))\n",
    "    feats = np.load('./{}/feats.npy'.format(prefix))\n",
    "    class_map = json.load(open('./{}/class_map.json'.format(prefix)))\n",
    "    class_map = {int(k):v for k,v in class_map.items()}\n",
    "    assert len(class_map) == feats.shape[0]\n",
    "    # ---- normalize feats ----\n",
    "    train_nodes = np.array(list(set(adj_train.nonzero()[0])))\n",
    "    train_feats = feats[train_nodes]\n",
    "    scaler = StandardScaler()\n",
    "    scaler.fit(train_feats)\n",
    "    feats = scaler.transform(feats)\n",
    "    # -------------------------\n",
    "    return adj_full, adj_train, feats, class_map, role"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "re_adj_full, re_adj_train, re_feats, re_class_map, re_role = gnn_re_load_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "re_role.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "re_adj_full"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "re_adj_train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "re_feats.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len(re_class_map.keys())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_nodes = re_feats.shape[0]\n",
    "\n",
    "train_mask = torch.zeros(num_nodes, dtype=torch.bool)\n",
    "train_mask[torch.tensor(re_role['tr'])] = True\n",
    "\n",
    "val_mask = torch.zeros(num_nodes, dtype=torch.bool)\n",
    "val_mask[torch.tensor(re_role['va'])] = True\n",
    "\n",
    "test_mask = torch.zeros(num_nodes, dtype=torch.bool)\n",
    "test_mask[torch.tensor(re_role['te'])] = True\n",
    "\n",
    "# train_mask = re_role['tr']\n",
    "# val_mask = re_role['va']\n",
    "# test_mask = re_role['te']\n",
    "row, col = re_adj_full.nonzero()\n",
    "edge_index = torch.tensor([row, col], dtype=torch.long)\n",
    "# Get the keys in sorted order\n",
    "keys = sorted(re_class_map.keys())\n",
    "\n",
    "# Create a numpy array with the class values in the same order\n",
    "re_class_map_np = np.array([re_class_map[key] for key in keys])\n",
    "\n",
    "# Now you can create your PyG Data object\n",
    "gnn_re = Data(edge_index=edge_index, x=torch.from_numpy(re_feats), y=torch.from_numpy(re_class_map_np).long())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "type(gnn_re)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Remove self-loops from GNN-RE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "is_undirected(gnn_re.edge_index)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "contains_self_loops(gnn_re.edge_index)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "newE, newA = remove_self_loops(gnn_re.edge_index, gnn_re.edge_attr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gnn_re.edge_index = newE\n",
    "gnn_re.edge_attr = newA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "contains_self_loops(gnn_re.edge_index)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(gnn_re) #the graph is not connected. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "numREFeatures = 34"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def check_symmetry(graphs):\n",
    "    \"\"\"\n",
    "    Checks whether each graph in a list of graphs (or a single graph) is symmetric.\n",
    "    If a graph is not symmetric, it prints the index of that graph.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    graphs : torch_geometric.data.Data or list of torch_geometric.data.Data\n",
    "        The graph or list of graphs to check.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    None\n",
    "    \"\"\"\n",
    "    def is_symmetric(matrix):\n",
    "        \"\"\"Checks whether a matrix is symmetric.\"\"\"\n",
    "        return torch.all(matrix == matrix.T)\n",
    "\n",
    "    # Handle the case where graphs is not a list\n",
    "    if not isinstance(graphs, list):\n",
    "        graphs = [graphs]\n",
    "\n",
    "    sym = True\n",
    "    for idx, graph in enumerate(graphs):\n",
    "        # Convert the edge_index tensor to an adjacency matrix\n",
    "        adj_matrix = torch.zeros((graph.num_nodes, graph.num_nodes))\n",
    "        edge_index = graph.edge_index.t().tolist()\n",
    "        for edge in edge_index:\n",
    "            adj_matrix[edge[0], edge[1]] = 1\n",
    "            adj_matrix[edge[1], edge[0]] = 1\n",
    "        \n",
    "        if not is_symmetric(adj_matrix):\n",
    "            sym = False\n",
    "            print(f\"Graph at index {idx} is not symmetric.\")\n",
    "    return sym"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# check_symmetry(gnn_re)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Get individual Graphs from GNN-RE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def count_individual_graphs(graph):\n",
    "    nx_graph = nx.Graph()\n",
    "    nx_graph.add_edges_from(graph.edge_index.T.tolist())\n",
    "\n",
    "    # Compute the connected components\n",
    "    components = nx.connected_components(nx_graph)\n",
    "\n",
    "    # Count the number of individual graphs\n",
    "    num_graphs = len(list(components))\n",
    "\n",
    "    return num_graphs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "count_individual_graphs(gnn_re) #there are 43 disconnected small graphs which make up this large graph "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_individual_graphs(graph):\n",
    "    # Create a NetworkX graph from the input graph\n",
    "    nx_graph = nx.Graph()\n",
    "    nx_graph.add_edges_from(graph.edge_index.T.tolist())\n",
    "\n",
    "    # Compute the connected components\n",
    "    components = list(nx.connected_components(nx_graph))\n",
    "\n",
    "    # Create PyTorch Geometric graphs for each component\n",
    "    individual_graphs = []\n",
    "    for component in tqdm(components):\n",
    "        # Create a mapping from the original node index to the new node index in the subgraph\n",
    "        index_mapping = {original_index: new_index for new_index, original_index in enumerate(component)}\n",
    "        \n",
    "        # Convert the component to a list of nodes\n",
    "        component_nodes = list(component)\n",
    "        \n",
    "        # Extract the subgraph for the current component\n",
    "        subgraph = nx_graph.subgraph(component)\n",
    "        \n",
    "        # Remap the edge indices to match the new node indices in the subgraph\n",
    "        remapped_edges = [[index_mapping[edge[0]], index_mapping[edge[1]]] for edge in subgraph.edges()]\n",
    "        component_edges = torch.tensor(remapped_edges).T\n",
    "\n",
    "        # Extract node features for the current component\n",
    "        component_feats = graph.x[torch.tensor(component_nodes)]\n",
    "        \n",
    "        # Create a PyTorch Geometric Data object for the subgraph\n",
    "        individual_graphs.append(Data(x=component_feats, edge_index=component_edges))\n",
    "\n",
    "    return individual_graphs\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gnn_re_43 = get_individual_graphs(gnn_re)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gnn_re_43 = sorted(gnn_re_43, key=lambda graph: graph.num_nodes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gnn_re_43"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "check_symmetry(gnn_re_43)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_graph_diameter(graph):\n",
    "    '''computes the diameter of a PYG graph. this function runs in polynomial time but can still take a long time to run'''\n",
    "    # Create a NetworkX graph from the input graph\n",
    "\n",
    "    nx_graph = nx.Graph()\n",
    "    nx_graph.add_edges_from(graph.edge_index.T.tolist())\n",
    "\n",
    "    # Compute the diameter using NetworkX\n",
    "    diameter = nx.diameter(nx_graph)\n",
    "\n",
    "    return diameter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Takes 14 min to run\n",
    "# diameters = []\n",
    "# j = 0 \n",
    "# for subG in tqdm(gnn_re_43):\n",
    "#     diameters.append([j, subG, compute_graph_diameter(subG)])\n",
    "#     j += 1\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#each graph and it's diameter\n",
    "\n",
    "\n",
    "diameters = [[0, Data(x=[3, 34], edge_index=[2, 2]), 2],\n",
    " [1, Data(x=[3, 34], edge_index=[2, 2]), 2],\n",
    " [2, Data(x=[6, 34], edge_index=[2, 5]), 5],\n",
    " [3, Data(x=[6, 34], edge_index=[2, 5]), 5],\n",
    " [4, Data(x=[10, 34], edge_index=[2, 9]), 9],\n",
    " [5, Data(x=[12, 34], edge_index=[2, 19]), 4],\n",
    " [6, Data(x=[13, 34], edge_index=[2, 13]), 6],\n",
    " [7, Data(x=[16, 34], edge_index=[2, 20]), 5],\n",
    " [8, Data(x=[35, 34], edge_index=[2, 53]), 7],\n",
    " [9, Data(x=[44, 34], edge_index=[2, 71]), 8],\n",
    " [10, Data(x=[61, 34], edge_index=[2, 109]), 7],\n",
    " [11, Data(x=[71, 34], edge_index=[2, 90]), 20],\n",
    " [12, Data(x=[79, 34], edge_index=[2, 133]), 10],\n",
    " [13, Data(x=[87, 34], edge_index=[2, 150]), 10],\n",
    " [14, Data(x=[119, 34], edge_index=[2, 208]), 11],\n",
    " [15, Data(x=[131, 34], edge_index=[2, 231]), 11],\n",
    " [16, Data(x=[149, 34], edge_index=[2, 242]), 10],\n",
    " [17, Data(x=[207, 34], edge_index=[2, 416]), 11],\n",
    " [18, Data(x=[222, 34], edge_index=[2, 317]), 22],\n",
    " [19, Data(x=[239, 34], edge_index=[2, 461]), 14],\n",
    " [20, Data(x=[255, 34], edge_index=[2, 494]), 15],\n",
    " [21, Data(x=[309, 34], edge_index=[2, 591]), 12],\n",
    " [22, Data(x=[342, 34], edge_index=[2, 649]), 13],\n",
    " [23, Data(x=[484, 34], edge_index=[2, 718]), 29],\n",
    " [24, Data(x=[576, 34], edge_index=[2, 1065]), 13],\n",
    " [25, Data(x=[735, 34], edge_index=[2, 1502]), 14],\n",
    " [26, Data(x=[863, 34], edge_index=[2, 1677]), 20],\n",
    " [27, Data(x=[926, 34], edge_index=[2, 1797]), 18],\n",
    " [28, Data(x=[1111, 34], edge_index=[2, 2138]), 16],\n",
    " [29, Data(x=[1161, 34], edge_index=[2, 2216]), 16],\n",
    " [30, Data(x=[1637, 34], edge_index=[2, 3267]), 17],\n",
    " [31, Data(x=[3215, 34], edge_index=[2, 6624]), 18],\n",
    " [32, Data(x=[3977, 34], edge_index=[2, 7808]), 20],\n",
    " [33, Data(x=[4278, 34], edge_index=[2, 8310]), 19],\n",
    " [34, Data(x=[4618, 34], edge_index=[2, 9031]), 19],\n",
    " [35, Data(x=[4953, 34], edge_index=[2, 9545]), 19],\n",
    " [36, Data(x=[5037, 34], edge_index=[2, 10750]), 19],\n",
    " [37, Data(x=[9268, 34], edge_index=[2, 20010]), 20],\n",
    " [38, Data(x=[10336, 34], edge_index=[2, 21728]), 21],\n",
    " [39, Data(x=[10408, 34], edge_index=[2, 21818]), 24],\n",
    " [40, Data(x=[12014, 34], edge_index=[2, 24871]), 19],\n",
    " [41, Data(x=[12520, 34], edge_index=[2, 25639]), 20],\n",
    " [42, Data(x=[14191, 34], edge_index=[2, 32765]), 22]]\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "diameters = np.array([np.array(lst) for lst in diameters])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "diameters.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gSize = []\n",
    "for i in range(43):\n",
    "   gSize.append(gnn_re_43[i].num_nodes)\n",
    "plt.stem(gSize);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.median(gSize)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sns.histplot(gSize, bins=20) #[20,50,100,200,500,800,1000,5000,10000,20000]);\n",
    "plt.xlabel('Number of Nodes');\n",
    "plt.ylabel('Number of Graphs');\n",
    "plt.title('Histogram of GNN-RE Dataset Graph Sizes');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.median(gSize)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "diameters[:,2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.average(diameters[:,2].astype(int))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sns.histplot(diameters[:,2].astype(int), bins = 20);\n",
    "plt.xlabel('Diameter');\n",
    "plt.ylabel('Number of Graphs');\n",
    "plt.title('Histogram of GNN-RE Dataset Diameters');"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generate Permutations of GNN_RE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install deeprobust"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from deeprobust.graph.data import Dataset\n",
    "from deeprobust.graph.global_attack import Random\n",
    "\n",
    "from deeprobust.graph.data import Dataset, Dpr2Pyg, Pyg2Dpr\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_isomorphic_permutations(G, num_samples=100): \n",
    "\n",
    "    def create_isomorph(G):\n",
    "        num_nodes = G.num_nodes\n",
    "\n",
    "        # Generate a permutation of the nodes\n",
    "        permutation = torch.randperm(num_nodes)\n",
    "\n",
    "        # Apply the permutation to the node features\n",
    "        G.x = G.x[permutation]\n",
    "\n",
    "        # Create a mapping tensor\n",
    "        mapping_tensor = torch.tensor(permutation)\n",
    "\n",
    "        # Apply the permutation to the edges\n",
    "        G.edge_index = mapping_tensor[G.edge_index]\n",
    "\n",
    "        return G\n",
    "    dataset = []\n",
    "\n",
    "    for _ in range(num_samples):\n",
    "        # Create a deep copy of the graph\n",
    "        G_copy = copy.deepcopy(G)\n",
    "\n",
    "        # Create an isomorphic graph and add it to the dataset\n",
    "        G_isomorph = create_isomorph(G_copy)\n",
    "        dataset.append(G_isomorph)\n",
    "\n",
    "    return dataset\n",
    "\n",
    "def generate_nonisomorphic_permutations(graph,num_samples):\n",
    "\n",
    "    model = Random()\n",
    "    non_isomorphic_graphs = []\n",
    "    while len(non_isomorphic_graphs) < num_samples:\n",
    "        # Create a deep copy of the graph\n",
    "        graph_copy = copy.deepcopy(graph)\n",
    "        num_edges = graph.num_edges\n",
    "\n",
    "        adj = pyg_to_csr(graph_copy)  # Convert to CSR matrix\n",
    "        n_perturbations = random.randint(max(1, num_edges//2), num_edges-1)\n",
    "\n",
    "        model.attack(adj, n_perturbations=n_perturbations, type='remove')  # remove random amount of edges\n",
    "        modified_adj1 = model.modified_adj  # Get the modified adjacency matrix\n",
    "\n",
    "        model.attack(modified_adj1, n_perturbations=n_perturbations, type='add')  # add (same) random amount of edges back \n",
    "        modified_adj2 = model.modified_adj  # Get the modified adjacency matrix\n",
    "\n",
    "        tri_mod = remove_redundant_edges_csr(modified_adj2)\n",
    "\n",
    "        non_isomorphic_graph = csr_to_pyg(tri_mod, graph_copy.x, graph_copy.y)  # Convert back to PyG\n",
    "\n",
    "        # Convert the graphs to networkx for WL hashing\n",
    "        graph_nx = to_networkx(graph, to_undirected=True) #although it says \"to undirected = true\", when we convert back to PYG, it becomes directed.\n",
    "        non_isomorphic_graph_nx = to_networkx(non_isomorphic_graph, to_undirected=True)  #although it says \"to undirected = true\", when we convert back to PYG, it becomes directed.\n",
    "\n",
    "        # Check if the WL hashes are different\n",
    "        if nx.weisfeiler_lehman_graph_hash(graph_nx) != nx.weisfeiler_lehman_graph_hash(non_isomorphic_graph_nx):\n",
    "            non_isomorphic_graphs.append(non_isomorphic_graph)\n",
    "\n",
    "    return non_isomorphic_graphs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_permutations(G, k):\n",
    "    '''G: A list of graphs you want to make permutations of\n",
    "        k: the number of permutations per graph\n",
    "        '''\n",
    "    data_list = []\n",
    "    for graph in tqdm(G):\n",
    "\n",
    "        if graph.num_nodes < 10 or graph.num_edges == (graph.num_nodes*(graph.num_nodes-1))/2: #we can not generate non-isomorphic graphs if less than 10 nodes or fully connected graphs\n",
    "\n",
    "            # Generate isomorphic graphs\n",
    "            isomorphic_graphs = generate_isomorphic_permutations(graph, num_samples=2*k)\n",
    "            # Pair each isomorphic graph with the original graph and assign label 1.0\n",
    "            data_list.extend([(graph, isomorphic_graph, 1.0) for isomorphic_graph in isomorphic_graphs])\n",
    "            \n",
    "        else:\n",
    "        \n",
    "            # Generate isomorphic graphs\n",
    "            isomorphic_graphs = generate_isomorphic_permutations(graph, num_samples=k)\n",
    "            # Pair each isomorphic graph with the original graph and assign label 1.0\n",
    "            data_list.extend([(graph, isomorphic_graph, 1.0) for isomorphic_graph in isomorphic_graphs])\n",
    "\n",
    "            # Generate non-isomorphic graphs\n",
    "            non_isomorphic_graphs = generate_nonisomorphic_permutations(graph, num_samples=k)\n",
    "            # Pair each non-isomorphic graph with the original graph and assign label 0.0\n",
    "            data_list.extend([(graph, non_isomorphic_graph, 0.0) for non_isomorphic_graph in non_isomorphic_graphs])\n",
    "            \n",
    "    return data_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "k = 16 #the number of permutations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#note that the '.dataset' function of PYG graphs SHOULD NOT be used as it returns the entirety of gnnre43"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#For reproducability, we hardcode the train and test set\n",
    "gnnRETrain = gnn_re_43[::2]  # Elements at even indices\n",
    "gnnRETest = gnn_re_43[1::2]  # Elements at odd indices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "largeD_IDX = [11, 23,37, 39, 41] #the largest diameter graphs in the test set\n",
    "largeDTestList = [gnn_re_43[g] for g in largeD_IDX]\n",
    "\n",
    "largestN_idx = [37,39,41]\n",
    "largestN_graphs_list = [gnn_re_43[g] for g in largestN_idx] #largest num nodes part of the test set \n",
    " \n",
    "print(len(gnnRETest), len(gnnRETrain), len(largeDTestList), len(largestN_graphs_list))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#takes ~12 min to run\n",
    "# gnnRETrain = generate_permutations(gnnRETrain, k)\n",
    "# gnnRETest = generate_permutations(gnnRETest, k)\n",
    "# largeDTest = generate_permutations(largeDTestList, k)\n",
    "# largestN_graphs = generate_permutations(largestN_graphs_list, k) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# with gzip.open('gnnRETrain.pkl.gz', 'wb') as f:\n",
    "#     pickle.dump(gnnRETrain, f)\n",
    "\n",
    "# with gzip.open('gnnRETest.pkl.gz', 'wb') as f:\n",
    "#     pickle.dump(gnnRETest, f)\n",
    "\n",
    "# with gzip.open('largeDTest.pkl.gz', 'wb') as f:\n",
    "#     pickle.dump(largeDTest, f)\n",
    "\n",
    "# with gzip.open('largestN_graphs.pkl.gz', 'wb') as f:\n",
    "#     pickle.dump(largestN_graphs, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with gzip.open('gnnRETrain.pkl.gz', 'rb') as f:\n",
    "    gnnRETrain = pickle.load(f)\n",
    "\n",
    "with gzip.open('gnnRETest.pkl.gz', 'rb') as f:\n",
    "    gnnRETest = pickle.load(f)\n",
    "\n",
    "with gzip.open('largeDTest.pkl.gz', 'rb') as f:\n",
    "    largeDTest = pickle.load(f)\n",
    "\n",
    "with gzip.open('largestN_graphs.pkl.gz', 'rb') as f:\n",
    "    largestN_graphs = pickle.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "largeDTestList"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "largestN_graphs_list"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generate Artificial Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_synthetic_CGI(num_graphs, num_iso_pairs, num_features, min_nodes=5, max_nodes=7):\n",
    "    \"\"\"\n",
    "    Generates a list of graph pairs and their corresponding labels for identifying circuit graph isomorphisms.\n",
    "    Takes 7 minutes to make 50 graphs of size 2000 nodes\n",
    "\n",
    "    Args:\n",
    "        num_graphs (int): The total number of connected graphs to generate.\n",
    "        num_iso_pairs (int): The number of isomorphic pairs to include in the data_list.\n",
    "        num_features (int): The number of features for each node in the graphs.\n",
    "        min_nodes (int): The minimum number of nodes in each graph.\n",
    "        max_nodes (int): The maximum number of nodes in each graph.\n",
    "\n",
    "    Returns:\n",
    "        data_list (list): A list of graph pairs and their corresponding labels, where each pair consists of two\n",
    "        Data objects and a label (1 for isomorphic pair, 0 for non-isomorphic pair).\n",
    "    \"\"\"\n",
    "\n",
    "    \n",
    "    def generate_isomorphic_pair(num_nodes):\n",
    "            \"\"\"\n",
    "            Generates a pair of isomorphic graphs with `num_nodes` nodes.\n",
    "\n",
    "            Args:\n",
    "                num_nodes (int): The number of nodes in each graph.\n",
    "\n",
    "            Returns:\n",
    "                A tuple containing two NetworkX graphs which are isomorphic.\n",
    "            \"\"\"\n",
    "            # Generate a random graph\n",
    "            #G1 = nx.barabasi_albert_graph(num_nodes,random.randint(2, num_nodes-1)) #this produces HARD graphs for vf2. \n",
    "            G1 = nx.gnm_random_graph(num_nodes,random.randint(num_nodes, num_nodes*(num_nodes-1)//2))\n",
    "\n",
    "            # Shuffle the nodes to generate a second graph\n",
    "            nodes = list(G1.nodes())\n",
    "            random.shuffle(nodes)\n",
    "            mapping = {n: nodes[i] for i, n in enumerate(G1.nodes())}\n",
    "            G2 = nx.relabel_nodes(G1, mapping)\n",
    "\n",
    "            #assert(nx.is_c              onnected(G1) == True and nx.is_connected(G2) == True)\n",
    "\n",
    "            return G1, G2\n",
    "\n",
    "    def generate_non_isomorphic_pair(num_nodes, num_edges):\n",
    "        while True:\n",
    "            #G1 = nx.barabasi_albert_graph(num_nodes, num_edges, seed=123) #if you want to make things really hard for vf2\n",
    "            G1 = nx.gnm_random_graph(num_nodes, num_edges, seed=random.randint(2, num_nodes-1) )\n",
    "            G2 = nx.gnm_random_graph(num_nodes, num_edges, seed=random.randint(2, num_nodes-1) )\n",
    "            if nx.weisfeiler_lehman_graph_hash(G1) != nx.weisfeiler_lehman_graph_hash(G2):\n",
    "                return G1, G2\n",
    "\n",
    "    data_list = []\n",
    "\n",
    "    for i in tqdm(range(num_iso_pairs), leave = False):\n",
    "        num_nodes = random.randint(min_nodes, max_nodes)\n",
    "        G1, G2 = generate_isomorphic_pair(num_nodes)\n",
    "        features = torch.randn((num_nodes, num_features))\n",
    "        graph1 = Data(edge_index=torch.tensor(list(G1.edges)).t().contiguous(), x=features)\n",
    "        graph2 = Data(edge_index=torch.tensor(list(G2.edges)).t().contiguous(), x=features)\n",
    "        data_list.append((graph1, graph2, 1.0))\n",
    "\n",
    "    assert(num_graphs - num_iso_pairs > 0)\n",
    "\n",
    "    for i in tqdm(range(num_graphs - num_iso_pairs), leave=False):\n",
    "        num_nodes = random.randint(min_nodes, max_nodes)\n",
    "        num_edges = random.randint(2, num_nodes-1)\n",
    "        G1, G2 = generate_non_isomorphic_pair(num_nodes, num_edges)\n",
    "        features1 = torch.randn((num_nodes, num_features))\n",
    "        features2 = torch.randn((num_nodes, num_features))\n",
    "        graph1 = Data(edge_index=torch.tensor(list(G1.edges)).t().contiguous(), x=features1)\n",
    "        graph2 = Data(edge_index=torch.tensor(list(G2.edges)).t().contiguous(), x=features2)\n",
    "        data_list.append((graph1, graph2, 0.0))\n",
    "\n",
    "    return data_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "synthetic_data = generate_synthetic_CGI(800, 400, 34, min_nodes=50, max_nodes=50)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Set Data Loaders"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Synthetic Dataloaders\n",
    "\n",
    "# Create your dataset and split it into training and test sets\n",
    "synDS = DataLoader(synthetic_data)\n",
    "train_dataset, val_dataset, test_dataset = torch.utils.data.random_split(synthetic_data, [0.8, 0.1, 0.1])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# with gzip.open('synTrain.pkl.gz', 'wb') as f:\n",
    "#     pickle.dump(train_dataset, f)\n",
    "\n",
    "# with gzip.open('synVal.pkl.gz', 'wb') as f:\n",
    "#     pickle.dump(val_dataset, f)\n",
    "\n",
    "# with gzip.open('synTest.pkl.gz', 'wb') as f:\n",
    "#     pickle.dump(test_dataset, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with gzip.open('synTrain.pkl.gz', 'rb') as f:\n",
    "    syn_train_dataset = pickle.load(f)\n",
    "\n",
    "with gzip.open('synVal.pkl.gz', 'rb') as f:\n",
    "    syn_val_dataset = pickle.load(f)\n",
    "\n",
    "with gzip.open('synTest.pkl.gz', 'rb') as f:\n",
    "    syn_test_dataset = pickle.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Create your data loaders\n",
    "train_loaderSyn = DataLoader(syn_train_dataset, batch_size=8, shuffle=True)\n",
    "val_loaderSyn = DataLoader(syn_val_dataset, batch_size=8, shuffle=False)\n",
    "test_loaderSyn = DataLoader(syn_test_dataset, batch_size=1, shuffle=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(len(train_loaderSyn.dataset), len(val_loaderSyn.dataset), len(test_loaderSyn))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "type(train_loaderSyn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# RE Dataloaders\n",
    "\n",
    "train_datasetRE =  DataLoader(gnnRETrain, batch_size = 8, shuffle = True)\n",
    "test_datasetRE = DataLoader(gnnRETest, batch_size = 1, shuffle = False)\n",
    "largeD_datasetRE = DataLoader(largeDTest, batch_size=1, shuffle = False)\n",
    "largestN_graphsDL = DataLoader(largestN_graphs, batch_size = 1, shuffle = False) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(len(train_datasetRE.dataset), len(test_datasetRE.dataset), len(largeD_datasetRE))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "type(train_datasetRE)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Full WL SGNN (AKA DEPTH GRAPHNET)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "class FullWLGraphSiameseNet(torch.nn.Module):\n",
    "    def __init__(self, num_node_features=numREFeatures, hidden_dim=36, encoding_size=32, num_layers=29): #num layers set to the max diameter\n",
    "        super(FullWLGraphSiameseNet, self).__init__()\n",
    "\n",
    "        # Define MLP for GINConv\n",
    "        nn1 = torch.nn.Sequential(torch.nn.Linear(num_node_features, hidden_dim),\n",
    "                                  torch.nn.ReLU(),\n",
    "                                  torch.nn.Linear(hidden_dim, hidden_dim))\n",
    "        \n",
    "        self.num_layers = num_layers\n",
    "\n",
    "        # Define the GINConv layers for the encoder network\n",
    "        self.convs = torch.nn.ModuleList()\n",
    "        self.convs.append(GINConv(nn1))\n",
    "        for _ in range(num_layers - 1):\n",
    "            nn2 = torch.nn.Sequential(torch.nn.Linear(hidden_dim, hidden_dim),\n",
    "                                      torch.nn.ReLU(),\n",
    "                                      torch.nn.Linear(hidden_dim, hidden_dim))\n",
    "            self.convs.append(GINConv(nn2))\n",
    "\n",
    "        self.encoder = torch.nn.Linear(hidden_dim, encoding_size)\n",
    "\n",
    "        # Define the layers for the classification network\n",
    "        self.classifier = torch.nn.Sequential(\n",
    "            torch.nn.Linear(encoding_size, encoding_size),\n",
    "            torch.nn.ReLU(),\n",
    "            torch.nn.Linear(encoding_size, encoding_size),\n",
    "            torch.nn.ReLU(),\n",
    "            torch.nn.Linear(encoding_size, 1),\n",
    "            torch.nn.Sigmoid()\n",
    "        )\n",
    "\n",
    "    def forward(self, graph1, graph2):\n",
    "        x1, edge_index1 = graph1.x, graph1.edge_index\n",
    "        x2, edge_index2 = graph2.x, graph2.edge_index\n",
    "\n",
    "        for i in range(self.num_layers):\n",
    "            x1 = self.convs[i](x1, edge_index1)\n",
    "            x2 = self.convs[i](x2, edge_index2)\n",
    "            x1 = F.relu(x1)\n",
    "            x2 = F.relu(x2)\n",
    "\n",
    "        # Encode the two graphs\n",
    "        encoding1 = self.encoder(x1)\n",
    "        encoding2 = self.encoder(x2)\n",
    "\n",
    "        # Apply add/mean pooling to obtain fixed-size representations\n",
    "        encoding1 = global_add_pool(encoding1, graph1.batch) #global_mean_pool(encoding1, graph1.batch)\n",
    "        encoding2 = global_add_pool(encoding2, graph2.batch)#global_mean_pool(encoding2, graph2.batch)\n",
    "\n",
    "        # Compute the difference between the encodings\n",
    "        diff = torch.abs(encoding1 - encoding2)\n",
    "\n",
    "        # Pass the difference through the classification network\n",
    "        out = self.classifier(diff.view(-1, diff.shape[1]))\n",
    "\n",
    "        return out.squeeze()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train 29 Layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "l = 29\n",
    "model_RE_Full = FullWLGraphSiameseNet(num_node_features=34,hidden_dim=36, encoding_size=32, num_layers=l)\n",
    "model_RE_Full = model_RE_Full.to(device)\n",
    "num_epochs = 30\n",
    "# Define the loss function and optimizer\n",
    "criterion = nn.BCELoss() #nn.MSELoss() #experiment with mse, nn.BCELoss()\n",
    "optimizer = torch.optim.Adam(model_RE_Full.parameters(), lr=1e-4)\n",
    "\n",
    "# Count the number of parameters\n",
    "total_params = sum(p.numel() for p in model_RE_Full.parameters())\n",
    "print(f\"Total number of parameters: {total_params}\")\n",
    "# num_params.append(total_params)\n",
    "\n",
    "# Train the model on your dataset. FULLWLSGNN takes a 5 min for 30 epochs\n",
    "for epoch in range(num_epochs):\n",
    "    for i, (g1, g2, label) in enumerate(train_datasetRE):\n",
    "        # print(g1, g2, label)\n",
    "        # assert(False)\n",
    "        g1 = g1.to(device)\n",
    "        g2 = g2.to(device)\n",
    "        label = label.to(device)\n",
    "        optimizer.zero_grad()\n",
    "        output = model_RE_Full(g1, g2)\n",
    "        # print(label.dtype, output.dtype)\n",
    "        loss = criterion(output, label)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        \n",
    "        if (epoch+1) % 2 == 0 and (i+1) % 35 == 0:\n",
    "            print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'.format(epoch+1, num_epochs, i+1, len(train_datasetRE), loss.item()))\n",
    "\n",
    "#num params 80513"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test the model on your test dataset\n",
    "correct = 0\n",
    "total = 0\n",
    "with torch.no_grad():\n",
    "    \n",
    "    for data1, data2, label in test_datasetRE:\n",
    "        data1 = data1.to(device)\n",
    "        data2 = data2.to(device)\n",
    "        label = label.to(device)\n",
    "        output = model_RE_Full(data1, data2)\n",
    "        # print('Num Nodes', data1.num_nodes)\n",
    "        # print(output.shape[0])\n",
    "        # print(label.shape)\n",
    "        # _, predicted = torch.max(output, 1)\n",
    "        predicted = torch.round(output)\n",
    "        total += label.size(0)\n",
    "        correct += (predicted == label).sum().item()\n",
    "accuracy = 100 * correct / total\n",
    "print(f'Test Accuracy: {accuracy:.2f}%', )#' using ', l, ' message passing layers')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Train 29 Layer GREED\n",
    "#Greed needs a slightly different dataloader than the others. (the y value must be duplicated) \n",
    "class GREEDDataLoader:\n",
    "    def __init__(self, original_dataloader):\n",
    "        self.original_dataloader = original_dataloader\n",
    "\n",
    "    def __iter__(self):\n",
    "        for g1, g2, y in self.original_dataloader:\n",
    "            yield g1, g2, y, y\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.original_dataloader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "from neuro import models, config, train\n",
    "greed_train_loader = GREEDDataLoader(train_datasetRE)\n",
    "greed_val_loader = GREEDDataLoader(test_datasetRE)\n",
    "config.device# = torch.device('cpu')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#I am using this for timing so i don't actually need to train this. epoch set to just 1\n",
    "greed = models.NeuralSiameseModel(n_layers= 20, input_dim=34, hidden_dim=36, output_dim= 32)\n",
    "ad = torch.optim.Adam(greed.parameters(), lr=2e-4)\n",
    "print(total_params := sum(p.numel() for p in greed.parameters()))\n",
    "train.train_loop(model = greed, opt=ad, \n",
    "                        loader=greed_train_loader, n_epochs=1)\n",
    "\n",
    "#num params: 120729"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "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.9.12"
  },
  "toc-showtags": false,
  "vscode": {
   "interpreter": {
    "hash": "ad2bdc8ecc057115af97d19610ffacc2b4e99fae6737bb82f5d7fb13d2f2c186"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}