COSIE.utils.compute_neighborhood_embedding

compute_neighborhood_embedding(edge_index_spatial, embedding_matrix, device)[source]

Compute the neighborhood embedding of each cell based on its spatial k-nearest neighbor graph.

Return type:

Tensor

Parameters

edge_index_spatialtorch.Tensor

A tensor of shape (2, num_edges), where each column represents a directed edge from source node to target node in the spatial graph.

embedding_matrixtorch.Tensor

A tensor of shape (n_cells, dim) representing the embedding of each cell.

devicetorch.device

The computation device, e.g., torch.device(‘cuda’) or torch.device(‘cpu’).

Returns

neighborhood_embedding_matrixtorch.Tensor

A tensor of shape (n_cells, dim) representing the average of each cell’s neighbors’ embeddings.