MolecularDiffusion.modules.models.gcdm.gcp_layers¶

Geometry-Complete Perceptron (GCP) blocks.

Near-verbatim port of src/models/components/gcpnet.py:35-931 from https://github.com/BioinfoMachineLearning/bio-diffusion (commit a328950c), which adapts https://github.com/BioinfoMachineLearning/GCPNet.

Changes from upstream, all mechanical – no module is renamed and no ``nn.Parameter`` shape changes, so the published GCDM checkpoints map onto these classes one-to-one:

torchtyping/typeguard decorators stripped.
the five omegaconf.DictConfig objects upstream threads through every constructor are replaced by two plain dataclasses, GCPModuleConfig and GCPLayerConfig, with the shipped QM9/GEOM values as defaults. copy.copy + attribute assignment (which upstream relies on for its soft_cfg / ff_cfg variants) still works on them.
module_cfg.selected_GCP, upstream a Hydra _partial_ pointing at GCP2, becomes a plain gcp_version: "GCP" | "GCP2" string.
upstream’s mp_cfg sub-config is flattened into GCPLayerConfig (it only ever carried three used fields).

Classes¶

`GCP`
`GCP2`
`GCPEmbedding`
`GCPInteractions`
`GCPLayerConfig`	Flattened `layer_cfg/*_gcp_interaction_layer.yaml` + its `mp_cfg`.
`GCPMessagePassing`
`GCPModuleConfig`	Flattened `configs/model/module_cfg/*_gcp_module.yaml`.

Functions¶

get_GCP_with_custom_cfg(input_dims, output_dims, cfg, ...)

Upstream gcpnet.py:606, with OmegaConf swapped for dataclasses.

Module Contents¶

class MolecularDiffusion.modules.models.gcdm.gcp_layers.GCP(input_dims: MolecularDiffusion.modules.models.gcdm.gcp_utils.ScalarVector, output_dims: MolecularDiffusion.modules.models.gcdm.gcp_utils.ScalarVector, nonlinearities: Tuple[str | None, str | None] = ('silu', 'silu'), scalar_out_nonlinearity: str | None = 'silu', scalar_gate: int = 0, vector_gate: bool = True, frame_gate: bool = False, sigma_frame_gate: bool = False, feedforward_out: bool = False, bottleneck: int = 1, vector_residual: bool = False, vector_frame_residual: bool = False, ablate_frame_updates: bool = False, ablate_scalars: bool = False, ablate_vectors: bool = False, scalarization_vectorization_output_dim: int = 3, **kwargs)¶

Bases: torch.nn.Module

create_zero_vector(scalar_rep)¶

forward(s_maybe_v, edge_index, frames, node_inputs: bool = False, node_mask=None)¶

process_vector(scalar_rep, v_pre, vector_hidden_rep)¶

process_vector_frames(scalar_rep, v_pre, edge_index, frames, node_inputs, node_mask=None)¶

ablate_frame_updates = False¶

class MolecularDiffusion.modules.models.gcdm.gcp_layers.GCP2(input_dims: MolecularDiffusion.modules.models.gcdm.gcp_utils.ScalarVector, output_dims: MolecularDiffusion.modules.models.gcdm.gcp_utils.ScalarVector, nonlinearities: Tuple[str | None, str | None] = ('silu', 'silu'), scalar_out_nonlinearity: str | None = 'silu', scalar_gate: int = 0, vector_gate: bool = True, frame_gate: bool = False, sigma_frame_gate: bool = False, feedforward_out: bool = False, bottleneck: int = 1, vector_residual: bool = False, vector_frame_residual: bool = False, ablate_frame_updates: bool = False, ablate_scalars: bool = False, ablate_vectors: bool = False, scalarization_vectorization_output_dim: int = 3, **kwargs)¶

Bases: torch.nn.Module

create_zero_vector(scalar_rep)¶

forward(s_maybe_v, edge_index, frames, node_inputs: bool = False, node_mask=None)¶

process_vector_with_frames(scalar_rep, v_pre, vector_hidden_rep, edge_index, frames, node_inputs, node_mask=None)¶

process_vector_without_frames(scalar_rep, v_pre, vector_hidden_rep)¶

ablate_frame_updates = False¶

