athena__neural_operator

Module containing implementation of a simple neural operator layer

This module implements a neural operator layer that approximates a discretized integral operator. The layer combines a standard affine transform (local component) with a mean-field integral operator (global/non-local component):

$\mathbf{v} = \sigma\!\left(\mathbf{W}\mathbf{u} + \mathbf{w}_k \langle\mathbf{u}\rangle + \mathbf{b}\right)$

where: - $\mathbf{u} \in \mathbb{R}^{n_{in}}$ is the input (discretised function) - $\mathbf{W} \in \mathbb{R}^{n_{out} \times n_{in}}$ are the local weights - $\mathbf{w}_k \in \mathbb{R}^{n_{out}}$ are the integral kernel weights - $\langle\mathbf{u}\rangle = \frac{1}{n_{in}}\sum_j u_j$ is the input mean - $\mathbf{b} \in \mathbb{R}^{n_{out}}$ is the bias - $\sigma$ is the activation function

The global mean $\langle\mathbf{u}\rangle$ acts as a rank-1 approximation to a continuous integral operator $\mathcal{K}[u](x) = \kappa(x)\int u\,\mathrm{d}y$ , where $\mathbf{w}_k$ discretises $\kappa$ . Using this layer stacked in sequence provides a resolution-invariant building block similar to the graph neural operator family.

Number of parameters: - with bias: $n_{out}(n_{in} + 1) + n_{out} = n_{out}(n_{in}+2)$ - without bias: $n_{out}(n_{in} + 1)$

Uses

- athena__onnx_nop_utils
- diffstruc
- athena__misc_types
- coreutils
- athena__initialiser_data
- athena__base_layer

Interfaces

public interface neural_operator_layer_type

Interface for setting up the neural operator layer

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

Setup a neural operator layer

Arguments

Type	Intent	Optional		Name
integer,	intent(in)		::	num_outputs	Number of outputs
integer,	intent(in),	optional	::	num_inputs	Number of inputs
logical,	intent(in),	optional	::	use_bias	Whether to use bias
class(*),	intent(in),	optional	::	activation	Activation function
class(*),	intent(in),	optional	::	kernel_initialiser	Kernel and bias initialisers
class(*),	intent(in),	optional	::	bias_initialiser	Kernel and bias initialisers
integer,	intent(in),	optional	::	verbose	Verbosity level

Return Value type(neural_operator_layer_type)

Instance of the neural operator layer

Derived Types

type, public, extends(learnable_layer_type) :: neural_operator_layer_type

Type for a 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
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
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		::	num_inputs			Number of inputs (discretisation points of the input function)
integer,	public		::	num_outputs			Number of outputs (discretisation points of the output function)
integer,	public		::	num_params	=	0	Number of learnable parameters
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
type(array_type),	public,	dimension(1)	::	z			Temporary array for pre-activation values (forward propagation)

Constructor

Interface for setting up the neural operator layer

private module function layer_setup (num_outputs, num_inputs, use_bias, activation, kernel_initialiser, bias_initialiser, verbose)	Setup a neural operator layer

Finalizations Procedures

final :: finalise_neural_operator	Finalise neural operator layer

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_neural_operator	Emit format-aware ONNX nodes for the layer
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_neural_operator	Forward propagation
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_neural_operator	Get layer attributes for ONNX export
procedure, public, pass(this) :: get_gradients	Get parameter gradients of layer
procedure, public, pass(this) :: get_num_params => get_num_params_neural_operator	Get the number of parameters for the neural operator layer
procedure, public, pass(this) :: get_params	Get learnable parameters of layer
procedure, public, pass(this) :: init => init_neural_operator	Initialise the neural operator layer
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_neural_operator	Print the layer to a file
procedure, public, pass(this) :: read => read_neural_operator	Read the layer from a file
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_base	Set the graph structure of the input data !! this is adjacency and edge weighting
procedure, public, pass(this) :: set_hyperparams => set_hyperparams_neural_operator	Set the hyperparameters for the neural operator layer
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

Functions

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

Read neural operator layer from file and return as base_layer_type

Arguments

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

Return Value class(base_layer_type), allocatable

Allocated layer instance

private function get_attributes_neural_operator(this) result(attributes)

Return list of neural operator attributes for ONNX export

Arguments

Type	Intent	Optional	Attributes		Name
class(neural_operator_layer_type),	intent(in)			::	this	Instance of the neural operator layer

Return Value type(onnx_attribute_type), allocatable, dimension(:)

List of attributes for ONNX export

private pure function get_num_params_neural_operator(this) result(num_params)

Get the number of parameters for the neural operator layer

Arguments

Type	Intent	Optional	Attributes		Name
class(neural_operator_layer_type),	intent(in)			::	this	Instance of the neural operator layer

Return Value integer

Number of parameters

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

Setup a neural operator layer

Arguments

Type	Intent	Optional		Name
integer,	intent(in)		::	num_outputs	Number of outputs
integer,	intent(in),	optional	::	num_inputs	Number of inputs
logical,	intent(in),	optional	::	use_bias	Whether to use bias
class(*),	intent(in),	optional	::	activation	Activation function
class(*),	intent(in),	optional	::	kernel_initialiser	Kernel and bias initialisers
class(*),	intent(in),	optional	::	bias_initialiser	Kernel and bias initialisers
integer,	intent(in),	optional	::	verbose	Verbosity level

