preprocess.py

import numpy as np
import torch
import dgl
import os
import sys
sys.path.append('../')
import utils
from scipy.spatial.distance import cosine
from sklearn.model_selection import train_test_split
from dgl.data import FraudYelpDataset, FraudAmazonDataset


def build_DGraphFin(data_path=None, save_path=None):
    # load data
    if data_path is None:
        data_path = './dataset/DGraphFin/raw/dgraphfin.npz'
    ds = np.load(data_path)
    
    x = ds['x']
    # # node features normalization
    # x = (x - x.mean(axis=0)) / x.std(axis=0)
    x = x.astype(dtype=np.float32)
    y = ds['y']  # label: 0, 1, 2, 3 for normal, fraud, background1 and background2
    y = y.astype(dtype=np.int64)
    trmask = np.zeros(x.shape[0]).astype(bool)
    trmask[ds['train_mask']] = True
    valmask = np.zeros(x.shape[0]).astype(bool)
    valmask[ds['valid_mask']] = True
    ttmask = np.zeros(x.shape[0]).astype(bool)
    ttmask[ds['test_mask']] = True
    edge_type = ds['edge_type']
    ets = ds['edge_timestamp']  # edge time stamp

    edge_idx = ds['edge_index']
    rev_edge_idx = np.vstack((edge_idx.transpose()[1], edge_idx.transpose()[0])).transpose()
    sl_idx = np.vstack((np.arange(x.shape[0]), np.arange(x.shape[0]))).transpose()  # self loop index
    sl_type = (edge_type.max() * np.ones(x.shape[0]) + 1).astype(int)  # self loop type
    sl_ets = np.zeros(x.shape[0])  # self loop time stamp

    # build bi-directional homogeneous graph
    edge_idx_set = set(map(tuple, edge_idx))
    rev_flg = np.array([False if tuple(e) in edge_idx_set else True for e in rev_edge_idx])  # non-repeated reverse edges' index
    homo_edge_idx = np.vstack((edge_idx, rev_edge_idx[rev_flg], sl_idx))
    homo_edge_type = np.hstack((edge_type, edge_type[rev_flg], sl_type))
    homo_ets = np.hstack((ets, ets[rev_flg], sl_ets))
    homo_rev = np.hstack((np.ones_like(edge_type), -1 * np.ones_like(edge_type)[rev_flg], -1 * np.ones_like(sl_type)))  # a flag to indicate whether it is added manually or not: 1 for original edges and -1 for manually-added reverse edges
    
    g = dgl.graph(tuple(homo_edge_idx.transpose()))
    g.ndata['feat'] = torch.from_numpy(x)
    g.ndata['label'] = torch.from_numpy(y)
    mask = torch.zeros(x.shape[0]).bool()
    mask[trmask] = True
    g.ndata['train_mask'] = mask
    mask = torch.zeros(x.shape[0]).bool()
    mask[valmask] = True
    g.ndata['val_mask'] = mask
    mask = torch.zeros(x.shape[0]).bool()
    mask[ttmask] = True
    g.ndata['test_mask'] = mask

    # edge features
    g.edata['ts'] = torch.from_numpy(homo_ets)
    g.edata['rev'] = torch.from_numpy(homo_rev)
    g.edata['type'] = torch.nn.functional.one_hot(torch.from_numpy(homo_edge_type))
        
    if save_path is not None:
        save_name = 'dgraphfin.bin'
        save_path = os.path.join(save_path, save_name)
        dgl.save_graphs(save_path, [g])
    
    return g

