Official Code for FedRule: Federated Rule Recommendation System with Graph Neural Networks
#!/usr/bin/env python
# coding: utf-8

from huggingface_hub import login
# enter your token of huggingface
# https://huggingface.co/docs/hub/security-tokens
login(token="XXXXXXXXXXXXXX")

from datasets import load_dataset
import collections
import torch
import torch.nn.functional as F
import dgl
import json

train_device = load_dataset("wyzelabs/RuleRecommendation", data_files="train_device.csv")
train_rule = load_dataset("wyzelabs/RuleRecommendation", data_files="train_rule.csv")
test_device = load_dataset("wyzelabs/RuleRecommendation", data_files="test_device.csv")
test_rule = load_dataset("wyzelabs/RuleRecommendation", data_files="test_rule.csv")

print("start preprocessing")
df = train_rule['train'].to_pandas()
all_devices = {'Camera': 0,
                 'ClimateSensor': 1,
                 'Cloud': 2,
                 'ContactSensor': 3,
                 'Irrigation': 4,
                 'LeakSensor': 5,
                 'Light': 6,
                 'LightStrip': 7,
                 'Lock': 8,
                 'MeshLight': 9,
                 'MotionSensor': 10,
                 'OutdoorPlug': 11,
                 'Plug': 12,
                 'RobotVacuum': 13,
                 'Switch': 14,
                 'Thermostat': 15}

all_trigger_actions = collections.defaultdict(int)
for index, row in df.iterrows():
                all_trigger_actions[str(row['trigger_state_id'])+' '+str(row['action_id'])] += 1

#filter some action 
for index in list(all_trigger_actions.keys()):
        if all_trigger_actions[index] <= 10:
            all_trigger_actions.pop(index)


for index, key in enumerate(all_trigger_actions):
            all_trigger_actions[key] = index

class_num = len(set(all_devices.values()))
one_hot_matrix = F.one_hot(torch.tensor(list(range(class_num))), num_classes=class_num)

user_device_id_to_node_id = collections.defaultdict(dict)
user_number_of_devices = collections.defaultdict(int)

user_graphs = collections.defaultdict(dgl.DGLGraph)

count = 0
for index, row in df.iterrows():
    trigger_action = str(row['trigger_state_id'])+' '+str(row['action_id'])
    if trigger_action in all_trigger_actions and row['trigger_device'] in all_devices and row['action_device'] in all_devices:

        if row['trigger_device_id'] not in user_device_id_to_node_id[row["user_id"]]:
            #assign id to the current device for supporting multiple devices with the same type
            user_device_id_to_node_id[row["user_id"]][row['trigger_device_id']] = user_number_of_devices[row["user_id"]]
            user_number_of_devices[row["user_id"]] += 1

            device = all_devices[row['trigger_device']]
            user_graphs[row["user_id"]].add_nodes(1, data = {'feat':one_hot_matrix[device].reshape(1,-1)})

        if row['action_device_id'] not in user_device_id_to_node_id[row["user_id"]]:
            user_device_id_to_node_id[row["user_id"]][row['action_device_id']] = user_number_of_devices[row["user_id"]]
            user_number_of_devices[row["user_id"]] += 1

            device = all_devices[row['action_device']]
            user_graphs[row["user_id"]].add_nodes(1, data = {'feat':one_hot_matrix[device].reshape(1,-1)})
        node1 = user_device_id_to_node_id[row["user_id"]][row['trigger_device_id']]
        node2 = user_device_id_to_node_id[row["user_id"]][row['action_device_id']]
        #the file contains same rules but with different devices
        user_graphs[row["user_id"]].add_edges(node1, node2, data = {'etype':torch.tensor([all_trigger_actions[trigger_action]])}) #directed
    count += 1
    if count > 1000:
        break

#filter, remove graph with devices < 3
for i in list(user_graphs.keys()):
    if user_graphs[i].number_of_nodes() <= 2:
        user_graphs.pop(i)

print("save processed data")

user_id_dict = dict()
for i in user_graphs.keys():
    user_id_dict[i] = i

with open('user_id_dict.json', 'w') as f:
    json.dump(user_id_dict, f)

dgl.save_graphs("usergraphs.bin", list(user_graphs.values()))