Return Value type(neural_operator_layer_type)

Instance of the neural operator layer

Subroutines

private subroutine emit_onnx_nodes_neural_operator(this, prefix, nodes, num_nodes, max_nodes, inits, num_inits, max_inits, input_name, is_last_layer, format)

Emit decomposed standard ONNX nodes for a Neural Operator layer.

Arguments

Type	Intent	Optional	Attributes		Name
class(neural_operator_layer_type),	intent(in)			::	this	Neural operator layer instance
character(len=*),	intent(in)			::	prefix	Layer name prefix
type(onnx_node_type),	intent(inout),		dimension(:)	::	nodes	Node accumulator
integer,	intent(inout)			::	num_nodes	Node counter
integer,	intent(in)			::	max_nodes	Node limit
type(onnx_initialiser_type),	intent(inout),		dimension(:)	::	inits	Initialiser accumulator
integer,	intent(inout)			::	num_inits	Initialiser counter
integer,	intent(in)			::	max_inits	Initialiser limit
character(len=*),	intent(in),	optional		::	input_name	Name of the input tensor
logical,	intent(in),	optional		::	is_last_layer	Whether this is the last layer
integer,	intent(in),	optional		::	format	Export format selector

private subroutine finalise_neural_operator(this)

Finalise neural operator layer

Arguments

Type	Intent	Optional	Attributes		Name
type(neural_operator_layer_type),	intent(inout)			::	this	Instance of the neural operator layer

private subroutine forward_neural_operator(this, input)

Forward propagation for the neural operator layer

Arguments

Type	Intent	Optional	Attributes		Name
class(neural_operator_layer_type),	intent(inout)			::	this	Instance of the neural operator layer
class(array_type),	intent(in),		dimension(:,:)	::	input	Input values

private subroutine init_neural_operator(this, input_shape, verbose)

Initialise neural operator layer

Arguments

Type	Intent	Optional	Attributes		Name
class(neural_operator_layer_type),	intent(inout)			::	this	Instance of the neural operator layer
integer,	intent(in),		dimension(:)	::	input_shape	Input shape
integer,	intent(in),	optional		::	verbose	Verbosity level

private subroutine print_to_unit_neural_operator(this, unit)

Print neural operator layer to unit

Arguments

Type	Intent	Optional	Attributes		Name
class(neural_operator_layer_type),	intent(in)			::	this	Instance of the neural operator layer
integer,	intent(in)			::	unit	File unit

private subroutine read_neural_operator(this, unit, verbose)

Read neural operator layer from file

Arguments

Type	Intent	Optional		Name
class(neural_operator_layer_type),	intent(inout)		::	this	Instance of the neural operator layer
integer,	intent(in)		::	unit	Unit number
integer,	intent(in),	optional	::	verbose	Verbosity level

private subroutine set_hyperparams_neural_operator(this, num_outputs, use_bias, activation, kernel_initialiser, bias_initialiser, verbose)

Set the hyperparameters for the neural operator layer

Arguments

Type	Intent	Optional	Attributes		Name
class(neural_operator_layer_type),	intent(inout)			::	this	Instance of the neural operator layer
integer,	intent(in)			::	num_outputs	Number of outputs
logical,	intent(in)			::	use_bias	Whether to use bias
class(base_actv_type),	intent(in),		allocatable	::	activation	Activation function
class(base_init_type),	intent(in),		allocatable	::	kernel_initialiser	Kernel and bias initialisers
class(base_init_type),	intent(in),		allocatable	::	bias_initialiser	Kernel and bias initialisers
integer,	intent(in),	optional		::	verbose	Verbosity level

athena__neural_operator_layer Module

Contents

Interfaces

Derived Types

Functions

Subroutines

Uses

Interfaces

public interface neural_operator_layer_type

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

Arguments

Return Value type(neural_operator_layer_type)

Derived Types

type, public, extends(learnable_layer_type) :: neural_operator_layer_type

Components

Constructor

Finalizations Procedures

Type-Bound Procedures

Functions

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

Arguments

Return Value class(base_layer_type), allocatable

private function get_attributes_neural_operator(this) result(attributes)

Arguments

Return Value type(onnx_attribute_type), allocatable, dimension(:)

private pure function get_num_params_neural_operator(this) result(num_params)

Arguments

Return Value integer

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

Arguments

Return Value type(neural_operator_layer_type)

Subroutines

private subroutine emit_onnx_nodes_neural_operator(this, prefix, nodes, num_nodes, max_nodes, inits, num_inits, max_inits, input_name, is_last_layer, format)

Arguments

private subroutine finalise_neural_operator(this)

Arguments

private subroutine forward_neural_operator(this, input)

Arguments

private subroutine init_neural_operator(this, input_shape, verbose)

Arguments

private subroutine print_to_unit_neural_operator(this, unit)

Arguments

private subroutine read_neural_operator(this, unit, verbose)

Arguments

private subroutine set_hyperparams_neural_operator(this, num_outputs, use_bias, activation, kernel_initialiser, bias_initialiser, verbose)

Arguments

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