MolecularDiffusion.modules.models.jodo.layers

Classes

CondCoorsNorm
CondGaussianLayer	Gaussian basis function layer for 3D distance features, with time embedding condition
CoorsNorm
EdgeGateTransLayer	The version of edge feature gating.
GaussianLayer	Gaussian basis function layer for 3D distance features
LearnedSinusodialposEmb	following @crowsonkb 's lead with learned sinusoidal pos emb
TransMixLayer	The version for involving the edge feature. Multiply Msg. Without FFN and norm.
Trans_Layer	The version for involving the edge feature. Multiply Msg. Without FFN and norm.

Functions

gaussian(x, mean, std)

Module Contents

class MolecularDiffusion.modules.models.jodo.layers.CondCoorsNorm(time_dim, eps=1e-08)

Bases: torch.nn.Module

forward(coors, time_emb)

eps = 1e-08

time_emb

class MolecularDiffusion.modules.models.jodo.layers.CondGaussianLayer(K, time_dim)

Bases: torch.nn.Module

Gaussian basis function layer for 3D distance features, with time embedding condition

forward(x, time_emb=None)

K

means

stds

time_mlp

class MolecularDiffusion.modules.models.jodo.layers.CoorsNorm(eps=1e-08, scale_init=1.0)

Bases: torch.nn.Module

forward(coors)

eps = 1e-08

scale

class MolecularDiffusion.modules.models.jodo.layers.EdgeGateTransLayer(x_channels: int, out_channels: int, heads: int = 1, dropout: float = 0.0, edge_dim: int | None = None, bias: bool = True, **kwargs)

Bases: torch_geometric.nn.conv.MessagePassing

The version of edge feature gating.

forward(x: torch_geometric.typing.OptTensor, edge_index: torch_geometric.typing.Adj, edge_attr: torch_geometric.typing.OptTensor = None) → torch.Tensor

message(query_i: torch.Tensor, key_j: torch.Tensor, value_j: torch.Tensor, edge_attr: torch_geometric.typing.OptTensor, index: torch.Tensor, ptr: torch_geometric.typing.OptTensor, size_i: int | None) → Tuple[torch.Tensor, torch.Tensor]

reset_parameters()

dropout = 0.0

edge_dim = None

heads = 1

lin_edge0

lin_edge1

lin_key

lin_query

lin_value

out_channels

x_channels

class MolecularDiffusion.modules.models.jodo.layers.GaussianLayer(K, *args, **kwargs)

Bases: torch.nn.Module

Gaussian basis function layer for 3D distance features

forward(x, *args, **kwargs)

K

means

stds

class MolecularDiffusion.modules.models.jodo.layers.LearnedSinusodialposEmb(dim)

Bases: torch.nn.Module

following @crowsonkb ‘s lead with learned sinusoidal pos emb https://github.com/crowsonkb/v-diffusion-jax/blob/master/diffusion/models/danbooru_128.py#L8

forward(x)

weights

class MolecularDiffusion.modules.models.jodo.layers.TransMixLayer(x_channels: int, out_channels: int, extra_heads: int = 2, heads: int = 4, dropout: float = 0.0, edge_dim: int | None = None, bias: bool = True, inf: bool = False, **kwargs)

Bases: torch_geometric.nn.conv.MessagePassing

The version for involving the edge feature. Multiply Msg. Without FFN and norm. Extra attention heads from adjacency matrix.

forward(x: torch_geometric.typing.OptTensor, edge_index: torch_geometric.typing.Adj, edge_attr: torch_geometric.typing.OptTensor, extra_heads: torch_geometric.typing.OptTensor) → torch.Tensor

message(query_i: torch.Tensor, key_j: torch.Tensor, value_j: torch.Tensor, edge_attr: torch_geometric.typing.OptTensor, extra_heads: torch_geometric.typing.OptTensor, index: torch.Tensor, ptr: torch_geometric.typing.OptTensor, size_i: int | None) → Tuple[torch.Tensor, torch.Tensor]

reset_parameters()

dropout = 0.0

edge_dim = None

extra_heads = 2

heads = 4

lin_edge0

lin_edge1

lin_key

lin_query

lin_value

out_channels

set_inf = False

x_channels

class MolecularDiffusion.modules.models.jodo.layers.Trans_Layer(x_channels: int, out_channels: int, heads: int = 1, dropout: float = 0.0, edge_dim: int | None = None, bias: bool = True, **kwargs)

Bases: torch_geometric.nn.conv.MessagePassing

The version for involving the edge feature. Multiply Msg. Without FFN and norm.

forward(x: torch_geometric.typing.OptTensor, edge_index: torch_geometric.typing.Adj, edge_attr: torch_geometric.typing.OptTensor = None) → torch.Tensor

message(query_i: torch.Tensor, key_j: torch.Tensor, value_j: torch.Tensor, edge_attr: torch_geometric.typing.OptTensor, index: torch.Tensor, ptr: torch_geometric.typing.OptTensor, size_i: int | None) → Tuple[torch.Tensor, torch.Tensor]

reset_parameters()

dropout = 0.0

edge_dim = None

heads = 1

lin_edge0

lin_edge1

lin_key

lin_query

lin_value

out_channels

x_channels

MolecularDiffusion.modules.models.jodo.layers.gaussian(x, mean, std)