GraphNeuralOperator#

class GraphNeuralOperator(lifting_operator, projection_operator, edge_features, n_layers=10, internal_n_layers=0, inner_size=None, internal_layers=None, internal_func=None, external_func=None, shared_weights=True)[source]#

Bases: KernelNeuralOperator

Graph Neural Operator model class.

The Graph Neural Operator is a general architecture for learning operators, which map functions to functions. It can be trained both with Supervised and Physics-Informed learning strategies. The Graph Neural Operator performs graph convolution by means of a Graph Neural Kernel.

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.

Initialization of the GraphNeuralOperator class.

Parameters:

  • lifting_operator (torch.nn.Module) – The lifting neural network mapping the input to its hidden dimension.

  • projection_operator (torch.nn.Module) – The projection neural network mapping the hidden representation to the output function.

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

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

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

  • inner_size (int) – The size of the hidden layers of the neural network inside each kernel layer. Default is None.

  • 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)[source]#

The forward pass of the Graph Neural Operator.

Parameters:

x (torch_geometric.data.Batch) – The input graph.

Returns:

The output tensor.

Return type:

torch.Tensor