MolecularDiffusion.modules.models.diffsmol.uni_transformer¶

DiffSMol’s shape-conditioned SE(3)-equivariant backbone (uni_o2).

Ported from DiffSMol source/models/uni_transformer.py, bond-free.

Dropped relative to upstream (per the approved INTEGRATION_PLAN gate resolution): the ligand_bond_index/ligand_bond_type arguments, the per-layer bond CE loss, and the cov_radius cutoff mode. None of these fed the geometry path – upstream already rebuilds a kNN graph from scratch inside every block (_connect_graph), never using the ground-truth bond graph.

What is kept is _build_edge_type: a purely distance-and-element based bond-order heuristic computed from the current (noisy) coordinates. It reads no ground-truth bonds, so it survives the bond drop intact and still supplies the 5-wide edge feature the attention layers expect.

``pred_bond_type`` (default ``False``, opt-in). Upstream does not merely predict bond types, it feeds them back: uni_transformer.py:504 does edge_type = F.softmax(next_edge_type, dim=-1) inside the layer loop, so layers 1..N-1 were trained on iteratively refined predicted bond types as their 5-dim edge feature, not on the distance heuristic. Loading upstream’s released diffusion.pt without the b_inference head would therefore feed every layer an input distribution it never saw – a silent correctness failure, not a crash. Switching pred_bond_type: true restores the head and that feedback loop, which is what makes the port checkpoint-compatible.

Bond supervision (the bond-type CE / bond-distance / bond-angle / torsion losses) still needs a ground-truth bond channel the data layer does not have, and stays out of scope. It is not needed here: upstream’s _pred_edge_type(..., if_test=True) path – the only one this port implements – produces the feedback features without any ground truth.

The shape latent enters at every layer, twice: as an invariant scalar (InvariantShapeEmbLayer) and as raw equivariant 3-vectors mixed into the vector channel. That is why the latent must be in the same rotational frame as the coordinates.

Attributes¶

`BOND_MARGINS`
`EPS`
`N_BOND_TYPES`

Classes¶

`AttentionLayerO2TwoUpdateNodeGeneral`	Thin wrapper: distance expansion + edge-feature outer product, then
`BaseGVPAttLayer`	One shape-conditioned GVP attention layer.
`InvariantShapeEmbLayer`	Collapse the equivariant `(shape_dim, 3)` latent to an invariant
`UniTransformerO2TwoUpdateGeneral`	`uni_o2` backbone; bond-free unless `pred_bond_type` is set.

Functions¶

`build_bond_length_tensors`(→ torch.Tensor)	Reference bond lengths per (element_i, element_j) and bond order.
`build_diffsmol_bond_tensors`(→ torch.Tensor)	Upstream `construct_bond_tensors('add_aromatic')`, replicated exactly.

Module Contents¶

class MolecularDiffusion.modules.models.diffsmol.uni_transformer.AttentionLayerO2TwoUpdateNodeGeneral(input_dim: tuple[int, int], hidden_dim: tuple[int, int], output_dim: tuple[int, int], n_heads: int, num_r_gaussian: int, edge_feat_dim: int, shape_dim: int, shape_latent_dim: int, act_fn: str = 'relu', norm: bool = True, r_min: float = 0.0, r_max: float = 10.0, output_norm: bool = True, use_shape_vec_mul: bool = True, use_residue: bool = True, pred_bond_type: bool = False)¶

Bases: torch.nn.Module

Thin wrapper: distance expansion + edge-feature outer product, then one BaseGVPAttLayer.

forward(h: torch.Tensor, vec: torch.Tensor, x: torch.Tensor, ligand_emb: torch.Tensor, edge_attr: torch.Tensor, edge_index: torch.Tensor, ligand_shape_emb: torch.Tensor, invar_ligand_shape: torch.Tensor)¶

distance_expansion¶

edge_feat_dim¶

gvp_layer¶

output_dim¶

pred_bond_type = False¶

class MolecularDiffusion.modules.models.diffsmol.uni_transformer.BaseGVPAttLayer(input_dim: tuple[int, int], hidden_dim: tuple[int, int], output_dim: tuple[int, int], shape_dim: int, shape_latent_dim: int, n_heads: int, edge_feat_dim: int, r_feat_dim: int, mess_gvp_layer_num: int = 3, node_gvp_layer_num: int = 3, act_fn: str = 'relu', use_shape_vec_mul: bool = True, use_residue: bool = True, mlp_norm: bool = True, output_norm: bool = True)¶

Bases: torch.nn.Module

One shape-conditioned GVP attention layer.

embed_ligand_shape(scalar_feat: torch.Tensor, vec_feat: torch.Tensor, ligand_shape: torch.Tensor, invar_ligand_shape: torch.Tensor) → torch.Tensor¶

embed_message_att_weight(mess_scalar_emb: torch.Tensor, mess_vec_emb: torch.Tensor, node_scalar_emb: torch.Tensor, node_vec_emb: torch.Tensor, edge_index: torch.Tensor) → torch.Tensor¶

