MolecularDiffusion.modules.models.kgdiff.common¶

Shared building blocks for the KGDiff backbone.

Ported from KGDiff’s models/common.py (commit ad893fc). Only the pieces uni_transformer.py / score_model.py actually reach are kept: the custom-offset GaussianSmearing, the MLP used inside every attention head, outer_product(), compose_context(), and ShiftedSoftplus.

Dropped on purpose (dead for this integration, verified by call site): AngleExpansion/Swish/get_h_dist/get_r_feat are referenced nowhere in the ported path, and hybrid_edge_connection / batch_hybrid_edge_connection only serve cutoff_mode='hybrid', which neither the released config nor the released checkpoint uses (both are cutoff_mode: knn).

Attributes¶

NONLINEARITIES

Classes¶

`GaussianSmearing`	Distance expansion on a hand-picked, non-uniform offset grid.
`MLP`	MLP with the same hidden dim across all layers.
`ShiftedSoftplus`	`softplus(x) - log(2)`, so `f(0) == 0`.

Functions¶

`compose_context`(h_protein, h_ligand, pos_protein, ...)	Interleave pocket and ligand tokens into one graph-ordered node set.
`outer_product`(→ torch.Tensor)	Flattened outer product, folded left to right.

Module Contents¶

class MolecularDiffusion.modules.models.kgdiff.common.GaussianSmearing(start: float = 0.0, stop: float = 5.0, num_gaussians: int = 50)¶

Bases: torch.nn.Module

Distance expansion on a hand-picked, non-uniform offset grid.

start/stop are recorded for __repr__ only – upstream overrides the linspace with a fixed 20-point table biased towards bond distances, and coeff is derived from the first gap (1.0 A), so the kernel width is constant even though the grid is not.

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

coeff¶

num_gaussians = 50¶

start = 0.0¶

stop = 5.0¶

class MolecularDiffusion.modules.models.kgdiff.common.MLP(in_dim: int, out_dim: int, hidden_dim: int, num_layer: int = 2, norm: bool = True, act_fn: str = 'relu', act_last: bool = False)¶

Bases: torch.nn.Module

MLP with the same hidden dim across all layers.

The net attribute name and the exact nn.Sequential ordering (Linear, LayerNorm, activation) are load-bearing: the released checkpoint’s keys are *.net.0/1/3.*.

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

net¶

class MolecularDiffusion.modules.models.kgdiff.common.ShiftedSoftplus¶

Bases: torch.nn.Module

softplus(x) - log(2), so f(0) == 0.

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

shift¶

MolecularDiffusion.modules.models.kgdiff.common.compose_context(h_protein: torch.Tensor, h_ligand: torch.Tensor, pos_protein: torch.Tensor, pos_ligand: torch.Tensor, batch_protein: torch.Tensor, batch_ligand: torch.Tensor)¶

Interleave pocket and ligand tokens into one graph-ordered node set.

The sort is stable, so ligand atoms keep their relative order within each complex – upstream’s comment records that an unstable sort broke fixed-atom-type runs.

Returns:: (h_ctx, pos_ctx, batch_ctx, mask_ligand) where mask_ligand is True on ligand rows.

MolecularDiffusion.modules.models.kgdiff.common.outer_product(*vectors: torch.Tensor) → torch.Tensor¶: Flattened outer product, folded left to right.

MolecularDiffusion.modules.models.kgdiff.common.NONLINEARITIES¶