diff --git a/topomodelx/nn/hypergraph/allset.py b/topomodelx/nn/hypergraph/allset.py index 219f9049..0219747f 100644 --- a/topomodelx/nn/hypergraph/allset.py +++ b/topomodelx/nn/hypergraph/allset.py @@ -16,18 +16,20 @@ class AllSet(torch.nn.Module): Dimension of the input features. hidden_channels : int Dimension of the hidden features. - n_layers : int, default: 2 + n_layers : int, default = 2 Number of AllSet layers in the network. - layer_dropout: float, default: 0.2 + layer_dropout : float, default = 0.2 Dropout probability for the AllSet layer. - mlp_num_layers : int, default: 2 + mlp_num_layers : int, default = 2 Number of layers in the MLP. - mlp_dropout : float, default: 0.0 - Dropout probability for the MLP. - mlp_activation : torch.nn.Module, default: None + mlp_activation : torch.nn.Module, default = None Activation function in the MLP. - mlp_norm : bool, default: False + mlp_dropout : float, default = 0.0 + Dropout probability for the MLP. + mlp_norm : bool, default = False Whether to apply input normalization in the MLP. + **kwargs : optional + Additional arguments for the inner layers. References ---------- @@ -47,6 +49,7 @@ def __init__( mlp_activation=None, mlp_dropout=0.0, mlp_norm=None, + **kwargs, ): super().__init__() @@ -59,6 +62,7 @@ def __init__( mlp_activation=mlp_activation, mlp_dropout=mlp_dropout, mlp_norm=mlp_norm, + **kwargs, ) for i in range(n_layers) ) diff --git a/topomodelx/nn/hypergraph/allset_transformer.py b/topomodelx/nn/hypergraph/allset_transformer.py index ffaa4279..67bfd882 100644 --- a/topomodelx/nn/hypergraph/allset_transformer.py +++ b/topomodelx/nn/hypergraph/allset_transformer.py @@ -26,6 +26,8 @@ class AllSetTransformer(torch.nn.Module): Number of layers in the MLP. mlp_dropout : float, default: 0.2 Dropout probability in the MLP. + **kwargs : optional + Additional arguments for the inner layers. References ---------- @@ -44,6 +46,7 @@ def __init__( dropout=0.2, mlp_num_layers=2, mlp_dropout=0.2, + **kwargs, ): super().__init__() @@ -55,6 +58,7 @@ def __init__( heads=heads, mlp_num_layers=mlp_num_layers, mlp_dropout=mlp_dropout, + **kwargs, ) for i in range(n_layers) ) diff --git a/topomodelx/nn/hypergraph/dhgcn.py b/topomodelx/nn/hypergraph/dhgcn.py index 846e78d9..e6f02e6c 100644 --- a/topomodelx/nn/hypergraph/dhgcn.py +++ b/topomodelx/nn/hypergraph/dhgcn.py @@ -16,8 +16,10 @@ class DHGCN(torch.nn.Module): Dimension of the input features. hidden_channels : int Dimension of the hidden features. - n_layer : int, default = 2 + n_layers : int, default = 2 Amount of message passing layers. + **kwargs : optional + Additional arguments for the inner layers. References ---------- @@ -31,6 +33,7 @@ def __init__( in_channels, hidden_channels, n_layers=1, + **kwargs, ): super().__init__() @@ -39,6 +42,7 @@ def __init__( in_channels=in_channels if i == 0 else hidden_channels, intermediate_channels=hidden_channels, out_channels=hidden_channels, + **kwargs, ) for i in range(n_layers) ) diff --git a/topomodelx/nn/hypergraph/hmpnn.py b/topomodelx/nn/hypergraph/hmpnn.py index d2ef5a6e..6dce8406 100644 --- a/topomodelx/nn/hypergraph/hmpnn.py +++ b/topomodelx/nn/hypergraph/hmpnn.py @@ -12,18 +12,18 @@ class HMPNN(torch.nn.Module): Parameters ---------- in_channels : int - Dimension of input features + Dimension of input features. hidden_channels : Tuple[int] A tuple of hidden feature dimensions to gradually reduce node/hyperedge representations feature dimension from in_features to the last item in the tuple. - num_classes: int - Number of classes n_layers : int, default = 2 Number of HMPNNLayer layers. - adjacency_dropout_rate: int, default = 0.7 + adjacency_dropout_rate : int, default = 0.7 Adjacency dropout rate. regular_dropout_rate : int, default = 0.5 Regular dropout rate applied on features. + **kwargs : optional + Additional arguments for the inner layers. References ---------- @@ -40,6 +40,7 @@ def __init__( n_layers=2, adjacency_dropout_rate=0.7, regular_dropout_rate=0.5, + **kwargs, ): super().__init__() @@ -52,6 +53,7 @@ def __init__( hidden_channels, adjacency_dropout=adjacency_dropout_rate, updating_dropout=regular_dropout_rate, + **kwargs, ) for _ in range(n_layers) ] @@ -66,7 +68,7 @@ def forward(self, x_0, x_1, incidence_1): Node features. x_1 : torch.Tensor, shape = (n_hyperedges, in_features) Hyperedge features. - incidence_1: torch.sparse.Tensor, shape = (n_nodes, n_hyperedges) + incidence_1 : torch.sparse.Tensor, shape = (n_nodes, n_hyperedges) Incidence matrix (B1). Returns diff --git a/topomodelx/nn/hypergraph/hnhn.py b/topomodelx/nn/hypergraph/hnhn.py index 6d8c0dc6..ba4fb878 100644 --- a/topomodelx/nn/hypergraph/hnhn.py +++ b/topomodelx/nn/hypergraph/hnhn.py @@ -20,6 +20,8 @@ class HNHN(torch.nn.Module): Number of HNHN message passing layers. layer_drop : float, default = 0.2 Dropout rate for the hidden features. + **kwargs : optional + Additional arguments for the inner layers. References ---------- @@ -30,7 +32,13 @@ class HNHN(torch.nn.Module): """ def __init__( - self, in_channels, hidden_channels, incidence_1, n_layers=2, layer_drop=0.2 + self, + in_channels, + hidden_channels, + incidence_1, + n_layers=2, + layer_drop=0.2, + **kwargs, ): super().__init__() @@ -39,6 +47,7 @@ def __init__( in_channels=in_channels if i == 0 else hidden_channels, hidden_channels=hidden_channels, incidence_1=incidence_1, + **kwargs, ) for i in range(n_layers) ) diff --git a/topomodelx/nn/hypergraph/hypergat.py b/topomodelx/nn/hypergraph/hypergat.py index e8449cc7..02d480c7 100644 --- a/topomodelx/nn/hypergraph/hypergat.py +++ b/topomodelx/nn/hypergraph/hypergat.py @@ -16,8 +16,10 @@ class HyperGAT(torch.nn.Module): Dimension of the hidden features. n_layers : int, default = 2 Amount of message passing layers. - layer_drop: float, default = 0.2 + layer_drop : float, default = 0.2 Dropout rate for the hidden features. + **kwargs : optional + Additional arguments for the inner layers. References ---------- @@ -32,6 +34,7 @@ def __init__( hidden_channels, n_layers=2, layer_drop=0.2, + **kwargs, ): super().__init__() @@ -39,6 +42,7 @@ def __init__( HyperGATLayer( in_channels=in_channels if i == 0 else hidden_channels, hidden_channels=hidden_channels, + **kwargs, ) for i in range(n_layers) ) @@ -49,8 +53,8 @@ def forward(self, x_0, incidence_1): Parameters ---------- - x_1 : torch.Tensor, shape = (n_edges, channels_edge) - Edge features. + x_0 : torch.Tensor, shape = (n_nodes, channels_nodes) + Node features. incidence_1 : torch.Tensor, shape = (n_nodes, n_edges) Boundary matrix of rank 1. diff --git a/topomodelx/nn/hypergraph/hypersage.py b/topomodelx/nn/hypergraph/hypersage.py index a9744225..588461ab 100644 --- a/topomodelx/nn/hypergraph/hypersage.py +++ b/topomodelx/nn/hypergraph/hypersage.py @@ -14,10 +14,12 @@ class HyperSAGE(torch.nn.Module): Dimension of the input features. hidden_channels : int Dimension of the hidden features. - n_layer : int, default = 2 + n_layers : int, default = 2 Amount of message passing layers. alpha : int, default = -1 - Max number of nodes in a neighborhood to consider. If -1 it considers all the nodes.รน + Max number of nodes in a neighborhood to consider. If -1 it considers all the nodes. + **kwargs : optional + Additional arguments for the inner layers. References ---------- @@ -39,14 +41,14 @@ def __init__(self, in_channels, hidden_channels, n_layers=2, alpha=-1, **kwargs) for i in range(n_layers) ) - def forward(self, x_0, incidence): + def forward(self, x_0, incidence_1): """Forward computation through layers, then linear layer, then global max pooling. Parameters ---------- - x : torch.Tensor, shape = (n_nodes, features_nodes) + x_0 : torch.Tensor, shape = (n_nodes, features_nodes) Edge features. - incidence : torch.Tensor, shape = (n_nodes, n_edges) + incidence_1 : torch.Tensor, shape = (n_nodes, n_edges) Boundary matrix of rank 1. Returns @@ -55,6 +57,6 @@ def forward(self, x_0, incidence): Label assigned to whole complex. """ for layer in self.layers: - x_0 = layer.forward(x_0, incidence) + x_0 = layer.forward(x_0, incidence_1) return x_0 diff --git a/topomodelx/nn/hypergraph/unigcn.py b/topomodelx/nn/hypergraph/unigcn.py index 0852fb23..58998943 100644 --- a/topomodelx/nn/hypergraph/unigcn.py +++ b/topomodelx/nn/hypergraph/unigcn.py @@ -16,6 +16,8 @@ class UniGCN(torch.nn.Module): Dimension of the hidden features. n_layers : int, default = 2 Amount of message passing layers. + **kwargs : optional + Additional arguments for the inner layers. References ---------- @@ -30,6 +32,7 @@ def __init__( in_channels, hidden_channels, n_layers=2, + **kwargs, ): super().__init__() @@ -37,6 +40,7 @@ def __init__( UniGCNLayer( in_channels=in_channels if i == 0 else hidden_channels, hidden_channels=hidden_channels, + **kwargs, ) for i in range(n_layers) ) diff --git a/topomodelx/nn/hypergraph/unigcnii.py b/topomodelx/nn/hypergraph/unigcnii.py index 47d3673c..6aed07c8 100644 --- a/topomodelx/nn/hypergraph/unigcnii.py +++ b/topomodelx/nn/hypergraph/unigcnii.py @@ -28,6 +28,8 @@ class UniGCNII(torch.nn.Module): Dropout rate for the hidden features. use_norm : bool, default=False Whether to apply row normalization after every layer. + **kwargs : optional + Additional arguments for the inner layers. References ---------- @@ -47,6 +49,7 @@ def __init__( input_drop=0.2, layer_drop=0.2, use_norm=False, + **kwargs, ): super().__init__() layers = [] @@ -65,6 +68,7 @@ def __init__( alpha=alpha, beta=beta, use_norm=use_norm, + **kwargs, ) ) diff --git a/topomodelx/nn/hypergraph/unigin.py b/topomodelx/nn/hypergraph/unigin.py index d9549ada..1c0cd0a3 100644 --- a/topomodelx/nn/hypergraph/unigin.py +++ b/topomodelx/nn/hypergraph/unigin.py @@ -26,7 +26,8 @@ class UniGIN(torch.nn.Module): Whether to make eps a trainable parameter. use_norm : bool, default=False Whether to apply row normalization after every layer. - + **kwargs : optional + Additional arguments for the inner layers. References ---------- @@ -46,6 +47,7 @@ def __init__( eps=0, train_eps=False, use_norm=False, + **kwargs, ): super().__init__() @@ -60,6 +62,7 @@ def __init__( eps=eps, train_eps=train_eps, use_norm=use_norm, + **kwargs, ) for _ in range(n_layers) ) diff --git a/topomodelx/nn/hypergraph/unisage.py b/topomodelx/nn/hypergraph/unisage.py index a43b89dc..4b2cce19 100644 --- a/topomodelx/nn/hypergraph/unisage.py +++ b/topomodelx/nn/hypergraph/unisage.py @@ -26,9 +26,10 @@ class UniSAGE(torch.nn.Module): Aggregator function for hyperedges. v_aggr : Literal["sum", "mean",], default="mean" Aggregator function for nodes. - use_norm : boolean + use_norm : bool Whether to apply row normalization after every layer. - + **kwargs : optional + Additional arguments for the inner layers. References ---------- @@ -54,6 +55,7 @@ def __init__( "mean", ] = "mean", use_norm: bool = False, + **kwargs, ): super().__init__() @@ -67,6 +69,7 @@ def __init__( e_aggr=e_aggr, v_aggr=v_aggr, use_norm=use_norm, + **kwargs, ) for i in range(n_layers) )