"""Layers for PELICAN architecture."""
import torch
from torch import nn
from .primitives import (
aggregate_0to2,
aggregate_1to2,
aggregate_2to0,
aggregate_2to1,
aggregate_2to2,
bell_number,
)
ACTIVATION = {
"leaky_relu": nn.LeakyReLU(),
"relu": nn.ReLU(),
"gelu": nn.GELU(),
"tanh": nn.Tanh(),
"sigmoid": nn.Sigmoid(),
"silu": nn.SiLU(),
}
[docs]
class GeneralAggregator(nn.Module):
"""General aggregator class."""
def __init__(
self,
aggregator: callable,
in_rank: int,
out_rank: int,
in_channels: int,
out_channels: int,
map_multipliers: bool = True,
factorize: bool = False,
aggr: str = "mean",
):
"""
Parameters
----------
aggregator : callable
Aggregation function to use.
in_rank : int
Input rank (0 for graph, 1 for nodes, 2 for edges).
out_rank : int
Output rank (0 for graph, 1 for nodes, 2 for edges).
in_channels : int
Number of input channels.
out_channels : int
Number of output channels.
map_multipliers : bool
Whether to use learnable multipliers for each aggregation map, by default True.
factorize : bool
Whether to use factorized coefficients, by default False.
Factorization reduces the number of parameters.
aggr : str
Aggregation method to use ('mean', 'sum', 'prod', 'amin', 'amax'), by default 'mean'.
"""
super().__init__()
self.num_maps = bell_number(in_rank + out_rank)
self.in_channels = in_channels
self.aggr = aggr
self.aggregator = aggregator
self.factorize = factorize
self.map_multipliers = nn.Parameter(torch.ones(self.num_maps)) if map_multipliers else None
if factorize:
self.coeffs00 = nn.Parameter(torch.empty(in_channels, self.num_maps))
self.coeffs01 = nn.Parameter(torch.empty(out_channels, self.num_maps))
self.coeffs10 = nn.Parameter(torch.empty(in_channels, out_channels))
self.coeffs11 = nn.Parameter(torch.empty(in_channels, out_channels))
nn.init.kaiming_uniform_(self.coeffs00, nonlinearity="linear")
nn.init.kaiming_uniform_(self.coeffs01, nonlinearity="linear")
nn.init.kaiming_uniform_(self.coeffs10, nonlinearity="linear")
nn.init.kaiming_uniform_(self.coeffs11, nonlinearity="linear")
else:
self.coeffs_direct = nn.Parameter(torch.empty(in_channels, out_channels, self.num_maps))
nn.init.kaiming_uniform_(self.coeffs_direct, nonlinearity="relu")
@property
def coeffs(self):
if self.factorize:
coeffs = self.coeffs00.unsqueeze(1) * self.coeffs10.unsqueeze(
2
) + self.coeffs01.unsqueeze(0) * self.coeffs11.unsqueeze(2)
else:
coeffs = self.coeffs_direct
if self.map_multipliers is not None:
coeffs = coeffs * self.map_multipliers.view(1, 1, self.num_maps)
return coeffs
[docs]
def forward(self, x, *args, **kwargs):
"""Forward pass of the aggregator.
Parameters
----------
x : torch.Tensor
Input features of shape (in_objects, in_channels).
*args
Additional arguments to pass to the aggregator.
**kwargs
Additional keyword arguments to pass to the aggregator.
Returns
-------
out : torch.Tensor
Output features of shape (out_objects, in_channels).
"""
x = self.aggregator(x, *args, reduce=self.aggr, **kwargs)
in_objects, in_channels, num_maps = x.shape
x_flat = x.reshape(in_objects, in_channels * num_maps)
coeffs_flat = self.coeffs.reshape(-1, in_channels * num_maps).contiguous()
out = torch.nn.functional.linear(x_flat, coeffs_flat)
return out
[docs]
class Aggregator2to2(GeneralAggregator):
"""Aggregator from edges (rank 2) to edges (rank 2)."""
def __init__(self, in_channels, out_channels, **kwargs):
super().__init__(aggregate_2to2, 2, 2, in_channels, out_channels, **kwargs)
[docs]
class Aggregator2to1(GeneralAggregator):
"""Aggregator from edges (rank 2) to nodes (rank 1)."""
def __init__(self, in_channels, out_channels, **kwargs):
super().__init__(aggregate_2to1, 2, 1, in_channels, out_channels, **kwargs)
[docs]
class Aggregator2to0(GeneralAggregator):
"""Aggregator from edges (rank 2) to graph (rank 0)."""
def __init__(self, in_channels, out_channels, **kwargs):
super().__init__(aggregate_2to0, 2, 0, in_channels, out_channels, **kwargs)
[docs]
class Aggregator1to2(GeneralAggregator):
"""Aggregator from nodes (rank 1) to edges (rank 2)."""
def __init__(self, in_channels, out_channels, **kwargs):
super().__init__(aggregate_1to2, 1, 2, in_channels, out_channels, **kwargs)
[docs]
class Aggregator0to2(GeneralAggregator):
"""Aggregator from graph (rank 0) to edges (rank 2)."""
def __init__(self, in_channels, out_channels, **kwargs):
super().__init__(aggregate_0to2, 0, 2, in_channels, out_channels, **kwargs)
[docs]
class PELICANBlock(nn.Module):
"""PELICAN edge-to-edge aggregation block.
A single PELICAN block consisting of a feedforward network and edge-to-edge aggregation."""
def __init__(
self,
hidden_channels: int,
increase_hidden_channels: float = 1.0,
activation: str = "leaky_relu",
dropout_prob: float | None = None,
**kwargs,
):
"""
Parameters
----------
hidden_channels : int
Number of hidden channels.
increase_hidden_channels : float
Factor to increase hidden channels in the feedforward network, by default 1.0.
activation : str
Activation function to use ('gelu', 'relu', 'leaky_relu', 'tanh', 'sigmoid', 'silu'), by default 'leaky_relu'.
dropout_prob: float
Dropout probability, by default None.
**kwargs
Additional keyword arguments to pass to the aggregator.
"""
super().__init__()
hidden_channels_2 = int(increase_hidden_channels * hidden_channels)
linear_in = nn.Linear(hidden_channels, hidden_channels_2)
self.activation = ACTIVATION[activation]
norm = nn.RMSNorm(normalized_shape=hidden_channels_2)
dropout = nn.Dropout(p=dropout_prob) if dropout_prob is not None else nn.Identity()
self.mlp = nn.ModuleList([linear_in, self.activation, norm, dropout])
self.aggregator = Aggregator2to2(
in_channels=hidden_channels_2,
out_channels=hidden_channels,
**kwargs,
)
[docs]
def forward(self, x, edge_index, batch, **kwargs):
"""Forward pass of the PELICAN block.
Parameters
----------
x : torch.Tensor
Edge-level features of shape (E, C).
edge_index : torch.Tensor
Edge index tensor of shape (2, E).
batch : torch.Tensor
Batch tensor of shape (N).
**kwargs
Returns
-------
x : torch.Tensor
Updated edge-level features of shape (E, C).
"""
for layer in self.mlp:
x = layer(x)
x = self.aggregator(x, edge_index=edge_index, batch=batch, **kwargs)
x = self.activation(x)
return x