MolecularDiffusion.modules.models.shepherd_arch.utils

Utility functions for ShEPhERD model architecture.

Functions

add_virtual_edges_to_edge_index(edge_index, ...)

Adds edges to edge_index that connect all (real) nodes to the virtual node(s).

positional_encoding(position, dim, device)

Module Contents

MolecularDiffusion.modules.models.shepherd_arch.utils.add_virtual_edges_to_edge_index(edge_index, virtual_node_mask, batch)

Adds edges to edge_index that connect all (real) nodes to the virtual node(s).

Parameters:

  • shape (batch -- torch.LongTensor with)

  • shape

  • shape

Returns:

new_edge_index – updated edge_index with additional virtual edges

MolecularDiffusion.modules.models.shepherd_arch.utils.positional_encoding(position, dim, device)