forward(scalar_feat: torch.Tensor, vec_feat: torch.Tensor, r_feat: torch.Tensor, rel_x: torch.Tensor, x: torch.Tensor, ligand_emb: torch.Tensor, edge_feat: torch.Tensor | None, edge_index: torch.Tensor, ligand_shape: torch.Tensor, invar_ligand_shape: torch.Tensor)¶

message_passing(scalar_emb: torch.Tensor, vec_emb: torch.Tensor, x: torch.Tensor, ligand_emb: torch.Tensor, edge_feat: torch.Tensor | None, edge_index: torch.Tensor, r_feat: torch.Tensor, rel_x: torch.Tensor)¶

edge_feat_dim¶

hidden_dim¶

hk_func¶

hq_func¶

input_dim¶

message_layer¶

n_heads¶

node_output_layer¶

output_dim¶

output_norm = True¶

shape_dim¶

shape_emb_layer¶

shape_in¶

shape_scalar_layer¶

use_residue = True¶

use_shape_vec_mul = True¶

class MolecularDiffusion.modules.models.diffsmol.uni_transformer.InvariantShapeEmbLayer(input_dim: int, output_dim: int, act_fn: str = 'relu', norm: bool = True)¶

Bases: torch.nn.Module

Collapse the equivariant (shape_dim, 3) latent to an invariant shape_latent_dim scalar by projecting onto its own mean direction.

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

hidden_layer¶

class MolecularDiffusion.modules.models.diffsmol.uni_transformer.UniTransformerO2TwoUpdateGeneral(num_blocks: int, num_layers: int, scalar_hidden_dim: int, vec_hidden_dim: int, shape_dim: int, shape_latent_dim: int, atom_vocab: list[str], n_heads: int = 1, k: int = 32, num_r_gaussian: int = 20, edge_feat_dim: int = 5, act_fn: str = 'relu', norm: bool = True, r_max: float = 10.0, use_shape_vec_mul: bool = False, use_residue: bool = True, pred_bond_type: bool = False, aromatic_flags: list[bool] | None = None)¶

Bases: torch.nn.Module

uni_o2 backbone; bond-free unless pred_bond_type is set.

forward(v: torch.Tensor, h: torch.Tensor, x: torch.Tensor, batch_ligand: torch.Tensor, ligand_shape: torch.Tensor, return_all: bool = False) → Dict[str, torch.Tensor]¶

act_fn = 'relu'¶

aromatic_flags¶

atom_vocab¶

base_block¶

edge_feat_dim = 5¶

invariant_shape_layer¶

k = 32¶

n_heads = 1¶

norm = True¶

num_blocks¶

num_layers¶

num_r_gaussian = 20¶

pred_bond_type = False¶

r_max = 10.0¶

scalar_hidden_dim¶

shape_dim¶

shape_latent_dim¶

use_residue = True¶

use_shape_vec_mul = False¶

vec_hidden_dim¶

MolecularDiffusion.modules.models.diffsmol.uni_transformer.build_bond_length_tensors(atom_vocab: list[str]) → torch.Tensor¶

Reference bond lengths per (element_i, element_j) and bond order.

Returns [4, V, V] in 0.01 A units: single / double / triple / aromatic. Missing pairs get -1, upstream’s sentinel for “this bond order does not exist for this pair” (the heuristic’s dist - ref < margin test then never fires for real distances).

The aromatic slice is all -1: aromaticity is a bond-graph property and was dropped with the bonds, so the vocab is plain elements. Reuses the platform’s own bond-length tables rather than re-vendoring DiffSMol’s copies of them.

MolecularDiffusion.modules.models.diffsmol.uni_transformer.build_diffsmol_bond_tensors(elements: list[str], aromatic_flags: list[bool]) → torch.Tensor¶

Upstream construct_bond_tensors('add_aromatic'), replicated exactly.

elements[i]/aromatic_flags[i] describe class i of the 15-class (element, is_aromatic) vocabulary. A pair that is aromatic on both ends gets single + aromatic reference lengths and no double/triple; every other pair gets single/double/triple and no aromatic.

Used only by the pred_bond_type / add_aromatic checkpoint- compatible path. build_bond_length_tensors above stays the element- vocab default and is untouched.

# ponytail: upstream’s loop writes triple_bond_tensor[i, j] twice and # never [j, i] (analyze.py:102-103), so its triple-bond table is # upper-triangular and a triple bond is only ever detected when # class(src) <= class(dst). Replicated deliberately: the released # checkpoint was trained against exactly that asymmetry. Do not “fix” it # without retraining.

MolecularDiffusion.modules.models.diffsmol.uni_transformer.BOND_MARGINS = (10, 5, 3, 5)¶

MolecularDiffusion.modules.models.diffsmol.uni_transformer.EPS = 1e-06¶

MolecularDiffusion.modules.models.diffsmol.uni_transformer.N_BOND_TYPES = 5¶