topobench.nn.backbones.combinatorial package

Submodules

topobench.nn.backbones.combinatorial.gccn module

Define the TopoTune class, which, given a choice of hyperparameters, instantiates a GCCN expecting a collection of strictly augmented Hasse graphs as input.

class topobench.nn.backbones.combinatorial.gccn.TopoTune(GNN, neighborhoods, layers, use_edge_attr, activation)

Bases: Module

Tunes a GNN model using higher-order relations.

This class takes a GNN and its kwargs as inputs, and tunes it with specified additional relations.

Parameters:

GNNtorch.nn.Module, a class not an object

The GNN class to use. ex: GAT, GCN.

neighborhoodslist of lists

The neighborhoods of interest.

layersint

The number of layers to use. Each layer contains one GNN.

use_edge_attrbool

Whether to use edge attributes.

activationstr

The activation function to use. ex: ‘relu’, ‘tanh’, ‘sigmoid’.

aggregate_inter_nbhd(x_out_per_route)

Aggregate the outputs of the GNN for each rank.

While the GNN takes care of intra-nbhd aggregation, this will take care of inter-nbhd aggregation. Default: sum.

Parameters:

x_out_per_routedict

The outputs of the GNN for each route.

Returns:

dict

The aggregated outputs of the GNN for each rank.

forward(batch)

Forward pass of the model.

Parameters:

batchComplex or ComplexBatch(Complex)

The input data.

Returns:

dict

The output hidden states of the model per rank.

generate_membership_vectors(batch: Data)

Generate membership vectors based on batch.cell_statistics.

Parameters:

batchtorch_geometric.data.Data

Batch object containing the batched domain data.

Returns:

dict

The batch membership of the graphs per rank.

get_nbhd_cache(params)

Cache the nbhd information into a dict for the complex at hand.

Parameters:

paramsdict

The parameters of the batch, containing the complex.

Returns:

dict

The neighborhood cache.

interrank_expand(params, src_rank, dst_rank, nbhd_cache, membership)

Expand the complex into an interrank Hasse graph.

Parameters:

paramsdict

The parameters of the batch, containting the complex.

src_rankint

The source rank.

dst_rankint

The destination rank.

nbhd_cachedict

The neighborhood cache containing the expanded boundary index and edge attributes.

membershipdict

The batch membership of the graphs per rank.

Returns:

torch_geometric.data.Data

The expanded batch of interrank Hasse graphs for this route.

interrank_gnn_forward(batch_route, layer_idx, route_index, n_dst_cells)

Forward pass of the GNN (one layer) for an interrank Hasse graph.

Parameters:

batch_routetorch_geometric.data.Data

The batch of interrank Hasse graphs for this route.

layer_idxint

The index of the layer.

route_indexint

The index of the route.

n_dst_cellsint

The number of destination cells in the whole batch.

Returns:

torch.tensor

The output of the GNN (updated features).

intrarank_expand(params, src_rank, nbhd)

Expand the complex into an intrarank Hasse graph.

Parameters:

paramsdict

The parameters of the batch, containting the complex.

src_rankint

The source rank.

nbhdstr

The neighborhood to use.

Returns:

torch_geometric.data.Data

The expanded batch of intrarank Hasse graphs for this route.

intrarank_gnn_forward(batch_route, layer_idx, route_index)

Forward pass of the GNN (one layer) for an intrarank Hasse graph.

Parameters:

batch_routetorch_geometric.data.Data

The batch of intrarank Hasse graphs for this route.

layer_idxint

The index of the TopoTune layer.

route_indexint

The index of the route.

Returns:

torch.tensor

The output of the GNN (updated features).

topobench.nn.backbones.combinatorial.gccn.get_activation(nonlinearity, return_module=False)

Activation resolver from CWN.

Parameters:

nonlinearitystr

The nonlinearity to use.

return_modulebool

Whether to return the module or the function.

Returns:

module or function

The module or the function.

topobench.nn.backbones.combinatorial.gccn.interrank_boundary_index(x_src, boundary_index, n_dst_nodes)

Recover lifted graph.

Edge-to-node boundary relationships of a graph with n_nodes and n_edges can be represented as up-adjacency node relations. There are n_nodes+n_edges nodes in this lifted graph. Desgiend to work for regular (edge-to-node and face-to-edge) boundary relationships.

Parameters:

x_srctorch.tensor

Source node features. Shape [n_src_nodes, n_features]. Should represent edge or face features.

