MolecularDiffusion.modules.models.loqi.fn_model¶

LoQI / Megalodon MegaFNV3Conf backbone.

Ported from others/LoQI/src/megalodon/dynamics/fn_model.py (NVIDIA, Apache-2.0). A 10-layer interleaved stack of DiTeBlock (DiT-with-edges, invariant) and XEGNNK (coordinate-only EGNN, equivariant). Only the coordinate head exists: this model denoises x against a fixed conditioning graph (atom types, formal charges, bond orders, stereo edges), so upstream’s MegaFNV3 atom/edge prediction heads and BondRefine are not ported – MegaFNV3Conf never instantiates them, and porting dead modules would only add state-dict keys the released weights do not have.

Two deliberate deviations from upstream, both behaviour-preserving:

einops is replaced by plain unflatten/transpose/flatten (upstream lines 512 and 541). einops is not installed in this environment, and the precedent for this substitution is modules/models/pmdm/encoders.py:15.
BondRefine/MegaFNV3/PredictionHead are not ported (dead here).

Module names and construction order are load-bearing: they are the state-dict keys the converted pretrained checkpoint is remapped onto.

Attributes¶

NONLINEARITIES

Classes¶

`DiTeBlock`	DiT block with an edge channel, PyG-batched (batch dim of 1 + attn mask).
`E3Norm`	Per-channel norm over the 3-vector axis, mean-normalized per molecule.
`MLP`	Upstream's `fn_model.MLP` (note: the `bias` flag, unlike the copy in
`MegaFNV3Conf`	The conformer backbone: coordinates out, conditioning graph in.
`SwiGLU`
`TimestepEmbedder`	Sinusoidal timestep embedding + 2-layer MLP.
`XEGNNK`	Coordinate-only EGNN update (`X` in, `X` out), with a cross-product

Functions¶

`coord2distfn`(x, edge_index[, scale_dist_features, batch])	Pairwise squared distances per vector channel, plus optional extras.
`modulate`(x, shift, scale)
`swiglu_correction_fn`(→ int)	Nearest multiple of 256 after the expansion ratio (from ESM3).
`swiglu_ffn`(→ torch.nn.Sequential)
`swiglu_ffn_edge`(→ torch.nn.Sequential)

Module Contents¶

class MolecularDiffusion.modules.models.loqi.fn_model.DiTeBlock(hidden_size: int, edge_hidden_size: int, num_heads: int, mlp_expansion_ratio: float = 4.0, use_z: bool = True, mask_z: bool = True, use_rotary: bool = False, n_vector_features: int = 128, dist_size: int = 128, **block_kwargs)¶

Bases: torch.nn.Module

DiT block with an edge channel, PyG-batched (batch dim of 1 + attn mask).

forward(batch: torch.Tensor, x: torch.Tensor, t_emb_h: torch.Tensor, edge_attr: torch.Tensor | None = None, edge_index: torch.Tensor | None = None, t_emb_e: torch.Tensor | None = None, dist: torch.Tensor | None = None, edge_batch: torch.Tensor | None = None, Z: torch.Tensor | None = None)¶

adaLN_edge_modulation¶

adaLN_modulation¶

d_head¶

edge_hidden_size¶

feature_embedder¶

ffn¶

ffn_edge¶

ffn_norm¶

ffn_norm_edge¶

hidden_size¶

lin_edge0¶

lin_edge1¶

norm1¶

norm1_edge¶

norm2¶

norm2_edge¶

norm_k¶

norm_q¶

num_heads¶

out_projection¶

qkv_proj¶

use_rotary = False¶

class MolecularDiffusion.modules.models.loqi.fn_model.E3Norm(n_vector_features: int = 1, eps: float = 1e-05)¶

Bases: torch.nn.Module

Per-channel norm over the 3-vector axis, mean-normalized per molecule.

forward(pos: torch.Tensor, batch: torch.Tensor) → torch.Tensor¶

reset_parameters() → None¶

eps = 1e-05¶

class MolecularDiffusion.modules.models.loqi.fn_model.MLP(input_dim: int, hidden_size: int, output_dim: int, num_hidden_layers: int = 0, activation: str = 'silu', dropout: float = 0.0, last_act: str | None = None, bias: bool = True)¶

Bases: torch.nn.Module

Upstream’s fn_model.MLP (note: the bias flag, unlike the copy in dynamics/utils.py).

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

act_layer¶

layers¶

class MolecularDiffusion.modules.models.loqi.fn_model.MegaFNV3Conf(num_layers: int = 8, equivariant_node_feature_dim: int = 3, invariant_node_feat_dim: int = 256, invariant_edge_feat_dim: int = 256, atom_classes: int = 16, edge_classes: int = 5, num_heads: int = 16, n_vector_features: int = 128, scale_dist_features: int = 4, dist_size: int = 4, prune_edges: bool = False)¶

Bases: torch.nn.Module

The conformer backbone: coordinates out, conditioning graph in.

forward(batch, X, H, E_idx, E, t) -> {"x_hat": (N,3), "H": (N,D)}, with x_hat centre-of-mass free.

forward(batch, X, H, E_idx, E, t)¶

atom_embedder¶

coord_emb¶

coord_pred¶

dist_projection¶

dit_layers¶

edge_embedder¶

edge_time_embedding¶

egnn_layers¶

n_vector_features = 128¶

node_time_embedding¶

num_atom_classes = 16¶

num_edge_classes = 5¶

scale_dist_features = 4¶

class MolecularDiffusion.modules.models.loqi.fn_model.SwiGLU¶

Bases: torch.nn.Module

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

class MolecularDiffusion.modules.models.loqi.fn_model.TimestepEmbedder(hidden_size: int, frequency_embedding_size: int = 256)¶

Bases: torch.nn.Module

Sinusoidal timestep embedding + 2-layer MLP.

forward(t: torch.Tensor, batch=None) → torch.Tensor¶

static timestep_embedding(t: torch.Tensor, dim: int, max_period: int = 10000) → torch.Tensor¶

frequency_embedding_size = 256¶

mlp¶

class MolecularDiffusion.modules.models.loqi.fn_model.XEGNNK(invariant_node_feat_dim: int = 64, invariant_edge_feat_dim: int = 64, n_vector_features: int = 128, dist_size: int = 4, prune_edges: bool = False)¶

Bases: torch.nn.Module

Coordinate-only EGNN update (X in, X out), with a cross-product term. Node features are read but never written.

forward(batch, X, H, edge_index, edge_attr=None, te=None)¶

coor_update_clamp_value = 10.0¶

coord_projection¶

h_norm¶

h_projection¶

message_input_size = 260¶

phi_message¶

phi_x¶

phi_x_cross¶

prune_edges = False¶

use_cross_product = True¶

x_norm¶

MolecularDiffusion.modules.models.loqi.fn_model.coord2distfn(x, edge_index, scale_dist_features=1, batch=None)¶: Pairwise squared distances per vector channel, plus optional extras.

MolecularDiffusion.modules.models.loqi.fn_model.modulate(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor)¶

MolecularDiffusion.modules.models.loqi.fn_model.swiglu_correction_fn(expansion_ratio: float, d_model: int) → int¶: Nearest multiple of 256 after the expansion ratio (from ESM3).

MolecularDiffusion.modules.models.loqi.fn_model.swiglu_ffn(d_model: int, expansion_ratio: float, bias: bool) → torch.nn.Sequential¶

MolecularDiffusion.modules.models.loqi.fn_model.swiglu_ffn_edge(d_model: int, bias: bool) → torch.nn.Sequential¶

MolecularDiffusion.modules.models.loqi.fn_model.NONLINEARITIES¶