GraphSL.GNN.SLVAE package

Submodules

GraphSL.GNN.SLVAE.main module

class GraphSL.GNN.SLVAE.main.SLVAE[source]

Bases: object

Implement the Source Localization Variational Autoencoder (SLVAE) model.

Ling C, Jiang J, Wang J, et al. Source localization of graph diffusion via variational autoencoders for graph inverse problems[C]//Proceedings of the 28th ACM SIGKDD conference on knowledge discovery and data mining. 2022: 1010-1020.

infer(test_dataset, slvae_model, seed_vae_train, thres, lr=0.0001, num_epoch=10, print_epoch=1)[source]

Infer using the SLVAE model.

Args:

  • test_dataset (torch.utils.data.dataset.Subset): the test dataset (number of simulations * number of graph nodes * 2 (the first column is seed vector and the second column is diffusion vector)).

  • slvae_model (SLVAE_model): Trained SLVAE model.

  • seed_vae_train (torch.Tensor): The latent representations of training seed vector from VAE, which is used to initialize seed vector in the test set.

  • thres (float): Threshold value.

  • lr (float): Learning rate.

  • num_epoch (int): Number of epochs.

  • print_epoch (int): Number of epochs every time to print loss.

Returns:

  • Metric: Evaluation metric containing accuracy, precision, recall, F1 score, and AUC.

Example:

import os

curr_dir = os.getcwd()

from GraphSL.utils import load_dataset, diffusion_generation, split_dataset

from GraphSL.GNN.SLVAE.main import SLVAE

data_name = ‘karate’

graph = load_dataset(data_name, data_dir=curr_dir)

dataset = diffusion_generation(graph=graph, infect_prob=0.3, diff_type=’IC’, sim_num=100, seed_ratio=0.1)

adj, train_dataset, test_dataset =split_dataset(dataset)

slave = SLVAE()

slvae_model, seed_vae_train, thres, auc, f1, pred = slave.train(adj, train_dataset)

print(“SLVAE:”)

print(f”train auc: {auc:.3f}, train f1: {f1:.3f}”)

metric = slave.infer(test_dataset, slvae_model, seed_vae_train, thres)

print(f”test acc: {metric.acc:.3f}, test pr: {metric.pr:.3f}, test re: {metric.re:.3f}, test f1: {metric.f1:.3f}, test auc: {metric.auc:.3f}”)

train(adj, train_dataset, num_thres=10, lr=0.0001, weight_decay=0.0001, num_epoch=100, print_epoch=10, random_seed=0)[source]

Train the SLVAE model.

Args:

  • adj (scipy.sparse.csr_matrix): The adjacency matrix of the graph.

  • train_dataset (torch.utils.data.dataset.Subset): the training dataset (number of simulations * number of graph nodes * 2 (the first column is seed vector and the second column is diffusion vector)).

  • num_thres (int): Number of threshold values to try.

  • lr (float): Learning rate.

  • weight_decay (float): Weight decay.

  • num_epoch (int): Number of training epochs.

  • print_epoch (int): Number of epochs every time to print loss.

  • random_seed (int): Random seed.

Returns:

  • slvae_model (SLVAE_model): Trained SLVAE model.

  • seed_vae_train (torch.Tensor): The latent representations of training seed vector from VAE, which is used to initialize seed vector in the test set.

  • opt_thres (float): Optimal threshold.

  • train_auc (float): Train AUC.

  • opt_f1 (float): Optimal F1 score.

  • pred (numpy.ndarray): Predicted seed vector of the training set, every column is the prediction of every simulation. It is used to adjust thres_list.

Example:

import os

curr_dir = os.getcwd()

from GraphSL.utils import load_dataset, diffusion_generation, split_dataset

from GraphSL.GNN.SLVAE.main import SLVAE

data_name = ‘karate’

graph = load_dataset(data_name, data_dir=curr_dir)

dataset = diffusion_generation(graph=graph, infect_prob=0.3, diff_type=’IC’, sim_num=100, seed_ratio=0.1)

adj, train_dataset, test_dataset =split_dataset(dataset)

slave = SLVAE()

slvae_model, seed_vae_train, thres, auc, f1, pred = slave.train(adj, train_dataset)

print(“SLVAE:”)

print(f”train auc: {auc:.3f}, train f1: {f1:.3f}”)

class GraphSL.GNN.SLVAE.main.SLVAE_model(vae: Module, gnn: Module)[source]

Bases: Module

Source Localization Variational Autoencoder (SLVAE) model combining VAE and GNN.

Attributes: - vae (nn.Module): Variational Autoencoder module.

  • gnn (nn.Module): Graph Neural Network module.

  • reg_params (list): List of parameters requiring gradients.

forward(seed_vec, train_mode)[source]

Forward pass method of the SLVAE model.

Args:

  • seed_vec (torch.Tensor): Seed vector.

  • train_mode (bool): Flag indicating whether in training mode.

Returns:

  • seed_hat (torch.Tensor): reconstructed seed vector.

  • mean (torch.Tensor): Mean of the VAE.

  • log_var (torch.Tensor): Log variance of the VAE.

  • predictions (torch.Tensor): Predictions made by the SLVAE model.

infer_loss(y_true, y_hat, x_hat, train_pred)[source]

Compute inference loss.