boundary_indexlist of lists or list of tensors

List boundary_index[0] stores node ids in the boundary of edge stored in boundary_index[1]. List boundary_index[1] stores list of edges.

n_dst_nodesint

Number of destination nodes.

Returns:

edge_indexlist of lists

The edge_index[0][i] and edge_index[1][i] are the two nodes of edge i.

edge_attrtensor

Edge features are given by feature of bounding node represnting an edge. Shape [n_edges, n_features].

topobench.nn.backbones.combinatorial.gccn_onehasse module

Define the TopoTune class, which, given a choice of hyperparameters, instantiates a GCCN expecting a single Hasse graph as input.

class topobench.nn.backbones.combinatorial.gccn_onehasse.TopoTune_OneHasse(GNN, neighborhoods, layers, use_edge_attr, activation)

Bases: Module

Tunes a GNN model using higher-order relations.

This class takes a GNN and its kwargs as inputs, and tunes it with specified additional relations. Unlike the case of TopoTune, this class expects a single Hasse graph as input, where all higher-order neighborhoods are represented as a single adjacency matrix.

Parameters:

GNNtorch.nn.Module, a class not an object

The GNN class to use. ex: GAT, GCN.

neighborhoodslist of lists

The neighborhoods of interest.

layersint

The number of layers to use. Each layer contains one GNN.

use_edge_attrbool

Whether to use edge attributes.

activationstr

The activation function to use. ex: ‘relu’, ‘tanh’, ‘sigmoid’.

aggregate_inter_nbhd(x_out)

Aggregate the outputs of the GNN for each rank.

While the GNN takes care of intra-nbhd aggregation, this will take care of inter-nbhd aggregation. Default: sum.

Parameters:

x_outtorch.tensor

The output of the GNN, concatenated features of each rank.

Returns:

dict

The aggregated outputs of the GNN for each rank.

all_nbhds_expand(params, membership)

Expand the complex into a single Hasse graph which contains all ranks and all nbhd.

Parameters:

paramsdict

The parameters of the batch, containing the complex.

membershipdict

The batch membership of the graphs per rank.

Returns:

torch_geometric.data.Data

The expanded Hasse graph.

all_nbhds_gnn_forward(batch_route, layer_idx)

Forward pass of the GNN (one layer) for an intrarank Hasse graph.

Parameters:

batch_routetorch_geometric.data.Data

The batch of intrarank Hasse graphs for this route.

layer_idxint

The index of the TopoTune layer.

Returns:

torch.tensor

The output of the GNN (updated features).

forward(batch)

Forward pass of the model.

Parameters:

batchComplex or ComplexBatch(Complex)

The input data.

Returns:

dict

The output hidden states of the model per rank.

generate_membership_vectors(batch: Data)

Generate membership vectors based on batch.cell_statistics.

Parameters:

batchtorch_geometric.data.Data

Batch object containing the batched domain data.

Returns:

dict

The batch membership of the graphs per rank.

topobench.nn.backbones.combinatorial.gccn_onehasse.get_activation(nonlinearity, return_module=False)

Activation resolver from CWN.

Parameters:

nonlinearitystr

The nonlinearity to use.

return_modulebool

Whether to return the module or the function.

Returns:

module or function

The module or the function.

Module contents

Combinatorial backbones with automated exports.

class topobench.nn.backbones.combinatorial.TopoTune(GNN, neighborhoods, layers, use_edge_attr, activation)

Bases: Module

Tunes a GNN model using higher-order relations.

This class takes a GNN and its kwargs as inputs, and tunes it with specified additional relations.

Parameters:

GNNtorch.nn.Module, a class not an object

The GNN class to use. ex: GAT, GCN.

neighborhoodslist of lists

The neighborhoods of interest.

layersint

The number of layers to use. Each layer contains one GNN.

use_edge_attrbool

Whether to use edge attributes.

activationstr

The activation function to use. ex: ‘relu’, ‘tanh’, ‘sigmoid’.

aggregate_inter_nbhd(x_out_per_route)

Aggregate the outputs of the GNN for each rank.

While the GNN takes care of intra-nbhd aggregation, this will take care of inter-nbhd aggregation. Default: sum.

Parameters:

x_out_per_routedict

The outputs of the GNN for each route.

Returns:

dict

The aggregated outputs of the GNN for each rank.

forward(batch)

Forward pass of the model.

Parameters:

batchComplex or ComplexBatch(Complex)

The input data.

Returns:

dict

