topobench.nn.wrappers.simplicial.scn_wrapper module#
Wrapper for the SCNW model.
- class AbstractWrapper(backbone, **kwargs)#
-
Abstract class that provides an interface to handle the network output.
- Parameters:
- backbonetorch.nn.Module
Backbone model.
- **kwargsdict
Additional arguments for the class. It should contain the following keys: - out_channels (int): Number of output channels. - num_cell_dimensions (int): Number of cell dimensions.
- __init__(backbone, **kwargs)#
- abstract forward(batch)#
Forward pass for the model.
- Parameters:
- batchtorch_geometric.data.Data
Batch object containing the batched data.
- residual_connection(model_out, batch)#
Residual connection for the model.
This method sums, for the embeddings of the cells of any rank, the output of the model with the input embeddings and applies layer normalization.
- Parameters:
- model_outdict
Dictionary containing the model output.
- batchtorch_geometric.data.Data
Batch object containing the batched data.
- Returns:
- dict
Dictionary containing the updated model output.
- class SCNWrapper(backbone, **kwargs)#
Bases:
AbstractWrapperWrapper for the SCNW model.
This wrapper defines the forward pass of the model. The SCNW model returns the embeddings of the cells of rank 0, 1, and 2.
- forward(batch)#
Forward pass for the SCNW wrapper.
- Parameters:
- batchtorch_geometric.data.Data
Batch object containing the batched data.
- Returns:
- dict
Dictionary containing the updated model output.
- normalize_matrix(matrix)#
Normalize the input matrix.
The normalization is performed using the diagonal matrix of the inverse square root of the sum of the absolute values of the rows.
- Parameters:
- matrixtorch.sparse.FloatTensor
Input matrix to be normalized.
- Returns:
- torch.sparse.FloatTensor
Normalized matrix.