GraphNeuralKernel#

class GraphNeuralKernel(width, edge_features, n_layers=2, internal_n_layers=0, internal_layers=None, inner_size=None, internal_func=None, external_func=None, shared_weights=False)[source]#

Bases: Module

Graph Neural Operator kernel model class.

This class implements the Graph Neural Operator kernel network.

See also

Original reference: Li, Z., Kovachki, N., Azizzadenesheli, K., Liu, B., Bhattacharya, K., Stuart, A., Anandkumar, A. (2020). Neural Operator: Graph Kernel Network for Partial Differential Equations. DOI: arXiv preprint arXiv:2003.03485

Example:
>>> import torch
>>> from pina.model._src.model.graph_neural_operator import (
...     GraphNeuralKernel
... )
>>> model = GraphNeuralKernel(width=16, edge_features=4)
>>> x = torch.randn(10, 16)
>>> edge_index = torch.randint(0, 10, (2, 20))
>>> edge_attr = torch.randn(20, 4)
>>> out = model(x, edge_index, edge_attr)

Initialization of the GraphNeuralKernel class.

Parameters:

  • width (int) – The width of the kernel.

  • edge_features (int) – The number of edge features.

  • n_layers (int) – The number of kernel layers. Default is 2.

  • internal_n_layers (int) – The number of layers of the neural network inside each kernel layer. Default is 0.

  • internal_layers (list[int] | tuple[int]) – The number of neurons for each layer of the neural network inside each kernel layer. Default is None.

  • internal_func (torch.nn.Module) – The activation function used inside each kernel layer. If None, it uses the torch.nn.Tanh activation. Default is None.

  • external_func (torch.nn.Module) – The activation function applied to the output of the each kernel layer. If None, it uses the torch.nn.Tanh activation. Default is None.

  • shared_weights (bool) – If True, the weights of each kernel layer are shared. Default is False.

forward(x, edge_index, edge_attr)[source]#

The forward pass of the Graph Neural Kernel.

Parameters:
Returns:

The output tensor.

Return type:

torch.Tensor