def split_DGraphFin(g, split):
    '''
    prepare data splits
    '''
    params = {'out_dim': 2}
    params['in_dim'] = g.ndata['feat'].shape[1]
    g.ndata['train_mask'] = g.ndata['train_mask'].bool()
    g.ndata['val_mask'] = g.ndata['val_mask'].bool()
    g.ndata['test_mask'] = g.ndata['test_mask'].bool()

    # add edge type information
    edgetypes = g.canonical_etypes
    for i, etype in enumerate(edgetypes):
        g.edges[etype].data['typeid'] = i * torch.ones(g.num_edges(etype)).long()

    if len(split) == 1:
        return g, params
    labels = g.ndata['label']
    idx = np.arange(len(labels))[labels <= 1]  # only include normal and fraud nodes
    if '.' in split[0]:
        trsize = float(split[0])
        valsize, testsize = list(map(float, split[1].split('_')))
        if testsize == '':
            testsize = 1 - valsize
    else:
        trsize = int(split[0])
        valsize, testsize = list(map(int, split[1].split('_')))
        if testsize == -1:
            testsize = len(idx) - trsize - valsize
    idx_tr, idx_rest, _, y_rest = train_test_split(idx, labels[idx], train_size=trsize, stratify=labels[idx])
    idx_val, idx_test, _, _ = train_test_split(idx_rest, y_rest, train_size=valsize, stratify=y_rest)
    g.ndata['train_mask'] = torch.zeros(len(labels)).bool()
    g.ndata['train_mask'][idx_tr] = True
    g.ndata['val_mask'] = torch.zeros(len(labels)).bool()
    g.ndata['val_mask'][idx_val] = True
    g.ndata['test_mask'] = torch.zeros(len(labels)).bool()
    g.ndata['test_mask'][idx_test] = True

    return g, params

def split_tfintsoc(g, split):
    '''
    prepare data splits for t-finance and t-social
    split: "X" or "X,Y_Z", X stands for the ratio of training set over the whole set, Y and Z stands for the ratio over data for evaluations (val + test)
    '''
    params = {'out_dim': 2}
    params['in_dim'] = g.ndata['feature'].shape[1]
    params['n_rel'] = 2
    g.edata['type'] = torch.zeros(g.num_edges(), 2).float()
    g.edata['type'][:, 0] = 1
    g.ndata['feat'] = g.ndata.pop('feature').float()  # rename the node feature attributes
    # normalize input features
    g.ndata['feat'] = (g.ndata['feat'] - g.ndata['feat'].mean(dim=0)) / g.ndata['feat'].std(dim=0)
    if len(split) <= 1:
        return g, params
    labels = g.ndata['label']
    idx = list(range(len(labels)))
    if '.' in split[0]:
        trsize = float(split[0])
        valsize, testsize = list(map(float, split[1].split('_')))
    else:
        trsize = int(split[0])
        valsize, testsize = list(map(int, split[1].split('_')))
        if testsize == -1:
            testsize = len(labels) - trsize - valsize
    idx_tr, idx_rest, _, y_rest = train_test_split(idx, labels, train_size=trsize, stratify=labels)
    idx_val, idx_test, _, _ = train_test_split(idx_rest, y_rest, train_size=valsize, stratify=y_rest)
    g.ndata['train_mask'] = torch.zeros(len(labels)).bool()
    g.ndata['train_mask'][idx_tr] = True
    g.ndata['val_mask'] = torch.zeros(len(labels)).bool()
    g.ndata['val_mask'][idx_val] = True
    g.ndata['test_mask'] = torch.zeros(len(labels)).bool()
    g.ndata['test_mask'][idx_test] = True

    return g, params  

