Neural Network Components

Low-level neural network layers and attention mechanisms.

Attention Mechanisms

class spatialvi.nn.SpatialAttention(embed_dim, n_heads=4, dropout=0.1, distance_decay=1.0)

Bases: Module

Spatial attention mechanism that weights neighbors by distance.

This attention mechanism uses spatial distances to modulate attention weights between cells.

Parameters:

• embed_dim (int) – Embedding dimension.

• n_heads (int) – Number of attention heads.

• dropout (float) – Dropout rate.

• distance_decay (float) – Decay factor for distance weighting.

• embed_dim (int)

• n_heads (int)

• dropout (float)

• distance_decay (float)

forward(query, key, value, distances=None, mask=None)

Forward pass.

Parameters:

• query (Tensor) – Query tensor of shape (batch, seq_q, embed_dim).

• key (Tensor) – Key tensor of shape (batch, seq_k, embed_dim).

• value (Tensor) – Value tensor of shape (batch, seq_k, embed_dim).

• distances (Tensor | None) – Spatial distances of shape (batch, seq_q, seq_k).

• mask (Tensor | None) – Attention mask of shape (batch, seq_q, seq_k).

• query (Tensor)

• key (Tensor)

• value (Tensor)

• distances (Tensor | None)

• mask (Tensor | None)

Return type:

tuple[Tensor, Tensor]

Returns:

Tuple of (output, attention_weights).

training: bool

class spatialvi.nn.CrossAttention(embed_dim, n_heads=4, dropout=0.1)

Bases: Module

Cross-attention between two modalities or cell types.

Useful for cell-cell interaction modeling where one cell type attends to another.

Parameters:

• embed_dim (int) – Embedding dimension.

• n_heads (int) – Number of attention heads.

• dropout (float) – Dropout rate.

• embed_dim (int)

• n_heads (int)

• dropout (float)

forward(query, context, context_mask=None)

Forward pass.

Parameters:

• query (Tensor) – Query tensor of shape (batch, seq_q, embed_dim).

• context (Tensor) – Context tensor of shape (batch, seq_c, embed_dim).

• context_mask (Tensor | None) – Mask for context of shape (batch, seq_c).

• query (Tensor)

• context (Tensor)

• context_mask (Tensor | None)

Return type:

Tensor

Returns:

Output tensor of shape (batch, seq_q, embed_dim).

class spatialvi.nn.NeighborAttention(input_dim, hidden_dim=64, n_heads=4, dropout=0.1, use_distance_weights=True)

Bases: Module

Attention mechanism for aggregating neighbor information.

Specifically designed for spatial neighbor aggregation where each cell attends to its k-nearest neighbors.

Parameters:

• input_dim (int) – Input feature dimension.

• hidden_dim (int) – Hidden dimension for attention computation.

• n_heads (int) – Number of attention heads.

• dropout (float) – Dropout rate.

• use_distance_weights (bool) – Whether to use distance-based attention weights.

• input_dim (int)

• hidden_dim (int)

• n_heads (int)

• dropout (float)

• use_distance_weights (bool)

forward(x, neighbor_x, neighbor_distances=None)

Forward pass.

Parameters:

• x (Tensor) – Cell features of shape (batch, input_dim).

• neighbor_x (Tensor) – Neighbor features of shape (batch, n_neighbors, input_dim).

• neighbor_distances (Tensor | None) – Distances to neighbors of shape (batch, n_neighbors).

• x (Tensor)

• neighbor_x (Tensor)

• neighbor_distances (Tensor | None)

Return type:

Tensor

Returns:

Aggregated features of shape (batch, input_dim).

training: bool

class spatialvi.nn.GATLayer(in_features, out_features, n_heads=4, dropout=0.1, alpha=0.2, concat=True)

Bases: Module

Graph Attention Layer.

Parameters:

• in_features (int) – Number of input features.

• out_features (int) – Number of output features.

• n_heads (int) – Number of attention heads.

• dropout (float) – Dropout rate.

• alpha (float) – Negative slope for LeakyReLU.

• concat (bool) – Whether to concatenate heads or average.

• in_features (int)

• out_features (int)

• n_heads (int)

• dropout (float)

• alpha (float)

• concat (bool)

forward(x, edge_index)

Forward pass.

Parameters:

• x (Tensor) – Node features of shape (n_nodes, in_features).

• edge_index (Tensor) – Edge indices of shape (2, n_edges).

• x (Tensor)

• edge_index (Tensor)

Return type:

Tensor

Returns:

Output features.

Encoders

Encoder architectures for spatial models.

