MolecularDiffusion.modules.tasks.diffusion_diffsmol

Shape-conditioned DiffSMol diffusion task.

Binds modules/models/diffsmol to the duck-typed Task contract (docs/adding_new_models.md §2.1) and adapts the platform’s padded PointCloud dict batch to DiffSMol’s flat/ragged (N_total, ...) + batch_ligand layout.

Scope (approved INTEGRATION_PLAN): bonds dropped, D3PM atom-type diffusion kept, shape conditioning + classifier-free guidance kept, heavy atoms only.

Frame. The shape latent is a (128, 3) SO(3)-equivariant feature that rotates with the molecule and is not translation invariant. Upstream puts coordinates in the surface-centroid frame – the mean of the 512 sampled mesh points – and this task does the same by subtracting the precomputed shape_center. It therefore deliberately does not call remove_mean_with_mask: COM-centering would put the coordinates in a different frame from the latent and silently break conditioning. If rotation augmentation is ever switched on, the same R must be applied to both coords and shape_emb (the latter is a stack of 128 3-vectors, so shape_emb @ R.T is the correct action).

Shape cache. The mesh -> point cloud -> VN-AE chain is precomputed offline by docs/model_integrations/diffsmol/scripts/precompute_shapes.py into a .pt keyed by each sample’s xyz identifier – the same key the PointCloud dataset already carries through the default collate. No custom DataModule and no collate change is needed: the platform’s collate passes unknown keys (including a list of strings) straight through.

Attributes

DIFFSMOL_AROMATIC_FLAGS
DIFFSMOL_AROMATIC_VOCAB
DIFFSMOL_ATOM_VOCAB

Classes

DiffSMolDiffusionTask	Task contract implementation for DiffSMol.
DiffSMolShapeGenerator	Shape-conditioned generator behind interference/gen_diffsmol_shape.
DiffSMolShapeNodeDistribution	Atom-count sampler, marginal and shape-volume-conditioned.
DiffSMolTaskFactory	Instantiated by cli/train.py.

Functions

load_shape_cache(→ Dict[str, Dict[str, Any]])

Load the precomputed {key: {shape_emb, shape_center, shape_volume}}

Module Contents

class MolecularDiffusion.modules.tasks.diffusion_diffsmol.DiffSMolDiffusionTask(model_config: dict, atom_vocab: List[str], shape_cache: Dict[str, Dict[str, Any]], dataset_stats: dict, default_num_steps: int, guide_stren: float = 0.0, threshold_type: str | None = None, threshold_p: float | None = None, task_type: str = 'diffusion_diffsmol', atom_mode: str = 'basic')

Bases: torch.nn.Module

Task contract implementation for DiffSMol.

evaluate(pred: torch.Tensor, target: torch.Tensor)

forward(batch: Dict[str, Any])

predict_and_target(batch: Dict[str, Any])

sample(batch_size: int | None = None, nodesxsample: torch.Tensor | None = None, num_steps: int | None = None, batch: Dict[str, Any] | None = None, shape_emb: torch.Tensor | None = None, guide_stren: float | None = None, mode=None, n_frames: int = 0, **kwargs: Any)

Reverse-diffuse a batch of molecules.

shape_emb=None uses the zero latent, i.e. the classifier-free unconditional branch – so the bundled gen_unconditional config works with no reference molecule. Returns the platform’s EDM 4-tuple (one_hot, charges, coords, node_mask).

Coordinates come out in the surface-centered frame of the reference shape, which is the frame that makes generated molecules land inside the reference envelope.

T

atom_mode = 'basic'

atom_vocab

default_num_steps

property device

guide_stren = 0.0

property model

n_atom_types

property n_node_dist: Dict[int, int]

property node_dist_model: DiffSMolShapeNodeDistribution

prop_dist_model = None

score_model

shape_dim

task_type = 'diffusion_diffsmol'

threshold_p = None

threshold_type = None

class MolecularDiffusion.modules.tasks.diffusion_diffsmol.DiffSMolShapeGenerator(task, reference_shape: str | None = None, reference_key: str | None = None, num_generate: int = 100, batch_size: int = 8, num_steps: int | None = None, guide_stren: float = 0.0, mol_size: list | None = None, output_path: str = 'generated_diffsmol', seed: int = 42, save_xyzrender_figures: bool = False, **kwargs: Any)

Shape-conditioned generator behind interference/gen_diffsmol_shape.