def split_yelpamz(g, split, normfeat=True):
    '''
    prepare data splits for yelp and amazon
    split: "X" or "X,Y_Z", X stands for the ratio of training set over the whole set, Y and Z stands for the ratio over data for evaluations (val + test)
    '''
    params = {'out_dim': 2}
    params['in_dim'] = g.ndata['feature'].shape[1]
    g.ndata['feat'] = g.ndata.pop('feature').float()  # rename the node feature attributes

    if normfeat:
        # normalize input features
        g.ndata['feat'] = (g.ndata['feat'] - g.ndata['feat'].min(dim=0)[0]) / (g.ndata['feat'].max(dim=0)[0] - g.ndata['feat'].min(dim=0)[0])

    # add edge type information
    edgetypes = g.canonical_etypes
    for i, etype in enumerate(edgetypes):
        g.edges[etype].data['typeid'] = i * torch.ones(g.num_edges(etype)).long()

    if len(split) <= 1:
        return g, params
    labels = g.ndata['label']
    idx = list(range(len(labels)))
    if '.' in split[0]:
        trsize = float(split[0])
        valsize, testsize = list(map(float, split[1].split('_')))
    else:
        trsize = int(split[0])
        valsize, testsize = list(map(int, split[1].split('_')))
        if testsize == -1:
            testsize = len(labels) - trsize - valsize
    idx_tr, idx_rest, _, y_rest = train_test_split(idx, labels, train_size=trsize, stratify=labels)
    idx_val, idx_test, _, _ = train_test_split(idx_rest, y_rest, train_size=valsize, stratify=y_rest)
    g.ndata['train_mask'] = torch.zeros(len(labels)).bool()
    g.ndata['train_mask'][idx_tr] = True
    g.ndata['val_mask'] = torch.zeros(len(labels)).bool()
    g.ndata['val_mask'][idx_val] = True
    g.ndata['test_mask'] = torch.zeros(len(labels)).bool()
    g.ndata['test_mask'][idx_test] = True

    return g, params

def load_dataset(dataset_name, split, seed=0, het=False, save_path=None):
    utils.set_random_seed(1000 + seed)
    split = split.strip().split(',')
    if dataset_name =='dgraphfin':
        gfilename = 'dgraphfin.bin'
        gpath = os.path.join('data_processing', gfilename)
        if os.path.exists(gpath):
            g, _ = dgl.load_graphs(gpath)
            g = g[0]
        else:
            g = build_DGraphFin(het=het, save_path='data_processing' if save_path is None else save_path)
        g, params = split_DGraphFin(g, split)
        params['n_rel'] = g.edata['type'].shape[1]
        g.edata['type'] = g.edata['type'].float()
        
    elif dataset_name == 'tfin':
        gpath = './dataset/T-Finance/tfinance'
        g, _ = dgl.load_graphs(gpath)
        g = g[0]
        g.ndata['label'] = g.ndata['label'].argmax(1)
        g, params = split_tfintsoc(g, split)
    elif dataset_name == 'tsoc':
        gpath = './dataset/T-Social/tsocial'
        g, _ = dgl.load_graphs(gpath)
        g = g[0]
        g, params = split_tfintsoc(g, split)
    elif dataset_name == 'yelp':
        gpath = './dataset/FraudYelp'
        g = FraudYelpDataset(gpath)
        g = g[0]
        
        g = dgl.to_homogeneous(g, ndata=['feature', 'label', 'train_mask', 'val_mask', 'test_mask'])
        g.edata['type'] = torch.nn.functional.one_hot(g.edata['_TYPE']).float()
        g.ndata['train_mask'] = g.ndata['train_mask'].bool()
        g.ndata['val_mask'] = g.ndata['val_mask'].bool()
        g.ndata['test_mask'] = g.ndata['test_mask'].bool()
        g, params = split_yelpamz(g, split)
        params['n_rel'] = g.edata['type'].shape[1]
    elif dataset_name == 'amazon':
        gpath = './dataset/FraudAmazon'
        g = FraudAmazonDataset(gpath)
        g = g[0]
        
        g = dgl.to_homogeneous(g, ndata=['feature', 'label', 'train_mask', 'val_mask', 'test_mask'])
        g.edata['type'] = torch.nn.functional.one_hot(g.edata['_TYPE']).float()
        g.ndata['train_mask'] = g.ndata['train_mask'].bool()
        g.ndata['val_mask'] = g.ndata['val_mask'].bool()
        g.ndata['test_mask'] = g.ndata['test_mask'].bool()
        g, params = split_yelpamz(g, split)
        params['n_rel'] = g.edata['type'].shape[1]
    
    else:
        raise NotImplementedError('dataset {} not implemented'.format(dataset_name))

    return g, params