with open('all_trigger_actions.json', 'w') as f:
    json.dump(all_trigger_actions, f)
    
with open('all_devices.json', 'w') as f:
    json.dump(all_devices, f)
    
with open('user_device_id_to_node_id.json', 'w') as f:
    json.dump(user_device_id_to_node_id, f)


data_preprocess.py
yh-yao/FedRule
import numpy as np
import dgl
import torch
import torch.nn.functional as F

def split_data(user_graphs, all_trigger_actions, more_neg = False):
    train_gs = dict()
    train_pos_gs = dict()
    train_neg_gs = dict()
    test_pos_gs = dict()
    test_neg_gs = dict()

    for user_index in user_graphs:
        g = user_graphs[user_index]

        # Split edge set for training and testing
        u, v = g.edges()

        # eids = np.arange(g.number_of_edges())
        eids = torch.randperm(g.number_of_edges()).to(g.device)
        test_size = int(len(eids) * 0.2) 
        train_size = g.number_of_edges() - test_size
        test_pos_u, test_pos_v = u[eids[:test_size]], v[eids[:test_size]]
        train_pos_u, train_pos_v = u[eids[test_size:]], v[eids[test_size:]]

        test_pos_edata = g.edata['etype'][eids[:test_size]]
        train_pos_edata = g.edata['etype'][eids[test_size:]]

        if len(test_pos_u) == 0:
            continue #remove some graph with few edges

        # Find all negative edges and split them for training and testing
        adj = g.adjacency_matrix().to(g.device)
        adj_neg = 1 - adj.to_dense() - torch.eye(g.number_of_nodes()).to(g.device) #no self loop, filtered before
        neg_inds = torch.nonzero(adj_neg)
        neg_edata = torch.randint(len(set(all_trigger_actions.values())), (neg_inds.shape[0],)).to(g.device)
        
        
        #add the same edge in dataset but different edge type for negative sampleing
        neg_u_different_type = u
        neg_v_different_type = v
        #neg_sampling
        if more_neg == True:
            for i in range(10):
                neg_u_different_type = torch.cat((neg_u_different_type, u), 0)
                neg_v_different_type = torch.cat((neg_v_different_type, v), 0)

        neg_edge_fetures = torch.randint(len(set(all_trigger_actions.values())), (len(neg_u_different_type),)).to(g.device)

        for i in range(len(u)):
            same_edges_with_different_type = g.edata['etype'][g.edge_ids(u[i],v[i], return_uv = True)[2]]
            while neg_edge_fetures[i] in same_edges_with_different_type:
                neg_edge_fetures[i] = np.random.choice(len(set(all_trigger_actions.values())), 1)[0]
                
        neg_u = torch.cat((neg_inds[:,0], neg_u_different_type), 0)
        neg_v = torch.cat((neg_inds[:,1], neg_v_different_type), 0)
        neg_edata = torch.cat((neg_edata, neg_edge_fetures), 0)
        
        #print(len(neg_u), len(neg_edata))
        
        if len(neg_u) == 0:
            continue #some graphs are too small, become complete graphs, skip it

        #neg_eids = list(range(len(neg_u)))#np.random.choice(len(neg_u), len(neg_u))   #### super negative sampling, many edges
        
        neg_eids = torch.randperm(len(neg_u))[:g.number_of_edges()]
        
        test_neg_edata = neg_edata[neg_eids[:test_size]]
        train_neg_edata = neg_edata[neg_eids[test_size:]]

        test_neg_u, test_neg_v = neg_u[neg_eids[:test_size]], neg_v[neg_eids[:test_size]]
        train_neg_u, train_neg_v = neg_u[neg_eids[test_size:]], neg_v[neg_eids[test_size:]]

        train_pos_g = dgl.graph((train_pos_u, train_pos_v), num_nodes=g.number_of_nodes())
        train_pos_g.edata['etype'] = train_pos_edata

        train_neg_g = dgl.graph((train_neg_u, train_neg_v), num_nodes=g.number_of_nodes())
        train_neg_g.edata['etype'] = train_neg_edata

        test_pos_g = dgl.graph((test_pos_u, test_pos_v), num_nodes=g.number_of_nodes())
        test_pos_g.edata['etype'] = test_pos_edata

        test_neg_g = dgl.graph((test_neg_u, test_neg_v), num_nodes=g.number_of_nodes())
        test_neg_g.edata['etype'] = test_neg_edata

        train_g = dgl.remove_edges(g, eids[:test_size])

        train_g.ndata['feat'] = train_g.ndata['feat'].float()
        
        
        #add efeat
        edge_types = len(set(all_trigger_actions.values()))
        one_hot_matrix = F.one_hot(torch.tensor(list(range(edge_types))), num_classes=edge_types).to(g.device)
        train_g.edata['efeat'] = one_hot_matrix[train_g.edata['etype']]
        train_pos_g.edata['efeat'] = one_hot_matrix[train_pos_g.edata['etype']]
        train_neg_g.edata['efeat'] = one_hot_matrix[train_neg_g.edata['etype']]
        test_pos_g.edata['efeat'] = one_hot_matrix[test_pos_g.edata['etype']]
        test_neg_g.edata['efeat'] = one_hot_matrix[test_neg_g.edata['etype']]


        train_gs[user_index] = train_g
        train_pos_gs[user_index] = train_pos_g
        train_neg_gs[user_index] = train_neg_g
        test_pos_gs[user_index] = test_pos_g
        test_neg_gs[user_index] = test_neg_g
    return train_gs, train_pos_gs, train_neg_gs, test_pos_gs, test_neg_gs

