MolecularDiffusion.modules.tasks.diffusion_syncogen

SynCoGen task: joint masked graph diffusion + coordinate flow matching.

SynCoGen generates molecules the way a chemist would make them: nodes are catalogue building blocks, edges are reaction templates, and the 3D coordinates are flow-matched alongside. Chemical bonds are never diffused – they are a consequence of the sampled reaction graph, recovered by RDKit assembly at the end (BBRxnGraph.build_rdkit -> RDKitMoleculeAssembly).

Two things follow from that, and both shape this file.

The batch is not a point cloud. Training consumes upstream’s own PyG Data/Batch end to end, produced by the task-owned modules/models/syncogen/datamodule.py. No platform data_type is involved and nothing under MolecularDiffusion/data/ is touched.

Generation is. GenerativeFactory’s unconditional loop wants (one_hot, charges, coords, node_mask) plus a dense bond matrix on last_bond_types. So the only adapters in this integration live in sample(), converting the sampled (BBRxnGraph, Coordinates) pair into those tensors – see _assemble and _to_pointcloud.

Two traps worth naming here rather than rediscovering:

• ``charges`` carries signed formal charges, not atomic numbers. The channel is overloaded in the platform: save_xyz_file would read it as atomic numbers, but only when the model sets use_unknown_fallback (which this one does not, so the value is unused there), while GenerativeFactory._write_molecule_sdf passes it straight through as formal_charge= (runmodes/generate/tasks_generate.py:344). Atomic numbers here would silently stamp every SDF atom with a +6/+7/+8 charge.

• Atom ordering is fixed by ``set_mol_coordinates``. Upstream assigns the masked coordinate rows positionally onto the assembled molecule’s atoms, so the point-cloud order and the RDKit atom order are the same order by construction. one_hot, charges, coords and the bond matrix are all built in that one order, or the .xyz and the .sdf would describe different molecules.

Vocabulary loading is process-global and must happen before any other import from the ported package – see modules/models/syncogen/vocab.py. That is why every syncogen import in this file is deferred into a function body.

Attributes

logger

Classes

ModelTaskFactory	Instantiated by cli/train.py / cli/generate.py from Hydra.
SyncogenDiffusionTask	The duck-typed Task object (docs/adding_new_models.md §2.1).
SyncogenSizePrior	Fragment-count prior, and the small surface core.Diffusion reads.

Module Contents

class MolecularDiffusion.modules.tasks.diffusion_syncogen.ModelTaskFactory(task_type: str = 'diffusion_syncogen', vocab_dir: str | None = None, backbone: str = 'semla', d_model: int = 384, d_message: int = 128, n_coord_sets: int = 64, n_layers: int = 12, n_head: int = 32, d_message_hidden: int = 128, d_edge: int = 128, coord_norm: str = 'length', size_emb: int = 64, pos_emb: int = 64, max_bbs: int = 5, self_conditioning: bool = True, pharm_subset: int = 7, noise_sigma_min: float = 0.001, noise_sigma_max: float = 7.0, num_sample_steps: int = 100, sampling_eps: float = 0.001, antithetic_sampling: bool = True, importance_sampling: bool = False, sampling_noise_removal: bool = True, use_compat: bool = True, time_conditioning: bool = True, train_rot_align: bool = True, augmentations: Sequence[str] = ('center', 'random_rotate', 'normalize'), scale_noise: bool = True, scale_noise_factor: float = 0.2, inference_annealing: bool = True, annealing_coef: float = 10.0, ema_decay: float = 0.0, nll_coef: float = 1.0, mse_coef: float = 1.0, bond_length_coef: float = 0.2, pairwise_distance_coef: float = 0.4, pairwise_distance_threshold: float = 5.0, smooth_lddt_coef: float = 0.4, loss_t_threshold: float = 0.25, load_bonds: bool = True, load_pharmacophores: bool = False, num_fragments_probs: dict | None = None, eval_batch_size: int = 4, reference_ligand: str | None = None, sdf_output_path: str | None = None, max_sample_attempts: int = 5, atom_vocab: list | None = None, **kwargs: Any)

Instantiated by cli/train.py / cli/generate.py from Hydra.

No train_set parameter is declared, so the §2.5 injection seam stays inert: SynCoGen’s only construction-time statistic is the fragment-count histogram, and that is config-owned (num_fragments_probs) precisely so generation needs no dataset.

