Models API Reference

The k3_node.models package contains complete, production-ready Graph Neural Network architectures implemented natively in Keras 3.

Standard GNN Architectures

BasicGNN

`k3_node.models.BasicGNN`

Bases: K3NodeHubMixin, Model

An abstract base class for implementing basic GNN models.

Parameters:

Name	Type	Description	Default
`in_channels`	`int or tuple`	Size of each input sample.	required
`hidden_channels`	`int`	Size of each hidden sample.	required
`num_layers`	`int`	Number of message passing layers.	required
`out_channels`	`int`	If not set to :obj:`None`, will apply a final linear transformation to convert hidden node embeddings to output size :obj:`out_channels`. (default: :obj:`None`)	`None`
`dropout`	`float`	Dropout probability. (default: :obj:`0.`)	`0.0`
`act`	`str or Callable`	The non-linear activation function to use. (default: :obj:`"relu"`)	`'relu'`
`act_first`	`bool`	If set to :obj:`True`, activation is applied before normalization. (default: :obj:`False`)	`False`
`act_kwargs`	`Dict[str, Any]`	Arguments passed to the respective activation function defined by :obj:`act`. (default: :obj:`None`)	`None`
`norm`	`str or Callable`	The normalization function to use. (default: :obj:`None`)	`None`
`norm_kwargs`	`Dict[str, Any]`	Arguments passed to the respective normalization function defined by :obj:`norm`. (default: :obj:`None`)	`None`
`jk`	`str`	The Jumping Knowledge mode. If specified, the model will additionally apply a final linear transformation to transform node embeddings to the expected output feature dimensionality. (:obj:`None`, :obj:`"last"`, :obj:`"cat"`, :obj:`"max"`, :obj:`"lstm"`). (default: :obj:`None`)	`None`
`**kwargs`	`optional`	Additional arguments of the underlying :class:`torch_geometric.nn.conv.MessagePassing` layers.	`{}`