data_splitter.py
yh-yao/FedRule
#!/usr/bin/env python
# coding: utf-8

import os
import yaml
from pathlib import Path
import torch
import argparse
import dgl
import json
import itertools
import numpy as np
from torch.utils.tensorboard import SummaryWriter

from data_splitter import split_data
from utils import compute_loss, increment_dir, test_global_model, get_recommendation_result
from model import GraphSAGE, GCN, GAT, HeteroMLPPredictor


if __name__=="__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('-d', '--dataset', default='wyze', type=str)
    parser.add_argument('-l', '--logdir', default='./runs', type=str)
    parser.add_argument('-lr', '--learning_rate', default=0.1, type=float)
    parser.add_argument('-c', '--num_comms',default=100, type=int) # num_iterations in centralized training
    parser.add_argument('-m', '--model_type', default='graphsage', type=str)
    parser.add_argument('-neg', '--more_negative', action='store_true')

    seed = 0
    dgl.seed(seed)
    torch.manual_seed(seed)
    np.random.seed(seed)
    
    args = parser.parse_args()

    args.log_dir =  increment_dir(Path(args.logdir) / 'exp')
    args.log_dir += args.dataset + "_" + 'center'
    os.makedirs(args.log_dir)
    yaml_file = str(Path(args.log_dir) / "args.yaml")
    with open(yaml_file, 'w') as out:
        yaml.dump(args.__dict__, out, default_flow_style=False)
    args.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    
    tb_writer = SummaryWriter(log_dir=args.log_dir)
    
    print("load data")
    user_graphs_list, _ = dgl.load_graphs("usergraphs.bin")
    user_id_dict = json.load(open('user_id_dict.json', 'r'))
    user_graphs = dict()
    user_ids = list(user_id_dict.keys())

    for i in range(len(user_ids)):
        user_graphs[user_ids[i]] = user_graphs_list[i]

    all_trigger_actions = json.load(open('all_trigger_actions.json', 'r'))
    all_devices = json.load(open('all_devices.json', 'r'))
    user_device_id_to_node_id = json.load(open('user_device_id_to_node_id.json', 'r'))
    
    train_gs, train_pos_gs, train_neg_gs, test_pos_gs, test_neg_gs = split_data(user_graphs, all_trigger_actions, args.more_negative)
    if args.model_type == 'graphsage':
        model = GraphSAGE(len(set(all_devices.values())), 32).to(args.device) #feature dim: len(set(all_devices.values()))
    elif args.model_type == 'gcn':
        model = GCN(len(set(all_devices.values())), 32).to(args.device) 
    elif args.model_type == 'gat':
        model = GAT(len(set(all_devices.values())), 32).to(args.device) 
        
    pred = HeteroMLPPredictor(32, len(set(all_trigger_actions.values()))).to(args.device)
    optimizer = torch.optim.Adam(itertools.chain(model.parameters(), pred.parameters()), lr=args.learning_rate)

    print("start training")
    for e in range(args.num_comms):
        model.train()
        pred.train()
        loss = None
        for user_index in train_gs:
            train_g = train_gs[user_index]
            train_pos_g = train_pos_gs[user_index]
            train_neg_g = train_neg_gs[user_index]
            
            h = model(train_g, train_g.ndata['feat'])
            pos_score = pred(train_pos_g, h)[list(range(len(train_pos_g.edata['etype']))), train_pos_g.edata['etype']]
            neg_score = pred(train_neg_g, h)[list(range(len(train_neg_g.edata['etype']))), train_neg_g.edata['etype']]
            if loss == None:
                loss = compute_loss(pos_score, neg_score)
            else:
                loss += compute_loss(pos_score, neg_score)
            
        tb_writer.add_scalar('Train/Loss',loss.item() / len(train_gs),e-1)  #-1 since it is before backward
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        print('In epoch {}, loss: {}'.format(e-1, loss.item() / len(train_gs)))
        
        if (e + 1) % 5 == 0:
            global_train_loss, global_train_AUC, global_train_MR, global_test_loss, global_test_AUC, global_test_MR = test_global_model(train_gs, train_pos_gs, train_neg_gs, test_pos_gs, test_neg_gs, model, pred)

            tb_writer.add_scalar('Global Train/Loss', global_train_loss, e)
            tb_writer.add_scalar('Global Train/AUC', global_train_AUC, e)
            tb_writer.add_scalar('Global Train/POS_MR', global_train_MR, e)

            tb_writer.add_scalar('Global Test/Loss', global_test_loss, e)
            tb_writer.add_scalar('Global Test/AUC', global_test_AUC, e)
            tb_writer.add_scalar('Global Test/POS_MR', global_test_MR, e)
            
    torch.save(model.state_dict(), args.dataset + "central_model_" + args.model_type)
    torch.save(pred.state_dict(), args.dataset + "central_pred_" + args.model_type)
