gnp.models module

gnp.models.gnp

class gnp.models.gnp.GNP(node_dim: int, edge_dim: int, out_dim: int, layers: list[int], conv_name: str, conv_args: dict, nonlinearity: str, skip_connection: bool, device: str)

Bases: Module

General purpose model for Geometric Neural Operators (GNPs).

This model consists of a lifting layer, multiple graph convolution blocks, and a final projection layer.

Parameters:

• node_dim (int) -- Input dimension of node features.

• edge_dim (int) -- Input dimension of edge features.

• out_dim (int) -- Output dimension.

• layers (list[int]) -- List of hidden dimensions.

• conv_name (str) -- Name of the convolution block class to use.

• conv_args (dict) -- Arguments for the convolution block.

• nonlinearity (str) -- The name of the activation function to use.

• skip_connection (bool) -- Whether to use a skip connection for each layer.

• device (str) -- The device to place the model on. Defaults to "cuda".

forward(data: Data)

Forward pass.

Parameters:

data (Data) -- PyG Data object containing x, edge_index, and edge_attr.

Returns:

Output features.

Return type:

torch.Tensor

class gnp.models.gnp.PatchGNP(node_dim: int, out_dim: int, layers: list[int], num_channels: int, neurons: int, nonlinearity: str, device: str = 'cuda')

Bases: Module

Geometric Neural Operator for processing point cloud patches.

Parameters:

• node_dim (int) -- The dimensionality of the input node features (e.g., 3 for xyz).

• out_dim (int) -- The dimensionality of the final output vector for each patch.

• layers (list[int]) -- A list of integers defining the width of each layer in the network. The first element is the width after the initial lifting layer.

• num_channels (int) -- The number of channels to use in the 'block' type convolution.

• neurons (int) -- The number of neurons in the hidden layers of the MLPs within the convolutional layers.

• nonlinearity (str) -- The name of the activation function to use.

• device (str) -- The device to place the model on. Defaults to "cuda".

forward(x: Tensor, batch: Tensor)

Perform a forward pass on a batch of patches.

Parameters:

• x (torch.Tensor) -- Input coordinates/features (N, node_dim).

• batch (torch.Tensor) -- Batch indices (N,).

Returns:

A tensor of shape (num_patches, out_dim) containing the output vector for each patch.

Return type:

torch.Tensor

gnp.models.layers

class gnp.models.layers.BlockFactorizedConvolution(edge_dim: int, in_dim: int, out_dim: int, num_channels: int, neurons: int, nonlinearity: str = 'ReLU')

Bases: MessagePassing

A graph convolution layer using a block-factorized kernel.

Parameters:

• edge_dim (int) -- Dimension of edge features.

• in_dim (int) -- Input node feature dimension.

• out_dim (int) -- Output node feature dimension.

• num_channels (int) -- Number of blocks/channels for factorization.

• neurons (int) -- Hidden neurons in the kernel MLP.

• nonlinearity (str) -- Activation function, by default 'ReLU'.

forward(x: Tensor, edge_index: Tensor, edge_attr: Tensor) → Tensor

Perform message passing.

Parameters:

• x (torch.Tensor) -- Node features.

• edge_index (torch.Tensor) -- Graph connectivity.

• edge_attr (torch.Tensor) -- Edge features.

Returns:

Updated node features.

Return type:

torch.Tensor

message(x_j: Tensor, edge_attr: Tensor) → Tensor

Construct messages using the block kernel.

update(aggr_out: Tensor) → Tensor

Update node embeddings.

class gnp.models.layers.BlockKernel(edge_dim: int, in_dim: int, out_dim: int, num_channels: int, neurons: int, nonlinearity: str = 'ReLU')

Bases: Module

An MLP that computes a block-factorized kernel for a graph convolution.

This module takes edge features as input and produces a weight matrix (kernel) that is factorized into smaller blocks. This is used to reduce the number of parameters in the convolution.

forward(x: Tensor) → Tensor

Compute the kernel weights from the input features.

Parameters:

x (torch.Tensor) -- The input features.

Returns:

The computed kernel, reshaped into a block-diagonal-like structure of shape (..., num_channels, head_in, head_out).

Return type:

torch.Tensor

class gnp.models.layers.ConvolutionBlock(in_dim: int, out_dim: int, edge_dim: int, conv_name: str, conv_args: dict, nonlinearity: str, skip: bool = True)

Bases: Module

A unified block combining a graph convolution, residual connection, and activation.

Parameters:

• in_dim (int) -- Input dimension.

• out_dim (int) -- Output dimension.

• edge_dim (int) -- Edge feature dimension.

• conv_name (str) -- Name of the convolution layer class.

• conv_args (dict) -- Arguments for the convolution layer.

• nonlinearity (str) -- Activation function name.

• skip (bool) -- Whether to use a skip connection, by default True.

forward(x: Tensor, edge_index: Tensor, edge_attr: Tensor, use_activation: bool = True) → Tensor

Forward pass.

Parameters:

• x (torch.Tensor) -- Node features.

• edge_index (torch.Tensor) -- Graph connectivity.

• edge_attr (torch.Tensor) -- Edge features.

• use_activation (bool) -- Whether to apply activation at the end, by default True.

Returns:

Updated features.

Return type:

torch.Tensor

