# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# +
# default_exp model
# -

# # DGL Model

#hide
from nbdev.showdoc import *

#export
import dgl.function as fn
from dgl.nn.pytorch import GraphConv
import torch
import torch.nn as nn
import torch.nn.functional as F

#export
gcn_msg = fn.copy_src(src='h', out='m')
gcn_reduce = fn.sum(msg='m', out='h')  # sum aggregation


#export
class GCN(nn.Module):
    def __init__(self, in_feats, out_feats, activation):
        super().__init__()
        self.linear = nn.Linear(in_feats, out_feats)
        self.activation = activation
        torch.nn.init.xavier_uniform_(self.linear.weight)
     
    def forward(self, g, feature):
        with g.local_scope():
            feature = feature * g.ndata['norm']
            g.ndata['h'] = feature
            g.update_all(gcn_msg, gcn_reduce)
            h = g.ndata['h']
            h = self.linear(h)
            h = h * g.ndata['norm']
            return self.activation(h)


#export
class InnerProductDecoder(nn.Module):
    def __init__(self, activation=torch.sigmoid, dropout=0.1):
        super().__init__()
        self.dropout = dropout
        self.activation = activation

    def forward(self, z):
        z = F.dropout(z, self.dropout)
        adj = self.activation(torch.matmul(z, z.t()))
        return adj


#export
class GAEN(nn.Module):
    def __init__(self, in_dim, hidden_dims, dropout=0.1, decode_activation=torch.sigmoid):
        super().__init__()
        if len(hidden_dims)>=2:
            layers = [GCN(in_dim, hidden_dims[0], F.relu)]
            for i in range(1,len(hidden_dims)):
                if i != len(hidden_dims)-1:
                    layers.append(GCN(hidden_dims[i-1], hidden_dims[i], F.relu))
                else:
                    layers.append(GCN(hidden_dims[i-1], hidden_dims[i], lambda x:x))
        else:
            layers = [GCN(in_dim, hidden_dims[0], lambda x:x)]
        self.layers = nn.ModuleList(layers)
        self.decoder = InnerProductDecoder(activation=decode_activation, dropout=dropout)
    
    def forward(self, g):
        h = self.encode(g)
        adj_rec = self.decoder(h)
        return adj_rec

    def encode(self, g):
        h = g.ndata['h']
        for conv in self.layers:
            h = conv(g, h)
        return h


#export
class GAE(nn.Module):
    def __init__(self, in_dim, hidden_dims, dropout=0.1, decode_activation=torch.sigmoid):
        super().__init__()
        if len(hidden_dims)>=2:
            layers = [GraphConv(in_dim, hidden_dims[0], activation=F.relu)]
            for i in range(1,len(hidden_dims)):
                if i != len(hidden_dims)-1:
                    layers.append(GraphConv(hidden_dims[i-1], hidden_dims[i], activation=F.relu))
                else:
                    layers.append(GraphConv(hidden_dims[i-1], hidden_dims[i], activation=lambda x: x))
        else:
            layers = [GraphConv(in_dim, hidden_dims[0], activation=lambda x: x)]
        self.layers = nn.ModuleList(layers)
        self.decoder = InnerProductDecoder(activation=decode_activation, dropout=dropout)
    
    def forward(self, g):
        h = self.encode(g)
        adj_rec = self.decoder(h)
        return adj_rec

    def encode(self, g):
        with g.local_scope():
            h = g.ndata['h']
            for conv in self.layers:
                h = conv(g, h)
            return h

StarcoderdataJupyterScriptsDedupFiltered