athena__graph_nop

Module containing implementation of a Graph Neural Operator (GNO) layer

This module implements a Graph Neural Operator layer that learns a continuous kernel on graph edges. It combines a learnable kernel network (small MLP) evaluated on relative coordinates with a linear transform of the node features:

$h_i^{(l+1)} = \sigma\!\left( \mathbf{W}\,h_i^{(l)} + \sum_{j \in \mathcal{N}(i)} \kappa_\theta(x_i - x_j)\,h_j^{(l)} \right)$

where: - $h_i^{(l)} \in \mathbb{R}^{F_{in}}$ is the node feature at layer l - $x_i \in \mathbb{R}^{d}$ is the node coordinate / attribute - $\kappa_\theta \colon \mathbb{R}^d \to \mathbb{R}^{F_{out} \times F_{in}}$ is a learnable kernel MLP - $\mathbf{W} \in \mathbb{R}^{F_{out} \times F_{in}}$ is a learnable linear (bypass) transform - $\sigma$ is the activation function

The kernel MLP has one hidden layer: $\kappa_\theta(\Delta x) = V\,\text{relu}(U\,\Delta x + b_u) + b_v$ where $U \in \mathbb{R}^{H \times d}$ , $V \in \mathbb{R}^{(F_{out} F_{in}) \times H}$ , and $H$ is the hidden width of the kernel network.

Input layout: input(1,s) = node features [F_in x num_vertices] input(2,s) = edge geometry / relative coords [d x num_edges]

Number of learnable parameters: Kernel MLP: $H d + H + (F_{out} F_{in}) H + F_{out} F_{in}$ Linear: $F_{out} F_{in}$ Bias: $F_{out}$ (optional)

This layer extends the message passing layer type and uses the diffstruc autodiff framework to support physics-informed neural networks (PINNs). The forward pass builds a computation graph through two differentiable operations: gno_kernel_eval (kernel MLP evaluation on every edge) and gno_aggregate (neighbour message aggregation), followed by the standard matmul, add_bias, and activation operations.

Uses

- athena__base_layer
- coreutils
- athena__diffstruc_extd
- diffstruc
- graphstruc
- athena__msgpass_layer
- athena__misc_types

Interfaces

public interface graph_nop_layer_type

private module function layer_setup(num_outputs, coord_dim, kernel_hidden, num_inputs, use_bias, activation, kernel_initialiser, bias_initialiser, verbose) result(layer)

Arguments

Type	Intent	Optional		Name
integer,	intent(in)		::	num_outputs	Number of output node features
integer,	intent(in)		::	coord_dim	Dimensionality of edge geometric features
integer,	intent(in),	optional	::	kernel_hidden	Hidden width of kernel MLP (default: num_outputs)
integer,	intent(in),	optional	::	num_inputs	Number of input node features (deferred if absent)
logical,	intent(in),	optional	::	use_bias	Whether to use bias (default: .true.)
class(*),	intent(in),	optional	::	activation	Activation function
class(*),	intent(in),	optional	::	kernel_initialiser	Parameter initialisers
class(*),	intent(in),	optional	::	bias_initialiser	Parameter initialisers
integer,	intent(in),	optional	::	verbose	Verbosity level

Return Value type(graph_nop_layer_type)

Derived Types

type, public, extends(msgpass_layer_type) :: graph_nop_layer_type

Type for a Graph Neural Operator layer

Components

