{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/RecoHut-Projects/recohut/blob/master/tutorials/modeling/siren_movielens_pytorch_cpu.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Sign-Aware Recommendation Using Graph Neural Networks (SiReN) on ML-1m Dataset in PyTorch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Description:** In this tutorial, we are building the SiReN recommender model on MovieLens-1m dataset. For the given user id, the model will recommend Top-K relevant movies to the user. We are using PyTorch and PyTorch Geometric libraries heavily. For graph embeddings, LightGCN is used and this can be replaced with LR-GCCF also.\n",
    "\n",
    "**Links:**\n",
    " 1. Refer to [this](https://recohut.notion.site/Item-Recommendations-with-Sign-aware-Graph-Neural-Network-3ba80ce0460246b2999408d431bd8d76) report for theoretical understanding. \n",
    " 2. Widen your understanding by applying this codebase on other datasets also. [Here](https://nbviewer.org/github/sparsh-ai/stanza/tree/S138006/nbs/) are some notebooks to start with."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install torch-scatter -f https://data.pyg.org/whl/torch-1.10.0+cu111.html\n",
    "!pip install torch-sparse -f https://data.pyg.org/whl/torch-1.10.0+cu111.html\n",
    "!pip install torch-geometric"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'1.10.0+cu111'"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "torch.__version__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch import nn\n",
    "from torch import Tensor\n",
    "from torch import optim\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from torch_geometric.data import Data\n",
    "import torch.nn.functional as F\n",
    "from torch_geometric.nn import MessagePassing\n",
    "from torch_geometric.utils import add_self_loops, degree\n",
    "\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import time\n",
    "from tqdm.notebook import tqdm\n",
    "\n",
    "import os\n",
    "import pickle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings('ignore')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "mappingproxy({'K': 40,\n",
       "              'MLP_layers': 2,\n",
       "              '__dict__': <attribute '__dict__' of 'Args' objects>,\n",
       "              '__doc__': None,\n",
       "              '__module__': '__main__',\n",
       "              '__weakref__': <attribute '__weakref__' of 'Args' objects>,\n",
       "              'batch_size': 1024,\n",
       "              'dataset': 'ML-1M',\n",
       "              'dim': 64,\n",
       "              'epoch': 5,\n",
       "              'lr': 0.005,\n",
       "              'num_layers': 4,\n",
       "              'offset': 3.5,\n",
       "              'reg': 0.05,\n",
       "              'version': 1})"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "class Args:\n",
    "    dataset = 'ML-1M' # Dataset\n",
    "    version = 1 # Dataset version\n",
    "    batch_size = 1024 # Batch size\n",
    "    dim = 64 # Dimension\n",
    "    lr = 5e-3 # Learning rate\n",
    "    offset = 3.5 # Criterion of likes/dislikes\n",
    "    K = 40 # The number of negative samples\n",
    "    num_layers = 4 # The number of layers of a GNN model for the graph with positive edges\n",
    "    MLP_layers = 2 # The number of layers of MLP for the graph with negative edges\n",
    "    epoch = 5 # The number of epochs\n",
    "    reg = 0.05 # Regularization coefficient\n",
    "\n",
    "args = Args()\n",
    "Args.__dict__"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!git clone -q --branch v2 https://github.com/RecoHut-Datasets/movielens_1m.git"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Data_loader():\n",
    "    def __init__(self,dataset,version):\n",
    "        self.dataset=dataset; self.version=version\n",
    "        self.sep='::'\n",
    "        self.names=['userId','movieId','rating','timestemp'];\n",
    "        self.path_for_whole='./movielens_1m/ratings.dat'\n",
    "        self.path_for_train='./movielens_1m/train_1m%s.dat'%(version)\n",
    "        self.path_for_test='./movielens_1m/test_1m%s.dat'%(version)\n",
    "        self.num_u=6040; self.num_v=3952;\n",
    "        \n",
    "    def data_load(self):\n",
    "        self.whole_=pd.read_csv(self.path_for_whole, names = self.names, sep=self.sep, engine='python').drop('timestemp',axis=1).sample(frac=1,replace=False,random_state=self.version)\n",
    "        self.train_set = pd.read_csv(self.path_for_train,engine='python',names=self.names).drop('timestemp',axis=1)\n",
    "        self.test_set = pd.read_csv(self.path_for_test,engine='python',names=self.names).drop('timestemp',axis=1)            \n",
    "        return self.train_set, self.test_set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class bipartite_dataset(Dataset): \n",
    "    def __init__(self, train,neg_dist,offset,num_u,num_v,K): \n",
    "        self.edge_1 = torch.tensor(train['userId'].values-1)\n",
    "        self.edge_2 = torch.tensor(train['movieId'].values-1) +num_u\n",
    "        self.edge_3 = torch.tensor(train['rating'].values) - offset\n",
    "        self.neg_dist = neg_dist\n",
    "        self.K = K;\n",
    "        self.num_u = num_u\n",
    "        self.num_v = num_v\n",
    "        self.tot = np.arange(num_v)\n",
    "        self.train = train\n",
    "        \n",
    "    def negs_gen_(self):\n",
    "        print('negative sampling...'); st=time.time()\n",
    "        self.edge_4 = torch.empty((len(self.edge_1),self.K),dtype=torch.long)\n",
    "        prog = tqdm(desc='negative sampling for each epoch...',total=len(set(self.train['userId'].values)),position=0)\n",
    "        for j in set(self.train['userId'].values):\n",
    "            pos=self.train[self.train['userId']==j]['movieId'].values-1\n",
    "            neg = np.setdiff1d(self.tot,pos)\n",
    "            temp = (torch.tensor(np.random.choice(neg,len(pos)*self.K,replace=True,p=self.neg_dist[neg]/self.neg_dist[neg].sum()))+self.num_u).long()\n",
    "            self.edge_4[self.edge_1==j-1]=temp.view(int(len(temp)/self.K),self.K)\n",
    "            prog.update(1)\n",
    "        prog.close()\n",
    "        self.edge_4 = torch.tensor(self.edge_4).long()\n",
    "        print('complete ! %s'%(time.time()-st))\n",
    "        \n",
    "    def negs_gen_EP(self,epoch):\n",
    "        print('negative sampling for next epochs...'); st=time.time()\n",
    "        self.edge_4_tot = torch.empty((len(self.edge_1),self.K,epoch),dtype=torch.long)\n",
    "        prog = tqdm(desc='negative sampling for next epochs...',total=len(set(self.train['userId'].values)),position=0)\n",
    "        for j in set(self.train['userId'].values):\n",
    "            pos=self.train[self.train['userId']==j]['movieId'].values-1\n",
    "            neg = np.setdiff1d(self.tot,pos)\n",
    "            temp = (torch.tensor(np.random.choice(neg,len(pos)*self.K*epoch,replace=True,p=self.neg_dist[neg]/self.neg_dist[neg].sum()))+self.num_u).long()\n",
    "            self.edge_4_tot[self.edge_1==j-1]=temp.view(int(len(temp)/self.K/epoch),self.K,epoch)\n",
    "            prog.update(1)\n",
    "        prog.close()\n",
    "        self.edge_4_tot = torch.tensor(self.edge_4_tot).long()\n",
    "        print('complete ! %s'%(time.time()-st))\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.edge_1)\n",
    "\n",
    "    def __getitem__(self,idx):\n",
    "        u = self.edge_1[idx]\n",
    "        v = self.edge_2[idx]\n",
    "        w = self.edge_3[idx]\n",
    "        negs = self.edge_4[idx]\n",
    "        return u,v,w,negs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def deg_dist(train, num_v):\n",
    "    uni, cou = np.unique(train['movieId'].values-1,return_counts=True)\n",
    "    cou = cou**(0.75)\n",
    "    deg = np.zeros(num_v)\n",
    "    deg[uni] = cou\n",
    "    return torch.tensor(deg)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Modules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def gen_top_K(data_class,emb,train,directory_):\n",
    "    no_items = np.array(list(set(np.arange(1,data_class.num_v+1))-set(train['movieId'])))\n",
    "    total_users = set(np.arange(1,data_class.num_u+1))\n",
    "    reco = dict()\n",
    "    pbar = tqdm(desc = 'top-k recommendation...',total=len(total_users),position=0)\n",
    "    for j in total_users:\n",
    "        pos = train[train['userId']==j]['movieId'].values-1\n",
    "        embedding_ = emb[j-1].view(1,len(emb[0])).mm(emb[data_class.num_u:].t()).detach();\n",
    "        embedding_[0][no_items-1]=-np.inf;\n",
    "        embedding_[0][pos]=-np.inf;\n",
    "        reco[j]=torch.topk(embedding_[0],300).indices.cpu().numpy()+1\n",
    "        pbar.update(1)\n",
    "    pbar.close()\n",
    "    return reco"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LightGConv(MessagePassing):\n",
    "    def __init__(self):\n",
    "        super().__init__(aggr='add')\n",
    "        \n",
    "    def forward(self,x,edge_index):\n",
    "        row, col = edge_index\n",
    "        deg = degree(col, x.size(0), dtype=x.dtype)\n",
    "        deg_inv_sqrt = deg.pow(-0.5)\n",
    "        deg_inv_sqrt[deg_inv_sqrt == float('inf')] = 0\n",
    "        norm = deg_inv_sqrt[row] * deg_inv_sqrt[col]\n",
    "        return self.propagate(edge_index, x=x, norm=norm)\n",
    "    \n",
    "    def message(self,x_j,norm):\n",
    "        return norm.view(-1,1) * x_j\n",
    "        \n",
    "    def update(self,inputs: Tensor) -> Tensor:\n",
    "        return inputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LRGCCF(MessagePassing):\n",
    "    def __init__(self, in_channels,out_channels):\n",
    "        super(LRGCCF,self).__init__(aggr='mean')\n",
    "        self.lin = torch.nn.Linear(in_channels, out_channels)\n",
    "\n",
    "    def forward(self,x,edge_index):\n",
    "        edge_index, _ = add_self_loops(edge_index, num_nodes = x.size(0));\n",
    "        return self.lin(self.propagate(edge_index,x=x))\n",
    "\n",
    "    def message(self,x_j):\n",
    "        return x_j\n",
    "        \n",
    "    def update(self,inputs: Tensor) -> Tensor:\n",
    "        return inputs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SiReN(nn.Module):\n",
    "    def __init__(self,train,num_u,num_v,offset,num_layers = 2,MLP_layers=2,dim = 64,reg=1e-4\n",
    "                 ,device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")):\n",
    "        super(SiReN,self).__init__()\n",
    "        self.M = num_u; self.N = num_v;\n",
    "        self.num_layers = num_layers\n",
    "        self.MLP_layers = MLP_layers\n",
    "        self.device = device\n",
    "        self.reg = reg\n",
    "        self.embed_dim = dim\n",
    "        edge_user = torch.tensor(train[train['rating']>offset]['userId'].values-1)\n",
    "        edge_item = torch.tensor(train[train['rating']>offset]['movieId'].values-1)+self.M\n",
    "        edge_ = torch.stack((torch.cat((edge_user,edge_item),0),torch.cat((edge_item,edge_user),0)),0)\n",
    "        self.data_p=Data(edge_index=edge_)\n",
    "        # For the graph with positive edges\n",
    "        self.E = nn.Parameter(torch.empty(self.M + self.N, dim))\n",
    "        nn.init.xavier_normal_(self.E.data)\n",
    "        self.convs = nn.ModuleList()\n",
    "        self.mlps = nn.ModuleList()\n",
    "        for _ in range(num_layers):\n",
    "            # self.convs.append(LRGCCF(dim,dim)) \n",
    "            self.convs.append(LightGConv()) \n",
    "        # For the graph with negative edges\n",
    "        self.E2 = nn.Parameter(torch.empty(self.M + self.N, dim))\n",
    "        nn.init.xavier_normal_(self.E2.data)\n",
    "        for _ in range(MLP_layers):\n",
    "            self.mlps.append(nn.Linear(dim,dim,bias=True))\n",
    "            nn.init.xavier_normal_(self.mlps[-1].weight.data)\n",
    "        # Attntion model\n",
    "        self.attn = nn.Linear(dim,dim,bias=True)\n",
    "        self.q = nn.Linear(dim,1,bias=False)\n",
    "        self.attn_softmax = nn.Softmax(dim=1)\n",
    "        \n",
    "    def aggregate(self):\n",
    "        # Generate embeddings z_p\n",
    "        B=[]; B.append(self.E)\n",
    "        x = self.convs[0](self.E,self.data_p.edge_index)\n",
    "        B.append(x)\n",
    "        for i in range(1,self.num_layers):\n",
    "            x = self.convs[i](x,self.data_p.edge_index)\n",
    "            B.append(x)\n",
    "        z_p = sum(B)/len(B) \n",
    "        # Generate embeddings z_n\n",
    "        C = []; C.append(self.E2)\n",
    "        x = F.dropout(F.relu(self.mlps[0](self.E2)),p=0.5,training=self.training)\n",
    "        for i in range(1,self.MLP_layers):\n",
    "            x = self.mlps[i](x);\n",
    "            x = F.relu(x)\n",
    "            x = F.dropout(x,p=0.5,training=self.training)\n",
    "            C.append(x)\n",
    "        z_n = C[-1]\n",
    "        # Attntion for final embeddings Z\n",
    "        w_p = self.q(F.dropout(torch.tanh((self.attn(z_p))),p=0.5,training=self.training))\n",
    "        w_n = self.q(F.dropout(torch.tanh((self.attn(z_n))),p=0.5,training=self.training))\n",
    "        alpha_ = self.attn_softmax(torch.cat([w_p,w_n],dim=1))\n",
    "        Z = alpha_[:,0].view(len(z_p),1) * z_p + alpha_[:,1].view(len(z_p),1) * z_n\n",
    "        return Z\n",
    "    \n",
    "    def forward(self,u,v,w,n,device):\n",
    "        emb = self.aggregate()\n",
    "        u_ = emb[u].to(device);\n",
    "        v_ = emb[v].to(device);\n",
    "        n_ = emb[n].to(device);\n",
    "        w_ = w.to(device)\n",
    "        positivebatch = torch.mul(u_ , v_ ); \n",
    "        negativebatch = torch.mul(u_.view(len(u_),1,self.embed_dim),n_)  \n",
    "        sBPR_loss =  F.logsigmoid((torch.sign(w_).view(len(u_),1) * (positivebatch.sum(dim=1).view(len(u_),1))) - negativebatch.sum(dim=2)).sum(dim=1)\n",
    "        reg_loss = u_.norm(dim=1).pow(2).sum() + v_.norm(dim=1).pow(2).sum() + n_.norm(dim=2).pow(2).sum();\n",
    "        return -torch.sum(sBPR_loss) + self.reg * reg_loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class evaluate():\n",
    "    def __init__(self,reco,train,test,threshold,num_u,num_v,N=[5,10,15,20,25],ratings=[20,50]):\n",
    "        '''\n",
    "        train : training set\n",
    "        test : test set\n",
    "        threshold : To generate ground truth set from test set\n",
    "        '''\n",
    "        self.reco = reco\n",
    "        self.num_u = num_u;\n",
    "        self.num_v = num_v;\n",
    "        self.N=N\n",
    "        self.p=[]\n",
    "        self.r=[]\n",
    "        self.NDCG=[]\n",
    "        self.p_c1=[]; self.p_c2=[]; self.p_c3=[]\n",
    "        self.r_c1=[]; self.r_c2=[]; self.r_c3=[]\n",
    "        self.NDCG_c1=[]; self.NDCG_c2=[]; self.NDCG_c3=[]\n",
    "        self.tr = train; self.te = test;\n",
    "        self.threshold = threshold;\n",
    "        self.gen_ground_truth_set()\n",
    "        self.ratings = ratings\n",
    "        self.partition_into_groups_(self.ratings)\n",
    "        print('\\nevaluating recommendation accuracy....')\n",
    "        self.precision_and_recall_G(self.group1,1)\n",
    "        self.precision_and_recall_G(self.group2,2)\n",
    "        self.precision_and_recall_G(self.group3,3)\n",
    "        self.Normalized_DCG_G(self.group1,1)\n",
    "        self.Normalized_DCG_G(self.group2,2)\n",
    "        self.Normalized_DCG_G(self.group3,3)\n",
    "        self.metric_total()\n",
    "\n",
    "    def gen_ground_truth_set(self):\n",
    "        result = dict()\n",
    "        GT = self.te[self.te['rating']>=self.threshold];\n",
    "        U = set(GT['userId'])\n",
    "        for i in U:\n",
    "            result[i] = list(set([j for j in GT[GT['userId']==i]['movieId']]))#-set(self.TOP))\n",
    "            if len(result[i])==0:\n",
    "                del(result[i])\n",
    "        self.GT = result\n",
    "\n",
    "    def precision_and_recall(self):\n",
    "        user_in_GT=[j for j in self.GT];\n",
    "        for n in self.N:\n",
    "            p=0; r=0;\n",
    "            for i in user_in_GT:\n",
    "                topn=self.reco[i][:n]\n",
    "                num_hit=len(set(topn).intersection(set(self.GT[i])));\n",
    "                p+=num_hit/n; r+=num_hit/len(self.GT[i]);\n",
    "            self.p.append(p/len(user_in_GT)); self.r.append(r/len(user_in_GT));\n",
    "                \n",
    "    def Normalized_DCG(self):\n",
    "        maxn=max(self.N);\n",
    "        user_in_GT=[j for j in self.GT];\n",
    "        ndcg=np.zeros(maxn);\n",
    "        for i in user_in_GT:\n",
    "            idcg_len = min(len(self.GT[i]), maxn)\n",
    "            temp_idcg = np.cumsum(1.0 / np.log2(np.arange(2, maxn + 2)))\n",
    "            temp_idcg[idcg_len:] = temp_idcg[idcg_len-1]\n",
    "            temp_dcg=np.cumsum([1.0/np.log2(idx+2) if item in self.GT[i] else 0.0 for idx, item in enumerate(self.reco[i][:maxn])])\n",
    "            ndcg+=temp_dcg/temp_idcg;\n",
    "        ndcg/=len(user_in_GT);\n",
    "        for n in self.N:\n",
    "            self.NDCG.append(ndcg[n-1])\n",
    "            \n",
    "    def metric_total(self):\n",
    "        self.p = self.len1 * np.array(self.p_c1) + self.len2 * np.array(self.p_c2) + self.len3 * np.array(self.p_c3);\n",
    "        self.p/= self.len1 + self.len2 + self.len3\n",
    "        self.p = list(self.p)\n",
    "        self.r = self.len1 * np.array(self.r_c1) + self.len2 * np.array(self.r_c2) + self.len3 * np.array(self.r_c3);\n",
    "        self.r/= self.len1 + self.len2 + self.len3\n",
    "        self.r = list(self.r)\n",
    "        self.NDCG = self.len1 * np.array(self.NDCG_c1) + self.len2 * np.array(self.NDCG_c2) + self.len3 * np.array(self.NDCG_c3);\n",
    "        self.NDCG/= self.len1 + self.len2 + self.len3\n",
    "        self.NDCG = list(self.NDCG)\n",
    "\n",
    "    def partition_into_groups_(self,ratings=[20,50]):\n",
    "        unique_u, counts_u = np.unique(self.tr['userId'].values,return_counts=True)\n",
    "        self.group1 = unique_u[np.argwhere(counts_u<ratings[0])]\n",
    "        temp = unique_u[np.argwhere(counts_u<ratings[1])]\n",
    "        self.group2 = np.setdiff1d(temp,self.group1)\n",
    "        self.group3 = np.setdiff1d(unique_u,temp)\n",
    "        self.cold_groups = ratings\n",
    "        self.group1 = list(self.group1.reshape(-1))\n",
    "        self.group2 = list(self.group2.reshape(-1))\n",
    "        self.group3 = list(self.group3.reshape(-1))\n",
    "    \n",
    "    def precision_and_recall_G(self,group,gn):\n",
    "        user_in_GT=[j for j in self.GT];\n",
    "        leng = 0 ; maxn = max(self.N) ; p = np.zeros(maxn); r = np.zeros(maxn);\n",
    "        for i in user_in_GT:\n",
    "            if i in group:\n",
    "                leng+=1\n",
    "                hit_ = np.cumsum([1.0 if item in self.GT[i] else 0.0 for idx, item in enumerate(self.reco[i][:maxn])])\n",
    "                p+=hit_ / np.arange(1,maxn+1); r+=hit_/len(self.GT[i])\n",
    "        p/= leng; r/=leng;\n",
    "        for n in self.N:\n",
    "            if gn == 1 :\n",
    "                self.p_c1.append(p[n-1])\n",
    "                self.r_c1.append(r[n-1])\n",
    "                self.len1 = leng;\n",
    "            elif gn == 2 :\n",
    "                self.p_c2.append(p[n-1])\n",
    "                self.r_c2.append(r[n-1])\n",
    "                self.len2 = leng;\n",
    "            elif gn == 3 :\n",
    "                self.p_c3.append(p[n-1])\n",
    "                self.r_c3.append(r[n-1])\n",
    "                self.len3 = leng;\n",
    "            \n",
    "    def Normalized_DCG_G(self,group,gn):\n",
    "        maxn=max(self.N);\n",
    "        user_in_GT=[j for j in self.GT];\n",
    "        ndcg=np.zeros(maxn);\n",
    "        leng = 0\n",
    "        for i in user_in_GT:\n",
    "            if i in group:\n",
    "                leng+=1\n",
    "                idcg_len = min(len(self.GT[i]), maxn)\n",
    "                temp_idcg = np.cumsum(1.0 / np.log2(np.arange(2, maxn + 2)))\n",
    "                temp_idcg[idcg_len:] = temp_idcg[idcg_len-1]\n",
    "                temp_dcg=np.cumsum([1.0/np.log2(idx+2) if item in self.GT[i] else 0.0 for idx, item in enumerate(self.reco[i][:maxn])])\n",
    "                ndcg+=temp_dcg/temp_idcg;\n",
    "        ndcg/=leng\n",
    "        for n in self.N:\n",
    "            if gn == 1 :\n",
    "                self.NDCG_c1.append(ndcg[n-1])\n",
    "            elif gn == 2 :\n",
    "                self.NDCG_c2.append(ndcg[n-1])\n",
    "            elif gn == 3 :\n",
    "                self.NDCG_c3.append(ndcg[n-1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Main"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "data loading...\n",
      "loading complete! time :: 9.332336664199829\n",
      "generate negative candidates...\n",
      "complete ! time : 0.07818031311035156\n"
     ]
    }
   ],
   "source": [
    "data_class=Data_loader(args.dataset,args.version)\n",
    "threshold = round(args.offset) # To generate ground truth set \n",
    "\n",
    "print('data loading...'); st=time.time()\n",
    "train,test = data_class.data_load();\n",
    "train = train.astype({'userId':'int64', 'movieId':'int64'})\n",
    "print('loading complete! time :: %s'%(time.time()-st))\n",
    "\n",
    "print('generate negative candidates...'); st=time.time()\n",
    "neg_dist = deg_dist(train,data_class.num_v)\n",
    "print('complete ! time : %s'%(time.time()-st))    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Training on cuda:0...\n",
      "\n",
      "negative sampling for next epochs...\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ba3e794fad9a4546af6c1dacd62a7f26",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "negative sampling for next epochs...:   0%|          | 0/6040 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "complete ! 45.37964224815369\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ef65bd04e71947e1a142870b8be5266e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Version : 1 Epoch 1/5:   0%|          | 0/782 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8a7e8b66c3134a428703613bb96c2a33",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Version : 1 Epoch 2/5:   0%|          | 0/782 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "af91aebbfe4a4024a1a94ef867e3aad6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "top-k recommendation...:   0%|          | 0/6040 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "evaluating recommendation accuracy....\n",
      "\n",
      "***************************************************************************************\n",
      " /* Recommendation Accuracy */\n",
      "Precision at [10, 15, 20] ::  [0.19966571953869444, 0.1753969580478047, 0.15879157613237457]\n",
      "Recall at [10, 15, 20] ::  [0.12008447726136233, 0.1551362865560481, 0.18508128357830886]\n",
      "NDCG at [10, 15, 20] ::  [0.23887732176542237, 0.23155000884988997, 0.23061049717781704]\n",
      "***************************************************************************************\n",
      "negative sampling for next epochs...\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b6fd47cce5ae4e37bca20745cb761a15",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "negative sampling for next epochs...:   0%|          | 0/6040 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "complete ! 47.99065804481506\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "375570507e10452187107ec4d5279253",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Version : 1 Epoch 3/5:   0%|          | 0/782 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6eafffdd1964412ba4b4daf6f24c0688",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Version : 1 Epoch 4/5:   0%|          | 0/782 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0b097cf1b7c84af883a98cbfb7607dab",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "top-k recommendation...:   0%|          | 0/6040 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "evaluating recommendation accuracy....\n",
      "\n",
      "***************************************************************************************\n",
      " /* Recommendation Accuracy */\n",
      "Precision at [10, 15, 20] ::  [0.22502089252883245, 0.19700261853028295, 0.17802941668059427]\n",
      "Recall at [10, 15, 20] ::  [0.1377630580296285, 0.17745219739496146, 0.20962261836939539]\n",
      "NDCG at [10, 15, 20] ::  [0.27053789795527644, 0.2620976765066821, 0.26082565591434304]\n",
      "***************************************************************************************\n",
      "negative sampling for next epochs...\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "70b4cf9f12fd4fb6827d2e84c8a01237",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "negative sampling for next epochs...:   0%|          | 0/6040 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "complete ! 43.91022324562073\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c53abef792e945a8a5d19b1484f47c72",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Version : 1 Epoch 5/5:   0%|          | 0/782 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "model= SiReN(train, data_class.num_u,data_class.num_v,offset=args.offset,num_layers = args.num_layers,MLP_layers=args.MLP_layers,dim=args.dim,device=device,reg=args.reg)#.to(device);\n",
    "model.data_p.to(device)\n",
    "model.to(device)\n",
    "optimizer = optim.Adam(model.parameters(), lr = args.lr)\n",
    "\n",
    "print(\"\\nTraining on {}...\\n\".format(device))\n",
    "model.train()\n",
    "training_dataset=bipartite_dataset(train,neg_dist,args.offset,data_class.num_u,data_class.num_v,args.K);\n",
    "\n",
    "for EPOCH in range(1,args.epoch+1):\n",
    "    if EPOCH%2-1==0:training_dataset.negs_gen_EP(2)\n",
    "    LOSS=0\n",
    "    training_dataset.edge_4 = training_dataset.edge_4_tot[:,:,EPOCH%2-1]\n",
    "    ds = DataLoader(training_dataset,batch_size=args.batch_size,shuffle=True)\n",
    "    q=0\n",
    "    pbar = tqdm(desc = 'Version : {} Epoch {}/{}'.format(args.version,EPOCH,args.epoch),total=len(ds),position=0)\n",
    "    for u,v,w,negs in ds:\n",
    "        q+=len(u)\n",
    "        st=time.time()\n",
    "        optimizer.zero_grad()\n",
    "        loss = model(u,v,w,negs,device) # original\n",
    "        loss.backward()                \n",
    "        optimizer.step()\n",
    "        LOSS+=loss.item() * len(ds)\n",
    "        pbar.update(1);\n",
    "        pbar.set_postfix({'loss':loss.item()})\n",
    "    pbar.close()\n",
    "\n",
    "    if EPOCH%2==0 :\n",
    "        directory = os.getcwd() + '/results/%s/SiReN/epoch%s_batch%s_dim%s_lr%s_offset%s_K%s_num_layers%s_MLP_layers%s_threshold%s_reg%s/'%(args.dataset,EPOCH,args.batch_size,args.dim,args.lr,args.offset,args.K,args.num_layers,args.MLP_layers,threshold,args.reg)\n",
    "        if not os.path.exists(directory):\n",
    "            os.makedirs(directory)\n",
    "        model.eval()\n",
    "        emb = model.aggregate();\n",
    "        top_k_list = gen_top_K(data_class,emb,train,directory+'r%s_reco.pickle'%(args.version)) \n",
    "        eval_ = evaluate(top_k_list,train,test,threshold,data_class.num_u,data_class.num_v,N=[10,15,20],ratings=[20,50])\n",
    "        print(\"\\n***************************************************************************************\")\n",
    "        print(\" /* Recommendation Accuracy */\")\n",
    "        print('Precision at [10, 15, 20] :: ',eval_.p)\n",
    "        print('Recall at [10, 15, 20] :: ',eval_.r)\n",
    "        print('NDCG at [10, 15, 20] :: ',eval_.NDCG)\n",
    "        print(\"***************************************************************************************\")\n",
    "        directory_ = directory+'r%s_reco.pickle'%(args.version)\n",
    "        with open(directory_,'wb') as fw:\n",
    "            pickle.dump(eval_,fw)\n",
    "        model.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Author: Sparsh A.\n",
      "\n",
      "Last updated: 2021-11-29 07:15:56\n",
      "\n",
      "torch_geometric: 2.0.2\n",
      "\n",
      "Compiler    : GCC 7.5.0\n",
      "OS          : Linux\n",
      "Release     : 5.4.104+\n",
      "Machine     : x86_64\n",
      "Processor   : x86_64\n",
      "CPU cores   : 2\n",
      "Architecture: 64bit\n",
      "\n",
      "pandas : 1.1.5\n",
      "torch  : 1.10.0+cu111\n",
      "numpy  : 1.19.5\n",
      "IPython: 5.5.0\n",
      "\n"
     ]
    }
   ],
   "source": [
    "!pip install -q watermark\n",
    "%reload_ext watermark\n",
    "%watermark -a \"Sparsh A.\" -m -iv -u -t -d -p torch_geometric"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**END**"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
