ugnn package

Submodules

ugnn.experiment_config module

ugnn.experiment_config.validate_params(params)

Validate the experiment parameters.

Parameters:

params (dict) – Dictionary of experiment parameters.

Raises:

ValueError – If any parameter is invalid.

ugnn.conformal module

ugnn.conformal.APS_scores(probs, label)

ugnn.conformal.RAPS_scores(probs, label, penalty, kreg)

ugnn.conformal.SAPS_scores(probs, label, weight)

ugnn.conformal.THR_scores(probs, label)

ugnn.conformal.get_prediction_sets(output: Tensor, data: Data, calib_mask: Tensor | ndarray, test_mask: Tensor | ndarray, score_function: Literal['APS', 'RAPS', 'SAPS', 'THR'] = 'APS', alpha=0.1, kreg=1)

Computes prediction sets for a given model’s output using conformal prediction.

This function uses data points from the calibration set to compute a non-conformity score, which measures how “strange” a data-label pair is. It then calculates the quantile of the non-conformity scores to form prediction sets for the test set.

Parameters:

• output – Tensor The model’s output, typically softmax probabilities or logits, for all nodes or samples.

• data – Data The dataset object containing features, labels, and other graph-related information.

• calib_mask – Tensor A boolean mask indicating which samples belong to the calibration set.

• test_mask – Tensor A boolean mask indicating which samples belong to the test set.

• score_function – Literal[“APS”, “RAPS”, “SAPS”, “THR”], optional The scoring function to use for conformal prediction. Options include: - “APS”: Adaptive Prediction Sets. - “RAPS”: Regularized Adaptive Prediction Sets. - “SAPS”: Smoothed Adaptive Prediction Sets. - “THR”: Threshold-based Prediction Sets. Default is “APS”.

• alpha – float, optional The miscoverage level for conformal prediction. Determines the confidence level (e.g., alpha=0.1 corresponds to 90% confidence). Default is 0.1.

• kreg – int, optional Regularization parameter used in certain scoring functions like “RAPS” and “SAPS”. Default is 1.

Returns:

Tensor

A tensor containing the prediction sets for the test samples. Each set contains the indices of the predicted classes for each test sample.

Return type:

prediction_sets

Notes

• For “RAPS” and “SAPS” scoring functions, the calibration set is further split into

a validation set (20%) and a calibration set (80%) to tune the hyperparameter kreg. - The function assumes that the calibration and test masks are disjoint.

References

• “Conformal Prediction for Reliable Machine Learning: Theory and Applications” (2023)

https://arxiv.org/pdf/2310.06430

ugnn.experiments module

class ugnn.experiments.Experiment(method: Literal['BD', 'UA'], GNN_model: Literal['GCN', 'GAT'], regime: Literal['transductive', 'semi-inductive', 'temporal transductive'], data: Data, masks: Masks, experiment_params: ExperimentParams, data_params: DataParams, conformal_method: Literal['APS', 'RAPS', 'SAPS'] = 'APS', seed: int = 123)

Bases: object

Initializes the experiment with the specified parameters.

This class trains a GNN on multiple networks (e.g., a discrete-time dynamic network) and performs conformal prediction with the GNN. The experiment evaluates the GNN’s performance using accuracy, average prediction set size, and coverage.

The data is split into train, validation, calibration, and test groups based on the specified regime:

• Transductive: Nodes are randomly assigned to train/valid/calib/test groups.

• Semi-inductive: Nodes after a certain time point are assigned to the test group, while earlier nodes are randomly assigned to train/valid/calib groups.

• Temporal transductive: Nodes after a certain time point are split between calib and test groups, while earlier nodes are assigned to train/valid groups.

Parameters:

• method (Literal["BD", "UA"]) – The method to represent multiple networks as a single network (“block diagonal” or “unfolded”).

• GNN_model (Literal["GCN", "GAT"]) – The GNN model to use.

• regime (Literal["transductive", "semi-inductive", "temporal transductive"]) – The experiment regime.

• data (Data) – The dataset object containing graph data and labels.

• masks (Masks) – A dictionary with train, validation, calibration, and test masks.

• experiment_params (ExperimentParams) – Parameters for the experiment (e.g., number of epochs, learning rate, etc.).

• data_params (DataParams) – Parameters for the dataset (e.g., number of nodes, time steps, and classes).

evaluate()

Evaluate the trained model and compute metrics.

initialise_model()

Initialise the GNN model based on the specified type.

train()

Train the GNN model.

ugnn.experiments.parse_args_load_data()

Organise experiment parameters and data parameters for running experiments.

Parameters for GNN fitting are defined in ugnn.experiment_config.py. Arguments can be modified in the command line using argparse.

Selected data is then loaded, with matching GNN fitting params for the data.

Returns:

Experiment parameters. DATA_PARAMS (DataParams): Data parameters.

Return type:

EXPERIMENT_PARAMS (ExperimentParams)

ugnn.gnns module

class ugnn.gnns.GAT(num_nodes, num_channels, num_classes, seed)

Bases: Module

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x, edge_index, edge_weight)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class ugnn.gnns.GCN(num_nodes, num_channels, num_classes, seed)

Bases: Module

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x, edge_index, edge_weight)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

ugnn.gnns.train(model, data, train_mask, optimizer)

ugnn.gnns.valid(model, data, valid_mask)

ugnn.networks module

class ugnn.networks.Block_Diagonal_Network(dataset: Dataset)

Bases: Dataset

A pytorch geometric dataset for the block diagonal version of a Dynamic Network object.

NOTE: Attributes are left empty (identity) for now. Functionality to take attributes will be added soon.

class ugnn.networks.Dynamic_Network(As: ndarray | List[csr_matrix], labels: ndarray)

Bases: Dataset

A pytorch geometric dataset for an n-node discrete-time dynamic network with T time points.

Here a dynamic network consists of:

As: A collection of T nxn adjacency matrices, shape (T, n, n) labels: a list of labels of shape (T, n)

As we will be using GNNs to embed these networks, we must supply some attributes. For this, we use the nxn identity matrix for each time point.

class ugnn.networks.Unfolded_Network(dataset: Dataset)

Bases: Dataset

A pytorch geometric dataset for the dilated unfolding of a Dynamic Network object.

NOTE: Attributes are left empty (identity) for now. Functionality to take attributes will be added soon.

ugnn.networks.block_diagonal_matrix_from_series(As: ndarray | List[csr_matrix])

Forms the block diagonal matrix from an adjacency series

ugnn.networks.form_empty_attributes(m, sparse=False)

Forms the appropriately sized identity matrix for the (empty) attributes of the network

ugnn.networks.get_edge_weights(A)

Returns the edge weights for the adjacency matrix A

ugnn.networks.unfolded_matrix_from_series(As: ndarray | List[csr_matrix])

Forms the general unfolded matrix from an adjacency series

ugnn.results_manager module

Module contents

UGNN: The Unfolded Graph Neural Networks package.