Type	Visibility	Attributes		Name		Initial
class(base_actv_type),	public,	allocatable	::	activation			Activation function
class(base_init_type),	public,	allocatable	::	bias_init			Initialisers for kernel and bias
character(len=14),	public		::	bias_initialiser	=	''	Initialisers for kernel and bias
integer,	public,	allocatable, dimension(:)	::	bias_shape			Shape of biases
integer,	public		::	coord_dim	=	0	Dimensionality of edge geometric features (d)
type(graph_type),	public,	allocatable, dimension(:)	::	graph			Graph structure of input data
integer,	public		::	id			Unique identifier
logical,	public		::	inference	=	.false.	Inference mode
integer,	public		::	input_rank	=	0	Rank of input data
integer,	public,	allocatable, dimension(:)	::	input_shape			Input shape
integer,	public		::	kernel_hidden	=	0	Hidden width of the kernel MLP (H)
class(base_init_type),	public,	allocatable	::	kernel_init			Initialisers for kernel and bias
character(len=14),	public		::	kernel_initialiser	=	''	Initialisers for kernel and bias
character(len=:),	public,	allocatable	::	name			Layer name
integer,	public,	dimension(:), allocatable	::	num_edge_features			Number of edge features for each time step
integer,	public		::	num_output_edge_features			Number of output edge features
integer,	public		::	num_output_vertex_features			Number of output vertex features
integer,	public		::	num_outputs			Number of outputs (if output is not graph structure)
integer,	public		::	num_params	=	0	Number of learnable parameters
integer,	public,	dimension(:), allocatable	::	num_params_msg			Number of learnable parameters for each message
integer,	public		::	num_params_readout			Number of learnable parameters for the readout
integer,	public		::	num_time_steps			Number of time steps
integer,	public,	dimension(:), allocatable	::	num_vertex_features			Number of vertex features for each time step
class(array_type),	public,	allocatable, dimension(:,:)	::	output			Output
integer,	public		::	output_rank	=	0	Rank of output data
integer,	public,	allocatable, dimension(:)	::	output_shape			Output shape
type(array_type),	public,	allocatable, dimension(:)	::	params			Learnable parameters
character(len=20),	public		::	subtype	=	repeat(" ", 20)
character(len=4),	public		::	type	=	'base'	Layer type
logical,	public		::	use_bias	=	.false.	Layer has bias
logical,	public		::	use_graph_input	=	.false.	Use graph input
logical,	public		::	use_graph_output	=	.false.	Use graph output
integer,	public,	allocatable, dimension(:,:)	::	weight_shape			Shape of weights

Constructor

private module function layer_setup (num_outputs, coord_dim, kernel_hidden, num_inputs, use_bias, activation, kernel_initialiser, bias_initialiser, verbose)

Type-Bound Procedures

procedure, public :: add_t_t => add_learnable	Add two layers
procedure, public, pass(this) :: build_from_onnx => build_from_onnx_base	Build layer from ONNX node and initialiser
procedure, public, pass(this) :: emit_onnx_graph_inputs => emit_onnx_graph_inputs_base	Emit graph input tensor declarations for this layer
procedure, public, pass(this) :: emit_onnx_nodes => emit_onnx_nodes_base	Emit ONNX JSON nodes for this layer (format-aware and polymorphic)
procedure, public, pass(this) :: extract_output => extract_output_base	Extract the output of the layer as a standard real array
procedure, public, pass(this) :: forward => forward_msgpass	Forward pass for message passing layer
procedure, public, pass(this) :: forward_eval => forward_eval_base	Forward pass of layer and return output for evaluation
procedure, public, pass(this) :: get_attributes => get_attributes_base	Get the attributes of the layer (for ONNX export)
procedure, public, pass(this) :: get_gradients	Get parameter gradients of layer
procedure, public, pass(this) :: get_num_params => get_num_params_gno
procedure, public, pass(this) :: get_params	Get learnable parameters of layer
procedure, public, pass(this) :: init => init_gno
procedure, public, pass(this) :: nullify_graph => nullify_graph_base	Nullify the forward pass data of the layer to free memory Read more…
generic, public :: operator(+) => add_t_t	Operator overloading for addition
procedure, public, pass(this) :: print => print_base	Print the layer to a file with additional information
procedure, public, pass(this) :: print_to_unit => print_to_unit_gno
procedure, public, pass(this) :: read => read_gno
procedure, public, pass(this) :: reduce => reduce_learnable	Merge another learnable layer into this one
procedure, public, pass(this) :: set_gradients	Set learnable parameters of layer
procedure, public, pass(this) :: set_graph => set_graph_msgpass
procedure, public, pass(this) :: set_hyperparams => set_hyperparams_gno
procedure, public, pass(this) :: set_params	Set learnable parameters of layer
procedure, public, pass(this) :: set_rank => set_rank_base	Set the input and output ranks of the layer
procedure, public, pass(this) :: set_shape => set_shape_base	Set the input shape of the layer
procedure, public, pass(this) :: update_message => update_message_gno
procedure, public, pass(this) :: update_readout => update_readout_gno

