MolecularDiffusion.modules.models.kgdiff.atom_num

Pocket-size-conditioned prior over ligand atom counts.

Ported from KGDiff utils/evaluation/atom_num.py (commit ad893fc).

The prior is a static table (atom_num_config, generated once by the authors from CrossDocked): pockets are binned by their spatial extent, and each bin carries an empirical distribution over ligand sizes. Nothing here depends on the training set, so it behaves identically at train and generate time and needs no checkpoint buffer – unlike the histogram priors the other in-tree pocket models carry.

Functions

get_space_size(→ float)	Pocket extent: the median of the 10 largest pairwise distances.
marginal_size_distribution(→ dict[int, float])	{n_atoms: probability} marginalised over all pocket-size bins.
sample_atom_num(→ int)	Draw one ligand atom count for a pocket of this extent.

Module Contents

MolecularDiffusion.modules.models.kgdiff.atom_num.get_space_size(pocket_3d_pos: numpy.ndarray) → float

Pocket extent: the median of the 10 largest pairwise distances.

MolecularDiffusion.modules.models.kgdiff.atom_num.marginal_size_distribution() → dict[int, float]

{n_atoms: probability} marginalised over all pocket-size bins.

Used for the task’s n_node_dist, which callers only read to derive max_atom; the real, pocket-conditioned draw goes through sample_atom_num().

MolecularDiffusion.modules.models.kgdiff.atom_num.sample_atom_num(space_size: float) → int

Draw one ligand atom count for a pocket of this extent.