main.py
yh-yao/FedRule
import torch
import torch.nn as nn
import torch.nn.functional as F
from dgl.nn import SAGEConv, GraphConv, GATConv
import dgl.function as fn

class GraphSAGE(nn.Module):
    def __init__(self, in_feats, h_feats, dropout = 0.5):
        super(GraphSAGE, self).__init__()
        self.conv1 = SAGEConv(in_feats, h_feats, 'mean')
        self.dropout = nn.Dropout(dropout)
        self.conv2 = SAGEConv(h_feats, h_feats, 'mean')

    def forward(self, g, in_feat):
        h = self.conv1(g, in_feat)
        h = self.dropout(h) #dropout before relu
        h = F.relu(h)
        h = self.conv2(g, h)
        return h

class GCN(nn.Module):
    def __init__(self, in_feats, h_feats, dropout = 0.5):
        super(GCN, self).__init__()
        self.conv1 = GraphConv(in_feats, h_feats, allow_zero_in_degree = True)
        self.dropout = nn.Dropout(dropout)
        self.conv2 = GraphConv(h_feats, h_feats, allow_zero_in_degree = True)

    def forward(self, g, in_feat):
        h = self.conv1(g, in_feat)
        h = self.dropout(h) #dropout before relu
        h = F.relu(h)
        h = self.conv2(g, h)
        return h

class GAT(nn.Module):
    def __init__(self, in_feats, h_feats, dropout = 0.5):
        super(GAT, self).__init__()
        self.conv1 = GATConv(in_feats, h_feats, 1, allow_zero_in_degree = True)
        self.dropout = nn.Dropout(dropout)
        self.conv2 = GATConv(h_feats, h_feats, 1, allow_zero_in_degree = True)

    def forward(self, g, in_feat):
        h = self.conv1(g, in_feat)
        h = h.reshape(h.shape[0],h.shape[2]) #1 attention head
        h = self.dropout(h) #dropout before relu
        h = F.relu(h)
        h = self.conv2(g, h)
        h = h.reshape(h.shape[0],h.shape[2]) #1 attention head
        return h
        
