gsnn.models.utils

Functions

apply_norm_and_nonlin(norm, nonlin, out, ...)

Apply normalization and nonlinearity to the input tensor.

corr_score(y, yhat[, multioutput, method, eps])

calculate the average pearson correlation score

edge2node(x, edge_index, output_node_mask)

Convert edge-level features back to node-level features, focusing on output nodes.

get_Win_indices(edge_index, channels, ...)

Build sparse COO indices for the input weight matrix \(W_{in}\).

get_Wout_indices(edge_index, function_nodes, ...)

Build sparse COO indices for the output weight matrix \(W_{out}\).

get_conv_indices(edge_index, channels, ...)

Compute indexing structures for convolutional (sparse linear) layers.

hetero2homo(edge_index_dict, node_names_dict)

Convert a heterogeneous GSNN graph into a homogeneous graph representation.

node2edge(x, edge_index)

Convert node-level features to edge-level features.

predict_gsnn(loader, model, device[, verbose])

Run model on loader; return stacked numpy y, yhat, and sig ids from batches.
gsnn.models.utils.apply_norm_and_nonlin(norm, nonlin, out, norm_first)[source]

Apply normalization and nonlinearity to the input tensor.

Parameters:

  • norm (callable) – Normalization layer or operation.

  • nonlin (callable) – Nonlinear activation function.

  • out (Tensor) – Input tensor to be normalized and activated.

  • norm_first (bool) – If True, apply normalization before nonlinearity.

Returns:

The transformed tensor.

Return type:

Tensor

Example

>>> norm = torch.nn.BatchNorm1d(32)
>>> nonlin = torch.nn.ReLU()
>>> x = torch.randn(16, 32)  # [batch_size, num_features]
>>> # Apply normalization first
>>> out = apply_norm_and_nonlin(norm, nonlin, x, norm_first=True)
>>> print(out.shape)  # [16, 32]
gsnn.models.utils.corr_score(y, yhat, multioutput='uniform_weighted', method='pearson', eps=1e-06)[source]

calculate the average pearson correlation score

y (n_samples, n_outputs): yhat (n_samples, n_outputs):

gsnn.models.utils.edge2node(x, edge_index, output_node_mask)[source]

Convert edge-level features back to node-level features, focusing on output nodes.

Typically, output nodes should be designed to have an in-degree of 1, however, in the case of multiple edges per output node, the output features are summed and normalized by the square root of the in-degree.

Parameters:

  • x (Tensor) – Edge features of shape [batch_size, num_edges].

  • edge_index (Tensor) – Edge indices of shape [2, num_edges].

  • output_node_mask (Tensor) – Boolean mask of shape [num_nodes] indicating output nodes.

Returns:

Node features of shape [batch_size, num_output_nodes].

Return type:

Tensor

Example

>>> x = torch.randn(32, 3)  # [batch_size, num_edges]
>>> edge_index = torch.tensor([[0, 1, 1], [2, 2, 3]])  # 3 edges
>>> output_mask = torch.tensor([0, 0, 1, 1])  # Nodes 2,3 are outputs
>>> node_features = edge2node(x, edge_index, output_mask)
>>> print(node_features.shape)  # [32, 2]
gsnn.models.utils.get_Win_indices(edge_index, channels, function_nodes)[source]

Build sparse COO indices for the input weight matrix \(W_{in}\).

Parameters:

  • edge_index (Tensor) – Homogeneous edge index of shape [2, num_edges].

  • channels (int or Tensor) – If int, every function node gets the same number of hidden channels. If 1-D tensor/array, it must contain the per-node channel count of length num_nodes.

  • function_nodes (Tensor) – Index list of nodes that represent functions.

Returns:

A tuple containing:

  • indices (Tensor): COO indices of shape [2, nnz] for sparse tensor construction

  • channel_count (numpy.ndarray): Per-node channel counts for later reuse

Return type:

tuple

Example

>>> edge_index = torch.tensor([[0, 1], [1, 0]])  # 2 edges
>>> channels = 3  # 3 channels per function node
>>> function_nodes = torch.tensor([0])  # Node 0 is a function node
>>> indices, counts = get_Win_indices(edge_index, channels, function_nodes)
>>> print(indices.shape)  # [2, 6] (2 edges * 3 channels)
>>> print(counts)  # [3, 0] (3 channels for node 0, 0 for node 1)
gsnn.models.utils.get_Wout_indices(edge_index, function_nodes, channels)[source]

Build sparse COO indices for the output weight matrix \(W_{out}\).

Parameters:

  • edge_index (Tensor) – Homogeneous edge index of shape [2, num_edges].

  • function_nodes (Tensor) – Index list of nodes that represent functions.

  • channels (numpy.ndarray) – Array indicating the number of channels for each node.

Returns:

COO indices of shape [2, nnz] for sparse tensor construction.

Return type:

Tensor

Example