Call arguments: x, edge_index and optionally edge_weight, edge_attr and batch. With mini-batches from :class:~k3_node.loader.NeighborLoader, also pass num_sampled_nodes_per_hop and num_sampled_edges_per_hop (the batch's num_sampled_nodes / num_sampled_edges) to skip the nodes and edges that no longer affect the seed nodes after each layer, as in PyG. This trimming needs concrete sizes, so compile the model with run_eagerly=True when using it.

`reset_parameters()`

Resets all learnable parameters of the module.

GCN

`k3_node.models.GCN`

Bases: BasicGNN

The Graph Neural Network from the "Semi-supervised Classification with Graph Convolutional Networks" <https://arxiv.org/abs/1609.02907>_ paper, using the :class:~k3_node.layers.conv.GCNConv operator for message passing.

Example

import numpy as np
from k3_node.models import GCN

x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
edge_index = np.random.randint(0, 10, size=(2, 30))  # 30 random edges

model = GCN(in_channels=8, hidden_channels=16, num_layers=2, out_channels=4)
out = model(x, edge_index)  # e.g. logits for 4 classes per node
print(tuple(out.shape))  # (10, 4)

GraphSAGE

`k3_node.models.GraphSAGE`

Bases: BasicGNN

The Graph Neural Network from the "Inductive Representation Learning on Large Graphs" <https://arxiv.org/abs/1706.02216>_ paper, using the :class:~k3_node.layers.conv.SAGEConv operator for message passing.

Example

import numpy as np
from k3_node.models import GraphSAGE

x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
edge_index = np.random.randint(0, 10, size=(2, 30))  # 30 random edges

model = GraphSAGE(in_channels=8, hidden_channels=16, num_layers=2, out_channels=4)
out = model(x, edge_index)  # e.g. logits for 4 classes per node
print(tuple(out.shape))  # (10, 4)

GIN

`k3_node.models.GIN`

Bases: BasicGNN

The Graph Neural Network from the "How Powerful are Graph Neural Networks?" <https://arxiv.org/abs/1810.00826>_ paper, using the :class:~k3_node.layers.conv.GINConv operator for message passing.

Example

import numpy as np
from k3_node.models import GIN

x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
edge_index = np.random.randint(0, 10, size=(2, 30))  # 30 random edges

model = GIN(in_channels=8, hidden_channels=16, num_layers=2, out_channels=4)
out = model(x, edge_index)  # e.g. logits for 4 classes per node
print(tuple(out.shape))  # (10, 4)

GAT

`k3_node.models.GAT`

Bases: BasicGNN

The Graph Neural Network from "Graph Attention Networks" <https://arxiv.org/abs/1710.10903> or "How Attentive are Graph Attention Networks?" <https://arxiv.org/abs/2105.14491> papers, using the :class:~k3_node.layers.conv.GATConv or :class:~k3_node.layers.conv.GATv2Conv operator for message passing.

Example

import numpy as np
from k3_node.models import GAT

x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
edge_index = np.random.randint(0, 10, size=(2, 30))  # 30 random edges

model = GAT(in_channels=8, hidden_channels=16, num_layers=2, out_channels=4, heads=2)
out = model(x, edge_index)  # e.g. logits for 4 classes per node
print(tuple(out.shape))  # (10, 4)

PNA

`k3_node.models.PNA`

Bases: BasicGNN

The Graph Neural Network from the "Principal Neighbourhood Aggregation for Graph Nets" <https://arxiv.org/abs/2004.05718>_ paper, using the :class:~k3_node.layers.conv.PNAConv operator for message passing.

Example

import numpy as np
from k3_node.models import PNA

x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
edge_index = np.random.randint(0, 10, size=(2, 30))  # 30 random edges

deg = np.array([0, 2, 4, 3, 1])  # in-degree histogram of the training graphs
model = PNA(in_channels=8, hidden_channels=16, num_layers=2, out_channels=4,
            aggregators=["mean", "min", "max", "std"],
            scalers=["identity", "amplification", "attenuation"], deg=deg)
out = model(x, edge_index)
print(tuple(out.shape))  # (10, 4)

EdgeCNN

`k3_node.models.EdgeCNN`

Bases: BasicGNN

The Graph Neural Network from the "Dynamic Graph CNN for Learning on Point Clouds" <https://arxiv.org/abs/1801.07829>_ paper, using the :class:~k3_node.layers.conv.EdgeConv operator for message passing.

Example

import numpy as np
from k3_node.models import EdgeCNN

x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
edge_index = np.random.randint(0, 10, size=(2, 30))  # 30 random edges

model = EdgeCNN(in_channels=8, hidden_channels=16, num_layers=2, out_channels=4)
out = model(x, edge_index)  # e.g. logits for 4 classes per node
print(tuple(out.shape))  # (10, 4)

MLP

`k3_node.models.MLP`

Bases: Model

A Multi-Layer Perceptron (MLP) model.

There exists two ways to instantiate an MLP:

By specifying explicit channel sizes, e.g., MLP([16, 32, 64, 128]) creates a three-layer MLP with differently sized hidden layers.
By specifying fixed hidden channel sizes over a number of layers, e.g., MLP(in_channels=16, hidden_channels=32, out_channels=128, num_layers=3) creates a three-layer MLP with equally sized hidden layers.

Parameters:

Name	Type	Description	Default
`channel_list`	`List[int] or int`	List of input, intermediate and output channels such that `len(channel_list) - 1` denotes the number of layers of the MLP. (default: `None`)	`None`
`in_channels`	`int`	Size of each input sample. Will override `channel_list`. (default: `None`)	`None`
`hidden_channels`	`int`	Size of each hidden sample. Will override `channel_list`. (default: `None`)	`None`
`out_channels`	`int`	Size of each output sample. Will override `channel_list`. (default: `None`)	`None`
`num_layers`	`int`	The number of layers. Will override `channel_list`. (default: `None`)	`None`
`dropout`	`float or List[float]`	Dropout probability of each hidden embedding. (default: `0.`)	`0.0`
`act`	`str or Callable`	The non-linear activation function to use. (default: `"relu"`)	`'relu'`
`act_first`	`bool`	If set to `True`, activation is applied before normalization. (default: `False`)	`False`
`act_kwargs`	`dict`	Arguments passed to the activation function, e.g. `{"negative_slope": 0.2}` for `"leaky_relu"`. (default: `None`)	`None`
`norm`	`str or Callable`	The normalization function to use. (default: `"batch_norm"`)	`'batch_norm'`
`norm_kwargs`	`dict`	Arguments passed to the respective normalization function. (default: `None`)	`None`
`plain_last`	`bool`	If set to `False`, will apply non-linearity, normalization and dropout to the last layer as well. (default: `True`)	`True`
`bias`	`bool or List[bool]`	If set to `False`, the module will not learn additive biases. (default: `True`)	`True`

Example

import numpy as np
from k3_node.models import MLP

x = np.random.rand(10, 16).astype("float32")
mlp = MLP([16, 32, 32, 4])  # channel sizes: input, hidden, hidden, output
print(tuple(mlp(x).shape))  # (10, 4)

mlp = MLP(in_channels=16, hidden_channels=32, out_channels=4, num_layers=3, dropout=0.1)
print(tuple(mlp(x).shape))  # (10, 4)

`num_layers` `property`

The number of layers.

`reset_parameters()`

Resets all learnable parameters of the module.

Graph Autoencoders & Self-Supervised Models

GAE

`k3_node.models.GAE`

The Graph Auto-Encoder model from the "Variational Graph Auto-Encoders" <https://arxiv.org/abs/1611.07308>_ paper based on user-defined encoder and decoder models.

Parameters:

Name	Type	Description	Default
`encoder`		The encoder module.	required
`decoder`	`optional`	The decoder module. If set to `None`, will default to `InnerProductDecoder`. (default: `None`)	`None`

`call(*args, **kwargs)`

Alias for encode.

`compile(optimizer, discriminator_optimizer=None, discriminator_steps=5)`

Sets the optimizers used by :meth:fit.

Parameters:

Name	Type	Description	Default
`optimizer`	`Optimizer`	Trains the encoder (and decoder).	required
`discriminator_optimizer`	`Optimizer`	Trains the discriminator of adversarial models (:class:`ARGA`, :class:`ARGVA`).	`None`
`discriminator_steps`	`int`	Discriminator updates per encoder update. (default: `5`)	`5`

`decode(*args, **kwargs)`

Runs the decoder and computes edge probabilities.

`embed(data)`

Returns the node embeddings of data (without sampling noise).

`encode(*args, **kwargs)`

Runs the encoder and computes node-wise latent variables.

`evaluate(data)`

Returns the link prediction AUC and average precision on data (which needs pos_edge_label_index and neg_edge_label_index).

`fit(data, epochs=1, validation_data=None, verbose=1)`

Trains the model on one graph for epochs full-graph steps.

Parameters:

Name	Type	Description	Default
`data`	`Data`	The training graph with node features `x`, the message passing edges `edge_index` and the edges to reconstruct `pos_edge_label_index`, as created by :class:`~k3_node.transforms.RandomLinkSplit` with `split_labels=True`.	required
`epochs`	`int`	The number of training steps. (default: `1`)	`1`
`validation_data`	`Data`	A graph with `pos_edge_label_index` and `neg_edge_label_index` on which AUC and average precision are reported.	`None`
`verbose`	`int`	`0` is silent, otherwise one line is printed per epoch.	`1`

Returns:

Name	Type	Description
`dict`		The loss (and validation metrics) of every epoch.

`recon_loss(z, pos_edge_index, neg_edge_index=None)`

Given latent variables z, computes the binary cross entropy loss for positive edges pos_edge_index and negative sampled edges.

`reset_parameters()`

Resets all learnable parameters of the module.

`test(z, pos_edge_index, neg_edge_index)`

Given latent variables z, positive edges pos_edge_index and negative edges neg_edge_index, computes area under the ROC curve (AUC) and average precision (AP) scores.

`train_step(data)`

Runs one training step on data and returns the loss.

VGAE

`k3_node.models.VGAE`

Bases: GAE

The Variational Graph Auto-Encoder model from the "Variational Graph Auto-Encoders" <https://arxiv.org/abs/1611.07308>_ paper.

Parameters:

Name	Type	Description	Default
`encoder`		The encoder module to compute :math:`\mu` and :math:`\log\sigma^2`.	required
`decoder`	`optional`	The decoder module. If set to `None`, will default to `InnerProductDecoder`. (default: `None`)	`None`

`kl_loss(mu=None, logstd=None)`

Computes the KL loss, either for the passed arguments mu and logstd, or based on latent variables from last encoding.

ARGA

`k3_node.models.ARGA`

Bases: GAE

The Adversarially Regularized Graph Auto-Encoder model from the "Adversarially Regularized Graph Autoencoder for Graph Embedding" <https://arxiv.org/abs/1802.04407>_ paper.

Parameters:

Name	Type	Description	Default
`encoder`		The encoder module.	required
`discriminator`		The discriminator module.	required
`decoder`	`optional`	The decoder module. If set to `None`, will default to `InnerProductDecoder`. (default: `None`)	`None`

`discriminator_loss(z)`

Computes the loss of the discriminator.

`reg_loss(z)`

Computes the regularization loss of the encoder.

ARGVA

`k3_node.models.ARGVA`

Bases: ARGA

The Adversarially Regularized Variational Graph Auto-Encoder model from the "Adversarially Regularized Graph Autoencoder for Graph Embedding" <https://arxiv.org/abs/1802.04407>_ paper.

Parameters:

Name	Type	Description	Default
`encoder`		The encoder module to compute :math:`\mu` and :math:`\log\sigma^2`.	required
`discriminator`		The discriminator module.	required
`decoder`	`optional`	The decoder module. If set to `None`, will default to `InnerProductDecoder`. (default: `None`)	`None`

InnerProductDecoder

`k3_node.models.InnerProductDecoder`

The inner product decoder from the "Variational Graph Auto-Encoders" <https://arxiv.org/abs/1611.07308>_ paper.

.. math:: \sigma(\mathbf{Z}\mathbf{Z}^{\top})

where :math:\mathbf{Z} \in \mathbb{R}^{N \times d} denotes the latent space produced by the encoder.

`call(z, edge_index, sigmoid=True)`

Decodes the latent variables z into edge probabilities for the given node-pairs edge_index.

`forward_all(z, sigmoid=True)`

Decodes the latent variables z into a probabilistic dense adjacency matrix.

DeepGraphInfomax

`k3_node.models.DeepGraphInfomax`

Bases: Model

The Deep Graph Infomax model from the "Deep Graph Infomax" <https://arxiv.org/abs/1809.10341>_ paper.

Example

import numpy as np
import keras
from k3_node.layers import GCNConv
from k3_node.models import DeepGraphInfomax

x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
edge_index = np.random.randint(0, 10, size=(2, 30))  # 30 random edges

encoder = GCNConv(8, 16)
model = DeepGraphInfomax(
    hidden_channels=16,
    encoder=lambda x, edge_index: encoder(x, edge_index),
    summary=lambda z, *args: keras.ops.sigmoid(keras.ops.mean(z, axis=0)),
    corruption=lambda x, edge_index: (keras.random.shuffle(x), edge_index),  # permute node features
)
pos_z, neg_z, summary = model(x, edge_index)
print(tuple(pos_z.shape), tuple(neg_z.shape), tuple(summary.shape))  # (10, 16) (10, 16) (16,)
print(tuple(model.loss(pos_z, neg_z, summary).shape))  # (): scalar contrastive loss

`call(*args, **kwargs)`

Returns the latent space for the input arguments, their corruptions and their summary representation.

`discriminate(z, summary, sigmoid=True)`

Given the patch-summary pair z and summary, computes the probability scores assigned to this patch-summary pair.

`loss_fn(pos_z, neg_z, summary)`

Computes the mutual information maximization objective.

`test(train_z, train_y, test_z, test_y, solver='lbfgs', *args, **kwargs)`

Evaluates latent space quality via a logistic regression downstream task.

Scalable Graph Transformers

Polynormer

`k3_node.models.Polynormer`

Bases: Layer

The Polynormer module from the "Polynormer: polynomial-expressive graph transformer in linear time" <https://arxiv.org/abs/2403.01232>_ paper.

Parameters:

Name	Type	Description	Default
`in_channels`	`int`	Input channels.	required
`hidden_channels`	`int`	Hidden channels.	required
`out_channels`	`int`	Output channels.	required
`local_layers`	`int`	The number of local attention layers. (default: :obj:`7`)	`7`
`global_layers`	`int`	The number of global attention layers. (default: :obj:`2`)	`2`
`in_dropout`	`float`	Input dropout rate. (default: :obj:`0.15`)	`0.15`
`dropout`	`float`	Dropout rate. (default: :obj:`0.5`)	`0.5`
`global_dropout`	`float`	Global dropout rate. (default: :obj:`0.5`)	`0.5`
`heads`	`int`	The number of heads. (default: :obj:`1`)	`1`
`beta`	`float`	Aggregate type. (default: :obj:`0.9`)	`0.9`
`qk_shared`	`bool`	Whether weight of query and key are shared. (default: :obj:`True`)	`False`
`pre_ln`	`bool`	Pre layer normalization. (default: :obj:`False`)	`False`
`post_bn`	`bool`	Post batch normalization. (default: :obj:`True`)	`True`
`local_attn`	`bool`	Whether use local attention (GATConv vs GCNConv). (default: :obj:`False`)	`False`

Example

import numpy as np
from k3_node.models import Polynormer

x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
edge_index = np.random.randint(0, 10, size=(2, 30))  # 30 random edges

batch = np.repeat([0, 1], 5)  # two graphs with 5 nodes each
model = Polynormer(in_channels=8, hidden_channels=32, out_channels=4,
                   local_layers=2, global_layers=1, heads=2)
out = model(x, edge_index, batch)
print(tuple(out.shape))  # (10, 4)

`call(x, edge_index, batch=None, training=None)`

Forward pass.

Parameters:

Name	Type	Description	Default
`x`	`Tensor`	The input node features.	required
`edge_index`	`Tensor`	The edge indices.	required
`batch`	`Tensor`	The batch vector assigning each node to a graph. (default: :obj:`None`)	`None`
`training`	`bool`	Whether in training mode. (default: :obj:`None`)	`None`

`reset_parameters()`

Resets all learnable parameters of the module.

SGFormer

`k3_node.models.SGFormer`

Bases: Model

The sgformer module from the "SGFormer: Simplifying and Empowering Transformers for Large-Graph Representations" <https://arxiv.org/abs/2306.10759>_ paper.

Parameters:

Name	Type	Description	Default
`in_channels`	`int`	Input channels.	required
`hidden_channels`	`int`	Hidden channels.	required
`out_channels`	`int`	Output channels.	required
`trans_num_layers`	`int`	The number of layers for all-pair attention. (default: :obj:`2`)	`2`
`trans_num_heads`	`int`	The number of heads for attention. (default: :obj:`1`)	`1`
`trans_dropout`	`float`	Global dropout rate. (default: :obj:`0.5`)	`0.5`
`gnn_num_layers`	`int`	The number of layers for GNN. (default: :obj:`3`)	`3`
`gnn_dropout`	`float`	GNN dropout rate. (default: :obj:`0.5`)	`0.5`
`graph_weight`	`float`	The weight balance global and gnn module. (default: :obj:`0.5`)	`0.5`
`aggregate`	`str`	Aggregate type (:obj:`'add'` or :obj:`'cat'`). (default: :obj:`'add'`)	`'add'`

Example

import numpy as np
from k3_node.models import SGFormer

x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
edge_index = np.random.randint(0, 10, size=(2, 30))  # 30 random edges

model = SGFormer(in_channels=8, hidden_channels=32, out_channels=4, trans_num_layers=1,
                 trans_num_heads=1, gnn_num_layers=1)
print(tuple(model(x, edge_index).shape))  # (10, 4)

LPFormer

`k3_node.models.LPFormer`

Bases: Model

The LPFormer model from the "LPFormer: An Adaptive Graph Transformer for Link Prediction" <https://arxiv.org/abs/2310.11009>_ paper, as in PyG.

For every candidate link it attends over the common neighbors, the one-hop neighbors and the other nodes with a high personalized PageRank (PPR) from both endpoints (ppr_thresholds for the three kinds), using their PPR scores as relative positional encodings, and combines this with counts of each kind of node and with the GCN embeddings of the two endpoints.

Parameters:

Name	Type	Description	Default
`in_channels`	`int`	Input feature dimension.	required
`hidden_channels`	`int`	Hidden dimension.	required
`num_gnn_layers`	`int`	Number of GCN layers. (default: `2`)	`2`
`gnn_dropout`	`float`	GCN dropout. (default: `0.1`)	`0.1`
`num_transformer_layers`	`int`	Number of attention layers. (default: `1`)	`1`
`num_heads`	`int`	Number of attention heads. (default: `1`)	`1`
`transformer_dropout`	`float`	Attention dropout; during training this share of the selected nodes is also dropped. (default: `0.1`)	`0.1`
`ppr_thresholds`	`list`	Minimum PPR of common neighbors, one-hop neighbors and other nodes. (default: `[0, 1e-4, 1e-2]`)	`None`

Call arguments: batch (the [2, num_links] candidate links), x (node features), edge_index (the graph) and the keyword ppr_matrix (from :meth:calc_sparse_ppr). Returns one logit per link. The node selection runs on the host, so train with run_eagerly=True or :func:~k3_node.training.gradient_step.

Example

import numpy as np
from k3_node.models import LPFormer

x = np.random.rand(10, 16).astype("float32")  # 10 nodes with 16 features each
edge_index = np.random.randint(0, 10, size=(2, 30))  # 30 random edges
model = LPFormer(in_channels=16, hidden_channels=16)
ppr = model.calc_sparse_ppr(edge_index, num_nodes=10)
target_links = np.array([[0, 1], [2, 3]])  # the (source, target) pairs to score
print(tuple(model(target_links, x, edge_index, ppr_matrix=ppr).shape))  # (2,): one logit per link

`calc_sparse_ppr(edge_index, num_nodes, alpha=0.15, eps=5e-05)` `staticmethod`

The personalized PageRank matrix LPFormer needs (see :func:get_ppr).

`compute_node_mask(u, v, adj, ppr, training)`

For the links (u, v): (pair index, node, PPR from u, PPR from v) of their common neighbors, one-hop neighbors and other high-PPR nodes.

GPSE

`k3_node.models.GPSE`

Bases: Layer

The Graph Positional and Structural Encoder (GPSE) model from the "Graph Positional and Structural Encoder" <https://arxiv.org/abs/2307.07107>_ paper.

Example

import numpy as np
from k3_node.models import GPSE

x = np.random.rand(6, 16).astype("float32")
edge_index = np.array([[0, 1, 2, 3, 4, 5], [1, 2, 0, 4, 5, 3]])

model = GPSE(dim_in=16, dim_inner=32, layers_pre_mp=1, layers_mp=2, layers_post_mp=1, use_repr=True)
print(tuple(model(x, edge_index).shape))  # (6, 32): positional/structural encodings per node

Molecular & Physical Models

AttentiveFP

`k3_node.models.AttentiveFP`

Bases: Model

The Attentive FP model for molecular representation learning from the "Pushing the Boundaries of Molecular Representation for Drug Discovery with the Graph Attention Mechanism" <https://pubs.acs.org/doi/10.1021/acs.jmedchem.9b00959>_ paper, based on graph attention mechanisms.

Parameters:

Name	Type	Description	Default
`in_channels`	`int`	Size of each input sample.	required
`hidden_channels`	`int`	Hidden node feature dimensionality.	required
`out_channels`	`int`	Size of each output sample.	required
`edge_dim`	`int`	Edge feature dimensionality.	required
`num_layers`	`int`	Number of GNN layers.	required
`num_timesteps`	`int`	Number of iterative refinement steps for global readout.	required
`dropout`	`float`	Dropout probability. (default: `0.0`)	`0.0`
`batch_size`	`int`	Fixed batch size (number of graphs) for JAX/XLA static shape compatibility. (default: `None`)	`None`

Example

import numpy as np
from k3_node.models import AttentiveFP

x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
edge_index = np.random.randint(0, 10, size=(2, 30))  # 30 random edges
edge_attr = np.random.rand(30, 3).astype("float32")  # bond features

batch = np.repeat([0, 1], 5)  # two molecules with 5 atoms each
model = AttentiveFP(in_channels=8, hidden_channels=16, out_channels=1, edge_dim=3,
                    num_layers=2, num_timesteps=2)
out = model(x, edge_index, edge_attr, batch)  # one prediction per molecule
print(tuple(out.shape))  # (2, 1)

NeuralFingerprint

`k3_node.models.NeuralFingerprint`

Bases: Layer

The Neural Fingerprint model from the "Convolutional Networks on Graphs for Learning Molecular Fingerprints" <https://arxiv.org/abs/1509.09292>__ paper to generate fingerprints of molecules.

Parameters:

Name	Type	Description	Default
`in_channels`	`int`	Size of each input sample.	required
`hidden_channels`	`int`	Size of each hidden sample.	required
`out_channels`	`int`	Size of each output fingerprint.	required
`num_layers`	`int`	Number of layers.	required
`**kwargs`	`optional`	Additional arguments of :class:`~k3_node.layers.conv.MFConv`.	`{}`

Example

import numpy as np
from k3_node.models import NeuralFingerprint

x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
edge_index = np.random.randint(0, 10, size=(2, 30))  # 30 random edges

batch = np.repeat([0, 1], 5)  # two molecules with 5 atoms each
model = NeuralFingerprint(in_channels=8, hidden_channels=32, out_channels=16, num_layers=3)
fingerprint = model(x, edge_index, batch)  # one learned fingerprint per molecule
print(tuple(fingerprint.shape))  # (2, 16)

`reset_parameters()`

Resets all learnable parameters of the module.

SchNet

`k3_node.models.SchNet`

Bases: K3NodeHubMixin, Model

The continuous-filter convolutional neural network SchNet from the "SchNet: A Continuous-filter Convolutional Neural Network for Modeling Quantum Interactions" <https://arxiv.org/abs/1706.08566>_ paper.

Parameters:

Name	Type	Description	Default
`hidden_channels`	`int`	Hidden embedding size. (default: 128)	`128`
`num_filters`	`int`	The number of filters to use. (default: 128)	`128`
`num_interactions`	`int`	The number of interaction blocks. (default: 6)	`6`
`num_gaussians`	`int`	The number of gaussians. (default: 50)	`50`
`cutoff`	`float`	Cutoff distance. (default: 10.0)	`10.0`
`interaction_graph`	`callable`	Interaction graph builder. (default: None)	`None`
`max_num_neighbors`	`int`	Maximum neighbors per atom. (default: 32)	`32`
`readout`	`str`	Readout pooling (add, sum, mean). (default: "add")	`'add'`
`dipole`	`bool`	Predict dipole moment magnitude. (default: False)	`False`
`mean`	`float`	Mean of target property. (default: None)	`None`
`std`	`float`	Standard deviation of target property. (default: None)	`None`
`atomref`	`tensor`	Reference atomic values. (default: None)	`None`

Example

import numpy as np
from k3_node.models import SchNet

z = np.array([6, 8, 1, 1, 1])  # atomic numbers of a small molecule
pos = np.random.rand(5, 3).astype("float32") * 2.0  # 3D coordinates (Angstrom)
batch = np.array([0, 0, 0, 1, 1])  # two molecules: atoms 0-2 and atoms 3-4

model = SchNet(hidden_channels=16, num_filters=16, num_interactions=2, num_gaussians=10, cutoff=5.0)
energy = model(z, pos, batch=batch)
print(tuple(energy.shape))  # (2, 1)

DimeNet

`k3_node.models.DimeNet`

Bases: Layer

The directional message passing neural network (DimeNet) from the "Directional Message Passing for Molecular Graphs" <https://arxiv.org/abs/2003.03123>_ paper.

Example

import numpy as np
from k3_node.models import DimeNet

z = np.array([6, 8, 1, 1, 1])  # atomic numbers of a small molecule
pos = np.random.rand(5, 3).astype("float32") * 2.0  # 3D coordinates (Angstrom)
batch = np.array([0, 0, 0, 1, 1])  # two molecules: atoms 0-2 and atoms 3-4

model = DimeNet(hidden_channels=16, out_channels=1, num_blocks=2, num_bilinear=8,
                num_spherical=3, num_radial=6, cutoff=5.0)
energy = model(z, pos, batch=batch)  # one prediction per molecule
print(tuple(energy.shape))  # (2, 1)

DimeNetPlusPlus

`k3_node.models.DimeNetPlusPlus`

Bases: DimeNet

The DimeNet++ from the "Fast and Uncertainty-Aware Directional Message Passing for Non-Equilibrium Molecules" <https://arxiv.org/abs/2011.14115>_ paper.

Example

import numpy as np
from k3_node.models import DimeNetPlusPlus

z = np.array([6, 8, 1, 1, 1])  # atomic numbers of a small molecule
pos = np.random.rand(5, 3).astype("float32") * 2.0  # 3D coordinates (Angstrom)
batch = np.array([0, 0, 0, 1, 1])  # two molecules: atoms 0-2 and atoms 3-4

model = DimeNetPlusPlus(hidden_channels=16, out_channels=1, num_blocks=2, int_emb_size=8,
                        basis_emb_size=8, out_emb_channels=16, num_spherical=3, num_radial=6, cutoff=5.0)
energy = model(z, pos, batch=batch)
print(tuple(energy.shape))  # (2, 1)

GNNFF

`k3_node.models.GNNFF`

Bases: Layer

The Graph Neural Network Force Field (GNNFF) from the "Accurate and scalable graph neural network force field and molecular dynamics with direct force architecture" <https://www.nature.com/articles/s41524-021-00543-3>_ paper. :class:GNNFF directly predicts atomic forces from automatically extracted features of the local atomic environment.

Parameters:

Name	Type	Description	Default
`hidden_node_channels`	`int`	Hidden node embedding size.	required
`hidden_edge_channels`	`int`	Hidden edge embedding size.	required
`num_layers`	`int`	Number of message passing blocks.	required
`cutoff`	`float`	Cutoff distance. (default: 5.0)	`5.0`
`max_num_neighbors`	`int`	Maximum neighbors per node. (default: 32)	`32`

Example

import numpy as np
from k3_node.models import GNNFF

z = np.array([6, 8, 1, 1, 1])  # atomic numbers of a small molecule
pos = np.random.rand(5, 3).astype("float32") * 2.0  # 3D coordinates (Angstrom)

model = GNNFF(hidden_node_channels=16, hidden_edge_channels=16, num_layers=2)
forces = model(z, pos)  # predicts a 3D force vector for every atom
print(tuple(forces.shape))  # (5, 3)

ViSNet

`k3_node.models.ViSNet`

Bases: Layer

The equivariant vector-scalar interactive graph neural network (ViSNet) from the "Enhancing Geometric Representations for Molecules with Equivariant Vector-Scalar Interactive Message Passing" <https://arxiv.org/abs/2210.16518>_ paper.

Example

import numpy as np
from k3_node.models import ViSNet

z = np.array([6, 8, 1, 1, 1])  # atomic numbers of a small molecule
pos = np.random.rand(5, 3).astype("float32") * 2.0  # 3D coordinates (Angstrom)
batch = np.array([0, 0, 0, 1, 1])  # two molecules: atoms 0-2 and atoms 3-4

model = ViSNet(lmax=1, num_heads=2, num_layers=2, hidden_channels=16, num_rbf=8, cutoff=5.0)
energy, forces = model(z, pos, batch=batch)  # forces are None unless derivative=True
print(tuple(energy.shape))  # (2, 1)

Semi-Supervised & Non-Homophilous Models

LabelPropagation

`k3_node.models.LabelPropagation`

Bases: MessagePassing

The label propagation operator, firstly introduced in the "Learning from Labeled and Unlabeled Data with Label Propagation" <http://mlg.eng.cam.ac.uk/zoubin/papers/CMU-CALD-02-107.pdf>_ paper.

.. math:: \mathbf{Y}^{\prime} = \alpha \cdot \mathbf{D}^{-1/2} \mathbf{A} \mathbf{D}^{-1/2} \mathbf{Y} + (1 - \alpha) \mathbf{Y},

where unlabeled data is inferred by labeled data via propagation. This concrete implementation here is derived from the "Combining Label Propagation And Simple Models Out-performs Graph Neural Networks" <https://arxiv.org/abs/2010.13993>_ paper.

Parameters:

Name	Type	Description	Default
`num_layers`	`int`	The number of propagations.	required
`alpha`	`float`	The :math:`\alpha` coefficient.	required

Example

import numpy as np
from k3_node.models import LabelPropagation

y = np.array([0, 1, 2, 0, 1, 2])  # node labels
train_mask = np.array([True, True, True, False, False, False])  # labels known for 3 nodes
edge_index = np.array([[0, 1, 2, 3, 4, 5], [3, 4, 5, 0, 1, 2]])
model = LabelPropagation(num_layers=3, alpha=0.9)
out = model(y, edge_index, train_mask)  # soft labels for every node
print(tuple(out.shape))  # (6, 3)

CorrectAndSmooth

`k3_node.models.CorrectAndSmooth`

Bases: Layer

The correct and smooth (C&S) post-processing model from the "Combining Label Propagation And Simple Models Out-performs Graph Neural Networks" <https://arxiv.org/abs/2010.13993>_ paper.

Parameters:

Name	Type	Description	Default
`num_correction_layers`	`int`	The number of propagations :math:`L_1`.	required
`correction_alpha`	`float`	The :math:`\alpha_1` coefficient.	required
`num_smoothing_layers`	`int`	The number of propagations :math:`L_2`.	required
`smoothing_alpha`	`float`	The :math:`\alpha_2` coefficient.	required
`autoscale`	`bool`	If set to :obj:`True`, will automatically determine the scaling factor :math:`\gamma`. (default: :obj:`True`)	`True`
`scale`	`float`	The scaling factor :math:`\gamma`, in case :obj:`autoscale = False`. (default: :obj:`1.0`)	`1.0`

Example

import numpy as np
from k3_node.models import CorrectAndSmooth

y_soft = np.random.rand(6, 3).astype("float32")  # base model's class probabilities
y_true = np.array([1, 0, 0, 2, 1, 1])
train_mask = np.array([True, False, True, False, True, False])
edge_index = np.array([[0, 1, 1, 2, 4, 5], [1, 0, 2, 1, 5, 4]])

model = CorrectAndSmooth(num_correction_layers=2, correction_alpha=0.5,
                         num_smoothing_layers=2, smoothing_alpha=0.5)
y_soft = model.correct(y_soft, y_true[train_mask], train_mask, edge_index)  # propagate residual errors
y_soft = model.smooth(y_soft, y_true[train_mask], train_mask, edge_index)  # propagate predictions
print(tuple(y_soft.shape))  # (6, 3)

RECT_L

`k3_node.models.RECT_L`

Bases: Layer

The RECT model, i.e. its supervised RECT-L part, from the "Network Embedding with Completely-imbalanced Labels" <https://arxiv.org/abs/2007.03545>_ paper.

Parameters:

Name	Type	Description	Default
`in_channels`	`int`	Size of each input sample.	required
`hidden_channels`	`int`	Intermediate size of each sample.	required
`normalize`	`bool`	Whether to add self-loops and compute symmetric normalization coefficients on-the-fly. (default: :obj:`True`)	`True`
`dropout`	`float`	The dropout probability. (default: :obj:`0.0`)	`0.0`

Example

import numpy as np
from k3_node.models import RECT_L

x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
edge_index = np.random.randint(0, 10, size=(2, 30))  # 30 random edges

model = RECT_L(in_channels=8, hidden_channels=16)
out = model(x, edge_index)  # reconstructs the (semantic) input features
print(tuple(out.shape))  # (10, 8)
print(tuple(model.embed(x, edge_index).shape))  # (10, 16): node embeddings

`get_semantic_labels(x, y, mask)`

Replaces the original labels by their class-centers.

`reset_parameters()`

Resets all learnable parameters of the module.

PMLP

`k3_node.models.PMLP`

Bases: Model

The P(ropagational)MLP model from the "Graph Neural Networks are Inherently Good Generalizers: Insights by Bridging GNNs and MLPs" <https://arxiv.org/abs/2212.09034>_ paper.

:class:PMLP is identical to a standard MLP during training, but then adopts a GNN architecture during testing.

Parameters:

Name	Type	Description	Default
`in_channels`	`int`	Size of each input sample.	required
`hidden_channels`	`int`	Size of each hidden sample.	required
`out_channels`	`int`	Size of each output sample.	required
`num_layers`	`int`	The number of layers.	required
`dropout`	`float`	Dropout probability of each hidden embedding. (default: :obj:`0.`)	`0.0`
`norm`	`bool`	If set to :obj:`False`, will not apply batch normalization. (default: :obj:`True`)	`True`
`bias`	`bool`	If set to :obj:`False`, the module will not learn additive biases. (default: :obj:`True`)	`True`

Example

import numpy as np
from k3_node.models import PMLP

x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
edge_index = np.random.randint(0, 10, size=(2, 30))  # 30 random edges

model = PMLP(in_channels=8, hidden_channels=32, out_channels=4, num_layers=2)
print(tuple(model(x, edge_index).shape))  # (10, 4): message passing is used at inference

`call(x, edge_index=None, training=None)`

Forward pass.

Parameters:

Name	Type	Description	Default
`x`	`Tensor`	The node features of shape `[N, in_channels]`.	required
`edge_index`	`Tensor`	The edge indices. Required during inference. (default: :obj:`None`)	`None`
`training`	`bool`	Override the instance-level `self.training` flag. (default: :obj:`None`)	`None`

`reset_parameters()`

Resets all learnable parameters of the module.

LINKX

`k3_node.models.LINKX`

Bases: Model

The LINKX model from the "Large Scale Learning on Non-Homophilous Graphs: New Benchmarks and Strong Simple Methods" <https://arxiv.org/abs/2110.14446>_ paper.

Parameters:

Name	Type	Description	Default
`num_nodes`	`int`	The number of nodes in the graph.	required
`in_channels`	`int`	Size of each input sample.	required
`hidden_channels`	`int`	Size of each hidden sample.	required
`out_channels`	`int`	Size of each output sample.	required
`num_layers`	`int`	Number of layers of :math:`\textrm{MLP}_{f}`.	required
`num_edge_layers`	`int`	Number of layers of :math:`\textrm{MLP}_{\mathbf{A}}`. (default: :obj:`1`)	`1`
`num_node_layers`	`int`	Number of layers of :math:`\textrm{MLP}_{\mathbf{X}}`. (default: :obj:`1`)	`1`
`dropout`	`float`	Dropout probability. (default: :obj:`0.0`)	`0.0`

Example

import numpy as np
from k3_node.models import LINKX

x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
edge_index = np.random.randint(0, 10, size=(2, 30))  # 30 random edges

# LINKX learns from node features and adjacency separately, which suits heterophilic graphs
model = LINKX(num_nodes=10, in_channels=8, hidden_channels=32, out_channels=4, num_layers=2)
out = model(x, edge_index)
print(tuple(out.shape))  # (10, 4)

`reset_parameters()`

Resets all learnable parameters of the module.

ARLinkPredictor

`k3_node.models.ARLinkPredictor`

Bases: Layer

Link predictor using Attract-Repel embeddings from the paper "Pseudo-Euclidean Attract-Repel Embeddings for Undirected Graphs" <https://arxiv.org/abs/2106.09671>_.

This model splits node embeddings into: attract and repel. The edge prediction score is computed as the dot product of attract components minus the dot product of repel components.

Parameters:

Name	Type	Description	Default
`in_channels`	`int`	Size of each input sample.	required
`hidden_channels`	`int`	Size of hidden embeddings.	required
`out_channels`	`int`	Size of output embeddings. If set to `None`, will default to `hidden_channels`. (default: `None`)	`None`
`num_layers`	`int`	Number of message passing layers. (default: `2`)	`2`
`dropout`	`float`	Dropout probability. (default: `0.0`)	`0.0`
`attract_ratio`	`float`	Ratio to use for attract component. Must be between 0 and 1. (default: `0.5`)	`0.5`

Example

import numpy as np
from k3_node.models import ARLinkPredictor

x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
edge_index = np.random.randint(0, 10, size=(2, 30))  # 30 random edges

model = ARLinkPredictor(in_channels=8, hidden_channels=16, num_layers=2)
scores = model(x, edge_index)  # one link score per edge in edge_index
print(tuple(scores.shape))  # (30,)

`calculate_r_fraction(attract_z, repel_z)`

Calculate the R-fraction (proportion of energy in repel space).

`decode(attract_z, repel_z, edge_index)`

Decode edge scores from attract-repel embeddings.

`encode(x, *args, training=None, **kwargs)`

Encode node features into attract-repel embeddings.

Recommendation & Temporal Models

LightGCN

`k3_node.models.LightGCN`

Bases: Model

The LightGCN model from the "LightGCN: Simplifying and Powering Graph Convolution Network for Recommendation" <https://arxiv.org/abs/2002.02126>_ paper.

Parameters:

Name	Type	Description	Default
`num_nodes`	`int`	The number of nodes in the graph.	required
`embedding_dim`	`int`	The dimensionality of node embeddings.	required
`num_layers`	`int`	The number of :class:`LGConv` layers.	required
`alpha`	`float or Tensor`	The scalar or vector specifying the re-weighting coefficients for aggregating the final embedding. (default: :obj:`None`)	`None`

Example

import numpy as np
from k3_node.models import LightGCN

edge_index = np.array([[0, 1, 2, 3, 4, 5, 6, 7], [1, 2, 3, 4, 5, 6, 7, 0]])  # user-item interactions
edge_label_index = np.array([[0, 1, 2, 3], [4, 5, 6, 7]])  # pairs to score
model = LightGCN(num_nodes=50, embedding_dim=16, num_layers=2)
scores = model(edge_index, edge_label_index)
print(tuple(scores.shape))  # (4,)
print(tuple(model.recommend(edge_index, k=2).shape))  # (50, 2): top-2 recommendations per node

`call(edge_index, edge_label_index=None, edge_weight=None)`

Computes rankings for pairs of nodes.

`get_embedding(edge_index, edge_weight=None)`

Returns the embedding of nodes in the graph.

`reset_parameters()`

Resets all learnable parameters of the module.

BPRLoss

`k3_node.models.BPRLoss`

The Bayesian Personalized Ranking (BPR) loss.

SignedGCN

`k3_node.models.SignedGCN`

Bases: Layer

The signed graph convolutional network model from the "Signed Graph Convolutional Network" <https://arxiv.org/abs/1808.06354>_ paper.

Parameters:

Name	Type	Description	Default
`in_channels`	`int`	Size of each input sample.	required
`hidden_channels`	`int`	Size of each hidden sample.	required
`num_layers`	`int`	Number of layers.	required
`lamb`	`float`	Balances the contributions of the overall objective. (default: :obj:`5`)	`5.0`
`bias`	`bool`	If set to :obj:`False`, all layers will not learn an additive bias. (default: :obj:`True`)	`True`

Example

import numpy as np
from k3_node.models import SignedGCN

x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
edge_index = np.random.randint(0, 10, size=(2, 30))  # 30 random edges

pos_edge_index, neg_edge_index = edge_index[:, :15], edge_index[:, 15:]  # positive / negative links
model = SignedGCN(in_channels=8, hidden_channels=16, num_layers=2)
z = model(x, pos_edge_index, neg_edge_index)
print(tuple(z.shape))  # (10, 16)
print(tuple(model.discriminate(z, edge_index).shape))  # (30, 3): scores for (positive, negative, no link)

`create_spectral_features(pos_edge_index, neg_edge_index, num_nodes=None)`

Creates in_channels spectral node features from the positive and negative edges: a truncated SVD of the signed, symmetric adjacency matrix.

TGNMemory

`k3_node.models.TGNMemory`

Bases: Layer

The Temporal Graph Network (TGN) memory model from the "Temporal Graph Networks for Deep Learning on Dynamic Graphs" <https://arxiv.org/abs/2006.10637>_ paper.

Parameters:

Name	Type	Description	Default
`num_nodes`	`int`	The number of nodes to save memories for.	required
`raw_msg_dim`	`int`	The raw message dimensionality.	required
`memory_dim`	`int`	The hidden memory dimensionality.	required
`time_dim`	`int`	The time encoding dimensionality.	required
`message_module`	`Callable`	Function combining source and destination node memory, raw message, and time encoding.	required
`aggregator_module`	`Callable`	Function aggregating messages to the same destination into a single representation.	required

Example

import numpy as np
from k3_node.models import IdentityMessage, LastAggregator, TGNMemory

memory = TGNMemory(
    num_nodes=5, raw_msg_dim=8, memory_dim=16, time_dim=16,
    message_module=IdentityMessage(raw_msg_dim=8, memory_dim=16, time_dim=16),
    aggregator_module=LastAggregator(),
)
src, dst = np.array([0, 1]), np.array([1, 2])  # two interaction events
t = np.array([1.0, 2.0], dtype="float32")
raw_msg = np.random.rand(2, 8).astype("float32")
memory.update_state(src, dst, t, raw_msg)  # update the memory of the involved nodes

mem, last_update = memory(np.array([0, 1, 2]))
print(tuple(mem.shape), tuple(last_update.shape))  # (3, 16) (3,)

`call(n_id)`

Returns the memory and last update time of the nodes n_id.

`detach()`

Kept for PyG compatibility: the stored memory is never differentiated.

`reset_state()`

Starts again from an empty memory and message stores.

`update_state(src, dst, t, raw_msg)`

Updates the memory with the new events (src, dst, t, raw_msg).

RENet

`k3_node.models.RENet`

Bases: Model

The Recurrent Event Network model from the "Recurrent Event Network for Reasoning over Temporal Knowledge Graphs" <https://arxiv.org/abs/1904.05530>_ paper.

Parameters:

Name	Type	Description	Default
`num_nodes`	`int`	The number of nodes in the knowledge graph.	required
`num_rels`	`int`	The number of relations in the knowledge graph.	required
`hidden_channels`	`int`	Hidden size of node and relation embeddings.	required
`seq_len`	`int`	The sequence length of past events.	required
`num_layers`	`int`	The number of recurrent layers. (default: :obj:`1`)	`1`
`dropout`	`float`	Dropout rate before final prediction. (default: :obj:`0.0`)	`0.0`
`bias`	`bool`	If set to :obj:`False`, all layers will not learn an additive bias. (default: :obj:`True`)	`True`

Example

import numpy as np
from k3_node.models import RENet

model = RENet(num_nodes=5, num_rels=4, hidden_channels=16, seq_len=3)
sub, rel, obj = np.array([0, 1]), np.array([0, 1]), np.array([2, 3])  # queries at time t
# Neighbor histories of subjects and objects: neighbor id, timestep and query index
h_sub, h_sub_t, h_sub_batch = np.array([0, 1, 2]), np.array([0, 1, 0]), np.array([0, 0, 1])
h_obj, h_obj_t, h_obj_batch = np.array([1, 2, 3]), np.array([1, 2, 0]), np.array([0, 0, 1])
log_prob_obj, log_prob_sub = model(sub, rel, obj, h_sub, h_sub_t, h_sub_batch, h_obj, h_obj_t, h_obj_batch)
print(tuple(log_prob_obj.shape))  # (2, 5): scores over all entities for each query

`pre_transform(seq_len)` `staticmethod`

Returns a pre-transform that adds to every event (processed in time order) the history of its subject and object: the entities they were linked to by the same relation in each of the last seq_len time steps (h_sub / h_obj, with the step in h_sub_t / h_obj_t), as in PyG.

`test(logits, y)`

Given ground-truth :obj:y, computes Mean Reciprocal Rank (MRR) and Hits at 1/3/10.

Graph Embedding Models

Node2Vec

`k3_node.models.Node2Vec`

Bases: Layer

The Node2Vec model from the "node2vec: Scalable Feature Learning for Networks" <https://arxiv.org/abs/1607.00653>_ paper where random walks of length :obj:walk_length are sampled in a given graph, and node embeddings are learned via negative sampling optimization.

Parameters:

Name	Type	Description	Default
`edge_index`		The edge indices.	required
`embedding_dim`	`int`	The size of each embedding vector.	required
`walk_length`	`int`	The walk length.	required
`context_size`	`int`	The actual context size which is considered for positive samples.	required
`walks_per_node`	`int`	The number of walks to sample for each node. (default: :obj:`1`)	`1`
`p`	`float`	Likelihood of immediately revisiting a node in the walk. (default: :obj:`1.0`)	`1.0`
`q`	`float`	Control parameter to interpolate between breadth-first strategy and depth-first strategy. (default: :obj:`1.0`)	`1.0`
`num_negative_samples`	`int`	The number of negative samples to use for each positive sample. (default: :obj:`1`)	`1`
`num_nodes`	`int`	The number of nodes. (default: :obj:`None`)	`None`

Example

import numpy as np
from k3_node.models import Node2Vec

edge_index = np.array([[0, 1, 2, 3, 0, 2], [1, 2, 3, 0, 2, 0]])
model = Node2Vec(edge_index, embedding_dim=16, walk_length=4, context_size=3, walks_per_node=2)
print(tuple(model().shape))  # (4, 16): embeddings of all nodes

batch = np.array([0, 1])
loss = model.loss(model.pos_sample(batch), model.neg_sample(batch))  # skip-gram loss on random walks
print(tuple(loss.shape))  # (): a scalar

`call(batch=None)`

Returns the embeddings for the nodes in :obj:batch.

`compile(optimizer)`

Sets the optimizer used by :meth:fit.

`fit(epochs=1, batch_size=128, verbose=1)`

Trains the embeddings on freshly sampled random walks for epochs passes over all nodes; returns the mean loss of every epoch.

`loader(batch_size=128, shuffle=True)`

Yields (pos_rw, neg_rw) random-walk batches, starting from batch_size nodes at a time, as PyG's Node2Vec.loader.

`loss(pos_rw, neg_rw)`

Computes the loss given positive and negative random walks.

`test(train_z, train_y, test_z, test_y, solver='lbfgs', *args, **kwargs)`

Evaluates the embeddings with a logistic regression classifier; returns its accuracy.

MetaPath2Vec

`k3_node.models.MetaPath2Vec`

Bases: Layer

The MetaPath2Vec model from the "metapath2vec: Scalable Representation Learning for Heterogeneous Networks" <https://ericdongyx.github.io/papers/ KDD17-dong-chawla-swami-metapath2vec.pdf>_ paper where random walks based on a given :obj:metapath are sampled in a heterogeneous graph, and node embeddings are learned via negative sampling optimization.

Parameters:

Name	Type	Description	Default
`edge_index_dict`	`Dict[Tuple[str, str, str], Tensor]`	Dictionary holding edge indices for each edge type.	required
`embedding_dim`	`int`	The size of each embedding vector.	required
`metapath`	`List[Tuple[str, str, str]]`	The sequence of edge types denoting the metapath.	required
`walk_length`	`int`	The walk length.	required
`context_size`	`int`	The context size considered for positive samples.	required
`walks_per_node`	`int`	The number of walks to sample for each node. (default: :obj:`1`)	`1`
`num_negative_samples`	`int`	The number of negative samples. (default: :obj:`1`)	`1`
`num_nodes_dict`	`Dict[str, int]`	The number of nodes for each node type. (default: :obj:`None`)	`None`

Example

import numpy as np
from k3_node.models import MetaPath2Vec

edge_index_dict = {
    ("author", "writes", "paper"): np.array([[0, 1, 1], [0, 0, 1]]),
    ("paper", "written_by", "author"): np.array([[0, 0, 1], [0, 1, 1]]),
}
metapath = [("author", "writes", "paper"), ("paper", "written_by", "author")]
model = MetaPath2Vec(edge_index_dict, embedding_dim=16, metapath=metapath,
                     walk_length=2, context_size=2, walks_per_node=2)
print(tuple(model("author").shape))  # (2, 16): embeddings of all authors
print(tuple(model("paper").shape))  # (2, 16)

`call(node_type, batch=None)`

Returns the embeddings for the nodes in :obj:batch of type :obj:node_type.

`loss(pos_rw, neg_rw)`

Computes the loss given positive and negative random walks.

Architectural Blocks

DeepGCNLayer

`k3_node.models.DeepGCNLayer`

Bases: Layer

The skip connection operations from the "DeepGCNs: Can GCNs Go as Deep as CNNs?" <https://arxiv.org/abs/1904.03751> and "All You Need to Train Deeper GCNs" <https://arxiv.org/abs/2006.07739> papers. The implemented skip connections includes the pre-activation residual connection ("res+"), the residual connection ("res"), the dense connection ("dense") and no connections ("plain").

Res+ ("res+"):

.. math:: \text{Normalization}\to\text{Activation}\to\text{Dropout}\to \text{GraphConv}\to\text{Res}

Res ("res") / Dense ("dense") / Plain ("plain"):

.. math:: \text{GraphConv}\to\text{Normalization}\to\text{Activation}\to \text{Res/Dense/Plain}\to\text{Dropout}

Parameters:

Name	Type	Description	Default
`conv`	`optional`	the GCN operator. (default: `None`)	`None`
`norm`	`optional`	the normalization layer. (default: `None`)	`None`
`act`	`optional`	the activation layer. (default: `None`)	`None`
`block`	`str`	The skip connection operation to use (`"res+"`, `"res"`, `"dense"` or `"plain"`). (default: `"res+"`)	`'res+'`
`dropout`	`float`	Whether to apply dropout. (default: `0.`)	`0.0`
`ckpt_grad`	`bool`	Kept for API compatibility; a no-op, since there is no gradient-checkpointing API shared uniformly across Keras backends. (default: `False`)	`False`

Example

import numpy as np
import keras
from k3_node.layers import GENConv, LayerNorm
from k3_node.models import DeepGCNLayer

x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
edge_index = np.random.randint(0, 10, size=(2, 30))  # 30 random edges

x = np.random.rand(10, 16).astype("float32")
layer = DeepGCNLayer(
    conv=GENConv(16, 16),
    norm=LayerNorm(16),
    act=keras.layers.ReLU(),
    block="res+",  # pre-activation residual block (DeeperGCN)
)
out = layer(x, edge_index)
print(tuple(out.shape))  # (10, 16)

`reset_parameters()`

Resets all learnable parameters of the module.

GroupAddRev

`k3_node.models.GroupAddRev`

Bases: Layer

The Grouped Reversible GNN module from the "Graph Neural Networks with 1000 Layers" <https://arxiv.org/abs/2106.07476>_ paper.

Parameters:

Name	Type	Description	Default
`conv`	`Layer or List[Layer]`	A seed GNN layer or list of GNN layers.	required
`split_dim`	`int`	The dimension across which to split groups. (default: :obj:`-1`)	`-1`
`num_groups`	`int`	The number of groups. (default: :obj:`None`)	`2`

Example

import numpy as np
from k3_node.layers import GCNConv
from k3_node.models import GroupAddRev

x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
edge_index = np.random.randint(0, 10, size=(2, 30))  # 30 random edges

x = np.random.rand(10, 32).astype("float32")
model = GroupAddRev([GCNConv(16, 16), GCNConv(16, 16)])  # channels are split into 2 groups
out = model(x, edge_index=edge_index)
print(tuple(out.shape))  # (10, 32)
print(tuple(model.inverse(out, edge_index=edge_index).shape))  # (10, 32): reversible: recover the input

JumpingKnowledge

`k3_node.models.JumpingKnowledge`

Bases: Layer

The Jumping Knowledge layer aggregation module from the "Representation Learning on Graphs with Jumping Knowledge Networks" <https://arxiv.org/abs/1806.03536>_ paper.

Jumping knowledge is performed based on either concatenation (:obj:"cat")

.. math::

\mathbf{x}_v^{(1)} \, \Vert \, \ldots \, \Vert \, \mathbf{x}_v^{(T)},

max pooling (:obj:"max")

.. math::

\max \left( \mathbf{x}_v^{(1)}, \ldots, \mathbf{x}_v^{(T)} \right),

or weighted summation

.. math::

\sum_{t=1}^T \alpha_v^{(t)} \mathbf{x}_v^{(t)}

with attention scores :math:\alpha_v^{(t)} obtained from a bi-directional LSTM (:obj:"lstm").

Parameters:

Name	Type	Description	Default
`mode`	`str`	The aggregation scheme to use (:obj:`"cat"`, :obj:`"max"` or :obj:`"lstm"`).	required
`channels`	`int`	The number of channels per representation. Needs to be only set for LSTM-style aggregation. (default: :obj:`None`)	`None`
`num_layers`	`int`	The number of layers to aggregate. Needs to be only set for LSTM-style aggregation. (default: :obj:`None`)	`None`

Example

import numpy as np
from k3_node.models import JumpingKnowledge

# Node representations from 3 GNN layers
xs = [np.random.rand(10, 16).astype("float32") for _ in range(3)]
print(tuple(JumpingKnowledge("cat")(xs).shape))  # (10, 48): concatenate all layers
print(tuple(JumpingKnowledge("max")(xs).shape))  # (10, 16): element-wise max
print(tuple(JumpingKnowledge("lstm", channels=16, num_layers=3)(xs).shape))  # (10, 16): attention over layers

`call(xs)`

Forward pass.

Parameters:

Name	Type	Description	Default
`xs`	`List[Tensor]`	List containing the layer-wise representations.	required

`reset_parameters()`

Resets all learnable parameters of the module.

MetaLayer

`k3_node.models.MetaLayer`

Bases: Layer

A meta layer for building any kind of graph network, inspired by the "Relational Inductive Biases, Deep Learning, and Graph Networks" <https://arxiv.org/abs/1806.01261>_ paper.

A graph network takes a graph as input and returns an updated graph as output (with same connectivity). The input graph has node features :obj:x, edge features :obj:edge_attr as well as graph-level features :obj:u. The output graph has the same structure, but updated features.

Edge features, node features as well as global features are updated by calling the modules :obj:edge_model, :obj:node_model and :obj:global_model, respectively.

To allow for batch-wise graph processing, all callable functions take an additional argument :obj:batch, which determines the assignment of edges or nodes to their specific graphs.

Parameters:

Name	Type	Description	Default
`edge_model`	`callable`	A callable which updates a graph's edge features based on its source and target node features, its current edge features and its global features. (default: :obj:`None`)	`None`
`node_model`	`callable`	A callable which updates a graph's node features based on its current node features, its graph connectivity, its edge features and its global features. (default: :obj:`None`)	`None`
`global_model`	`callable`	A callable which updates a graph's global features based on its node features, its graph connectivity, its edge features and its current global features. (default: :obj:`None`)	`None`

Example

import numpy as np
import keras
from k3_node.models import MetaLayer

x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
edge_index = np.random.randint(0, 10, size=(2, 30))  # 30 random edges
edge_attr = np.random.rand(30, 4).astype("float32")

edge_mlp = keras.layers.Dense(8)
node_mlp = keras.layers.Dense(8)

def edge_model(src, dst, edge_attr, u, batch):  # update every edge from its endpoints
    return edge_mlp(keras.ops.concatenate([src, dst, edge_attr], axis=-1))

def node_model(x, edge_index, edge_attr, u, batch):  # update nodes from incoming edges
    incoming = keras.ops.segment_sum(edge_attr, edge_index[1], num_segments=10)
    return node_mlp(keras.ops.concatenate([x, incoming], axis=-1))

model = MetaLayer(edge_model=edge_model, node_model=node_model)
x_out, edge_attr_out, u_out = model(x, edge_index, edge_attr=edge_attr)
print(tuple(x_out.shape), tuple(edge_attr_out.shape))  # (10, 8) (30, 8)

`call(x, edge_index, edge_attr=None, u=None, batch=None)`

Forward pass.

Parameters:

Name	Type	Description	Default
`x`	`Tensor`	The node features of shape `[N, F_x]`.	required
`edge_index`	`Tensor`	The edge indices of shape `[2, E]`.	required
`edge_attr`	`Tensor`	The edge features of shape `[E, F_e]`. (default: :obj:`None`)	`None`
`u`	`Tensor`	The global graph features of shape `[B, F_u]`. (default: :obj:`None`)	`None`
`batch`	`Tensor`	The batch vector :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N`. (default: :obj:`None`)	`None`

`reset_parameters()`

Resets all learnable parameters of the module.

Models API Reference

Standard GNN Architectures

BasicGNN

k3_node.models.BasicGNN

reset_parameters()

GCN

k3_node.models.GCN

GraphSAGE

k3_node.models.GraphSAGE

GIN

k3_node.models.GIN

GAT

k3_node.models.GAT

PNA

k3_node.models.PNA

EdgeCNN

k3_node.models.EdgeCNN

MLP

k3_node.models.MLP

num_layers property

reset_parameters()

Graph Autoencoders & Self-Supervised Models

GAE

k3_node.models.GAE

__call__(*args, **kwargs)

compile(optimizer, discriminator_optimizer=None, discriminator_steps=5)

decode(*args, **kwargs)

embed(data)

encode(*args, **kwargs)

evaluate(data)

fit(data, epochs=1, validation_data=None, verbose=1)

recon_loss(z, pos_edge_index, neg_edge_index=None)

reset_parameters()

test(z, pos_edge_index, neg_edge_index)

train_step(data)

VGAE

k3_node.models.VGAE

kl_loss(mu=None, logstd=None)

ARGA

k3_node.models.ARGA

discriminator_loss(z)

reg_loss(z)

ARGVA

k3_node.models.ARGVA

InnerProductDecoder

k3_node.models.InnerProductDecoder

__call__(z, edge_index, sigmoid=True)

forward_all(z, sigmoid=True)

DeepGraphInfomax

k3_node.models.DeepGraphInfomax

call(*args, **kwargs)

discriminate(z, summary, sigmoid=True)

loss_fn(pos_z, neg_z, summary)

test(train_z, train_y, test_z, test_y, solver='lbfgs', *args, **kwargs)

Scalable Graph Transformers

Polynormer

k3_node.models.Polynormer

call(x, edge_index, batch=None, training=None)

reset_parameters()

SGFormer

k3_node.models.SGFormer

LPFormer

k3_node.models.LPFormer

calc_sparse_ppr(edge_index, num_nodes, alpha=0.15, eps=5e-05) staticmethod

compute_node_mask(u, v, adj, ppr, training)

GPSE

k3_node.models.GPSE

Molecular & Physical Models

AttentiveFP

k3_node.models.AttentiveFP

NeuralFingerprint

k3_node.models.NeuralFingerprint

reset_parameters()

SchNet

k3_node.models.SchNet

DimeNet

k3_node.models.DimeNet

DimeNetPlusPlus

k3_node.models.DimeNetPlusPlus

GNNFF

`k3_node.models.BasicGNN`

`reset_parameters()`

`k3_node.models.GCN`

`k3_node.models.GraphSAGE`

`k3_node.models.GIN`

`k3_node.models.GAT`

`k3_node.models.PNA`

`k3_node.models.EdgeCNN`

`k3_node.models.MLP`

`num_layers` `property`

`reset_parameters()`

`k3_node.models.GAE`

`call(*args, **kwargs)`

`compile(optimizer, discriminator_optimizer=None, discriminator_steps=5)`

`decode(*args, **kwargs)`

`embed(data)`

`encode(*args, **kwargs)`

`evaluate(data)`

`fit(data, epochs=1, validation_data=None, verbose=1)`

`recon_loss(z, pos_edge_index, neg_edge_index=None)`

`reset_parameters()`

`test(z, pos_edge_index, neg_edge_index)`

`train_step(data)`

`k3_node.models.VGAE`

`kl_loss(mu=None, logstd=None)`

`k3_node.models.ARGA`

`discriminator_loss(z)`

`reg_loss(z)`

`k3_node.models.ARGVA`

`k3_node.models.InnerProductDecoder`

`call(z, edge_index, sigmoid=True)`

`forward_all(z, sigmoid=True)`

`k3_node.models.DeepGraphInfomax`

`call(*args, **kwargs)`

`discriminate(z, summary, sigmoid=True)`

`loss_fn(pos_z, neg_z, summary)`

`test(train_z, train_y, test_z, test_y, solver='lbfgs', *args, **kwargs)`

`k3_node.models.Polynormer`

`call(x, edge_index, batch=None, training=None)`

`reset_parameters()`

`k3_node.models.SGFormer`

`k3_node.models.LPFormer`

`calc_sparse_ppr(edge_index, num_nodes, alpha=0.15, eps=5e-05)` `staticmethod`

`compute_node_mask(u, v, adj, ppr, training)`

`k3_node.models.GPSE`

`k3_node.models.AttentiveFP`

`k3_node.models.NeuralFingerprint`

`reset_parameters()`

`k3_node.models.SchNet`

`k3_node.models.DimeNet`

`k3_node.models.DimeNetPlusPlus`

`k3_node.models.GNNFF`

`k3_node.models.ViSNet`

`k3_node.models.LabelPropagation`

`k3_node.models.CorrectAndSmooth`

`k3_node.models.RECT_L`

`get_semantic_labels(x, y, mask)`

`reset_parameters()`

`k3_node.models.PMLP`

`call(x, edge_index=None, training=None)`

`reset_parameters()`

`k3_node.models.LINKX`

`reset_parameters()`

`k3_node.models.ARLinkPredictor`

`calculate_r_fraction(attract_z, repel_z)`

`decode(attract_z, repel_z, edge_index)`

`encode(x, *args, training=None, **kwargs)`

`k3_node.models.LightGCN`

`call(edge_index, edge_label_index=None, edge_weight=None)`

`get_embedding(edge_index, edge_weight=None)`

`reset_parameters()`

`k3_node.models.BPRLoss`

`k3_node.models.SignedGCN`

`create_spectral_features(pos_edge_index, neg_edge_index, num_nodes=None)`

`k3_node.models.TGNMemory`

`call(n_id)`

`detach()`

`reset_state()`

`update_state(src, dst, t, raw_msg)`

`k3_node.models.RENet`