class MolecularDiffusion.modules.models.gcdm.gcp_layers.GCPEmbedding(edge_input_dims: MolecularDiffusion.modules.models.gcdm.gcp_utils.ScalarVector, node_input_dims: MolecularDiffusion.modules.models.gcdm.gcp_utils.ScalarVector, edge_hidden_dims: MolecularDiffusion.modules.models.gcdm.gcp_utils.ScalarVector, node_hidden_dims: MolecularDiffusion.modules.models.gcdm.gcp_utils.ScalarVector, num_atom_types: int, nonlinearities: Tuple[str | None, str | None] = ('silu', 'silu'), cfg: GCPModuleConfig | None = None, pre_norm: bool = True, use_gcp_norm: bool = True)¶

Bases: torch.nn.Module

forward(batch: Any)¶

edge_embedding¶

node_embedding¶

pre_norm = True¶

class MolecularDiffusion.modules.models.gcdm.gcp_layers.GCPInteractions(node_dims: MolecularDiffusion.modules.models.gcdm.gcp_utils.ScalarVector, edge_dims: MolecularDiffusion.modules.models.gcdm.gcp_utils.ScalarVector, cfg: GCPModuleConfig, layer_cfg: GCPLayerConfig, dropout: float = 0.0, nonlinearities: Tuple[Any, Any] | None = None, update_node_positions: bool = False)¶

Bases: torch.nn.Module

derive_x_update(node_rep, edge_index, f_ij, node_mask=None)¶

forward(node_rep, edge_rep, edge_index, frames, node_mask=None, node_pos=None)¶

feedforward_network¶

gcp_dropout¶

gcp_norm¶

interaction¶

node_positions_weight¶

pre_norm¶

update_node_positions = False¶

update_positions_with_vector_sum¶

class MolecularDiffusion.modules.models.gcdm.gcp_layers.GCPLayerConfig¶

Flattened layer_cfg/*_gcp_interaction_layer.yaml + its mp_cfg.

Defaults are the shipped QM9/GEOM values (the two files are identical).

dropout: float = 0.0¶

nonlinearity_slope: float = 0.01¶

num_feedforward_layers: int = 1¶

num_message_layers: int = 4¶

pre_norm: bool = False¶

self_message: bool = True¶

use_gcp_dropout: bool = False¶

use_gcp_norm: bool = False¶

use_residual_message_gcp: bool = True¶

use_scalar_message_attention: bool = True¶

class MolecularDiffusion.modules.models.gcdm.gcp_layers.GCPMessagePassing(input_dims: MolecularDiffusion.modules.models.gcdm.gcp_utils.ScalarVector, output_dims: MolecularDiffusion.modules.models.gcdm.gcp_utils.ScalarVector, edge_dims: MolecularDiffusion.modules.models.gcdm.gcp_utils.ScalarVector, cfg: GCPModuleConfig, mp_cfg: GCPLayerConfig, reduce_function: str = 'sum', use_scalar_message_attention: bool = True)¶

Bases: torch.nn.Module

aggregate(message, edge_index, dim_size: int)¶

forward(node_rep, edge_rep, edge_index, frames, node_mask=None)¶

message(node_rep, edge_rep, edge_index, frames, node_mask=None)¶

conv_cfg¶

message_fusion¶

reduce_function = 'sum'¶

self_message¶

use_residual_message_gcp¶

use_scalar_message_attention = True¶

class MolecularDiffusion.modules.models.gcdm.gcp_layers.GCPModuleConfig¶

Flattened configs/model/module_cfg/*_gcp_module.yaml.

Defaults are the values shipped for both QM9 and GEOM (the two upstream files are identical apart from their conditioning list).

ablate_frame_updates: bool = False¶

ablate_scalars: bool = False¶

ablate_vectors: bool = False¶

bottleneck: int = 4¶

default_bottleneck: int = 4¶

default_vector_residual: bool = False¶

frame_gate: bool = False¶

gcp_version: str = 'GCP2'¶

node_positions_weight: float = 1.0¶

nonlinearities: Tuple[str | None, str | None] | None = None¶

norm_x_diff: bool = True¶

scalar_gate: int = 0¶

scalar_nonlinearity: str | None = 'silu'¶

property selected_gcp¶

sigma_frame_gate: bool = False¶

update_positions_with_vector_sum: bool = False¶

vector_frame_residual: bool = False¶

vector_gate: bool = True¶

vector_nonlinearity: str | None = 'silu'¶

vector_residual: bool = False¶

MolecularDiffusion.modules.models.gcdm.gcp_layers.get_GCP_with_custom_cfg(input_dims, output_dims, cfg: GCPModuleConfig, **kwargs)¶: Upstream gcpnet.py:606, with OmegaConf swapped for dataclasses.