topobench.nn.backbones.graph.nsd_utils.laplacian_builders module#

Laplacian builders for Neural Sheaf Diffusion.

This module provides builders for constructing different types of sheaf Laplacians: diagonal, bundle (with orthogonal maps), and general (full matrices).

class DiagLaplacianBuilder(size, edge_index, d)#

Bases: LaplacianBuilder

Builder for sheaf Laplacian with diagonal restriction maps.

This builder constructs a sheaf Laplacian where the restriction maps are diagonal matrices, parameterized by d values per edge.

Parameters:

sizeint: Number of nodes in the graph.
edge_indextorch.Tensor: Edge indices of shape [2, num_edges].
dint: Dimension of the diagonal stalk space.

__init__(size, edge_index, d)#: Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(maps)#

Build the sheaf Laplacian from diagonal restriction maps.

Parameters:

mapstorch.Tensor: Diagonal restriction map parameters of shape [num_edges, d].

Returns:

Ltuple of torch.Tensor: Sparse Laplacian representation as (indices, values).
saved_tril_mapstorch.Tensor: Saved lower triangular restriction maps for analysis.

class GeneralLaplacianBuilder(size, edge_index, d, augmented=True)#

Bases: LaplacianBuilder

Builder for general sheaf Laplacian with full matrix restriction maps.

This builder constructs a sheaf Laplacian where the restriction maps are arbitrary d x d matrices learned from data.

Parameters:

sizeint: Number of nodes in the graph.
edge_indextorch.Tensor: Edge indices of shape [2, num_edges].
dint: Dimension of the stalk space.
augmentedbool, optional: Whether to use augmented representation (not currently used). Default is True.

__init__(size, edge_index, d, augmented=True)#: Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(maps)#

Build the sheaf Laplacian from general restriction maps.

Parameters:

mapstorch.Tensor: General restriction map matrices of shape [num_edges, d, d].

Returns:

Ltuple of torch.Tensor: Sparse Laplacian representation as (indices, values).
saved_tril_mapstorch.Tensor: Saved lower triangular transport maps for analysis.

class LaplacianBuilder(size, edge_index, d, normalised=False, deg_normalised=False)#

Bases: Module

Base class for building sheaf Laplacians.

This class provides common functionality for all Laplacian builders, including preprocessing edge indices and computing normalization.

Parameters:

sizeint: Number of nodes in the graph.
edge_indextorch.Tensor: Edge indices of shape [2, num_edges].
dint: Dimension of the stalk space.
normalisedbool, optional: Whether to use normalized Laplacian. Default is False.
deg_normalisedbool, optional: Whether to use degree normalization (not used). Default is False.

__init__(size, edge_index, d, normalised=False, deg_normalised=False)#: Initialize internal Module state, shared by both nn.Module and ScriptModule.

scalar_normalise(diag, tril, row, col)#

Apply scalar normalization to Laplacian entries.

Normalizes diagonal and off-diagonal entries by node degrees, similar to symmetric normalization in standard graph Laplacians.

Parameters:

diagtorch.Tensor: Diagonal block values.
triltorch.Tensor: Lower triangular block values.
rowtorch.Tensor: Row indices of edges.
coltorch.Tensor: Column indices of edges.

Returns:

diag_mapstorch.Tensor: Normalized diagonal block values.
non_diag_mapstorch.Tensor: Normalized off-diagonal block values.

class NormConnectionLaplacianBuilder(size, edge_index, d, orth_map=None)#

Bases: LaplacianBuilder

Builder for normalized bundle sheaf Laplacian with orthogonal restriction maps.

This builder constructs a normalized sheaf Laplacian where the restriction maps are orthogonal matrices parameterized via Cayley transform or matrix exponential. Used for bundle sheaf models.

Parameters:

sizeint: Number of nodes in the graph.
edge_indextorch.Tensor: Edge indices of shape [2, num_edges].
dint: Dimension of the stalk space.
orth_mapstr or None, optional: Method for orthogonalization (‘cayley’ or ‘matrix_exp’). Default is None.

__init__(size, edge_index, d, orth_map=None)#: Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(map_params)#

Build the normalized sheaf Laplacian from orthogonal restriction maps.

Parameters:

map_paramstorch.Tensor: Orthogonal map parameters of shape [num_edges, d*(d+1)/2].

Returns:

Ltuple of torch.Tensor: Sparse normalized Laplacian representation as (indices, values).
saved_tril_mapstorch.Tensor: Saved lower triangular transport maps for analysis.

class Orthogonal(d, orthogonal_map)#

Bases: Module

Orthogonal transformation module for sheaf restriction maps.