GenerativeFactory cannot express this conditioning – its cfg branch carries a run-wide list of scalar target_values validated against len(task.condition), which a per-sample (128, 3) latent has no expression in. A separate generator behind its own _target_ is the in-tree pattern for exactly this (cf. PharmacophoreConditionGenerator) and needs no core change: cli/generate.py only calls hydra.utils.instantiate(cfg.interference, task=task) then run().

reference_shape is a .pt written by precompute_shapes.py – the same code path the training cache came from, so the reference latent is guaranteed to be in the same frame and produced by the same weights. Embedding a reference molecule live is deliberately not supported here.

# ponytail: reference must be pre-embedded. Keeps skimage/trimesh out of # the generation path entirely; precompute is one command you already ran.

run() → None

batch_size = 8

guide_stren = 0.0

mol_size

num_generate = 100

num_steps = None

output_path = 'generated_diffsmol'

save_xyzrender_figures = False

seed = 42

shape_emb = None

shape_volume = None

task

class MolecularDiffusion.modules.tasks.diffusion_diffsmol.DiffSMolShapeNodeDistribution(n_node_dist: Dict[int, int], volume_bins: List[float] | None = None, volume_hist: Dict[int, Dict[int, int]] | None = None, bin_window: int = 1)

Atom-count sampler, marginal and shape-volume-conditioned.

sample(n) draws from the marginal atom-count histogram (the duck-typed contract every generic caller uses). sample_for_volume(volume, n) draws from the sub-histogram of training molecules whose mesh volume is near volume, mirroring upstream’s volume-conditioned draw (sample_diffusion_no_pocket.py:242-283).

# ponytail: keyed on trimesh mesh volume, not upstream’s 0.5 A # occupied-voxel count – the precompute already has mesh.volume for # free, both are monotone volume proxies, and the histogram is binned # anyway. Port get_atom_stamp/get_voxel_shape/make_grid only if the # generated size distribution measurably degrades.

sample(n_samples: int = 1) → torch.Tensor

sample_for_volume(volume: float, n_samples: int = 1) → torch.Tensor

Draw sizes from molecules of comparable mesh volume.

Falls back to the marginal histogram if the neighbourhood around volume is empty (or if no volume histogram was built).

bin_window = 1

property max_n_nodes: int

n_node_dist

volume_bins = []

volume_hist

class MolecularDiffusion.modules.tasks.diffusion_diffsmol.DiffSMolTaskFactory(task_type: str = 'diffusion_diffsmol', model_config: dict | None = None, shape_cache_path: str | None = None, num_diffusion_timesteps: int = 1000, default_num_steps: int | None = None, guide_stren: float = 0.0, threshold_type: str | None = None, threshold_p: float | None = None, n_volume_bins: int = 20, atom_mode: str = 'basic', dataset_stats: dict | None = None, atom_vocab: list | None = None, train_set: torch.utils.data.Dataset | None = None, **kwargs: Any)

Instantiated by cli/train.py.

Declares train_set so the declarative seam at cli/train.py:621 injects the training dataset; that is where the atom-count / mesh-volume histograms come from. No core change is involved.

build() → DiffSMolDiffusionTask

compute_dataset_stats(dataset, shape_cache: dict) → None

Build the marginal atom-count histogram and, if the shape cache covers the training set, the mesh-volume-conditioned one.

atom_mode = 'basic'

atom_vocab

dataset_stats

default_num_steps = 1000

guide_stren = 0.0

kwargs

model_config

n_volume_bins = 20

num_diffusion_timesteps = 1000

shape_cache_path = None

task_type = 'diffusion_diffsmol'

threshold_p = None

threshold_type = None

train_set = None

MolecularDiffusion.modules.tasks.diffusion_diffsmol.load_shape_cache(path: str | None) → Dict[str, Dict[str, Any]]

Load the precomputed {key: {shape_emb, shape_center, shape_volume}} map. An absent path yields an empty cache – every molecule then falls back to the unconditional (zero-latent) branch, which is exactly what gen_unconditional wants and is enough to smoke-test the plumbing.

MolecularDiffusion.modules.tasks.diffusion_diffsmol.DIFFSMOL_AROMATIC_FLAGS = [False, False, True, False, True, False, True, False, False, True, False, True, False, False, False]

MolecularDiffusion.modules.tasks.diffusion_diffsmol.DIFFSMOL_AROMATIC_VOCAB = ['H', 'C', 'C', 'N', 'N', 'O', 'O', 'F', 'P', 'P', 'S', 'S', 'Cl', 'Br', 'I']

MolecularDiffusion.modules.tasks.diffusion_diffsmol.DIFFSMOL_ATOM_VOCAB = ['C', 'N', 'O', 'F', 'P', 'S', 'Cl', 'Br', 'I']