class MLPPredictor(nn.Module):
    def __init__(self, h_feats):
        super().__init__()
        self.W1 = nn.Linear(h_feats * 2, h_feats)
        self.W2 = nn.Linear(h_feats, 1)
        self.sig = nn.Sigmoid()

    def apply_edges(self, edges):
        h = torch.cat([edges.src['h'], edges.dst['h']], 1)
        return {'score': self.sig(self.W2(F.relu(self.W1(h)))).squeeze(1)}

    def forward(self, g, h):
        with g.local_scope():
            g.ndata['h'] = h
            g.apply_edges(self.apply_edges)
            return g.edata['score']
        
class HeteroDotProductPredictor(nn.Module):
    def forward(self, graph, h, etype):
        with graph.local_scope():
            graph.ndata['h'] = h   # assigns 'h' of all node types in one shot
            graph.apply_edges(fn.u_dot_v('h', 'h', 'score'), etype=etype)
            return graph.edges[etype].data['score']
        
        
class HeteroMLPPredictor(nn.Module):
    def __init__(self, h_feats, edge_types, dropout = 0.5):
        super().__init__()
        self.W1 = nn.Linear(h_feats * 2, h_feats)
        self.dropout = nn.Dropout(dropout)
        self.W2 = nn.Linear(h_feats, edge_types)
        self.sig = nn.Sigmoid()
    def apply_edges(self, edges):
        h = torch.cat([edges.src['h'], edges.dst['h']], 1)
        return {'score': self.sig(self.W2(F.relu(self.dropout(self.W1(h)))))} # dim: edge_types

    def forward(self, g, h):
        with g.local_scope():
            g.ndata['h'] = h
            g.apply_edges(self.apply_edges)
            return g.edata['score']
model.py
yh-yao/FedRule
#!/usr/bin/env python
# coding: utf-8

# In[1]:


from huggingface_hub import login
login(token="XXXXXXXXXXXXXX")
from datasets import load_dataset
import collections
import torch
import torch.nn.functional as F
import dgl
import json
import numpy as np
import networkx as nx
from model import GraphSAGE, GCN, GAT, HeteroMLPPredictor

test_rule = load_dataset("wyzelabs/RuleRecommendation", data_files="test_rule.csv")
test_rule_df = test_rule['train'].to_pandas()


all_trigger_actions = json.load(open('all_trigger_actions.json', 'r'))
all_devices = json.load(open('all_devices.json', 'r'))


class_num = len(set(all_devices.values()))
one_hot_matrix = F.one_hot(torch.tensor(list(range(class_num))), num_classes=class_num)

test_device_df = load_dataset("wyzelabs/RuleRecommendation", data_files="test_device.csv")['train'].to_pandas()


user_device_id_to_node_id = collections.defaultdict(dict)
user_number_of_devices = collections.defaultdict(int)
test_user_graphs = collections.defaultdict(dgl.DGLGraph)

'''
# add devices that have no rules
for index, row in test_device_df.iterrows():
    user_device_id_to_node_id[row["user_id"]][row['device_id']] = user_number_of_devices[row["user_id"]]
    user_number_of_devices[row["user_id"]] += 1
    
    device_type = all_devices[row['device_model']]
    test_user_graphs[row["user_id"]].add_nodes(1, data = {'feat':one_hot_matrix[device_type].reshape(1,-1)})
'''

for index, row in test_rule_df.iterrows():
    trigger_action = str(row['trigger_state_id'])+' '+str(row['action_id'])
    if trigger_action in all_trigger_actions and row['trigger_device'] in all_devices and row['action_device'] in all_devices:
        if row['trigger_device_id'] not in user_device_id_to_node_id[row["user_id"]]:
            #assign id to the current device for supporting multiple devices with the same type
            user_device_id_to_node_id[row["user_id"]][row['trigger_device_id']] = user_number_of_devices[row["user_id"]]
            user_number_of_devices[row["user_id"]] += 1

            device = all_devices[row['trigger_device']]
            test_user_graphs[row["user_id"]].add_nodes(1, data = {'feat':one_hot_matrix[device].reshape(1,-1)})

        if row['action_device_id'] not in user_device_id_to_node_id[row["user_id"]]:
            user_device_id_to_node_id[row["user_id"]][row['action_device_id']] = user_number_of_devices[row["user_id"]]
            user_number_of_devices[row["user_id"]] += 1

            device = all_devices[row['action_device']]
            test_user_graphs[row["user_id"]].add_nodes(1, data = {'feat':one_hot_matrix[device].reshape(1,-1)})
        node1 = user_device_id_to_node_id[row["user_id"]][row['trigger_device_id']]
        node2 = user_device_id_to_node_id[row["user_id"]][row['action_device_id']]
        #the file contains same rules but with different devices
        test_user_graphs[row["user_id"]].add_edges(node1, node2, data = {'etype':torch.tensor([all_trigger_actions[trigger_action]])}) #directed