class spatialvi.nn._encoders.SpatialEncoder(n_input, n_output, n_cat_list=None, n_layers=1, n_hidden=128, dropout_rate=0.1, use_batch_norm=True, use_layer_norm=False, spatial_hidden=32, n_neighbors=20, aggregation='mean')

Bases: Module

Encoder that incorporates spatial context.

This encoder combines cell-intrinsic features with information from spatially neighboring cells.

Parameters:

• n_input (int) – Number of input features.

• n_output (int) – Dimensionality of the output (latent space).

• n_cat_list (list[int] | None) – List of number of categories for categorical variables.

• n_layers (int) – Number of hidden layers.

• n_hidden (int) – Number of nodes per hidden layer.

• dropout_rate (float) – Dropout rate.

• use_batch_norm (bool) – Whether to use batch normalization.

• use_layer_norm (bool) – Whether to use layer normalization.

• spatial_hidden (int) – Hidden dimension for spatial encoding.

• n_neighbors (int) – Number of neighbors to use.

• aggregation (Literal['mean', 'attention', 'gat']) – Neighbor aggregation method.

• n_input (int)

• n_output (int)

• n_cat_list (list[int] | None)

• n_layers (int)

• n_hidden (int)

• dropout_rate (float)

• use_batch_norm (bool)

• use_layer_norm (bool)

• spatial_hidden (int)

• n_neighbors (int)

• aggregation (Literal['mean', 'attention', 'gat'])

forward(x, *cat_list)

Forward pass without spatial context.

Parameters:

• x (Tensor) – Input tensor.

• cat_list (Tensor) – Categorical covariates.

• x (Tensor)

• cat_list (Tensor)

Return type:

Normal

Returns:

Normal distribution for latent representation.

forward_spatial(x, cat_list=None, spatial_coords=None, neighbor_indices=None, neighbor_expr=None)

Forward pass with spatial context.

Parameters:

• x (Tensor) – Input tensor of shape (n_cells, n_input).

• cat_list (list[Tensor] | None) – List of categorical covariates.

• spatial_coords (Tensor | None) – Spatial coordinates (not used in current implementation).

• neighbor_indices (Tensor | None) – Indices of neighboring cells.

• neighbor_expr (Tensor | None) – Expression of neighboring cells.

• x (Tensor)

• cat_list (list[Tensor] | None)

• spatial_coords (Tensor | None)

• neighbor_indices (Tensor | None)

• neighbor_expr (Tensor | None)

Return type:

tuple[Tensor, Tensor, Tensor]

Returns:

Tuple of (mean, variance, sample) for latent representation.

training: bool

class spatialvi.nn._encoders.GraphEncoder(n_input, n_output, n_hidden=128, n_layers=2, dropout_rate=0.1)

Bases: Module

Graph-based encoder using message passing.

Parameters:

• n_input (int) – Number of input features.

• n_output (int) – Dimensionality of the output.

• n_hidden (int) – Number of hidden nodes.

• n_layers (int) – Number of graph conv layers.

• dropout_rate (float) – Dropout rate.

• n_input (int)

• n_output (int)

• n_hidden (int)

• n_layers (int)

• dropout_rate (float)

forward(x, edge_index)

Forward pass.

Parameters:

• x (Tensor) – Node features of shape (n_nodes, n_input).

• edge_index (Tensor) – Edge indices of shape (2, n_edges).

• x (Tensor)

• edge_index (Tensor)

Return type:

tuple[Tensor, Tensor]

Returns:

Tuple of (mean, variance) for latent representation.

class spatialvi.nn._encoders.GraphConvLayer(n_input, n_output, dropout_rate=0.1)

Bases: Module

Simple graph convolution layer.

Parameters:

• n_input (int)

• n_output (int)

• dropout_rate (float)

forward(x, edge_index)

Forward pass.

Return type:

Tensor

Parameters:

• x (Tensor)

• edge_index (Tensor)

class spatialvi.nn._encoders.AttentionEncoder(n_input, n_output, n_hidden=128, n_heads=4, n_layers=2, dropout_rate=0.1)

Bases: Module

Transformer-based encoder with self-attention.

Parameters:

• n_input (int) – Number of input features.

• n_output (int) – Dimensionality of the output.

• n_hidden (int) – Number of hidden nodes.

• n_heads (int) – Number of attention heads.

• n_layers (int) – Number of transformer layers.

• dropout_rate (float) – Dropout rate.

• n_input (int)

• n_output (int)

• n_hidden (int)

• n_heads (int)

• n_layers (int)

• dropout_rate (float)

