IGMC.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "## IGMC CODE IMPLEMENTATION\n"
      ],
      "metadata": {
        "id": "ipZWiijdeF42"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xr5JI9jfZFhv"
      },
      "outputs": [],
      "source": [
        "!pip install torch-sparse\n",
        "!pip install torch-scatter\n",
        "!pip install torch-geometric"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import numpy as np\n",
        "import pandas as pd\n",
        "import multiprocessing as mp\n",
        "import argparse\n",
        "import scipy.io as sio\n",
        "import random\n",
        "import pdb\n",
        "import scipy.sparse as sp\n",
        "import torch\n",
        "from torch_geometric.data import Data, Dataset, InMemoryDataset\n",
        "import warnings\n",
        "from urllib.request import urlopen\n",
        "from zipfile import ZipFile\n",
        "import shutil\n",
        "import os.path\n",
        "from tqdm import tqdm\n",
        "import networkx as nx\n",
        "import pickle as pkl\n",
        "import os, sys, pdb, math, time\n",
        "from copy import deepcopy\n",
        "\n",
        "warnings.simplefilter('ignore', sp.SparseEfficiencyWarning)\n",
        "import torch.multiprocessing\n",
        "torch.multiprocessing.set_sharing_strategy('file_system')"
      ],
      "metadata": {
        "id": "DdrBQna6eeB_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def map_data(data):\n",
        "    unique_elements = list(set(data))\n",
        "    id_dictionary = {old: new for new, old in enumerate(sorted(unique_elements))}\n",
        "    data = np.array([id_dictionary[d] for d in data])\n",
        "    n = len(unique_elements)\n",
        "\n",
        "    return data, id_dictionary, n"
      ],
      "metadata": {
        "id": "h5U8RfMte_jM"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def load_data(fname, seed=1234, verbose=True):\n",
        "    u_features = None\n",
        "    v_features = None\n",
        "    sep = '\\t'\n",
        "    filename = './u.data'\n",
        "    dtypes = {\n",
        "        'u_nodes': np.int32, 'v_nodes': np.int32,\n",
        "        'ratings': np.float32, 'timestamp': np.float64}\n",
        "\n",
        "    data = pd.read_csv(\n",
        "        filename, sep=sep, header=None,\n",
        "        names=['u_nodes', 'v_nodes', 'ratings', 'timestamp'], dtype=dtypes)\n",
        "    data_array = data.values.tolist()\n",
        "    random.seed(seed)\n",
        "    random.shuffle(data_array)\n",
        "    data_array = np.array(data_array)\n",
        "\n",
        "    u_nodes_ratings = data_array[:, 0].astype(dtypes['u_nodes'])\n",
        "    v_nodes_ratings = data_array[:, 1].astype(dtypes['v_nodes'])\n",
        "    ratings = data_array[:, 2].astype(dtypes['ratings'])\n",
        "\n",
        "    u_nodes_ratings, u_dict, num_users = map_data(u_nodes_ratings)\n",
        "    v_nodes_ratings, v_dict, num_movies = map_data(v_nodes_ratings)\n",
        "\n",
        "    u_nodes_ratings, v_nodes_ratings = u_nodes_ratings.astype(np.int64), v_nodes_ratings.astype(np.int32)\n",
        "    ratings = ratings.astype(np.float64)\n",
        "\n",
        "    sep = r'|'\n",
        "    movie_file = './u.item'\n",
        "    movie_headers = ['movie id', 'movie title', 'release date', 'video release date',\n",
        "                      'IMDb URL', 'unknown', 'Action', 'Adventure', 'Animation',\n",
        "                      'Childrens', 'Comedy', 'Crime', 'Documentary', 'Drama', 'Fantasy',\n",
        "                      'Film-Noir', 'Horror', 'Musical', 'Mystery', 'Romance', 'Sci-Fi',\n",
        "                      'Thriller', 'War', 'Western']\n",
        "    movie_df = pd.read_csv(movie_file, sep=sep, header=None,\n",
        "                            names=movie_headers, engine='python')\n",
        "\n",
        "    genre_headers = movie_df.columns.values[6:]\n",
        "    num_genres = genre_headers.shape[0]\n",
        "\n",
        "    v_features = np.zeros((num_movies, num_genres), dtype=np.float32)\n",
        "    for movie_id, g_vec in zip(movie_df['movie id'].values.tolist(), movie_df[genre_headers].values.tolist()):\n",
        "        # Check if movie_id was listed in ratings file and therefore in mapping dictionary\n",
        "        if movie_id in v_dict.keys():\n",
        "            v_features[v_dict[movie_id], :] = g_vec\n",
        "\n",
        "    # User features\n",
        "\n",
        "    sep = r'|'\n",
        "    users_file = './u.user'\n",
        "    users_headers = ['user id', 'age', 'gender', 'occupation', 'zip code']\n",
        "    users_df = pd.read_csv(users_file, sep=sep, header=None,\n",
        "                            names=users_headers, engine='python')\n",
        "\n",
        "    occupation = set(users_df['occupation'].values.tolist())\n",
        "\n",
        "    gender_dict = {'M': 0., 'F': 1.}\n",
        "    occupation_dict = {f: i for i, f in enumerate(occupation, start=2)}\n",
        "\n",
        "    num_feats = 2 + len(occupation_dict)\n",
        "\n",
        "    u_features = np.zeros((num_users, num_feats), dtype=np.float32)\n",
        "    for _, row in users_df.iterrows():\n",
        "        u_id = row['user id']\n",
        "        if u_id in u_dict.keys():\n",
        "            u_features[u_dict[u_id], 0] = row['age']\n",
        "            u_features[u_dict[u_id], 1] = gender_dict[row['gender']]\n",
        "            u_features[u_dict[u_id], occupation_dict[row['occupation']]] = 1.\n",
        "\n",
        "    u_features = sp.csr_matrix(u_features)\n",
        "    v_features = sp.csr_matrix(v_features)\n",
        "\n",
        "    if verbose:\n",
        "        print('Number of users = %d' % num_users)\n",
        "        print('Number of items = %d' % num_movies)\n",
        "        print('Number of links = %d' % ratings.shape[0])\n",
        "        print('Fraction of positive links = %.4f' % (float(ratings.shape[0]) / (num_users * num_movies),))\n",
        "\n",
        "    return num_users, num_movies, u_nodes_ratings, v_nodes_ratings, ratings, u_features, v_features"
      ],
      "metadata": {
        "id": "xmPlnkbzfVGs"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class SparseRowIndex:\n",
        "    def __init__(self, csr_matrix):\n",
        "        data = []\n",
        "        indices = []\n",
        "        indptr = []\n",
        "\n",
        "        for row_start, row_end in zip(csr_matrix.indptr[:-1], csr_matrix.indptr[1:]):\n",
        "            data.append(csr_matrix.data[row_start:row_end])\n",
        "            indices.append(csr_matrix.indices[row_start:row_end])\n",
        "            indptr.append(row_end - row_start)  # nnz of the row\n",
        "\n",
        "        self.data = np.array(data, dtype=object)\n",
        "        self.indices = np.array(indices, dtype=object)\n",
        "        self.indptr = np.array(indptr, dtype=object)\n",
        "        self.shape = csr_matrix.shape\n",
        "\n",
        "    def __getitem__(self, row_selector):\n",
        "        indices = np.concatenate(self.indices[row_selector])\n",
        "        data = np.concatenate(self.data[row_selector])\n",
        "        indptr = np.append(0, np.cumsum(self.indptr[row_selector]))\n",
        "        shape = [indptr.shape[0] - 1, self.shape[1]]\n",
        "        return sp.csr_matrix((data, indices, indptr), shape=shape)"
      ],
      "metadata": {
        "id": "_KjAHwoZfoIe"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class SparseColIndex:\n",
        "    def __init__(self, csc_matrix):\n",
        "        data = []\n",
        "        indices = []\n",
        "        indptr = []\n",
        "\n",
        "        for col_start, col_end in zip(csc_matrix.indptr[:-1], csc_matrix.indptr[1:]):\n",
        "            data.append(csc_matrix.data[col_start:col_end])\n",
        "            indices.append(csc_matrix.indices[col_start:col_end])\n",
        "            indptr.append(col_end - col_start)\n",
        "\n",
        "        self.data = np.array(data, dtype=object)\n",
        "        self.indices = np.array(indices, dtype=object)\n",
        "        self.indptr = np.array(indptr, dtype=object)\n",
        "        self.shape = csc_matrix.shape\n",
        "\n",
        "    def __getitem__(self, col_selector):\n",
        "        indices = np.concatenate(self.indices[col_selector])\n",
        "        data = np.concatenate(self.data[col_selector])\n",
        "        indptr = np.append(0, np.cumsum(self.indptr[col_selector]))\n",
        "\n",
        "        shape = [self.shape[0], indptr.shape[0] - 1]\n",
        "        return sp.csc_matrix((data, indices, indptr), shape=shape)\n"
      ],
      "metadata": {
        "id": "sBShWlGcgCsl"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class GraphDataset(InMemoryDataset):\n",
        "    def __init__(self, root, A, links, labels, h, sample_ratio, max_nodes_per_hop,\n",
        "                 u_features, v_features, class_values, max_num=None, parallel=True):\n",
        "        self.Arow = SparseRowIndex(A)\n",
        "        self.Acol = SparseColIndex(A.tocsc())\n",
        "        self.links = links\n",
        "        self.labels = labels\n",
        "        self.h = h\n",
        "        self.sample_ratio = sample_ratio\n",
        "        self.max_nodes_per_hop = max_nodes_per_hop\n",
        "        self.u_features = u_features\n",
        "        self.v_features = v_features\n",
        "        self.class_values = class_values\n",
        "        self.parallel = parallel\n",
        "        self.max_num = max_num\n",
        "        if max_num is not None:\n",
        "            np.random.seed(123)\n",
        "            num_links = len(links[0])\n",
        "            perm = np.random.permutation(num_links)\n",
        "            perm = perm[:max_num]\n",
        "            self.links = (links[0][perm], links[1][perm])\n",
        "            self.labels = labels[perm]\n",
        "        super(GraphDataset, self).__init__(root)\n",
        "        self.data, self.slices = torch.load(self.processed_paths[0])\n",
        "\n",
        "    @property\n",
        "    def processed_file_names(self):\n",
        "        name = 'data.pt'\n",
        "        if self.max_num is not None:\n",
        "            name = 'data_{}.pt'.format(self.max_num)\n",
        "        return [name]\n",
        "\n",
        "    def process(self):\n",
        "        data_list = links2subgraphs(self.Arow, self.Acol, self.links, self.labels, self.h,\n",
        "                                    self.sample_ratio, self.max_nodes_per_hop,\n",
        "                                    self.u_features, self.v_features,\n",
        "                                    self.class_values, self.parallel)\n",
        "\n",
        "        data, slices = self.collate(data_list)\n",
        "        torch.save((data, slices), self.processed_paths[0])\n",
        "        del data_list"
      ],
      "metadata": {
        "id": "lzzr-Eo-gHhj"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class DynamicDataset(Dataset):\n",
        "    def __init__(self, root, A, links, labels, h, sample_ratio, max_nodes_per_hop,\n",
        "                 u_features, v_features, class_values, max_num=None):\n",
        "        super(DynamicDataset, self).__init__(root)\n",
        "        self.Arow = SparseRowIndex(A)\n",
        "        self.Acol = SparseColIndex(A.tocsc())\n",
        "        self.links = links\n",
        "        self.labels = labels\n",
        "        self.h = h\n",
        "        self.sample_ratio = sample_ratio\n",
        "        self.max_nodes_per_hop = max_nodes_per_hop\n",
        "        self.u_features = u_features\n",
        "        self.v_features = v_features\n",
        "        self.class_values = class_values\n",
        "        if max_num is not None:\n",
        "            np.random.seed(123)\n",
        "            num_links = len(links[0])\n",
        "            perm = np.random.permutation(num_links)\n",
        "            perm = perm[:max_num]\n",
        "            self.links = (links[0][perm], links[1][perm])\n",
        "            self.labels = labels[perm]\n",
        "\n",
        "    def len(self):\n",
        "        return len(self.links[0])\n",
        "\n",
        "    def get(self, idx):\n",
        "        i, j = self.links[0][idx], self.links[1][idx]\n",
        "        g_label = self.labels[idx]\n",
        "        tmp = subgraph_extraction_labeling(\n",
        "            (i, j), self.Arow, self.Acol, self.h, self.sample_ratio, self.max_nodes_per_hop,\n",
        "            self.u_features, self.v_features, self.class_values, g_label\n",
        "        )\n",
        "        return construct_pyg_graph(*tmp)\n"
      ],
      "metadata": {
        "id": "i5xZ5-ingLsn"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def links2subgraphs(Arow,\n",
        "                    Acol,\n",
        "                    links,\n",
        "                    labels,\n",
        "                    h=1,\n",
        "                    sample_ratio=1.0,\n",
        "                    max_nodes_per_hop=None,\n",
        "                    u_features=None,\n",
        "                    v_features=None,\n",
        "                    class_values=None,\n",
        "                    parallel=True):\n",
        "    # extract enclosing subgraphs\n",
        "    print('Enclosing subgraph extraction begins...')\n",
        "    g_list = []\n",
        "    if not parallel:\n",
        "        with tqdm(total=len(links[0])) as pbar:\n",
        "            for i, j, g_label in zip(links[0], links[1], labels):\n",
        "                tmp = subgraph_extraction_labeling(\n",
        "                    (i, j), Arow, Acol, h, sample_ratio, max_nodes_per_hop, u_features,\n",
        "                    v_features, class_values, g_label\n",
        "                )\n",
        "                data = construct_pyg_graph(*tmp)\n",
        "                g_list.append(data)\n",
        "                pbar.update(1)\n",
        "    else:\n",
        "        start = time.time()\n",
        "        pool = mp.Pool(mp.cpu_count())\n",
        "        results = pool.starmap_async(\n",
        "            subgraph_extraction_labeling,\n",
        "            [\n",
        "                ((i, j), Arow, Acol, h, sample_ratio, max_nodes_per_hop, u_features,\n",
        "                v_features, class_values, g_label)\n",
        "                for i, j, g_label in zip(links[0], links[1], labels)\n",
        "            ]\n",
        "        )\n",
        "        remaining = results._number_left\n",
        "        pbar = tqdm(total=remaining)\n",
        "        while True:\n",
        "            pbar.update(remaining - results._number_left)\n",
        "            if results.ready(): break\n",
        "            remaining = results._number_left\n",
        "            time.sleep(1)\n",
        "        results = results.get()\n",
        "        pool.close()\n",
        "        pbar.close()\n",
        "        end = time.time()\n",
        "        print(\"Time elapsed for subgraph extraction: {}s\".format(end-start))\n",
        "        print(\"Transforming to pytorch_geometric graphs...\")\n",
        "        g_list = []\n",
        "        pbar = tqdm(total=len(results))\n",
        "        while results:\n",
        "            tmp = results.pop()\n",
        "            g_list.append(construct_pyg_graph(*tmp))\n",
        "            pbar.update(1)\n",
        "        pbar.close()\n",
        "        end2 = time.time()\n",
        "        print(\"Time elapsed for transforming to pytorch_geometric graphs: {}s\".format(end2-end))\n",
        "    return g_list"
      ],
      "metadata": {
        "id": "xEFPcFFZgPyB"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def subgraph_extraction_labeling(ind, Arow, Acol, h=1, sample_ratio=1.0, max_nodes_per_hop=None,\n",
        "                                 u_features=None, v_features=None, class_values=None,\n",
        "                                 y=1):\n",
        "    # extract the h-hop enclosing subgraph around link 'ind'\n",
        "    u_nodes, v_nodes = [ind[0]], [ind[1]]\n",
        "    u_dist, v_dist = [0], [0]\n",
        "    u_visited, v_visited = set([ind[0]]), set([ind[1]])\n",
        "    u_fringe, v_fringe = set([ind[0]]), set([ind[1]])\n",
        "    for dist in range(1, h+1):\n",
        "        v_fringe, u_fringe = neighbors(u_fringe, Arow), neighbors(v_fringe, Acol)\n",
        "        u_fringe = u_fringe - u_visited\n",
        "        v_fringe = v_fringe - v_visited\n",
        "        u_visited = u_visited.union(u_fringe)\n",
        "        v_visited = v_visited.union(v_fringe)\n",
        "        if sample_ratio < 1.0:\n",
        "            u_fringe = random.sample(u_fringe, int(sample_ratio*len(u_fringe)))\n",
        "            v_fringe = random.sample(v_fringe, int(sample_ratio*len(v_fringe)))\n",
        "        if max_nodes_per_hop is not None:\n",
        "            if max_nodes_per_hop < len(u_fringe):\n",
        "                u_fringe = random.sample(u_fringe, max_nodes_per_hop)\n",
        "            if max_nodes_per_hop < len(v_fringe):\n",
        "                v_fringe = random.sample(v_fringe, max_nodes_per_hop)\n",
        "        if len(u_fringe) == 0 and len(v_fringe) == 0:\n",
        "            break\n",
        "        u_nodes = u_nodes + list(u_fringe)\n",
        "        v_nodes = v_nodes + list(v_fringe)\n",
        "        u_dist = u_dist + [dist] * len(u_fringe)\n",
        "        v_dist = v_dist + [dist] * len(v_fringe)\n",
        "    subgraph = Arow[u_nodes][:, v_nodes]\n",
        "    subgraph[0, 0] = 0\n",
        "\n",
        "    # prepare pyg graph constructor input\n",
        "    u, v, r = sp.find(subgraph)  # r is 1, 2... (rating labels + 1)\n",
        "    v += len(u_nodes)\n",
        "    r = r - 1  # transform r back to rating label\n",
        "    num_nodes = len(u_nodes) + len(v_nodes)\n",
        "    node_labels = [x*2 for x in u_dist] + [x*2+1 for x in v_dist]\n",
        "    max_node_label = 2*h + 1\n",
        "    y = class_values[y]\n",
        "\n",
        "    # get node features\n",
        "    if u_features is not None:\n",
        "        u_features = u_features[u_nodes]\n",
        "    if v_features is not None:\n",
        "        v_features = v_features[v_nodes]\n",
        "    node_features = None\n",
        "    if False:\n",
        "        if u_features is not None and v_features is not None:\n",
        "            u_extended = np.concatenate(\n",
        "                [u_features, np.zeros([u_features.shape[0], v_features.shape[1]])], 1\n",
        "            )\n",
        "            v_extended = np.concatenate(\n",
        "                [np.zeros([v_features.shape[0], u_features.shape[1]]), v_features], 1\n",
        "            )\n",
        "            node_features = np.concatenate([u_extended, v_extended], 0)\n",
        "    if False:\n",
        "        # use identity features (one-hot encodings of node idxes)\n",
        "        u_ids = one_hot(u_nodes, Arow.shape[0] + Arow.shape[1])\n",
        "        v_ids = one_hot([x+Arow.shape[0] for x in v_nodes], Arow.shape[0] + Arow.shape[1])\n",
        "        node_ids = np.concatenate([u_ids, v_ids], 0)\n",
        "        #node_features = np.concatenate([node_features, node_ids], 1)\n",
        "        node_features = node_ids\n",
        "    if True:\n",
        "        # only output node features for the target user and item\n",
        "        if u_features is not None and v_features is not None:\n",
        "            node_features = [u_features[0], v_features[0]]\n",
        "\n",
        "    return u, v, r, node_labels, max_node_label, y, node_features\n"
      ],
      "metadata": {
        "id": "3EfQl-DMgcNz"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def construct_pyg_graph(u, v, r, node_labels, max_node_label, y, node_features):\n",
        "    u, v = torch.LongTensor(u), torch.LongTensor(v)\n",
        "    r = torch.LongTensor(r)\n",
        "    edge_index = torch.stack([torch.cat([u, v]), torch.cat([v, u])], 0)\n",
        "    edge_type = torch.cat([r, r])\n",
        "    x = torch.FloatTensor(one_hot(node_labels, max_node_label+1))\n",
        "    y = torch.FloatTensor([y])\n",
        "    data = Data(x, edge_index, edge_type=edge_type, y=y)\n",
        "\n",
        "    if node_features is not None:\n",
        "        if type(node_features) == list:  # a list of u_feature and v_feature\n",
        "            u_feature, v_feature = node_features\n",
        "            data.u_feature = torch.FloatTensor(u_feature).unsqueeze(0)\n",
        "            data.v_feature = torch.FloatTensor(v_feature).unsqueeze(0)\n",
        "        else:\n",
        "            x2 = torch.FloatTensor(node_features)\n",
        "            data.x = torch.cat([data.x, x2], 1)\n",
        "    return data\n"
      ],
      "metadata": {
        "id": "So2cqfD3gkuL"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def neighbors(fringe, A):\n",
        "    # find all 1-hop neighbors of nodes in fringe from A\n",
        "    if not fringe:\n",
        "        return set([])\n",
        "    return set(A[list(fringe)].indices)"
      ],
      "metadata": {
        "id": "qlL0pjkJgpFE"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def one_hot(idx, length):\n",
        "    idx = np.array(idx)\n",
        "    x = np.zeros([len(idx), length])\n",
        "    x[np.arange(len(idx)), idx] = 1.0\n",
        "    return x"
      ],
      "metadata": {
        "id": "UHUtPM8Sgqws"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#icluded preprocessing steps as well\n",
        "def load_official_trainvaltest_split(testing=False, rating_map=None, post_rating_map=None, ratio=1.0):\n",
        "    sep = '\\t'\n",
        "    # Check if files exist and download otherwise\n",
        "    files = ['/u1.base', '/u1.test', '/u.item', '/u.user']\n",
        "    data_dir = '.'\n",
        "    dtypes = {\n",
        "        'u_nodes': np.int32, 'v_nodes': np.int32,\n",
        "        'ratings': np.float32, 'timestamp': np.float64}\n",
        "\n",
        "    filename_train = './u1.base'\n",
        "    filename_test = './u1.test'\n",
        "\n",
        "    data_train = pd.read_csv(\n",
        "        filename_train, sep=sep, header=None,\n",
        "        names=['u_nodes', 'v_nodes', 'ratings', 'timestamp'], dtype=dtypes)\n",
        "\n",
        "    data_test = pd.read_csv(\n",
        "        filename_test, sep=sep, header=None,\n",
        "        names=['u_nodes', 'v_nodes', 'ratings', 'timestamp'], dtype=dtypes)\n",
        "\n",
        "    data_array_train = data_train.values.tolist()\n",
        "    data_array_train = np.array(data_array_train)\n",
        "    data_array_test = data_test.values.tolist()\n",
        "    data_array_test = np.array(data_array_test)\n",
        "\n",
        "    if ratio < 1.0:\n",
        "        data_array_train = data_array_train[data_array_train[:, -1].argsort()[:int(ratio*len(data_array_train))]]\n",
        "\n",
        "    data_array = np.concatenate([data_array_train, data_array_test], axis=0)\n",
        "\n",
        "    u_nodes_ratings = data_array[:, 0].astype(dtypes['u_nodes'])\n",
        "    v_nodes_ratings = data_array[:, 1].astype(dtypes['v_nodes'])\n",
        "    ratings = data_array[:, 2].astype(dtypes['ratings'])\n",
        "    if rating_map is not None:\n",
        "        for i, x in enumerate(ratings):\n",
        "            ratings[i] = rating_map[x]\n",
        "\n",
        "    u_nodes_ratings, u_dict, num_users = map_data(u_nodes_ratings)\n",
        "    v_nodes_ratings, v_dict, num_movies = map_data(v_nodes_ratings)\n",
        "\n",
        "    u_nodes_ratings, v_nodes_ratings = u_nodes_ratings.astype(np.int64), v_nodes_ratings.astype(np.int32)\n",
        "    ratings = ratings.astype(np.float64)\n",
        "\n",
        "    u_nodes = u_nodes_ratings\n",
        "    v_nodes = v_nodes_ratings\n",
        "\n",
        "    neutral_rating = -1  # int(np.ceil(np.float(num_classes)/2.)) - 1\n",
        "\n",
        "    # assumes that ratings_train contains at least one example of every rating type\n",
        "    rating_dict = {r: i for i, r in enumerate(np.sort(np.unique(ratings)).tolist())}\n",
        "\n",
        "    labels = np.full((num_users, num_movies), neutral_rating, dtype=np.int32)\n",
        "    labels[u_nodes, v_nodes] = np.array([rating_dict[r] for r in ratings])\n",
        "\n",
        "    for i in range(len(u_nodes)):\n",
        "        assert(labels[u_nodes[i], v_nodes[i]] == rating_dict[ratings[i]])\n",
        "\n",
        "    labels = labels.reshape([-1])\n",
        "\n",
        "    # number of test and validation edges, see cf-nade code\n",
        "\n",
        "    num_train = data_array_train.shape[0]\n",
        "    num_test = data_array_test.shape[0]\n",
        "    num_val = int(np.ceil(num_train * 0.2))\n",
        "    num_train = num_train - num_val\n",
        "\n",
        "    pairs_nonzero = np.array([[u, v] for u, v in zip(u_nodes, v_nodes)])\n",
        "    idx_nonzero = np.array([u * num_movies + v for u, v in pairs_nonzero])\n",
        "\n",
        "    for i in range(len(ratings)):\n",
        "        assert(labels[idx_nonzero[i]] == rating_dict[ratings[i]])\n",
        "\n",
        "    idx_nonzero_train = idx_nonzero[0:num_train+num_val]\n",
        "    idx_nonzero_test = idx_nonzero[num_train+num_val:]\n",
        "\n",
        "    pairs_nonzero_train = pairs_nonzero[0:num_train+num_val]\n",
        "    pairs_nonzero_test = pairs_nonzero[num_train+num_val:]\n",
        "\n",
        "    # Internally shuffle training set (before splitting off validation set)\n",
        "    rand_idx = list(range(len(idx_nonzero_train)))\n",
        "    np.random.seed(42)\n",
        "    np.random.shuffle(rand_idx)\n",
        "    idx_nonzero_train = idx_nonzero_train[rand_idx]\n",
        "    pairs_nonzero_train = pairs_nonzero_train[rand_idx]\n",
        "\n",
        "    idx_nonzero = np.concatenate([idx_nonzero_train, idx_nonzero_test], axis=0)\n",
        "    pairs_nonzero = np.concatenate([pairs_nonzero_train, pairs_nonzero_test], axis=0)\n",
        "\n",
        "    val_idx = idx_nonzero[0:num_val]\n",
        "    train_idx = idx_nonzero[num_val:num_train + num_val]\n",
        "    test_idx = idx_nonzero[num_train + num_val:]\n",
        "\n",
        "    assert(len(test_idx) == num_test)\n",
        "\n",
        "    val_pairs_idx = pairs_nonzero[0:num_val]\n",
        "    train_pairs_idx = pairs_nonzero[num_val:num_train + num_val]\n",
        "    test_pairs_idx = pairs_nonzero[num_train + num_val:]\n",
        "\n",
        "    u_test_idx, v_test_idx = test_pairs_idx.transpose()\n",
        "    u_val_idx, v_val_idx = val_pairs_idx.transpose()\n",
        "    u_train_idx, v_train_idx = train_pairs_idx.transpose()\n",
        "\n",
        "    # create labels\n",
        "    train_labels = labels[train_idx]\n",
        "    val_labels = labels[val_idx]\n",
        "    test_labels = labels[test_idx]\n",
        "\n",
        "    if testing:\n",
        "        u_train_idx = np.hstack([u_train_idx, u_val_idx])\n",
        "        v_train_idx = np.hstack([v_train_idx, v_val_idx])\n",
        "        train_labels = np.hstack([train_labels, val_labels])\n",
        "        # for adjacency matrix construction\n",
        "        train_idx = np.hstack([train_idx, val_idx])\n",
        "\n",
        "    class_values = np.sort(np.unique(ratings))\n",
        "\n",
        "    # make training adjacency matrix\n",
        "    rating_mx_train = np.zeros(num_users * num_movies, dtype=np.float32)\n",
        "    if post_rating_map is None:\n",
        "        rating_mx_train[train_idx] = labels[train_idx].astype(np.float32) + 1.\n",
        "    else:\n",
        "        rating_mx_train[train_idx] = np.array([post_rating_map[r] for r in class_values[labels[train_idx]]]) + 1.\n",
        "    rating_mx_train = sp.csr_matrix(rating_mx_train.reshape(num_users, num_movies))\n",
        "\n",
        "\n",
        "    # movie features (genres)\n",
        "    sep = r'|'\n",
        "    movie_file = './u.item'\n",
        "    movie_headers = ['movie id', 'movie title', 'release date', 'video release date',\n",
        "                      'IMDb URL', 'unknown', 'Action', 'Adventure', 'Animation',\n",
        "                      'Childrens', 'Comedy', 'Crime', 'Documentary', 'Drama', 'Fantasy',\n",
        "                      'Film-Noir', 'Horror', 'Musical', 'Mystery', 'Romance', 'Sci-Fi',\n",
        "                      'Thriller', 'War', 'Western']\n",
        "    movie_df = pd.read_csv(movie_file, sep=sep, header=None,\n",
        "                            names=movie_headers, engine='python',encoding=\"latin-1\")\n",
        "\n",
        "    genre_headers = movie_df.columns.values[6:]\n",
        "    num_genres = genre_headers.shape[0]\n",
        "\n",
        "    v_features = np.zeros((num_movies, num_genres), dtype=np.float32)\n",
        "    for movie_id, g_vec in zip(movie_df['movie id'].values.tolist(), movie_df[genre_headers].values.tolist()):\n",
        "        # check if movie_id was listed in ratings file and therefore in mapping dictionary\n",
        "        if movie_id in v_dict.keys():\n",
        "            v_features[v_dict[movie_id], :] = g_vec\n",
        "\n",
        "    # user features\n",
        "\n",
        "    sep = r'|'\n",
        "    users_file = './u.user'\n",
        "    users_headers = ['user id', 'age', 'gender', 'occupation', 'zip code']\n",
        "    users_df = pd.read_csv(users_file, sep=sep, header=None,\n",
        "                            names=users_headers, engine='python')\n",
        "\n",
        "    occupation = set(users_df['occupation'].values.tolist())\n",
        "\n",
        "    age = users_df['age'].values\n",
        "    age_max = age.max()\n",
        "\n",
        "    gender_dict = {'M': 0., 'F': 1.}\n",
        "    occupation_dict = {f: i for i, f in enumerate(occupation, start=2)}\n",
        "\n",
        "    num_feats = 2 + len(occupation_dict)\n",
        "\n",
        "    u_features = np.zeros((num_users, num_feats), dtype=np.float32)\n",
        "    for _, row in users_df.iterrows():\n",
        "        u_id = row['user id']\n",
        "        if u_id in u_dict.keys():\n",
        "            # age\n",
        "            u_features[u_dict[u_id], 0] = row['age'] / np.float(age_max)\n",
        "            # gender\n",
        "            u_features[u_dict[u_id], 1] = gender_dict[row['gender']]\n",
        "            # occupation\n",
        "            u_features[u_dict[u_id], occupation_dict[row['occupation']]] = 1.\n",
        "\n",
        "\n",
        "\n",
        "    u_features = sp.csr_matrix(u_features)\n",
        "    v_features = sp.csr_matrix(v_features)\n",
        "\n",
        "    print(\"User features shape: \"+str(u_features.shape))\n",
        "    print(\"Item features shape: \"+str(v_features.shape))\n",
        "\n",
        "    return u_features, v_features, rating_mx_train, train_labels, u_train_idx, v_train_idx, \\\n",
        "        val_labels, u_val_idx, v_val_idx, test_labels, u_test_idx, v_test_idx, class_values"
      ],
      "metadata": {
        "id": "1N5ZQYVoh5fU"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "from torch.nn import Linear\n",
        "import torch.nn.functional as F\n",
        "from torch.optim import Adam\n",
        "from torch_geometric.data import DataLoader\n",
        "from torch_geometric.nn import RGCNConv\n",
        "from torch_geometric.utils import dropout_adj"
      ],
      "metadata": {
        "id": "bIp6CA6yiaNW"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Arguments\n",
        "EPOCHS=30\n",
        "BATCH_SIZE=128\n",
        "LR=1e-3\n",
        "LR_DECAY_STEP = 20\n",
        "LR_DECAY_VALUE = 10"
      ],
      "metadata": {
        "id": "wmD7MQWgighP"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "torch.manual_seed(1234)\n",
        "device = torch.device('cpu')\n",
        "if torch.cuda.is_available():\n",
        "    torch.cuda.manual_seed(123)\n",
        "    torch.cuda.synchronize()\n",
        "    device = torch.device('cuda')\n",
        "device"
      ],
      "metadata": {
        "id": "xuq0j9f8ijRj"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "(u_features, v_features, adj_train, train_labels, train_u_indices, train_v_indices, val_labels,\n",
        "val_u_indices, val_v_indices, test_labels, test_u_indices, test_v_indices, class_values\n",
        ") = load_official_trainvaltest_split(testing=True)"
      ],
      "metadata": {
        "id": "G3IlRpzLil-5"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "train_dataset = eval('DynamicDataset')(root='data/ml_100k/testmode/train', A=adj_train,\n",
        "    links=(train_u_indices, train_v_indices), labels=train_labels, h=1, sample_ratio=1.0,\n",
        "    max_nodes_per_hop=200, u_features=None, v_features=None, class_values=class_values)\n",
        "test_dataset = eval('GraphDataset')(root='data/ml_100k/testmode/test', A=adj_train,\n",
        "    links=(test_u_indices, test_v_indices), labels=test_labels, h=1, sample_ratio=1.0,\n",
        "    max_nodes_per_hop=200, u_features=None, v_features=None, class_values=class_values)\n",
        "\n",
        "len(train_dataset), len(test_dataset)"
      ],
      "metadata": {
        "id": "FL5VGBtJioDE"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class IGMC(torch.nn.Module):\n",
        "    def __init__(self):\n",
        "        super(IGMC, self).__init__()\n",
        "        self.rel_graph_convs = torch.nn.ModuleList()\n",
        "        self.rel_graph_convs.append(RGCNConv(in_channels=4, out_channels=32, num_relations=5, num_bases=4))\n",
        "        self.rel_graph_convs.append(RGCNConv(in_channels=32, out_channels=32, num_relations=5, num_bases=4))\n",
        "        self.rel_graph_convs.append(RGCNConv(in_channels=32, out_channels=32, num_relations=5, num_bases=4))\n",
        "        self.rel_graph_convs.append(RGCNConv(in_channels=32, out_channels=32, num_relations=5, num_bases=4))\n",
        "        self.linear_layer1 = Linear(256, 128)\n",
        "        self.linear_layer2 = Linear(128, 1)\n",
        "\n",
        "    def reset_parameters(self):\n",
        "        self.linear_layer1.reset_parameters()\n",
        "        self.linear_layer2.reset_parameters()\n",
        "        for i in self.rel_graph_convs:\n",
        "            i.reset_parameters()\n",
        "\n",
        "    def forward(self, data):\n",
        "        num_nodes = len(data.x)\n",
        "        edge_index_dr, edge_type_dr = dropout_adj(data.edge_index, data.edge_type, p=0.2, num_nodes=num_nodes, training=self.training)\n",
        "\n",
        "        out = data.x\n",
        "        h = []\n",
        "        for conv in self.rel_graph_convs:\n",
        "            out = conv(out, edge_index_dr, edge_type_dr)\n",
        "            out = torch.tanh(out)\n",
        "            h.append(out)\n",
        "        h = torch.cat(h, 1)\n",
        "        h = [h[data.x[:, 0] == True], h[data.x[:, 1] == True]]\n",
        "        g = torch.cat(h, 1)\n",
        "        out = self.linear_layer1(g)\n",
        "        out = F.relu(out)\n",
        "        out = F.dropout(out, p=0.2, training=self.training)\n",
        "        out = self.linear_layer2(out)\n",
        "        out = out[:,0]\n",
        "        return out\n",
        "\n",
        "model = IGMC()"
      ],
      "metadata": {
        "id": "I3YwLz2Xise0"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "train_dataset"
      ],
      "metadata": {
        "id": "TpsHfqusi-Pi"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "train_loader = DataLoader(train_dataset, 200, shuffle=True)\n",
        "test_loader = DataLoader(test_dataset, 200, shuffle=False)"
      ],
      "metadata": {
        "id": "H2NqXk-b2deK"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "model.to(device)\n",
        "model.reset_parameters()\n",
        "optimizer = Adam(model.parameters(), lr=LR, weight_decay=0)"
      ],
      "metadata": {
        "id": "hIowq0LHjAMj"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "train_losslist =[]\n",
        "test_losslist = []\n",
        "for epoch in range(1, EPOCHS+1):\n",
        "    model.train()\n",
        "    train_loss_all = 0\n",
        "\n",
        "    for train_batch in train_loader:\n",
        "        optimizer.zero_grad()\n",
        "        train_batch = train_batch.to(device)\n",
        "        y_pred = model(train_batch)\n",
        "        y_true = train_batch.y\n",
        "        train_loss = F.mse_loss(y_pred, y_true)\n",
        "        train_loss.backward()\n",
        "        train_loss_all += BATCH_SIZE * float(train_loss)\n",
        "        optimizer.step()\n",
        "        torch.cuda.empty_cache()\n",
        "    train_loss_all = train_loss_all / len(train_loader.dataset)\n",
        "    train_losslist.append(train_loss_all)\n",
        "    test_loss = 0\n",
        "    for test_batch in test_loader:\n",
        "          test_batch = test_batch.to(device)\n",
        "          with torch.no_grad():\n",
        "              y_pred = model(test_batch)\n",
        "          y_true = test_batch.y\n",
        "          test_loss += F.mse_loss(y_pred, y_true, reduction='sum')\n",
        "          # torch.cuda.empty_cache()\n",
        "    mse_loss = float(test_loss) / len(test_loader.dataset)\n",
        "    test_losslist.append(mse_loss)\n",
        "    print('epoch', epoch,'; train loss', train_loss_all,'; test loss',mse_loss)\n",
        "\n",
        "    if epoch % LR_DECAY_STEP == 0:\n",
        "      for param_group in optimizer.param_groups:\n",
        "          param_group['lr'] = param_group['lr'] / LR_DECAY_VALUE"
      ],
      "metadata": {
        "id": "C5tSEXC5jCbj"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "model.eval()\n",
        "test_loss = 0\n",
        "for test_batch in test_loader:\n",
        "    test_batch = test_batch.to(device)\n",
        "    with torch.no_grad():\n",
        "        y_pred = model(test_batch)\n",
        "    y_true = test_batch.y\n",
        "    test_loss += F.mse_loss(y_pred, y_true, reduction='sum')\n",
        "    # torch.cuda.empty_cache()\n",
        "mse_loss = float(test_loss) / len(test_loader.dataset)\n",
        "\n",
        "print('test MSE loss', mse_loss)\n",
        "print('test RMSE loss', math.sqrt(mse_loss))"
      ],
      "metadata": {
        "id": "11tOBo6KjGLD"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import matplotlib.pyplot as plt\n",
        "iters = [iter for iter in range(len(train_losslist))]\n",
        "plt.plot(iters, train_losslist, label='train')\n",
        "plt.plot(iters, test_losslist, label='test')\n",
        "plt.xlabel('epochs')\n",
        "plt.ylabel('loss')\n",
        "plt.title('training and test loss curves')\n",
        "plt.legend()\n",
        "plt.show()"
      ],
      "metadata": {
        "id": "cWWvQm5zjJLU"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}