nichecompass.nn.OneHopAttentionNodeLabelAggregator

class nichecompass.nn.OneHopAttentionNodeLabelAggregator(modality, n_input, n_heads=4, leaky_relu_negative_slope=0.2, dropout_rate=0.0)

One-hop Attention Node Label Aggregator class that uses a weighted sum of the omics features of a node’s 1-hop neighbors to build an aggregated neighbor omics feature vector for a node. The weights are determined by an additivite attention mechanism with learnable weights. It returns a concatenation of the node’s own omics feature vector and the attention-aggregated neighbor omics feature vector as node labels for the omics reconstruction task.

Parts of the implementation are inspired by https://github.com/pyg-team/pytorch_geometric/blob/master/torch_geometric/nn/conv/gatv2_conv.py#L16 (01.10.2022).

Parameters:

  • modality (Literal['rna', 'atac']) – Omics modality that is aggregated. Can be either rna or atac.

  • n_input (int) – Number of omics features used for the Node Label Aggregation.

  • n_heads (int (default: 4)) – Number of attention heads for multi-head attention.

  • leaky_relu_negative_slope (float (default: 0.2)) – Slope of the leaky relu activation function.

  • dropout_rate (float (default: 0.0)) – Dropout probability of the normalized attention coefficients which exposes each node to a stochastically sampled neighborhood during training.

Methods table

forward(x, edge_index[, return_agg_weights])

Forward pass of the One-hop Attention Node Label Aggregator.

message(x_j, x_i, g_j, g_i, index)

Message method of the MessagePassing parent class.

reset_parameters()

Reset weight parameters.

Methods

OneHopAttentionNodeLabelAggregator.forward(x, edge_index, return_agg_weights=False)

Forward pass of the One-hop Attention Node Label Aggregator.

Parameters:

  • x (Tensor) – Tensor containing the omics features of the nodes in the current node batch including sampled neighbors. (Size: n_nodes_batch_and_sampled_neighbors x n_node_features)

  • edge_index (Tensor) – Tensor containing the node indices of edges in the current node batch including sampled neighbors. (Size: 2 x n_edges_batch_and_sampled_neighbors)

  • return_agg_weights (bool (default: False)) – If ´True´, also return the aggregation weights (attention weights).

Return type:

Tensor

Returns:

x_neighbors:

Tensor containing the node labels of the nodes in the current node batch excluding sampled neighbors. These labels are used for the omics feature reconstruction task. (Size: n_nodes_batch x (2 x n_node_features))

alpha:

Aggregation weights for edges in ´edge_index´.

OneHopAttentionNodeLabelAggregator.message(x_j, x_i, g_j, g_i, index)

Message method of the MessagePassing parent class. Variables with “_i” suffix refer to the central nodes that aggregate information. Variables with “_j” suffix refer to the neigboring nodes.

Parameters:

  • x_j (Tensor) – Gene expression of neighboring nodes (dim: n_index x n_heads x n_node_features).

  • g_i (Tensor) – Key vector of central nodes (dim: n_index x n_heads x n_node_features).

  • g_j (Tensor) – Query vector of neighboring nodes (dim: n_index x n_heads x n_node_features).

Return type:

Tensor

OneHopAttentionNodeLabelAggregator.reset_parameters()

Reset weight parameters.