The output hidden states of the model per rank.

generate_membership_vectors(batch: Data)

Generate membership vectors based on batch.cell_statistics.

Parameters:

batchtorch_geometric.data.Data

Batch object containing the batched domain data.

Returns:

dict

The batch membership of the graphs per rank.

get_nbhd_cache(params)

Cache the nbhd information into a dict for the complex at hand.

Parameters:

paramsdict

The parameters of the batch, containing the complex.

Returns:

dict

The neighborhood cache.

interrank_expand(params, src_rank, dst_rank, nbhd_cache, membership)

Expand the complex into an interrank Hasse graph.

Parameters:

paramsdict

The parameters of the batch, containting the complex.

src_rankint

The source rank.

dst_rankint

The destination rank.

nbhd_cachedict

The neighborhood cache containing the expanded boundary index and edge attributes.

membershipdict

The batch membership of the graphs per rank.

Returns:

torch_geometric.data.Data

The expanded batch of interrank Hasse graphs for this route.

interrank_gnn_forward(batch_route, layer_idx, route_index, n_dst_cells)

Forward pass of the GNN (one layer) for an interrank Hasse graph.

Parameters:

batch_routetorch_geometric.data.Data

The batch of interrank Hasse graphs for this route.

layer_idxint

The index of the layer.

route_indexint

The index of the route.

n_dst_cellsint

The number of destination cells in the whole batch.

Returns:

torch.tensor

The output of the GNN (updated features).

intrarank_expand(params, src_rank, nbhd)

Expand the complex into an intrarank Hasse graph.

Parameters:

paramsdict

The parameters of the batch, containting the complex.

src_rankint

The source rank.

nbhdstr

The neighborhood to use.

Returns:

torch_geometric.data.Data

The expanded batch of intrarank Hasse graphs for this route.

intrarank_gnn_forward(batch_route, layer_idx, route_index)

Forward pass of the GNN (one layer) for an intrarank Hasse graph.

Parameters:

batch_routetorch_geometric.data.Data

The batch of intrarank Hasse graphs for this route.

layer_idxint

The index of the TopoTune layer.

route_indexint

The index of the route.

Returns:

torch.tensor

The output of the GNN (updated features).

class topobench.nn.backbones.combinatorial.TopoTune_OneHasse(GNN, neighborhoods, layers, use_edge_attr, activation)

Bases: Module

Tunes a GNN model using higher-order relations.

This class takes a GNN and its kwargs as inputs, and tunes it with specified additional relations. Unlike the case of TopoTune, this class expects a single Hasse graph as input, where all higher-order neighborhoods are represented as a single adjacency matrix.

Parameters:

GNNtorch.nn.Module, a class not an object

The GNN class to use. ex: GAT, GCN.

neighborhoodslist of lists

The neighborhoods of interest.

layersint

The number of layers to use. Each layer contains one GNN.

use_edge_attrbool

Whether to use edge attributes.

activationstr

The activation function to use. ex: ‘relu’, ‘tanh’, ‘sigmoid’.

aggregate_inter_nbhd(x_out)

Aggregate the outputs of the GNN for each rank.

While the GNN takes care of intra-nbhd aggregation, this will take care of inter-nbhd aggregation. Default: sum.

Parameters:

x_outtorch.tensor

The output of the GNN, concatenated features of each rank.

Returns:

dict

The aggregated outputs of the GNN for each rank.

all_nbhds_expand(params, membership)

Expand the complex into a single Hasse graph which contains all ranks and all nbhd.

Parameters:

paramsdict

The parameters of the batch, containing the complex.

membershipdict

The batch membership of the graphs per rank.

Returns:

torch_geometric.data.Data

The expanded Hasse graph.

all_nbhds_gnn_forward(batch_route, layer_idx)

Forward pass of the GNN (one layer) for an intrarank Hasse graph.

Parameters:

batch_routetorch_geometric.data.Data

The batch of intrarank Hasse graphs for this route.

layer_idxint

The index of the TopoTune layer.

Returns:

torch.tensor

The output of the GNN (updated features).

forward(batch)

Forward pass of the model.

Parameters:

batchComplex or ComplexBatch(Complex)

The input data.

Returns:

dict

The output hidden states of the model per rank.

generate_membership_vectors(batch: Data)

Generate membership vectors based on batch.cell_statistics.

Parameters:

batchtorch_geometric.data.Data

Batch object containing the batched domain data.

Returns:

dict

The batch membership of the graphs per rank.