MolecularDiffusion.modules.models.pmdm.common¶

Shared building blocks for the PMDM epsilon network.

Ported from the PMDM reference implementation (models/common.py, models/geometry.py, models/epsnet/diffusion.py, utils/misc.py::get_adj_matrix). Only the pieces MDM_full_pocket_coor_shared actually calls are kept; the upstream readouts, cluster helpers and losses have no consumer here.

Three deliberate deviations from upstream:

BOND_TYPES was imported from utils/chem.py, which imports openbabel at module scope. Only len(BOND_TYPES) is ever used, so the count is inlined as NUM_BOND_TYPES.
torch.sparse.LongTensor (removed in modern torch) is replaced by torch.sparse_coo_tensor().
get_edges upstream takes separate protein/ligand cutoffs but both call sites pass the same value twice, which makes the second (ligand-block) pass a no-op. It takes one cutoff here.

# ponytail: the vendored surface is the reachable subset only. If a future # pass turns on vae_context, port that branch then.

Attributes¶

NUM_BOND_TYPES

Classes¶

`GaussianSmearing`	Radial basis expansion of a distance.
`MultiLayerPerceptron`	MLP with no activation/dropout after the last layer.
`ShiftedSoftplus`	`softplus(x) - log(2)`.

Functions¶

`center_pos_lp`(→ Tuple[torch.Tensor, torch.Tensor])	Move ligand and pocket so the per-complex POCKET centroid is at 0.
`center_pos_pl`(→ Tuple[torch.Tensor, torch.Tensor])	Move ligand and pocket so the per-complex ligand centroid is at 0.
`clip_norm`(→ torch.Tensor)	Rescale rows of `vec` whose norm exceeds `limit`.
`eq_transform`(→ torch.Tensor)	Turn a per-edge distance score into an equivariant per-node score.
`extend_graph_order_radius`(→ Tuple[torch.Tensor, ...)	Bond graph -> (k-hop extended) -> (unioned with a radius graph).
`get_adj_matrix`(→ torch.Tensor)	Fully-connected, self-loop-free edge index for one molecule, `(2, n*(n-1))`.
`get_distance`(→ torch.Tensor)	Euclidean length of every edge, `(E,)`.
`get_edges`(→ torch.Tensor)	Dense within-graph radius graph over the joined ligand+pocket cloud.
`get_num_embedding`(→ torch.Tensor)	Sinusoidal embedding of a 1-D integer tensor, `(G, embedding_dim)`.

Module Contents¶

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

Bases: torch.nn.Module

Radial basis expansion of a distance.

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

coeff¶

class MolecularDiffusion.modules.models.pmdm.common.MultiLayerPerceptron(input_dim: int, hidden_dims: list, activation: str = 'relu', dropout: float = 0)¶

Bases: torch.nn.Module

MLP with no activation/dropout after the last layer.

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

activation¶

dims¶

dropout¶

layers¶

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

Bases: torch.nn.Module

softplus(x) - log(2).

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

shift¶

MolecularDiffusion.modules.models.pmdm.common.center_pos_lp(ligand_pos: torch.Tensor, pocket_pos: torch.Tensor, ligand_batch: torch.Tensor, pocket_batch: torch.Tensor) → Tuple[torch.Tensor, torch.Tensor]¶

Move ligand and pocket so the per-complex POCKET centroid is at 0.

Unlike center_pos_pl(), this does not shift the ligand relative to the pocket – needed when the ligand’s positions are real, already-placed coordinates (a starting fragment from a user’s SDF), not noise that can be freely re-centred (PMDMEpsNet.inpainting_sample/linker_sample).

MolecularDiffusion.modules.models.pmdm.common.center_pos_pl(ligand_pos: torch.Tensor, pocket_pos: torch.Tensor, ligand_batch: torch.Tensor, pocket_batch: torch.Tensor) → Tuple[torch.Tensor, torch.Tensor]¶: Move ligand and pocket so the per-complex ligand centroid is at 0.

MolecularDiffusion.modules.models.pmdm.common.clip_norm(vec: torch.Tensor, limit: float, p: int = 2) → torch.Tensor¶: Rescale rows of vec whose norm exceeds limit.

MolecularDiffusion.modules.models.pmdm.common.eq_transform(score_d: torch.Tensor, pos: torch.Tensor, edge_index: torch.Tensor, edge_length: torch.Tensor) → torch.Tensor¶: Turn a per-edge distance score into an equivariant per-node score.

MolecularDiffusion.modules.models.pmdm.common.extend_graph_order_radius(num_nodes: int, pos: torch.Tensor, edge_index: torch.Tensor, edge_type: torch.Tensor, batch: torch.Tensor, order: int = 3, cutoff: float = 10.0, extend_order: bool = True, extend_radius: bool = True) → Tuple[torch.Tensor, torch.Tensor]¶

Bond graph -> (k-hop extended) -> (unioned with a radius graph).

With PMDM’s training transform the input bond graph is already fully connected, so extend_order is a no-op and extend_radius only re-types edges: an edge inside cutoff ends up as type 0 (1 + (-1)), one outside stays type 1.

MolecularDiffusion.modules.models.pmdm.common.get_adj_matrix(n_particles: int, device=None) → torch.Tensor¶

Fully-connected, self-loop-free edge index for one molecule, (2, n*(n-1)).

Upstream builds this with a double python loop (utils/misc.py::get_adj_matrix); the edge set is identical here.

MolecularDiffusion.modules.models.pmdm.common.get_distance(pos: torch.Tensor, edge_index: torch.Tensor) → torch.Tensor¶: Euclidean length of every edge, (E,).

MolecularDiffusion.modules.models.pmdm.common.get_edges(pos: torch.Tensor, batch_mask: torch.Tensor, cutoff: float, max_pairs: int | None = None) → torch.Tensor¶

Dense within-graph radius graph over the joined ligand+pocket cloud.

Kept dense (torch.cdist) exactly as upstream: the pocket is a few hundred atoms, so the pairwise matrix is small, and a sparse rebuild would change which edges tie on the cutoff boundary.

# ponytail: O(N^2) in the joined point count, fine to ~5k atoms/batch. # Swap in torch_cluster.radius_graph if a batch ever gets bigger.

MolecularDiffusion.modules.models.pmdm.common.get_num_embedding(timesteps: torch.Tensor, embedding_dim: int) → torch.Tensor¶: Sinusoidal embedding of a 1-D integer tensor, (G, embedding_dim).

MolecularDiffusion.modules.models.pmdm.common.NUM_BOND_TYPES = 22¶