>>> edge_index = torch.tensor([[0, 1], [1, 0]])  # 2 edges
>>> function_nodes = torch.tensor([0])  # Node 0 is a function node
>>> channels = np.array([3, 0])  # 3 channels for node 0, 0 for node 1
>>> indices = get_Wout_indices(edge_index, function_nodes, channels)
>>> print(indices.shape)  # [2, 6] (3 channels * 2 edges)
gsnn.models.utils.get_conv_indices(edge_index, channels, function_nodes)[source]

Compute indexing structures for convolutional (sparse linear) layers.

Parameters:

  • edge_index (Tensor) – Homogeneous edge indices of shape [2, num_edges].

  • channels (int) – Number of channels per function node.

  • function_nodes (Tensor) – Indices of function nodes.

Returns:

A tuple containing:

  • w_in_indices (Tensor): Indexing for \(W_{in}\)

  • w_out_indices (Tensor): Indexing for \(W_{out}\)

  • w_in_size (tuple): Size specification for \(W_{in}\)

  • w_out_size (tuple): Size specification for \(W_{out}\)

  • channel_groups (List[int]): List mapping each channel to its node

Return type:

tuple

Example

>>> edge_index = torch.tensor([[0, 1], [1, 0]])  # 2 edges
>>> channels = 3  # 3 channels per function node
>>> function_nodes = torch.tensor([0])  # Node 0 is a function node
>>> indices = get_conv_indices(edge_index, channels, function_nodes)
>>> print(len(indices))  # 5 (w_in_indices, w_out_indices, sizes, groups)
gsnn.models.utils.hetero2homo(edge_index_dict, node_names_dict, edge_weight_dict=None)[source]

Convert a heterogeneous GSNN graph into a homogeneous graph representation.

The GSNN pipeline distinguishes three edge types:

  1. (‘input’, ‘to’, ‘function’)

  2. (‘function’, ‘to’, ‘function’)

  3. (‘function’, ‘to’, ‘output’)

This function stacks these edge sets into one homogeneous graph and returns boolean masks that let you recover the original node semantics.

Parameters:

  • edge_index_dict (Dict[Tuple[str, str, str], Tensor]) – Edge-type mapping where each value is a LongTensor with shape [2, num_edges_of_type].

  • node_names_dict (Dict[str, List[str]]) – Mapping of node types (‘input’, ‘function’, ‘output’) to their respective node names.

  • edge_weight_dict (Dict[Tuple[str, str, str], Tensor]) – Dictionary mapping edge types to edge weights. Expected keys are (‘input’, ‘to’, ‘function’), (‘function’, ‘to’, ‘function’), and (‘function’, ‘to’, ‘output’). Values should be tensors of shape [num_edges]. (default: None)

Returns:

A tuple containing:

  • edge_index (Tensor): Homogeneous edge indices of shape [2, num_edges]

  • input_mask (Tensor): Boolean mask for input nodes of shape [num_nodes]

  • output_mask (Tensor): Boolean mask for output nodes of shape [num_nodes]

  • num_nodes (int): Total number of nodes in the homogeneous graph

  • homo_names (List[str]): Node names in homogeneous ordering

  • edge_weight (Optional[Tensor]): Homogeneous edge weights of shape [num_edges], or None if edge_weight_dict was None.

Return type:

tuple

Example

>>> edge_index_dict = {
...     ('input', 'to', 'function'): torch.tensor([[0, 1], [0, 0]]),
...     ('function', 'to', 'function'): torch.tensor([[0], [0]]),
...     ('function', 'to', 'output'): torch.tensor([[0], [0]])
... }
>>> node_names_dict = {
...     'input': ['in1', 'in2'],
...     'function': ['func1'],
...     'output': ['out1']
... }
>>> edge_index, in_mask, out_mask, n_nodes, names = hetero2homo(
...     edge_index_dict, node_names_dict
... )
>>> print(edge_index.shape)  # [2, 4]
>>> print(in_mask.sum())  # 2 (number of input nodes)
>>> print(out_mask.sum())  # 1 (number of output nodes)
gsnn.models.utils.node2edge(x, edge_index)[source]

Convert node-level features to edge-level features. Every out-going edge receives the feature of the source node.

Parameters:

  • x (Tensor) – Node features of shape [batch_size, num_nodes].

  • edge_index (Tensor) – Edge indices of shape [2, num_edges].

Returns:

Edge features of shape [batch_size, num_edges].

Return type:

Tensor

Example

>>> x = torch.randn(32, 4)  # [batch_size, num_nodes]
>>> edge_index = torch.tensor([[0, 1], [1, 2]])  # 2 edges
>>> edge_features = node2edge(x, edge_index)
>>> print(edge_features.shape)  # [32, 2]
gsnn.models.utils.predict_gsnn(loader, model, device, verbose=True)[source]

Run model on loader; return stacked numpy y, yhat, and sig ids from batches.