build() → SyncogenDiffusionTask

annealing_coef

antithetic_sampling = True

atom_vocab

augmentations

backbone = 'semla'

backbone_kwargs

ema_decay

eval_batch_size = 4

generation_time_keys = ('vocab_dir', 'reference_ligand', 'sdf_output_path', 'num_fragments_probs', 'eval_batch_size',...

importance_sampling = False

inference_annealing = True

load_bonds = True

load_pharmacophores = False

loss_coefs

max_bbs = 5

max_sample_attempts = 5

noise_sigma_max

noise_sigma_min

num_fragments_probs

num_sample_steps = 100

pharm_subset = 7

reference_ligand = None

sampling_eps

sampling_noise_removal = True

scale_noise = True

scale_noise_factor

sdf_output_path = None

self_conditioning = True

task: SyncogenDiffusionTask | None = None

task_type = 'diffusion_syncogen'

time_conditioning = True

train_rot_align = True

use_compat = True

vocab_dir = None

class MolecularDiffusion.modules.tasks.diffusion_syncogen.SyncogenDiffusionTask(core: torch.nn.Module, size_prior: SyncogenSizePrior, atom_vocab: list, task_type: str, num_sample_steps: int, pharm_subset: int, max_bbs: int, reference_ligand: str | None = None, sdf_output_path: str | None = None, max_sample_attempts: int = 5)

Bases: torch.nn.Module

The duck-typed Task object (docs/adding_new_models.md §2.1).

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

forward(batch) → tuple[torch.Tensor, dict]

predict_and_target(batch) → tuple[torch.Tensor, torch.Tensor]

Pure-generative stub: the loss is both prediction and target.

sample(batch_size: int | None = None, nodesxsample: torch.Tensor | None = None, num_steps: int | None = None, mode: str | None = None, n_frames: int = 0, **kwargs: Any) → tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]

Return (one_hot, charges, coords, node_mask) for batch_size molecules.

nodesxsample is a fragment count per molecule, not an atom count (see SyncogenSizePrior). charges carries signed formal charges.

A sampled graph can come back with a reaction the assembler rejects, so a round is retried up to max_sample_attempts times rather than emitting an empty .xyz for the failed rows. All attempts exhausted -> raise, so the generation loop counts the batch as failed instead of writing nothing and exiting 0.

T

atom_vocab

core

property device: torch.device

last_bond_types: torch.Tensor | None = None

max_bbs

max_sample_attempts = 5

property model: SyncogenDiffusionTask

property n_node_dist: dict

property node_dist_model: SyncogenSizePrior

pharm_subset

prop_dist_model = None

reference_ligand = None

sdf_output_path = None

size_prior

task_type

class MolecularDiffusion.modules.tasks.diffusion_syncogen.SyncogenSizePrior(num_fragments_probs: dict, max_bbs: int = 5, batch_size: int = 4, eval_batch_size: int = 4, load_bonds: bool = True, load_pharmacophores: bool = False)

Fragment-count prior, and the small surface core.Diffusion reads.

Serves three roles at once, which is why it is one object and not three:

1. task.node_dist_model – .sample(n) -> LongTensor of per-molecule sizes. For SynCoGen a “size” is a fragment count (typically 2-5), not an atom count, so mol_size in a generate config must be [0, 0] or a range inside [min, max_bbs]. GenerativeFactory.__init__ clamps against max(n_node_dist) and raises with a clear message otherwise.

2. task.n_node_dist – the {n_fragments: count} histogram that clamp reads.

3. The stand-in for upstream’s SyncogenDataManager inside core.Diffusion, which reaches for max_bbs, batch_size, eval_batch_size, load_bonds, load_pharmacophores and sample_n_nodes. Keeping it config-owned rather than dataset-owned is what lets generation run with no dataset present at all – which is the whole point for the released checkpoints, where the 6.3 GB SynSpace download is not needed.

sample(n: int) → torch.Tensor

sample_n_nodes(batch_size: int) → torch.Tensor

batch_size = 4

eval_batch_size = 4

forced_n_nodes: torch.Tensor | None = None

load_bonds = True

load_pharmacophores = False

max_bbs = 5

property n_node_dist: dict

{n_fragments: count}. Counts are a scaled histogram; only the keys and their relative sizes matter to GenerativeFactory.

train_length_probs

train_length_values

MolecularDiffusion.modules.tasks.diffusion_syncogen.logger