forward(x, mask=None)

Forward pass.

Parameters:

• x (Tensor) – Input tensor of shape (batch, seq_len, n_input).

• mask (Tensor | None) – Attention mask.

• x (Tensor)

• mask (Tensor | None)

Return type:

tuple[Tensor, Tensor]

Returns:

Tuple of (mean, variance) for latent representation.

Decoders

Decoder architectures for spatial models.

class spatialvi.nn._decoders.SpatialDecoder(n_input, n_output, n_cat_list=None, n_layers=1, n_hidden=128, dropout_rate=0.1, use_batch_norm=True, use_layer_norm=False, use_spatial_context=False)

Bases: Module

Decoder that can incorporate spatial context for reconstruction.

Parameters:

• n_input (int) – Dimensionality of the latent space (input to decoder).

• n_output (int) – Number of output features (genes).

• n_cat_list (list[int] | None) – List of number of categories for categorical variables.

• n_layers (int) – Number of hidden layers.

• n_hidden (int) – Number of nodes per hidden layer.

• dropout_rate (float) – Dropout rate.

• use_batch_norm (bool) – Whether to use batch normalization.

• use_layer_norm (bool) – Whether to use layer normalization.

• use_spatial_context (bool) – Whether to use spatial context in decoding.

• n_input (int)

• n_output (int)

• n_cat_list (list[int] | None)

• n_layers (int)

• n_hidden (int)

• dropout_rate (float)

• use_batch_norm (bool)

• use_layer_norm (bool)

• use_spatial_context (bool)

forward(z, library, *cat_list, neighbor_z=None)

Forward pass.

Parameters:

• z (Tensor) – Latent representation.

• library (Tensor) – Library size.

• cat_list (Tensor) – Categorical covariates.

• neighbor_z (Tensor | None) – Latent representations of neighbors (for spatial context).

• z (Tensor)

• library (Tensor)

• cat_list (Tensor)

• neighbor_z (Tensor | None)

Return type:

tuple[Tensor, Tensor, Tensor, Tensor]

Returns:

Tuple of (px_scale, px_r, px_rate, px_dropout).

training: bool

class spatialvi.nn._decoders.ConditionalDecoder(n_input, n_output, n_conditions, n_hidden=128, n_layers=2, dropout_rate=0.1)

Bases: Module

Decoder conditioned on additional context (e.g., cell type).

Parameters:

• n_input (int) – Dimensionality of the latent space.

• n_output (int) – Number of output features.

• n_conditions (int) – Number of conditioning variables (e.g., cell types).

• n_hidden (int) – Number of hidden nodes.

• n_layers (int) – Number of hidden layers.

• dropout_rate (float) – Dropout rate.

• n_input (int)

• n_output (int)

• n_conditions (int)

• n_hidden (int)

• n_layers (int)

• dropout_rate (float)

forward(z, condition)

Forward pass.

Parameters:

• z (Tensor) – Latent representation.

• condition (Tensor) – Condition indices.

• z (Tensor)

• condition (Tensor)

Return type:

Tensor

Returns:

Decoded output.

class spatialvi.nn._decoders.FactorizedDecoder(n_input, n_output, n_factors, n_hidden=128, n_layers=2, dropout_rate=0.1)

Bases: Module

Decoder with factorized output for cell type deconvolution.

Parameters:

• n_input (int) – Dimensionality of the latent space.

• n_output (int) – Number of output features (genes).

• n_factors (int) – Number of factors (cell types).

• n_hidden (int) – Number of hidden nodes.

• n_layers (int) – Number of hidden layers.

• dropout_rate (float) – Dropout rate.

• n_input (int)

• n_output (int)

• n_factors (int)

• n_hidden (int)

• n_layers (int)

• dropout_rate (float)

forward(z)

Forward pass.

Parameters:

• z (Tensor) – Latent representation.

• z (Tensor)

Return type:

tuple[Tensor, Tensor]

Returns:

Tuple of (factor_outputs, mixing_weights).

get_mixed_output(z)

Get the weighted mixture of factor outputs.

Parameters:

• z (Tensor) – Latent representation.

• z (Tensor)

Return type:

Tensor

Returns:

Mixed output of shape (batch, n_output).

Layers

Custom neural network layers for spatial models.

class spatialvi.nn._layers.SpatialConv(in_channels, out_channels, n_neighbors=20, aggregation='mean', use_edge_features=True)

Bases: Module

Spatial convolution layer using k-nearest neighbors.

This layer performs convolution over spatial neighbors, aggregating information from nearby cells.

Parameters:

