MolecularDiffusion.modules.models.gcdm.gcp_utils¶

Geometry helpers and small containers used by the GCP blocks.

Near-verbatim port of the pieces of the GCDM repo that gcpnet.py imports:

src/models/components/__init__.py -> centralize, localize, scalarize, vectorize, safe_norm, norm_no_nan, is_identity, ScalarVector, VectorDropout, GCPDropout, GCPLayerNorm
src/models/__init__.py -> get_nonlinearity
src/datamodules/components/helper.py -> _normalize
src/datamodules/components/protein_graph_dataset.py -> _orientations
src/datamodules/components/edm_dataset.py -> _node_features / _edge_features

Changes from upstream, all mechanical: - torchtyping / typeguard decorators stripped (neither is a platform

dependency); the shapes they documented are kept as docstrings.

the plotting / PyMOL / wandb / ProDy helpers that shared the upstream module are not ported – nothing in the denoiser touches them.
_node_features / _edge_features are reduced to the edm_sampling path, the only one the denoiser calls.

Classes¶

`GCPDropout`	From https://github.com/drorlab/gvp-pytorch.
`GCPLayerNorm`	From https://github.com/drorlab/gvp-pytorch.
`ScalarVector`	From https://github.com/sarpaykent/GBPNet.
`VectorDropout`	From https://github.com/drorlab/gvp-pytorch.

Functions¶

`centralize`(→ Tuple[torch.Tensor, torch.Tensor])	Return `(centroid, centered)` for `batch[key]`.
`edge_features`(→ Tuple[torch.Tensor, torch.Tensor])	`(e, xi)` -- squared distance and unit displacement per edge.
`get_nonlinearity`(→ Any)
`is_identity`(→ bool)
`localize`(→ torch.Tensor)	Build the complete local frame of every edge.
`node_vector_features`(→ torch.Tensor)	`chi` -- the 2 equivariant node channels GCPNet embeds.
`norm_no_nan`(→ torch.Tensor)	From https://github.com/drorlab/gvp-pytorch.
`safe_norm`(→ torch.Tensor)
`scalarize`(→ torch.Tensor)	Project `(N, 3, 3)` equivariant values onto local frames -> `(N, 9)`.
`vectorize`(→ torch.Tensor)	Turn `(N, 9)` frame-space gates back into `(N, 3, 3)` vectors.

Module Contents¶

class MolecularDiffusion.modules.models.gcdm.gcp_utils.GCPDropout(drop_rate: float, use_gcp_dropout: bool = True)¶

Bases: torch.nn.Module

From https://github.com/drorlab/gvp-pytorch.

forward(x: torch.Tensor | ScalarVector)¶

scalar_dropout¶

vector_dropout¶

class MolecularDiffusion.modules.models.gcdm.gcp_utils.GCPLayerNorm(dims: ScalarVector, eps: float = 1e-08, use_gcp_norm: bool = True)¶

Bases: torch.nn.Module

From https://github.com/drorlab/gvp-pytorch.

forward(x: torch.Tensor | ScalarVector)¶

static norm_vector(v: torch.Tensor, use_gcp_norm: bool = True, eps: float = 1e-08) → torch.Tensor¶

eps = 1e-08¶

scalar_norm¶

use_gcp_norm = True¶

class MolecularDiffusion.modules.models.gcdm.gcp_utils.ScalarVector¶

Bases: tuple

From https://github.com/sarpaykent/GBPNet.

Initialize self. See help(type(self)) for accurate signature.

clone()¶

concat(others, dim=-1)¶

flatten()¶

idx(idx)¶

mask(node_mask: torch.Tensor)¶

static recover(x, vector_dim)¶

repeat(n, c=1, y=1)¶

vs()¶

property scalar¶

property vector¶

class MolecularDiffusion.modules.models.gcdm.gcp_utils.VectorDropout(drop_rate: float)¶

Bases: torch.nn.Module

From https://github.com/drorlab/gvp-pytorch.

forward(x)¶

drop_rate¶

MolecularDiffusion.modules.models.gcdm.gcp_utils.centralize(batch: Any, key: str, batch_index: torch.Tensor, node_mask: torch.Tensor | None = None, edm: bool = False) → Tuple[torch.Tensor, torch.Tensor]¶

Return (centroid, centered) for batch[key].

node_mask: (batch_num_nodes,) bool.

MolecularDiffusion.modules.models.gcdm.gcp_utils.edge_features(coords: torch.Tensor, edge_index: torch.Tensor) → Tuple[torch.Tensor, torch.Tensor]¶: (e, xi) – squared distance and unit displacement per edge.

MolecularDiffusion.modules.models.gcdm.gcp_utils.get_nonlinearity(nonlinearity: str | None = None, slope: float = 0.01, return_functional: bool = False) → Any¶

MolecularDiffusion.modules.models.gcdm.gcp_utils.is_identity(nonlinearity: Callable | torch.nn.Module | None = None) → bool¶

MolecularDiffusion.modules.models.gcdm.gcp_utils.localize(x: torch.Tensor, edge_index: torch.Tensor, norm_x_diff: bool = True, node_mask: torch.Tensor | None = None) → torch.Tensor¶

Build the complete local frame of every edge.

x: (batch_num_nodes, 3); edge_index: (2, batch_num_edges). Returns (batch_num_edges, 3, 3).

MolecularDiffusion.modules.models.gcdm.gcp_utils.node_vector_features(coords: torch.Tensor) → torch.Tensor¶: chi – the 2 equivariant node channels GCPNet embeds.

MolecularDiffusion.modules.models.gcdm.gcp_utils.norm_no_nan(x: torch.Tensor, dim: int = -1, keepdim: bool = False, eps: float = 1e-08, sqrt: bool = True) → torch.Tensor¶: From https://github.com/drorlab/gvp-pytorch.

MolecularDiffusion.modules.models.gcdm.gcp_utils.safe_norm(x: torch.Tensor, dim: int = -1, eps: float = 1e-08, keepdim: bool = False, sqrt: bool = True) → torch.Tensor¶

MolecularDiffusion.modules.models.gcdm.gcp_utils.scalarize(vector_rep: torch.Tensor, edge_index: torch.Tensor, frames: torch.Tensor, node_inputs: bool, dim_size: int, node_mask: torch.Tensor | None = None) → torch.Tensor¶: Project (N, 3, 3) equivariant values onto local frames -> (N, 9).

MolecularDiffusion.modules.models.gcdm.gcp_utils.vectorize(gate: torch.Tensor, edge_index: torch.Tensor, frames: torch.Tensor, node_inputs: bool, dim_size: int, node_mask: torch.Tensor | None = None) → torch.Tensor¶: Turn (N, 9) frame-space gates back into (N, 3, 3) vectors.