MolecularDiffusion.modules.models.etflow.torchmd¶

TorchMD-ET equivariant transformer, time-conditioned (ET-Flow’s backbone).

Ported near-verbatim from ET-Flow’s etflow/networks/torchmd_net/ (MIT, (c) 2024 Majdi Hassan, Nikhil Shenoy, Jungyoon Lee), which is itself a fork of torchmd-net’s TorchMD_ET. Module names and attribute names are load bearing: they are exactly the ones in the released Zenodo checkpoints (network.representation_model.* / network.output_model.*), so the conversion in docs/model_integrations/etflow/scripts/convert_checkpoint.py is an identity remap and can assert a strict bijection. Renaming anything here breaks that.

What differs from upstream’s file, and why:

Scalar, EquivariantVectorAndScalarOutput, Distance, GaussianSmearing, ShiftedSoftplus and CoorsNorm are dropped. None of them is reachable from ETFlowTask: the output head is fixed to EquivariantVectorOutput, edge_index is always supplied, and upstream’s own checkpoint schema types rbf_type/activation as Literal["expnorm"]/Literal["silu"] (commons/configs.py:102,104). norm_coors was never settable from BaseFlow either (models/model.py:74-95 does not pass it), so coors_norm was always nn.Identity. Dropping them changes no state_dict key.
The abstract OutputModel base is flattened into EquivariantVectorOutput – it carried no parameters.

Two upstream quirks are kept deliberately, because the published weights were trained with them and “fixing” either would silently invalidate the checkpoints:

edge_weight is the SQUARED interatomic distance, not the distance (r_ij below). It is what feeds the RBF expansion and the cosine cutoff.
EquivariantVectorOutput.pre_reduce adds pos and TorchMDDynamics.forward immediately subtracts it again.

Attributes¶

act_class_mapping

Classes¶

`CosineCutoff`	Smooth 1 -> 0 envelope. Parameter-free.
`EquivariantMultiHeadAttention`	One time-conditioned equivariant attention block.
`EquivariantVectorOutput`	Two gated blocks mapping (scalar, vector) features to one 3-vector.
`ExpNormalSmearing`	PhysNet-style exponential-normal radial basis.
`GatedEquivariantBlock`	Gated equivariant block (Schuett et al. 2021), TorchMD-ET variant.
`NeighborEmbedding`	One distance-weighted neighbour aggregation before the transformer.
`TorchMDDynamics`	ET-Flow's vector field: `(z, t, pos, edge_index) -> (N, 3)`.
`TorchMD_ET_dynamics`	The stack of time-conditioned equivariant attention blocks.

Functions¶

center(→ torch.Tensor)

Subtract each graph's centre of mass.

Module Contents¶

class MolecularDiffusion.modules.models.etflow.torchmd.CosineCutoff(cutoff_lower: float = 0.0, cutoff_upper: float = 5.0)¶

Bases: torch.nn.Module

Smooth 1 -> 0 envelope. Parameter-free.

forward(distances: torch.Tensor) → torch.Tensor¶

cutoff_lower = 0.0¶

cutoff_upper = 5.0¶

class MolecularDiffusion.modules.models.etflow.torchmd.EquivariantMultiHeadAttention(hidden_channels: int, num_rbf: int, distance_influence: str, num_heads: int, activation, attn_activation: str, cutoff_lower: float, cutoff_upper: float, node_attr_dim: int = 0, qk_norm: bool = False, so3_equivariant: bool = False)¶

Bases: torch_geometric.nn.MessagePassing

One time-conditioned equivariant attention block.

Time does not get its own embedding: it is concatenated with the node features and the projected node_attr and mixed by mixing_mlp INSIDE every block (upstream model_dynamics.py:116-118).

so3_equivariant splits the value projection into 4 parts instead of 3, adding a cross-product term – SO(3)- rather than O(3)-equivariant, i.e. chirality-aware by construction. False in both o3 checkpoints.

aggregate(features, index, ptr, dim_size)¶

forward(x, vec, edge_index, r_ij, f_ij, d_ij, t, node_attr)¶

message(q_i, k_j, v_j, vec_j, dk, dv, r_ij, d_ij)¶

reset_parameters() → None¶

update(inputs)¶

act¶

attn_activation¶

cutoff¶

distance_influence¶

dk_proj¶

dv_proj¶

head_dim¶

hidden_channels¶

layernorm¶

mixing_mlp¶

node_attr_dim = 0¶

num_heads¶

o_proj¶

qk_norm = False¶

so3_equivariant = False¶

v_proj¶

vec_proj¶

class MolecularDiffusion.modules.models.etflow.torchmd.EquivariantVectorOutput(hidden_channels: int, activation: str = 'silu', reduce_op: str = 'sum', layer_norm: bool = False)¶

Bases: torch.nn.Module

Two gated blocks mapping (scalar, vector) features to one 3-vector.

pre_reduce(x, v, z, pos, batch)¶

reset_parameters() → None¶

output_network¶