class gnp.models.layers.FullKernel(edge_dim: int, in_dim: int, out_dim: int, neurons: int, nonlinearity: str = 'ReLU')

Bases: Module

An MLP for a graph convolution.

This module maps edge features to a dense weight matrix for graph convolution.

Parameters:

• edge_dim (int) -- Dimension of the input edge features (e.g., 3 for xyz).

• in_dim (int) -- Dimension of the input node features for the convolution.

• out_dim (int) -- Dimension of the output node features for the convolution.

• neurons (int) -- Number of neurons in the hidden layers of the MLP.

• nonlinearity (str) -- Nonlinearity to use in the hidden layers, by default 'ReLU'.

forward(x: Tensor) → Tensor

Compute the kernel weights from the input features.

Parameters:

x (torch.Tensor) -- The input features.

Returns:

The computed kernel of shape (..., in_dim, out_dim).

Return type:

torch.Tensor

class gnp.models.layers.GraphConvolution(edge_dim: int, in_dim: int, out_dim: int, neurons: int, nonlinearity: str = 'ReLU')

Bases: MessagePassing

Standard graph convolution layer with a continuous kernel using mean aggregation.

Parameters:

• edge_dim (int) -- Dimension of edge features.

• in_dim (int) -- Input node feature dimension.

• out_dim (int) -- Output node feature dimension.

• neurons (int) -- Hidden neurons in the kernel MLP.

• nonlinearity (str) -- Activation function, by default 'ReLU'.

forward(x: Tensor, edge_index: Tensor, edge_attr: Tensor) → Tensor

Perform message passing.

Parameters:

• x (torch.Tensor) -- Node features.

• edge_index (torch.Tensor) -- Graph connectivity (2, num_edges).

• edge_attr (torch.Tensor) -- Edge features.

Returns:

Updated node features.

Return type:

torch.Tensor

message(x_j: Tensor, edge_attr: Tensor) → Tensor

Construct messages using the kernel.

update(aggr_out: Tensor) → Tensor

Update node embeddings.

class gnp.models.layers.PatchSeparableBlockFactorizedConvolutionBlock(in_dim: int, out_dim: int, dim_x: int, num_channels: int, neurons: int, nonlinearity: str = 'ReLU', skip: bool = True)

Bases: Module

A separable graph convolution using a block-factorized kernel specifically meant for the PatchGNP.

This convolutional layer performs a two-step message passing operation. It first aggregates messages from neighbors and then applies a second transformation based on the point's own features.

Parameters:

• in_dim (int) -- Input node feature dimension.

• out_dim (int) -- Output node feature dimension.

• dim_x (int) -- Dimension of spatial features.

• num_channels (int) -- Number of channels for block factorization.

• neurons (int) -- Hidden neurons in MLPs.

• nonlinearity (str) -- Activation function, by default 'ReLU'.

• skip (bool) -- Whether to use a residual skip connection, by default True.

forward(x: Tensor, v: Tensor, batch: Tensor) → Tensor

Perform the forward pass for the separable convolution.

Parameters:

• x (torch.Tensor) -- The spatial/edge features (e.g., local coordinates).

• v (torch.Tensor) -- The input node features.

• batch (torch.Tensor) -- A tensor mapping each node to its corresponding patch index.

Returns:

The updated node features after convolution.

Return type:

torch.Tensor

class gnp.models.layers.SeparableConvolution(in_dim: int, out_dim: int, edge_dim: int, neurons: int, kernel_name: str, kernel_args: dict, nonlinearity: str = 'ReLU')

Bases: Module

A separable convolution layer that evaluates a separable kernel on nodes instead of edges.

This layer applies convolution in two steps: convolution using a kernel that only depends on the source nodes followed by multiplication by a kernel that only depends on the target node.

Parameters:

• in_dim (int) -- Input node feature dimension.

• out_dim (int) -- Output node feature dimension.

• edge_dim (int) -- Dimension of spatial/edge features.

• neurons (int) -- Number of neurons in the kernel MLP.

• kernel_name (str) -- Name of the kernel class to use (e.g., 'BlockKernel').

• kernel_args (dict) -- Additional arguments for the kernel class.

• nonlinearity (str) -- Activation function, by default 'ReLU'.

block_matmul(x: Tensor, weights: Tensor) → Tensor

Perform matrix multiplication with weigths from the BlockKernel

Parameters:

• x (torch.Tensor) -- Features to be multiplied against kernel

• weights (torch.Tensor) -- Kernel weights

Returns:

Matrix multiplication product of x and weights.

Return type:

torch.Tensor

forward(x: Tensor, edge_index: Tensor, edge_attr: Tensor) → Tensor

Forward pass.

Parameters:

• x (torch.Tensor) -- Node features of shape (N, in_dim).

• edge_index (torch.Tensor) -- Graph connectivity.

• edge_attr (torch.Tensor) -- Spatial node features of shape (N, edge_dim).

Returns:

Updated node features.

Return type:

torch.Tensor

full_matmul(x: Tensor, weights: Tensor) → Tensor

Perform matrix multiplication with weigths from the FullKernel

Parameters:

• x (torch.Tensor) -- Features to be multiplied against kernel

• weights (torch.Tensor) -- Kernel weights

Returns:

Matrix multiplication product of x and weights.

Return type:

torch.Tensor

gnp.models.layers.get_activation(name: str) → Module

Helper to retrieve activation functions.

Module contents