MolecularDiffusion.modules.models.diffpharma.distributions

2D (n_ligand x n_pocket) size prior for DiffPharma.

Port of DistributionNodes from others/DiffPharma/equivariant_diffusion/conditional_model.py.

One deliberate deviation: upstream sample(n) returns a pair (n_lig, n_pocket), which does not satisfy the platform’s node_dist_model.sample(n) -> LongTensor of atom counts contract. Here sample returns only the ligand component (i.e. the marginal over pocket size) and the joint draw is available as sample_joint(). Generation uses sample_conditional(), which is what DiffPharma itself calls.

Classes

DistributionNodes

Module Contents

class MolecularDiffusion.modules.models.diffpharma.distributions.DistributionNodes(histogram)

log_prob(batch_n_nodes_1, batch_n_nodes_2)

log_prob_n1_given_n2(n1, n2)

log_prob_n2_given_n1(n2, n1)

sample(n_samples=1)

Ligand sizes only (pocket size marginalised out).

sample_conditional(n1=None, n2=None)

sample_joint(n_samples=1)

idx_to_n_nodes

m

n1_given_n2

n2_given_n1

property n_node_dist

{n_ligand_atoms: count} – the ligand-size marginal.

n_nodes_to_idx

prob