MolecularDiffusion.modules.models.flowmol_graph3d.vector_field

Bond-carrying FlowMol endpoint vector field (GVP, SE(3)-equivariant).

Ported from FlowMol (flowmol/models/vector_field.py) with the bond (``e``) modality intact, which is what makes this a separate package from the platform’s coordinate-only modules/models/flowmol/: that one has no n_bond_types, no token_embeddings, no to_edge_logits and an edge_embedding re-sourced to an RBF of interatomic distance, so restoring bonds there would mean changing its constructor and every public method (and breaking the existing diffusion_flowmol task, which pins canonical_feat_order = ["x", "a", "c"]).

Reused unchanged rather than re-ported:

• modules/layers/gvp – GVPConv, NodePositionUpdate, EdgeUpdate, _rbf, _norm_no_nan, get_time_embedding. Diffed against upstream: the platform copy already supports every FlowMol3 setting, including message_norm: 'sum', n_expansion_gvps, n_cp_feats and the edge_feat_size path.

• modules/models/flowmol.interpolant_scheduler.InterpolantScheduler – feature-agnostic (driven by canonical_feat_order) and linear-capable without cosine_params.

Not ported: VectorField and DirichletVectorField (the vector-field and dirichlet parameterizations, both unreachable at parameterization: ctmc), and trajectory visualization (n_frames is out of scope).

Symmetry is structural, not asserted. Loss and integration run on the upper triangle only and are written back to both halves, and the backbone pools both directions (ue_feats + le_feats) before the bond head. That makes the whole thing depend on the edge ordering laid down by build_edge_idxs, which is why get_upper_edge_mask may only ever be applied to graphs built with it.

Classes

EndpointVectorField

Predicts the trajectory endpoint x_1 for all four modalities.

Module Contents

class MolecularDiffusion.modules.models.flowmol_graph3d.vector_field.EndpointVectorField(n_atom_types: int, canonical_feat_order: list, interpolant_scheduler: MolecularDiffusion.modules.models.flowmol.interpolant_scheduler.InterpolantScheduler, n_charges: int = 6, n_bond_types: int = 4, n_vec_channels: int = 16, n_cp_feats: int = 0, n_hidden_scalars: int = 64, n_hidden_edge_feats: int = 64, n_recycles: int = 1, n_molecule_updates: int = 2, convs_per_update: int = 2, n_message_gvps: int = 3, n_update_gvps: int = 3, n_expansion_gvps: int = 3, separate_mol_updaters: bool = False, message_norm: float | str = 100, update_edge_w_distance: bool = False, rbf_dmax: float = 20, rbf_dim: int = 16, exclude_charges: bool = False, continuous_inv_temp_schedule=None, continuous_inv_temp_max: float = 10.0, time_embedding_dim: int = 1, a_token_dim: int = 0, c_token_dim: int = 0, e_token_dim: int = 0, attention: bool = False, n_heads: int = 1, s_message_dim: int = None, v_message_dim: int = None, dropout: float = 0.0, has_mask: bool = False, self_conditioning: bool = False, use_dst_feats: bool = False, dst_feat_msg_reduction_factor: float = 4, scprop: float = 0.5, adapter_indices: list | None = None, concat_indices: list | None = None)

Bases: torch.nn.Module

Predicts the trajectory endpoint x_1 for all four modalities.

forward consumes g.ndata['x_t','a_t','c_t'] and g.edata['e_t'] and returns a dict keyed {'x','a','c','e'}. The e entry covers the upper-triangle edges only (that is what to_edge_logits is fed).

static build_continuous_inv_temp_func(schedule, max_inv_temp: float = None) → collections.abc.Callable

Inverse-temperature schedule for the continuous (x) vector field.

denoise_graph(g: dgl.DGLGraph, node_scalar_features: torch.Tensor, node_vec_features: torch.Tensor, node_positions: torch.Tensor, edge_features: torch.Tensor, node_batch_idx: torch.Tensor, upper_edge_mask: torch.Tensor, apply_softmax: bool = False, remove_com: bool = False, adapter_ctx: torch.Tensor | None = None) → dict

The GVP message-passing stack plus the four output heads.

forward(g: dgl.DGLGraph, t: torch.Tensor, node_batch_idx: torch.Tensor, upper_edge_mask: torch.Tensor, apply_softmax: bool = False, remove_com: bool = False, prev_dst_dict: dict = None) → dict

Predict x_1 given x_t (and optionally a previous endpoint).

integrate(g: dgl.DGLGraph, node_batch_idx: torch.Tensor, upper_edge_mask: torch.Tensor, n_timesteps: int, **kwargs) → dgl.DGLGraph

Euler-integrate all four modalities from the prior to the data.

Overridden by CTMCVectorField, which is the only parameterization in scope; this stays for the sake of a complete, reusable bond-carrying parent.

precompute_distances(g: dgl.DGLGraph, node_positions: torch.Tensor = None)

Unit displacement vectors and RBF-embedded lengths for every edge.

sample_conditional_path(g: dgl.DGLGraph, t: torch.Tensor, node_batch_idx: torch.Tensor, edge_batch_idx: torch.Tensor, upper_edge_mask: torch.Tensor) → dgl.DGLGraph

Linearly interpolate between the prior and the data at time t.

step(g: dgl.DGLGraph, s_i: torch.Tensor, t_i: torch.Tensor, alpha_t_i: torch.Tensor, alpha_s_i: torch.Tensor, alpha_t_prime_i: torch.Tensor, node_batch_idx: torch.Tensor, upper_edge_mask: torch.Tensor, prev_dst_dict: dict = None, inv_temp_func: collections.abc.Callable = None, **kwargs)

One Euler step of the endpoint parameterization.

static vector_field(x_t, x_1, alpha_t, alpha_t_prime)

The endpoint-parameterized conditional vector field.

adapter_indices = []

canonical_feat_order

concat_indices = []

continuous_inv_temp_func

continuous_inv_temp_max = 10.0

continuous_inv_temp_schedule = None

conv_layers

convs_per_update = 2

edge_embedding

edge_updaters

has_mask = False

interpolant_scheduler

message_norm = 100

n_atom_types

n_bond_types = 4

n_cat_feats

n_charges = 6

n_hidden_edge_feats = 64

n_hidden_scalars = 64

n_molecule_updates = 2

n_recycles = 1

n_vec_channels = 16

node_output_head

node_position_updaters

rbf_dim = 16

rbf_dmax = 20

scalar_embedding

scprop = 0.5

self_conditioning = False

separate_mol_updaters = False

time_embedding_dim = 1

to_edge_logits

token_dims

token_embeddings