Functions

public function read_graph_nop_layer(unit, verbose) result(layer)

Read a graph NOP layer from file and return

Arguments

Type	Intent	Optional	Attributes		Name
integer,	intent(in)			::	unit	Input unit number
integer,	intent(in),	optional		::	verbose	Verbosity level

Return Value class(base_layer_type), allocatable

Allocated base-layer instance containing the result

private pure function get_num_params_gno(this) result(num_params)

Get the number of learnable parameters

Arguments

Type	Intent	Optional	Attributes		Name
class(graph_nop_layer_type),	intent(in)			::	this	Layer instance

Return Value integer

Total number of learnable parameters

private module function layer_setup(num_outputs, coord_dim, kernel_hidden, num_inputs, use_bias, activation, kernel_initialiser, bias_initialiser, verbose) result(layer)

Arguments

Type	Intent	Optional		Name
integer,	intent(in)		::	num_outputs	Number of output node features
integer,	intent(in)		::	coord_dim	Dimension of edge coordinate features
integer,	intent(in),	optional	::	kernel_hidden	Hidden width of the kernel MLP
integer,	intent(in),	optional	::	num_inputs	Number of input node features when known at construction time
logical,	intent(in),	optional	::	use_bias	Whether to allocate an output bias
class(*),	intent(in),	optional	::	activation	Activation function specification
class(*),	intent(in),	optional	::	kernel_initialiser	Kernel and bias initialiser specifications
class(*),	intent(in),	optional	::	bias_initialiser	Kernel and bias initialiser specifications
integer,	intent(in),	optional	::	verbose	Verbosity level

Return Value type(graph_nop_layer_type)

Constructed graph neural operator layer

Subroutines

private subroutine init_gno(this, input_shape, verbose)

Initialise the Graph Neural Operator layer

Arguments

Type	Intent	Optional	Attributes		Name
class(graph_nop_layer_type),	intent(inout)			::	this	Layer instance to initialise
integer,	intent(in),		dimension(:)	::	input_shape	Input feature/vertex shape
integer,	intent(in),	optional		::	verbose	Verbosity level

private subroutine print_to_unit_gno(this, unit)

Print graph neural operator settings and parameters to a unit

Arguments

Type	Intent	Optional	Attributes		Name
class(graph_nop_layer_type),	intent(in)			::	this	Layer instance to print
integer,	intent(in)			::	unit	Output unit number

private subroutine read_gno(this, unit, verbose)

Arguments

Type	Intent	Optional		Name
class(graph_nop_layer_type),	intent(inout)		::	this	Layer instance to populate from file data
integer,	intent(in)		::	unit	Input unit number
integer,	intent(in),	optional	::	verbose	Verbosity level

private subroutine set_hyperparams_gno(this, num_outputs, coord_dim, kernel_hidden, use_bias, activation, kernel_initialiser, bias_initialiser, verbose)

Arguments

Type	Intent	Optional	Attributes		Name
class(graph_nop_layer_type),	intent(inout)			::	this	Layer instance to configure
integer,	intent(in)			::	num_outputs	Number of output node features
integer,	intent(in)			::	coord_dim	Dimension of edge coordinate features
integer,	intent(in)			::	kernel_hidden	Hidden width of the kernel MLP
logical,	intent(in)			::	use_bias	Whether to use a bias term
class(base_actv_type),	intent(in),		allocatable	::	activation	Activation function object
class(base_init_type),	intent(in),		allocatable	::	kernel_initialiser	Kernel and bias initialiser objects
class(base_init_type),	intent(in),		allocatable	::	bias_initialiser	Kernel and bias initialiser objects
integer,	intent(in),	optional		::	verbose	Verbosity level

