Source code for spateo.external.CAST.model.aug

import numpy as np
import torch as th


# __func is original, func is GPU optimized

[docs]def __random_aug(graph, x, feat_drop_rate, edge_mask_rate):
    import dgl

    n_node = graph.number_of_nodes()

    edge_mask = mask_edge(graph, edge_mask_rate)
    feat = drop_feature(x, feat_drop_rate)

    ng = dgl.graph([])
    ng.add_nodes(n_node)
    src = graph.edges()[0]
    dst = graph.edges()[1]

    nsrc = src[edge_mask]
    ndst = dst[edge_mask]
    ng.add_edges(nsrc, ndst)

    return ng, feat




[docs]def __drop_feature(x, drop_prob):
    drop_mask = th.empty((x.size(1),), dtype=th.float32, device=x.device).uniform_(0, 1) < drop_prob
    x = x.clone()
    x[:, drop_mask] = 0

    return x




[docs]def __mask_edge(graph, mask_prob):
    E = graph.number_of_edges()

    mask_rates = th.FloatTensor(np.ones(E) * mask_prob)
    masks = th.bernoulli(1 - mask_rates)
    mask_idx = masks.nonzero().squeeze(1)
    return mask_idx




[docs]def random_aug(graph, x, feat_drop_rate, edge_mask_rate):
    import dgl

    n_node = graph.number_of_nodes()

    edge_mask = mask_edge(graph, edge_mask_rate)
    feat = x.clone()
    feat = drop_feature(feat, feat_drop_rate)

    ng = dgl.graph([], device=graph.device)
    ng.add_nodes(n_node)
    src = graph.edges()[0]
    dst = graph.edges()[1]

    nsrc = src[edge_mask]
    ndst = dst[edge_mask]
    ng.add_edges(nsrc, ndst)

    return ng, feat




[docs]def drop_feature(x, drop_prob):
    drop_mask = th.empty((x.size(1),), dtype=th.float32, device=x.device).uniform_(0, 1) < drop_prob
    # x = x.clone()
    x[:, drop_mask] = 0

    return x




[docs]def mask_edge(graph, mask_prob):
    E = graph.number_of_edges()
    mask_rates = th.ones(E, device=graph.device) * mask_prob
    masks = th.bernoulli(1 - mask_rates)
    mask_idx = masks.nonzero().squeeze(1)
    return mask_idx