Conv

Convolutional layer for message passing.

class topomodelx.base.conv.Conv(in_channels, out_channels, aggr_norm: bool = False, update_func: Literal['relu', 'sigmoid', None] = None, att: bool = False, initialization: Literal['xavier_uniform', 'xavier_normal'] = 'xavier_uniform', initialization_gain: float = 1.414, with_linear_transform: bool = True)

Message passing: steps 1, 2, and 3.

Builds the message passing route given by one neighborhood matrix. Includes an option for an x-specific update function.

Parameters:

in_channelsint

Dimension of input features.

out_channelsint

Dimension of output features.

aggr_normbool, default=False

Whether to normalize the aggregated message by the neighborhood size.

update_func{“relu”, “sigmoid”}, optional

Update method to apply to message.

attbool, default=False

Whether to use attention.

initialization{“xavier_uniform”, “xavier_normal”}, default=”xavier_uniform”

Initialization method.

initialization_gainfloat, default=1.414

Initialization gain.

with_linear_transformbool, default=True

Whether to apply a learnable linear transform. NB: if False in_channels has to be equal to out_channels.

forward(x_source, neighborhood, x_target=None) → Tensor

Forward pass.

This implements message passing: - from source cells with input features x_source, - via neighborhood defining where messages can pass, - to target cells with input features x_target.

In practice, this will update the features on the target cells.

If not provided, x_target is assumed to be x_source, i.e. source cells send messages to themselves.

Parameters:

x_sourceTensor, shape = (…, n_source_cells, in_channels)

Input features on source cells. Assumes that all source cells have the same rank r.

neighborhoodtorch.sparse, shape = (n_target_cells, n_source_cells)

Neighborhood matrix.

x_targetTensor, shape = (…, n_target_cells, in_channels)

Input features on target cells. Assumes that all target cells have the same rank s. Optional. If not provided, x_target is assumed to be x_source, i.e. source cells send messages to themselves.

Returns:

torch.Tensor, shape = (…, n_target_cells, out_channels)

Output features on target cells. Assumes that all target cells have the same rank s.

update(x_message_on_target) → Tensor

Update embeddings on each cell (step 4).

Parameters:

x_message_on_targettorch.Tensor, shape = (n_target_cells, out_channels)

Output features on target cells.

Returns:

torch.Tensor, shape = (n_target_cells, out_channels)

Updated output features on target cells.