MolecularDiffusion.modules.models.gcdm.dynamics

Adapter binding GCDM’s GCPNet denoiser to the platform’s dense point-cloud diffusion contract.

EnVariationalDiffusion.phi (modules/models/en_diffusion.py:234) calls dynamics._forward(t, xh, node_mask, edge_mask, context) with dense, padded tensors:

xh        (B, N, 3 + in_node_nf)
node_mask (B, N, 1)
edge_mask (B, N*N, 1)      -- discarded here, see below
context   (B, N, C) or None
t         (B, 1) or a 0-d/1-element tensor

GCDM’s GCPNetDynamics.atom_types_and_coords_forward (src/models/components/gcpnet.py:1069) instead consumes the flat, unpadded node list of a PyG Batch. GCDMDynamics bridges the two by compacting the padded batch down to its real atoms, running the ported network, and scattering the prediction back into the padded layout.

Compaction rather than a plain reshape(B*N, ...) is deliberate and load-bearing for checkpoint fidelity: GCDM’s equivariant node channel chi is _orientations – forward/backward displacements between consecutive rows of the flat node list – so padding rows interleaved between molecules would change those features. Row-major compaction of a dense (B, N, ...) batch reproduces exactly the concatenated node ordering Batch.from_data_list gives upstream.

edge_mask is discarded because GCPNet rebuilds its own fully-connected intra-molecule edge list every step from the batch index (get_fully_connected_edge_index, gcpnet.py:1056); the platform’s edge_mask is the identical mask outer-product.

Classes

GCDMDynamics

GCPNet denoiser, ported from GCPNetDynamics (gcpnet.py:933).

Module Contents

class MolecularDiffusion.modules.models.gcdm.dynamics.GCDMDynamics(in_node_nf: int, context_node_nf: int = 0, n_dims: int = 3, num_encoder_layers: int = 9, h_hidden_dim: int = 256, chi_hidden_dim: int = 32, e_hidden_dim: int = 64, xi_hidden_dim: int = 16, chi_input_dim: int = 2, e_input_dim: int = 1, xi_input_dim: int = 1, dropout: float = 0.0, condition_on_time: bool = True, self_condition: bool = False, module_cfg: MolecularDiffusion.modules.models.gcdm.gcp_layers.GCPModuleConfig | None = None, layer_cfg: MolecularDiffusion.modules.models.gcdm.gcp_layers.GCPLayerConfig | None = None)

Bases: torch.nn.Module

GCPNet denoiser, ported from GCPNetDynamics (gcpnet.py:933).

All five upstream DictConfig arguments are flattened into named keyword arguments; the defaults below are the shipped QM9 preset (configs/model/model_cfg/qm9_mol_gen_ddpm_gcp_model.yaml).

Parameters:

  • in_node_nf – Number of node scalar channels the diffusion latent carries – i.e. len(atom_vocab) + int(include_charges). Upstream’s num_atom_types + include_charges.

  • context_node_nf – Number of property-conditioning channels (upstream len(module_cfg.conditioning)). 0 disables context conditioning entirely.

  • n_dims – Spatial dimensionality (3).

  • num_encoder_layers – Number of GCPInteractions blocks (QM9: 9, GEOM: 4).

  • xi_hidden_dim (h_hidden_dim / chi_hidden_dim / e_hidden_dim /) – Hidden scalar/vector widths for nodes and edges.

  • xi_input_dim (chi_input_dim / e_input_dim /) – Input widths of the geometric features GCPNet builds itself.

  • self_condition – Ported for completeness; False in every shipped upstream config and doubles the input widths when on.

abstractmethod forward(t, xh, node_mask, edge_mask, context=None)
static get_fully_connected_edge_index(batch_index: torch.Tensor, node_mask: torch.Tensor | None = None) → torch.Tensor

gcpnet.py:1056. Every intra-molecule ordered pair, self loops included.

unwrap_forward()
wrap_forward(node_mask, edge_mask, context)
condition_on_context
condition_on_time = True
edge_dims
edge_input_dims
gcp_embedding
in_node_nf
interaction_layers
layer_cfg
module_cfg
n_dims = 3
node_dims
node_input_dims
norm_x_diff = True
num_context_node_features = 0
num_x_dims = 3
scalar_node_projection_gcp
self_condition = False