model = GraphSAGE(len(set(all_devices.values())), 32)
pred = HeteroMLPPredictor(32, len(set(all_trigger_actions.values())))
model.load_state_dict(torch.load("wyzecentral_model_graphsage"))
pred.load_state_dict(torch.load("wyzecentral_pred_graphsage"))


def get_recommendation_result(G, model, pred, topk, rule_dict=None):
    Complete_G = nx.complete_graph(list(G.nodes()), nx.MultiDiGraph())
    Complete_G = dgl.from_networkx(Complete_G)
    Complete_G = dgl.add_self_loop(Complete_G)
    
    model.eval()
    pred.eval()
    with torch.no_grad():
        h = model(G, G.ndata['feat'])
        #predictor, use node embeddings of source node and target node as input, predict the link probability of current edge
        #need a complete graph as input
        scores = pred(Complete_G, h)
    L = []
    edges = Complete_G.edges()
    for i in range(scores.shape[0]):
        for j in range(scores.shape[1]):
            # trigger_device, action_device, trigger_action_pair, score
            L.append([int(edges[0][i]), int(edges[1][i]), j, float(scores[i][j])])

    # rule filter
    if rule_dict is not None:
        node_transfer = np.argmax(G.ndata['feat'].numpy(), axis = 1)
        for i in L:
            trigger_device_type = node_transfer[int(i[0])]
            action_device_type = node_transfer[int(i[1])]
            trigger_action_pair = int(i[2])
            rule = f"{trigger_device_type} {action_device_type} {trigger_action_pair}"
            # give more score if the rule in rule_dict
            if rule in rule_dict:
                i[3] += 1
    L = torch.tensor(sorted(L, key= lambda e:e[3], reverse = True))[:,:-1]

    return L[:topk]


train_rule = load_dataset("wyzelabs/RuleRecommendation", data_files="train_rule.csv")
train_df = train_rule['train'].to_pandas()

rule_dict = dict()
for index, row in train_df.iterrows():
    trigger_action = str(row['trigger_state_id'])+' '+str(row['action_id'])
    if trigger_action in all_trigger_actions:
        current_rule = f"{all_devices[row['trigger_device']]} {all_devices[row['action_device']]} {all_trigger_actions[trigger_action]}"
        rule_dict[current_rule] = 1


inv_all_trigger_actions_with_id = dict()
for key in all_trigger_actions:
    inv_all_trigger_actions_with_id[all_trigger_actions[key]] = key
    
a = open("recommended_result.csv", "w")
a.write("user_id,rule,rank\n")
for user_index in test_user_graphs:
    test_user_graphs[user_index].ndata['feat'] = test_user_graphs[user_index].ndata['feat'].float()
    L = get_recommendation_result(test_user_graphs[user_index], model, pred, topk = 50, rule_dict=rule_dict)
    current_user_devices = list(user_device_id_to_node_id[user_index].keys())
    for rank in range(len(L)):
        current_rule = L[rank]
        trigger_device = current_user_devices[int(current_rule[0])]
        action_device = current_user_devices[int(current_rule[1])]
        trigger_state, action = inv_all_trigger_actions_with_id[int(current_rule[2])].split(' ')
        a.write(f"{user_index},{trigger_device}_{trigger_state}_{action}_{action_device},{rank + 1}\n")
a.close()




output_result.py
yh-yao/FedRule
import torch
import torch.nn.functional as F
from sklearn.metrics import roc_auc_score
import numpy as np
import networkx as nx
import dgl
from pathlib import Path
import glob
import re

def compute_loss(pos_score, neg_score):
    scores = torch.cat([pos_score, neg_score])
    labels = torch.cat([torch.ones(pos_score.shape[0]), torch.zeros(neg_score.shape[0])]).to(pos_score.device)
    return F.binary_cross_entropy(scores, labels)