Converts lower-triangular parameters into orthogonal matrices using either the matrix exponential or Cayley transform. Based on PyTorch’s parametrization utilities.

Reference: https://pytorch.org/docs/stable/_modules/torch/nn/utils/parametrizations.html#orthogonal

Parameters:

dint: Dimension of the square orthogonal matrices to generate.
orthogonal_mapstr: Method for generating orthogonal matrices. Options are ‘matrix_exp’ or ‘cayley’.

Raises:

AssertionError: If orthogonal_map is not ‘matrix_exp’ or ‘cayley’.

__init__(d, orthogonal_map)#: Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(params)#

Convert parameters to orthogonal matrices.

Parameters:

paramstorch.Tensor: Lower-triangular parameters of shape [batch_size, d*(d+1)/2].

Returns:

torch.Tensor: Orthogonal matrices of shape [batch_size, d, d].

Raises:

ValueError: If an unsupported orthogonal_map method is specified.

compute_learnable_diag_laplacian_indices(size, edge_index, learned_d, total_d)#

Compute sparse indices for a learnable Laplacian with diagonal matrices.

Generates indices for both diagonal and non-diagonal blocks of a block-structured Laplacian matrix where each block is diagonal (learned_d x learned_d).

Parameters:

sizeint: Number of nodes in the graph.
edge_indextorch.Tensor: Edge indices of shape [2, num_edges] where edge_index[0] < edge_index[1].
learned_dint: Dimension of each learned diagonal block.
total_dint: Total dimension per node (typically equal to learned_d).

Returns:

diag_indicestorch.Tensor: Sparse indices for diagonal blocks, shape [2, size * learned_d].
non_diag_indicestorch.Tensor: Sparse indices for non-diagonal blocks, shape [2, num_edges * learned_d].

compute_learnable_laplacian_indices(size, edge_index, learned_d, total_d)#

Compute sparse indices for a learnable Laplacian with full matrices.

Generates indices for both diagonal and non-diagonal blocks of a block-structured Laplacian matrix where each block is learned_d x learned_d.

Parameters:

sizeint: Number of nodes in the graph.
edge_indextorch.Tensor: Edge indices of shape [2, num_edges] where edge_index[0] < edge_index[1].
learned_dint: Dimension of each learned block matrix.
total_dint: Total dimension per node (typically equal to learned_d).

Returns:

diag_indicestorch.Tensor: Sparse indices for diagonal blocks, shape [2, size * learned_d * learned_d].
non_diag_indicestorch.Tensor: Sparse indices for non-diagonal blocks, shape [2, num_edges * learned_d * learned_d].

compute_left_right_map_index(edge_index, full_matrix=False)#

Compute indices for mapping edges to their reverse edges.

For each edge in the lower triangular part (source < target) or all edges, find the index of its corresponding reverse edge in the original edge list.

Parameters:

edge_indextorch.Tensor: Edge indices of shape [2, num_edges] representing directed edges.
full_matrixbool, optional: If True, use all edges. If False, only use edges where source < target. Default is False.

Returns:

left_right_indextorch.Tensor: Indices of shape [2, num_selected_edges] where first row contains indices of selected edges and second row contains indices of their reverse edges.
new_edge_indextorch.Tensor: Edge indices of selected edges, shape [2, num_selected_edges].

degree(index, num_nodes=None, dtype=None)#

Computes the (unweighted) degree of a given one-dimensional index tensor.

Parameters:

index (LongTensor) – Index tensor.
num_nodes (int, optional) – The number of nodes, i.e. max_val + 1 of index. (default: None)
dtype (torch.dtype, optional) – The desired data type of the returned tensor.

Return type:

Tensor

Example

>>> row = torch.tensor([0, 1, 0, 2, 0])
>>> degree(row, dtype=torch.long)
tensor([3, 1, 1])

mergesp(index1, value1, index2, value2)#

Merge two sparse matrices with disjoint indices into one.

Concatenates two sets of sparse matrix indices and values, assuming the indices are disjoint (no overlapping entries).

Parameters:

index1torch.Tensor: First set of sparse indices, shape [2, num_entries1].
value1torch.Tensor: First set of values, shape [num_entries1].
index2torch.Tensor: Second set of sparse indices, shape [2, num_entries2].
value2torch.Tensor: Second set of values, shape [num_entries2].

Returns:

indextorch.Tensor: Merged sparse indices, shape [2, num_entries1 + num_entries2].
valtorch.Tensor: Merged values, shape [num_entries1 + num_entries2].

scatter_add(src, index, dim=-1, out=None, dim_size=None)#