• in_channels (int) – Number of input channels.

• out_channels (int) – Number of output channels.

• n_neighbors (int) – Number of neighbors to use.

• aggregation (Literal['mean', 'max', 'sum', 'attention']) – Aggregation method (“mean”, “max”, “sum”, “attention”).

• use_edge_features (bool) – Whether to use edge features (distances).

• in_channels (int)

• out_channels (int)

• n_neighbors (int)

• aggregation (Literal['mean', 'max', 'sum', 'attention'])

• use_edge_features (bool)

forward(x, neighbor_indices, neighbor_distances=None)

Forward pass.

Parameters:

• x (Tensor) – Node features of shape (n_nodes, in_channels).

• neighbor_indices (Tensor) – Neighbor indices of shape (n_nodes, n_neighbors).

• neighbor_distances (Tensor | None) – Distances to neighbors of shape (n_nodes, n_neighbors).

• x (Tensor)

• neighbor_indices (Tensor)

• neighbor_distances (Tensor | None)

Return type:

Tensor

Returns:

Output features of shape (n_nodes, out_channels).

class spatialvi.nn._layers.GraphConv(in_channels, out_channels, bias=True, normalize=True)

Bases: Module

Graph convolution layer.

Parameters:

• in_channels (int) – Number of input channels.

• out_channels (int) – Number of output channels.

• bias (bool) – Whether to use bias.

• normalize (bool) – Whether to normalize by neighbor count.

• in_channels (int)

• out_channels (int)

• bias (bool)

• normalize (bool)

reset_parameters()

forward(x, edge_index, edge_weight=None)

Forward pass.

Parameters:

• x (Tensor) – Node features of shape (n_nodes, in_channels).

• edge_index (Tensor) – Edge indices of shape (2, n_edges).

• edge_weight (Tensor | None) – Optional edge weights of shape (n_edges,).

• x (Tensor)

• edge_index (Tensor)

• edge_weight (Tensor | None)

Return type:

Tensor

Returns:

Output features of shape (n_nodes, out_channels).

class spatialvi.nn._layers.PositionalEncoding(d_model, n_dims=2, max_len=10000, encoding_type='sinusoidal')

Bases: Module

Positional encoding for spatial coordinates.

Encodes 2D or 3D spatial coordinates using sinusoidal or learnable positional encodings.

Parameters:

• d_model (int) – Dimension of the positional encoding.

• n_dims (int) – Number of spatial dimensions (2 or 3).

• max_len (int) – Maximum length for each dimension.

• encoding_type (Literal['sinusoidal', 'learned']) – Type of encoding (“sinusoidal” or “learned”).

• d_model (int)

• n_dims (int)

• max_len (int)

• encoding_type (Literal['sinusoidal', 'learned'])

forward(coords, normalize=True)

Forward pass.

Parameters:

• coords (Tensor) – Spatial coordinates of shape (batch, n_dims).

• normalize (bool) – Whether to normalize coordinates to [0, max_len).

• coords (Tensor)

• normalize (bool)

Return type:

Tensor

Returns:

Positional encodings of shape (batch, d_model).

class spatialvi.nn._layers.FiLMLayer(feature_dim, condition_dim, hidden_dim=None)

Bases: Module

Feature-wise Linear Modulation layer.

Modulates features based on conditioning information.

Parameters:

• feature_dim (int) – Dimension of features to modulate.

• condition_dim (int) – Dimension of conditioning information.

• hidden_dim (int | None) – Hidden dimension for FiLM generators.

• feature_dim (int)

• condition_dim (int)

• hidden_dim (int | None)

forward(x, condition)

Forward pass.

Parameters:

• x (Tensor) – Features to modulate of shape (…, feature_dim).

• condition (Tensor) – Conditioning information of shape (…, condition_dim).

• x (Tensor)

• condition (Tensor)

Return type:

Tensor

Returns:

Modulated features of shape (…, feature_dim).

class spatialvi.nn._layers.MixedLinear(in_features, out_features, n_experts=4, bias=True)

Bases: Module

Linear layer with mixture of experts.

Parameters:

• in_features (int) – Number of input features.

• out_features (int) – Number of output features.

• n_experts (int) – Number of expert networks.

• bias (bool) – Whether to use bias.

• in_features (int)

• out_features (int)

• n_experts (int)

• bias (bool)

reset_parameters()

forward(x)

Forward pass.

Parameters:

• x (Tensor) – Input tensor of shape (batch, in_features).

• x (Tensor)

Return type:

Tensor

Returns:

Output tensor of shape (batch, out_features).