class MolecularDiffusion.modules.models.etflow.torchmd.ExpNormalSmearing(cutoff_lower: float = 0.0, cutoff_upper: float = 5.0, num_rbf: int = 50, trainable: bool = True)¶

Bases: torch.nn.Module

PhysNet-style exponential-normal radial basis.

trainable=True (upstream’s setting) registers means/betas as parameters, which is why they appear in the released checkpoints.

forward(dist: torch.Tensor) → torch.Tensor¶

reset_parameters() → None¶

alpha = 1.0¶

cutoff_fn¶

cutoff_lower = 0.0¶

cutoff_upper = 5.0¶

num_rbf = 50¶

trainable = True¶

class MolecularDiffusion.modules.models.etflow.torchmd.GatedEquivariantBlock(hidden_channels: int, out_channels: int, intermediate_channels: int | None = None, activation: str = 'silu', scalar_activation: bool = False, vector_output: bool = False, layer_norm: bool = True)¶

Bases: torch.nn.Module

Gated equivariant block (Schuett et al. 2021), TorchMD-ET variant.

forward(x: torch.Tensor, v: torch.Tensor) → tuple[torch.Tensor, torch.Tensor]¶

reset_parameters() → None¶

act¶

layer_norm = True¶

out_channels¶

update_net¶

vec1_proj¶

vec2_proj¶

vector_output = False¶

class MolecularDiffusion.modules.models.etflow.torchmd.NeighborEmbedding(hidden_channels: int, num_rbf: int, cutoff_lower: float, cutoff_upper: float, max_z: int = 100)¶

Bases: torch_geometric.nn.MessagePassing

One distance-weighted neighbour aggregation before the transformer.

forward(z, x, edge_index, edge_weight, edge_attr)¶

message(x_j, W)¶

reset_parameters() → None¶

combine¶

cutoff¶

distance_proj¶

embedding¶

class MolecularDiffusion.modules.models.etflow.torchmd.TorchMDDynamics(hidden_channels: int = 128, num_layers: int = 8, num_rbf: int = 64, rbf_type: str = 'expnorm', trainable_rbf: bool = False, activation: str = 'silu', neighbor_embedding: bool = True, cutoff_lower: float = 0.0, cutoff_upper: float = 10.0, max_z: int = 100, node_attr_dim: int = 0, edge_attr_dim: int = 0, attn_activation: str = 'silu', num_heads: int = 8, distance_influence: str = 'both', reduce_op: str = 'sum', qk_norm: bool = False, output_layer_norm: bool = True, clip_during_norm: bool = False, so3_equivariant: bool = False)¶

Bases: torch.nn.Module

ET-Flow’s vector field: (z, t, pos, edge_index) -> (N, 3).

The returned field is centre-of-mass free per graph, which is what makes the flow stay in the zero-COM subspace the harmonic prior lives in.

forward(z: torch.Tensor, t: torch.Tensor, pos: torch.Tensor, edge_index: torch.Tensor, batch: torch.Tensor, edge_attr: torch.Tensor | None = None, node_attr: torch.Tensor | None = None) → torch.Tensor¶: Args mirror upstream’s; t arrives already broadcast per atom.

reset_parameters() → None¶

output_model¶

representation_model¶

class MolecularDiffusion.modules.models.etflow.torchmd.TorchMD_ET_dynamics(hidden_channels: int = 128, num_layers: int = 6, num_rbf: int = 50, rbf_type: str = 'expnorm', trainable_rbf: bool = True, activation: str = 'silu', attn_activation: str = 'silu', neighbor_embedding: bool = True, num_heads: int = 8, distance_influence: str = 'both', cutoff_lower: float = 0.0, cutoff_upper: float = 10.0, max_z: int = 100, node_attr_dim: int = 0, edge_attr_dim: int = 0, qk_norm: bool = False, clip_during_norm: bool = False, so3_equivariant: bool = False)¶

Bases: torch.nn.Module

The stack of time-conditioned equivariant attention blocks.

forward(z: torch.Tensor, t: torch.Tensor, pos: torch.Tensor, batch: torch.Tensor, edge_index: torch.Tensor | None = None, node_attr: torch.Tensor | None = None, edge_attr: torch.Tensor | None = None)¶

reset_parameters() → None¶

activation = 'silu'¶

attention_layers¶

attn_activation = 'silu'¶

clip_during_norm = False¶

cutoff_lower = 0.0¶

cutoff_upper = 10.0¶

distance_expansion¶

distance_influence = 'both'¶

edge_attr_dim = 0¶

embedding¶

hidden_channels = 128¶

max_z = 100¶

neighbor_embedding¶

node_attr_dim = 0¶

num_heads = 8¶

num_layers = 6¶

num_rbf = 50¶

out_norm¶

rbf_type = 'expnorm'¶

trainable_rbf = True¶

MolecularDiffusion.modules.models.etflow.torchmd.center(pos: torch.Tensor, batch: torch.Tensor) → torch.Tensor¶: Subtract each graph’s centre of mass.

MolecularDiffusion.modules.models.etflow.torchmd.act_class_mapping¶