Args:

  • y_true (torch.Tensor): True label tensor.

  • y_hat (torch.Tensor): Predicted label tensor.

  • x_hat (torch.Tensor): Reconstructed input tensor.

  • train_pred (torch.Tensor): Predicted tensor during training.

Returns:

  • total_loss (torch.Tensor): Total loss tensor.

train_loss(x, x_hat, mean, log_var, y, y_hat)[source]

Compute training loss.

Args:

  • x (torch.Tensor): Seed vector.

  • x_hat (torch.Tensor): Reconstructed seed tensor.

  • mean (torch.Tensor): Mean of the VAE.

  • log_var (torch.Tensor): Log variance of the VAE.

  • y (torch.Tensor): Diffusion vector.

  • y_hat (torch.Tensor): Predicted Diffusion vector.

Returns:

  • total_loss (torch.Tensor): Total loss is the sum of prediction loss, reconstruction loss and KL divergence.

GraphSL.GNN.SLVAE.model module

class GraphSL.GNN.SLVAE.model.Decoder(output_dim=784, hidden_dim=512, latent_dim=256)[source]

Bases: Module

Decoder module for a variational autoencoder (VAE).

Attributes:

  • output_dim (int): Dimension of the output.

  • hidden_dim (int): Dimension of the hidden layer.

  • latent_dim (int): Dimension of the latent space.

forward(x)[source]

Forward pass of the Decoder.

Args:

  • x (torch.Tensor): Input tensor.

Returns:

  • x_hat (torch.Tensor): Decoded output tensor.

class GraphSL.GNN.SLVAE.model.Encoder(input_dim=784, hidden_dim=512, latent_dim=256)[source]

Bases: Module

Encoder module for a variational autoencoder (VAE).

Attributes:

  • input_dim (int): Dimension of the input.

  • hidden_dim (int): Dimension of the hidden layer.

  • latent_dim (int): Dimension of the latent space.

forward(x)[source]

Forward pass of the Encoder.

Args:

  • x (torch.Tensor): Input tensor.

Returns:

  • mean (torch.Tensor): The mean of the latent space.

  • log_var (torch.Tensor): The log variance of the latent space.

class GraphSL.GNN.SLVAE.model.GCNLayer(in_features, out_features, bias=True)[source]

Bases: Module

A single layer of a Graph Convolutional Network (GCN).

Attributes:

  • in_features (int): Number of input features for each node.

  • out_features (int): Number of output features for each node.

  • bias (bool): Whether to include a bias term in the layer.

forward(x, adj)[source]

Forward pass of the GCN layer.

Args:

  • x (torch.Tensor): Input feature matrix of shape (num_nodes, in_features).

  • adj (torch.Tensor): Adjacency matrix of shape (num_nodes, num_nodes).

Returns:

  • x(torch.Tensor): Output feature matrix of shape (num_nodes, out_features).

class GraphSL.GNN.SLVAE.model.GNN(adj_matrix, input_dim=1, hiddenunits=[64, 64], out_dim=1, bias=True, drop_prob=0.5)[source]

Bases: Module

Graph Neural Network (GNN) model using GCN layers.

Attributes:

  • adj_matrix (torch.Tensor): Adjacency matrix representing graph connectivity.

  • input_dim (int): Dimension of the input.

  • hiddenunits (List[int]): List of hidden units for each layer.

  • num_classes (int): Number of output classes.

  • bias (bool): Whether to include bias in linear layers.

  • drop_prob (float): Dropout probability.

forward(seed_vec)[source]

Forward pass of the GNN.

Args:

  • seed_vec (torch.Tensor): Input seed vector.

Returns:

  • x (torch.Tensor): Predicted output.

loss(y, y_hat)[source]

Calculate loss.

Args:

  • y (torch.Tensor): Ground truth.

  • y_hat (torch.Tensor): Predicted output.

Returns:

  • forward_loss (torch.Tensor): Forward loss.

class GraphSL.GNN.SLVAE.model.VAE(input_dim=1, hidden_dim=512, latent_dim=256)[source]

Bases: Module

Variational Autoencoder (VAE) model.

Attributes:

  • input_dim (int): Dimension of the input.

  • hidden_dim (int): Dimension of the hidden layer.

  • latent_dim (int): Dimension of the latent space.

decode(x)[source]

Decode latent vector into output space.

Args:

  • x (torch.Tensor): Latent vector.

Returns:

  • x (torch.Tensor): Decoded output tensor.

encode(x)[source]

Encode input data into latent space.

Args:

  • x (torch.Tensor): Input tensor.

Returns:

  • mean (torch.Tensor): The mean of latent space

  • logvar (torch.Tensor): Log variance of latent space.

forward(x)[source]

Forward pass of the VAE.

Args:

  • x (torch.Tensor): Input tensor.

Returns:

  • x_hat (tensor.Tensor): The reconstructed input.

  • mean (tensor.Tensor): The mean of the latent space.

  • log_var (tensor.Tensor): The log variance of the latent space.

reparameterization(mean, var, device)[source]

Reparameterization trick to sample from the latent space.

Args:

  • mean (torch.Tensor): Mean of the latent space.

  • var (torch.Tensor): Variance of the latent space.

  • device (torch.device): Device to be used for computation, cpu or cuda.

Returns:

  • z (torch.Tensor): Sampled latent vector.

Module contents