Model_I_0820_int.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_cell_guid": "b1076dfc-b9ad-4769-8c92-a6c4dae69d19",
    "_uuid": "8f2839f25d086af736a60e9eeb907d3b93b6e0e5"
   },
   "outputs": [],
   "source": [
    "class GNN_multiHead_noEdge_Int(torch.nn.Module):\n",
    "    # no edge update\n",
    "    def __init__(self,reuse,block,head,head_mol,head_atom,head_edge,dim,layer1,layer2,factor,\\\n",
    "                 node_in,edge_in,edge_in4,edge_in3=8,mol_shape=4,atom_shape=10,edge_shape=4,aggr='mean',interleave=False):\n",
    "        # block,head are nn.Module\n",
    "        # node_in,edge_in are dim for bonding and edge_in4,edge_in3 for coupling\n",
    "        super(GNN_multiHead_noEdge_Int, self).__init__()\n",
    "        if interleave:\n",
    "            assert layer1==layer2,'layer1 needs to be the same as layer2'\n",
    "        self.interleave = interleave\n",
    "        self.lin_node = Sequential(BatchNorm1d(node_in),Linear(node_in, dim*factor),LeakyReLU(), \\\n",
    "                                   BatchNorm1d(dim*factor),Linear(dim*factor, dim),LeakyReLU())\n",
    "\n",
    "        self.edge1 = Sequential(BatchNorm1d(edge_in),Linear(edge_in, dim*factor),LeakyReLU(), \\\n",
    "                                   BatchNorm1d(dim*factor),Linear(dim*factor, dim),LeakyReLU())\n",
    "\n",
    "        self.edge2 = Sequential(BatchNorm1d(edge_in4+edge_in3),Linear(edge_in4+edge_in3, dim*factor),LeakyReLU(), \\\n",
    "                                   BatchNorm1d(dim*factor),Linear(dim*factor, dim),LeakyReLU())        \n",
    "        if reuse:\n",
    "            self.conv1 = schnet_block(dim=dim,edge_dim=dim,aggr=aggr)\n",
    "            self.conv2 = block(dim=dim,edge_dim=dim,aggr=aggr)\n",
    "        else:\n",
    "            self.conv1 = nn.ModuleList([schnet_block(dim=dim,edge_dim=dim,aggr=aggr) for _ in range(layer1)])\n",
    "            self.conv2 = nn.ModuleList([block(dim=dim,edge_dim=dim,aggr=aggr) for _ in range(layer2)])            \n",
    "        \n",
    "        self.head = head(dim)\n",
    "        self.head_mol = head_mol(dim,mol_shape)\n",
    "        self.head_atom = head_atom(dim,atom_shape)\n",
    "        self.head_edge = head_edge(dim,edge_shape)\n",
    "        \n",
    "    def forward(self, data,IsTrain=False,typeTrain=False,logLoss=True,weight=None):\n",
    "        out = self.lin_node(data.x)\n",
    "        # edge_*3 only does not repeat for undirected graph. Hence need to add (j,i) to (i,j) in edges\n",
    "        edge_index3 = torch.cat([data.edge_index3,data.edge_index3[[1,0]]],1)\n",
    "        n = data.edge_attr3.shape[0]\n",
    "        temp_ = self.edge2(torch.cat([data.edge_attr3,data.edge_attr4],1))\n",
    "        edge_attr3 = torch.cat([temp_,temp_],0)\n",
    "        int_types = torch.cat([data.edge_attr3,data.edge_attr3],0)\n",
    "        edge_attr = self.edge1(data.edge_attr)\n",
    "        \n",
    "        if self.interleave:\n",
    "            for conv1,conv2 in zip(self.conv1,self.conv2):\n",
    "                out = conv1(out,data.edge_index,edge_attr)\n",
    "                out = conv2(out,edge_index3,edge_attr3,int_types)\n",
    "        else:\n",
    "            for conv in self.conv1:\n",
    "                out = conv(out,data.edge_index,edge_attr)\n",
    "            for conv in self.conv2:\n",
    "                out = conv(out,edge_index3,edge_attr3,int_types)    \n",
    "        \n",
    "        edge_attr3 = edge_attr3[:n]\n",
    "        if typeTrain:\n",
    "            if IsTrain:\n",
    "                y = data.y[data.type_attr]\n",
    "            edge_attr3 = edge_attr3[data.type_attr]\n",
    "            edge_index3 = data.edge_index3[:,data.type_attr]\n",
    "            edge_attr3_old = data.edge_attr3[data.type_attr]\n",
    "        else:\n",
    "            if IsTrain:\n",
    "                y = data.y\n",
    "            edge_index3 = data.edge_index3\n",
    "            edge_attr3_old = data.edge_attr3\n",
    "            \n",
    "        yhat = self.head(out,edge_index3,edge_attr3,edge_attr3_old)\n",
    "        \n",
    "        if IsTrain:\n",
    "            if weight is None:\n",
    "                loss_other = 0\n",
    "            else:\n",
    "                y_mol = self.head_mol(out,data.batch)\n",
    "                y_atom = self.head_atom(out)\n",
    "                y_edge = self.head_edge(out,edge_index3)\n",
    "                loss_other = weight * (torch.mean(torch.abs(data.y_mol - y_mol)) + \\\n",
    "                                       torch.mean(torch.abs(data.y_atom - y_atom)) + \\\n",
    "                                       torch.mean(torch.abs(data.y_coupling - y_edge)))\n",
    "\n",
    "            k = torch.sum(edge_attr3_old,0)\n",
    "            nonzeroIndex = torch.nonzero(k).squeeze(1)\n",
    "            abs_ = torch.abs(y-yhat).unsqueeze(1)\n",
    "            loss_perType = torch.zeros(8,device='cuda:0')\n",
    "            if logLoss:\n",
    "                loss_perType[nonzeroIndex] = torch.log(torch.sum(abs_ * edge_attr3_old[:,nonzeroIndex],0)/k[nonzeroIndex])\n",
    "                loss = torch.sum(loss_perType)/nonzeroIndex.shape[0]\n",
    "                return loss+loss_other,loss_perType         \n",
    "            else:\n",
    "                loss_perType[nonzeroIndex] = torch.sum(abs_ * edge_attr3_old[:,nonzeroIndex],0)/k[nonzeroIndex]\n",
    "                loss = torch.sum(loss_perType)/nonzeroIndex.shape[0]\n",
    "                loss_perType[nonzeroIndex] = torch.log(loss_perType[nonzeroIndex])\n",
    "                return loss+loss_other,loss_perType\n",
    "        else:\n",
    "            return yhat\n",
    "        \n",
    "\n",
    "class NNConv2_int(MessagePassing):\n",
    "    r\"\"\" use element-wise multiplication as in schnet instead of matrix multiplication\n",
    "    \"\"\"\n",
    "    def __init__(self,\n",
    "                 dim,\n",
    "                 nn,\n",
    "                 aggr='mean'):\n",
    "        super(NNConv2_int, self).__init__(aggr=aggr)\n",
    "        cat_factor = 2\n",
    "        multiple_factor = 3\n",
    "        self.dim = dim\n",
    "        self.type_factor = 8\n",
    "        self.nn = nn\n",
    "        self.aggr = aggr\n",
    "        self.v_update = Sequential(BatchNorm1d(dim*cat_factor),\n",
    "                                    Linear(dim*cat_factor,dim*cat_factor*multiple_factor),\n",
    "                                    LeakyReLU(inplace=True),\n",
    "                                    BatchNorm1d(dim*cat_factor*multiple_factor),\n",
    "                                    Linear(dim*cat_factor*multiple_factor,dim),\n",
    "                                    LeakyReLU(inplace=True))\n",
    "        \n",
    "    def forward(self, x, edge_index, edge_attr,int_types):\n",
    "        x = x.unsqueeze(-1) if x.dim() == 1 else x\n",
    "        pseudo = edge_attr.unsqueeze(-1) if edge_attr.dim() == 1 else edge_attr\n",
    "        return self.propagate(edge_index, x=x, pseudo=pseudo,int_types=int_types)\n",
    "\n",
    "    def message(self, x_j, pseudo,int_types):\n",
    "        weight = self.nn(pseudo).reshape(-1,self.dim,self.type_factor) # (n,d,k)\n",
    "        int_types=int_types.unsqueeze(1).to(torch.bool)\n",
    "        _,int_types = torch.broadcast_tensors(weight,int_types)\n",
    "        weight = weight[int_types].reshape(-1,self.dim)\n",
    "        return x_j * weight\n",
    "\n",
    "    def update(self, aggr_out, x):\n",
    "        return self.v_update(torch.cat([aggr_out,x],1))\n",
    "\n",
    "    def __repr__(self):\n",
    "        return 'NNConv2_int'   \n",
    "\n",
    "class schnet_block_int(torch.nn.Module):\n",
    "    # use both types of edges\n",
    "    def __init__(self,dim=64,edge_dim=12,aggr='mean'):\n",
    "        super(schnet_block_int, self).__init__()\n",
    "        multiple_factor = 3\n",
    "        type_factor = 8\n",
    "        nn = Sequential(BatchNorm1d(edge_dim),Linear(edge_dim, dim*multiple_factor),LeakyReLU(inplace=True), \\\n",
    "                        BatchNorm1d(dim*multiple_factor),Linear(dim*multiple_factor, dim*type_factor))\n",
    "        self.conv = NNConv2_int(dim, nn, aggr=aggr)\n",
    "        self.lin_covert = Sequential(BatchNorm1d(dim),Linear(dim, dim*multiple_factor),LeakyReLU(inplace=True), \\\n",
    "                                     BatchNorm1d(dim*multiple_factor),Linear(dim*multiple_factor, dim))\n",
    "        \n",
    "    def forward(self, x, edge_index, edge_attr,int_types):\n",
    "        m = self.conv(x, edge_index, edge_attr,int_types)\n",
    "        m = self.lin_covert(m)\n",
    "        return x + m"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "import torch\n",
    "from torch_geometric.data import Data,DataLoader\n",
    "from functions_refactor import *\n",
    "from pytorch_util import *\n",
    "#from torch.optim import Adam\n",
    "from torch.optim.lr_scheduler import ReduceLROnPlateau"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "_cell_guid": "79c7e3d0-c299-4dcb-8224-4455121ee9b0",
    "_uuid": "d629ff2d2480ee46fbb7e2d37f6b5fab8052498a"
   },
   "outputs": [],
   "source": [
    "# fixed parameters\n",
    "head_mol,head_atom,head_edge = head_mol,head_atom,head_edge2\n",
    "clip = 2\n",
    "batch_size = 32\n",
    "threshold = -1.3\n",
    "reuse = False\n",
    "lr = 1e-4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# changing parameters\n",
    "block = schnet_block_int\n",
    "head = cat3HeadInteraction_noEdge\n",
    "data = '../Data/{}_data_ACSF_SOAP_atomInfo_otherInfo.pickle'\n",
    "dim = 512\n",
    "logLoss = True\n",
    "weight = None\n",
    "layer1 = 3\n",
    "layer2 = 4\n",
    "factor = 2\n",
    "epochs = 150\n",
    "aggr = 'mean'\n",
    "interleave = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "prefix = '_'.join([str(i).split('}')[1] if '}' in str(i) else str(i) \\\n",
    "                                        for i in [block,head,data,dim,logLoss,weight,layer1,layer2,factor,epochs,aggr,interleave]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_df = pd.read_csv('../Data/train.csv')\n",
    "test_df = pd.read_csv('../Data/test.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "folds = []\n",
    "for f in range(5):\n",
    "    with open(data.format('train').split('pickle')[0][:-1]+'_f'+str(f)+'.pickle', 'rb') as handle:\n",
    "        folds.append(pickle.load(handle))\n",
    "folds = [[Data(**d) for d in fold] for fold in folds]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(5):\n",
    "    print('\\nstart fold '+str(i))\n",
    "    # parpare data\n",
    "    train_list = []\n",
    "    val_list = []\n",
    "    for j in range(5):\n",
    "        if i == j:\n",
    "            val_list.extend(folds[j])\n",
    "        else:\n",
    "            train_list.extend(folds[j])\n",
    "    \n",
    "    train_dl = DataLoader(train_list,batch_size,shuffle=True)\n",
    "    val_dl = DataLoader(val_list,batch_size,shuffle=False)\n",
    "    \n",
    "    # train model\n",
    "    model = GNN_multiHead_noEdge_Int(reuse,block,head,head_mol,head_atom,head_edge,\\\n",
    "                          dim,layer1,layer2,factor,**data_dict[data],aggr=aggr,interleave=interleave).to('cuda:0')    \n",
    "    paras = trainable_parameter(model)\n",
    "    opt = RAdam(paras,lr=lr,weight_decay=1e-2)\n",
    "    scheduler = ReduceLROnPlateau(opt, 'min',factor=0.5,patience=5,min_lr=1e-05)\n",
    "    model,train_loss_perType,val_loss_perType,bestWeight = train_type_earlyStop(opt,model,epochs,train_dl,val_dl,paras,clip,\\\n",
    "                                                                    scheduler=scheduler,logLoss=logLoss,weight=weight,threshold=threshold)\n",
    "    torch.save({'model_state_dict_type_'+str(j_):w for j_,w in enumerate(bestWeight)},\\\n",
    "                '../Model/'+prefix+'_fold'+str(i)+'.tar')\n",
    "    # predict oof for each type\n",
    "    for type_i in range(8):\n",
    "        # load val data and type_id\n",
    "        with open(data.format('train').split('pickle')[0][:-1]+'_f'+str(i)+'_type_'+str(type_i)+'.pickle', 'rb') as handle:\n",
    "            test_data = pickle.load(handle)\n",
    "        test_list = [Data(**d) for d in test_data]\n",
    "        test_dl = DataLoader(test_list,batch_size,shuffle=False)\n",
    "        \n",
    "        with open(data.format('train').split('pickle')[0][:-1]+'_f'+str(i)+'_type_'+str(type_i)+'_id.pickle', 'rb') as handle:\n",
    "            test_id = pickle.load(handle)\n",
    "    \n",
    "        # load model\n",
    "        model.load_state_dict(bestWeight[type_i])\n",
    "    \n",
    "        # predict\n",
    "        model.eval()\n",
    "        yhat_list = []\n",
    "        with torch.no_grad():\n",
    "            for data_torch in test_dl:\n",
    "                data_torch = data_torch.to('cuda:0')\n",
    "                yhat_list.append(model(data_torch,False,True))\n",
    "        yhat = torch.cat(yhat_list).cpu().detach().numpy()        \n",
    "    \n",
    "        # join\n",
    "        assert yhat.shape[0]==test_id.shape[0],'yhat and test_id should have same shape'\n",
    "        submit_ = dict(zip(test_id,yhat))\n",
    "        train_df['fold'+str(i)+'_type'+str(type_i)] = train_df.id.map(submit_)\n",
    "    \n",
    "    # predict test\n",
    "    for type_i in range(8):\n",
    "        # load val data and type_id\n",
    "        with open(data.format('test').split('pickle')[0][:-1]+'_type_'+str(type_i)+'.pickle', 'rb') as handle:\n",
    "            test_data = pickle.load(handle)\n",
    "        test_list = [Data(**d) for d in test_data]\n",
    "        test_dl = DataLoader(test_list,batch_size,shuffle=False)\n",
    "        \n",
    "        with open(data.format('test').split('pickle')[0][:-1]+'_id_type_'+str(type_i)+'.pickle', 'rb') as handle:\n",
    "            test_id = pickle.load(handle)\n",
    "    \n",
    "        # load model\n",
    "        model.load_state_dict(bestWeight[type_i])\n",
    "    \n",
    "        # predict\n",
    "        model.eval()\n",
    "        yhat_list = []\n",
    "        with torch.no_grad():\n",
    "            for data_torch in test_dl:\n",
    "                data_torch = data_torch.to('cuda:0')\n",
    "                yhat_list.append(model(data_torch,False,True))\n",
    "        yhat = torch.cat(yhat_list).cpu().detach().numpy()        \n",
    "    \n",
    "        # join\n",
    "        assert yhat.shape[0]==test_id.shape[0],'yhat and test_id should have same shape'\n",
    "        submit_ = dict(zip(test_id,yhat))\n",
    "        test_df['fold'+str(i)+'_type'+str(type_i)] = test_df.id.map(submit_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#assert set(test.iloc[:,5:].isnull().sum(1)) == set([7*5])\n",
    "test_df['yhat'] = np.nanmean(test_df.iloc[:,5:],1)\n",
    "#test = test[['id','yhat']]\n",
    "test_df.to_csv('../Data/test_oof_0820_'+prefix,index=False)\n",
    "\n",
    "#assert set(train.iloc[:,6:].isnull().sum(1)) == set([train.iloc[:,6:].shape[1]-1])\n",
    "train_df['yhat'] = np.nanmean(train_df.iloc[:,6:],1)\n",
    "#train = train[['id','yhat']]\n",
    "train_df.to_csv('../Data/train_oof_0820_'+prefix,index=False)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}