private subroutine update_message_gno(this, input)

Update message for the Graph Neural Operator layer

Arguments

Type	Intent	Optional	Attributes		Name
class(graph_nop_layer_type),	intent(inout),		target	::	this	Layer instance to execute
class(array_type),	intent(in),		dimension(:,:), target	::	input	Input node-feature and edge-feature tensors

private subroutine update_readout_gno(this)

No graph-level readout needed — GNO produces node-level output

Arguments

Type	Intent	Optional	Attributes		Name
class(graph_nop_layer_type),	intent(inout),		target	::	this	Layer instance retained for interface compatibility

athena__graph_nop_layer Module

Contents

Interfaces

Derived Types

Functions

Subroutines

Uses

Interfaces

public interface graph_nop_layer_type

private module function layer_setup(num_outputs, coord_dim, kernel_hidden, num_inputs, use_bias, activation, kernel_initialiser, bias_initialiser, verbose) result(layer)

Arguments

Return Value type(graph_nop_layer_type)

Derived Types

type, public, extends(msgpass_layer_type) :: graph_nop_layer_type

Components

Constructor

Type-Bound Procedures

Functions

public function read_graph_nop_layer(unit, verbose) result(layer)

Arguments

Return Value class(base_layer_type), allocatable

private pure function get_num_params_gno(this) result(num_params)

Arguments

Return Value integer

private module function layer_setup(num_outputs, coord_dim, kernel_hidden, num_inputs, use_bias, activation, kernel_initialiser, bias_initialiser, verbose) result(layer)

Arguments

Return Value type(graph_nop_layer_type)

Subroutines

private subroutine init_gno(this, input_shape, verbose)

Arguments

private subroutine print_to_unit_gno(this, unit)

Arguments

private subroutine read_gno(this, unit, verbose)

Arguments

private subroutine set_hyperparams_gno(this, num_outputs, coord_dim, kernel_hidden, use_bias, activation, kernel_initialiser, bias_initialiser, verbose)

Arguments

private subroutine update_message_gno(this, input)

Arguments

private subroutine update_readout_gno(this)

Arguments

athena__graph_nop_layer Module

Contents

Interfaces

Derived Types

Functions

Subroutines

Uses

Interfaces

public interface graph_nop_layer_type

private module function layer_setup(num_outputs, coord_dim, kernel_hidden, num_inputs, use_bias, activation, kernel_initialiser, bias_initialiser, verbose) result(layer) Implementation →

Arguments

Return Value type(graph_nop_layer_type)

Derived Types

type, public, extends(msgpass_layer_type) :: graph_nop_layer_type

Components

Constructor

Type-Bound Procedures

Functions

public function read_graph_nop_layer(unit, verbose) result(layer)

Arguments

Return Value class(base_layer_type), allocatable

private pure function get_num_params_gno(this) result(num_params)

Arguments

Return Value integer

private module function layer_setup(num_outputs, coord_dim, kernel_hidden, num_inputs, use_bias, activation, kernel_initialiser, bias_initialiser, verbose) result(layer)

Arguments

Return Value type(graph_nop_layer_type)

Subroutines

private subroutine init_gno(this, input_shape, verbose)

Arguments

private subroutine print_to_unit_gno(this, unit)

Arguments

private subroutine read_gno(this, unit, verbose)

Arguments

private subroutine set_hyperparams_gno(this, num_outputs, coord_dim, kernel_hidden, use_bias, activation, kernel_initialiser, bias_initialiser, verbose)

Arguments

private subroutine update_message_gno(this, input)

Arguments

private subroutine update_readout_gno(this)

Arguments

private module function layer_setup(num_outputs, coord_dim, kernel_hidden, num_inputs, use_bias, activation, kernel_initialiser, bias_initialiser, verbose) result(layer)