MolecularDiffusion.modules.models.equiformer_v2_dynamics

EquiformerV2_dynamics — wraps EquiformerV2 encoder as a drop-in denoising network for EnVariationalDiffusion and EnVariationalDiffusionPyG in en_diffusion.py.

The EquiformerV2 encoder and all layer files are NOT modified.

Classes

EquiformerV2_dynamics

Dynamics wrapper that adapts EquiformerV2 for use with EnVariationalDiffusion(PyG).

Module Contents

class MolecularDiffusion.modules.models.equiformer_v2_dynamics.EquiformerV2_dynamics(equiformer: torch.nn.Module, in_node_nf: int, n_dims: int = 3, condition_time: bool = True, context_node_nf: int = 0, adapter_indices=None, concat_indices=None, sphere_channels: int = 128, lmax_list=None)

Bases: torch.nn.Module, MolecularDiffusion.core.Configurable

Dynamics wrapper that adapts EquiformerV2 for use with EnVariationalDiffusion(PyG).

Parameters:

  • equiformer – Instantiated EquiformerV2 model (from shepherd_arch).

  • in_node_nf (int) – Number of output node features expected by the diffusion model (atomic_number_dim + extra_features + optional_charge). Does NOT include timestep — that is handled internally.

  • n_dims (int) – Spatial dimensions (default 3).

  • condition_time (bool) – Whether to condition on timestep (default True).

  • context_node_nf (int) – Total context feature dimension (default 0).

  • adapter_indices (list[int]) – Indices of context columns routed through additive adapter MLPs.

  • concat_indices (list[int]) – Indices of context columns concatenated to input features.

  • sphere_channels (int) – Must match equiformer.sphere_channels.

  • lmax_list (list[int]) – Must match equiformer.lmax_list.

forward(*args, **kwargs)
condition_time = True
context_node_nf = 0
equiformer
head_h
head_vel_ffn
in_node_nf
input_proj
lmax_list = None
n_adapter_context
n_concat_context
n_dims = 3
sphere_channels = 128
use_adapter_module
vel_SO3_grid