def compute_auc(pos_score, neg_score):
    scores = torch.cat([pos_score, neg_score]).cpu().numpy()
    scores = np.nan_to_num(scores, nan=0, posinf=0, neginf=0)
    #deal with overflow
    labels = torch.cat(
        [torch.ones(pos_score.shape[0]), torch.zeros(neg_score.shape[0])]).numpy()
    return roc_auc_score(labels, scores)

def compute_mr(predictions, test_g):
    _, indices = torch.sort(predictions, descending = True)
    hits = (test_g.edata['etype'].view(-1, 1) == indices).nonzero()[:,-1]
    return hits.float().mean()

def test_global_model(train_gs, train_pos_gs, train_neg_gs, test_pos_gs, test_neg_gs, global_model, global_pred):
    global_model.eval()
    global_pred.eval()
    total_loss = 0
    total_AUC = 0
    total_pos_MR = 0

    total_train_loss = 0
    total_train_AUC = 0
    total_train_pos_MR = 0
    with torch.no_grad():
        for user_index in test_pos_gs:
            train_g = train_gs[user_index]
            train_pos_g = train_pos_gs[user_index]
            train_neg_g = train_neg_gs[user_index]

            h = global_model(train_g, train_g.ndata['feat'])
            pos_score = global_pred(train_pos_g, h)[
                list(range(len(train_pos_g.edata['etype']))), train_pos_g.edata['etype']]
            neg_score = global_pred(train_neg_g, h)[
                list(range(len(train_neg_g.edata['etype']))), train_neg_g.edata['etype']]

            total_train_loss += compute_loss(pos_score, neg_score)
            total_train_AUC += compute_auc(pos_score, neg_score)
            total_train_pos_MR += compute_mr(global_pred(train_pos_g, h), train_pos_g)

            test_pos_g = test_pos_gs[user_index]
            test_neg_g = test_neg_gs[user_index]

            pos_score = global_pred(test_pos_g, h)[
                list(range(len(test_pos_g.edata['etype']))), test_pos_g.edata['etype']]

            neg_score = global_pred(test_neg_g, h)[
                list(range(len(test_neg_g.edata['etype']))), test_neg_g.edata['etype']]

            total_loss += compute_loss(pos_score, neg_score)
            total_pos_MR += compute_mr(global_pred(test_pos_g, h), test_pos_g)
            total_AUC += compute_auc(pos_score, neg_score)

    print('Global Test Loss', total_loss / len(test_pos_gs))
    print('Global Test AUC', total_AUC / len(test_pos_gs))
    print('Global Test Positive MR', float(total_pos_MR / len(test_pos_gs)))

    return float(total_train_loss / len(train_pos_gs)), total_train_AUC / len(train_pos_gs), float(
        total_train_pos_MR / len(train_pos_gs)), float(total_loss / len(test_pos_gs)), total_AUC / len(
        test_pos_gs), float(total_pos_MR / len(test_pos_gs))


def get_recommendation_result(G, model, pred, topk):
    Complete_G = nx.complete_graph(list(G.nodes()), nx.MultiDiGraph())
    Complete_G = dgl.from_networkx(Complete_G)

    model.eval()
    pred.eval()
    with torch.no_grad():
        h = model(G, G.ndata['feat'])
        # predictor, use node embeddings of source node and target node as input, predict the link probability of current edge
        # need a complete graph as input
        scores = pred(Complete_G, h)
    L = []
    edges = Complete_G.edges()
    for i in range(scores.shape[0]):
        for j in range(scores.shape[1]):
            L.append([int(edges[0][i]), int(edges[1][i]), j, float(scores[i][j])])
    L = torch.tensor(sorted(L, key=lambda e: e[3], reverse=True))[:, :-1]

    return L[:topk]

def increment_dir(dir, comment=''):
    # Increments a directory runs/exp1 --> runs/exp2_comment
    n = 0  # number
    dir = str(Path(dir))  # os-agnostic
    dirs = sorted(glob.glob(dir + '*'))  # directories
    if dirs:
        matches = [re.search(r"exp(\d+)", d) for d in dirs]
        idxs = [int(m.groups()[0]) for m in matches if m]
        if idxs:
            n = max(idxs) + 1  # increment
    return dir + str(n) + ('_' + comment if comment else '')
utils.py
yh-yao/FedRule