MolecularDiffusion.modules.tasks.diffusion_flowmol_graph3d¶

FlowMol3 (bond-generating, CTMC discrete flow matching) task.

The GEOM-Drugs-scale, bond-generating successor to the platform’s coordinate-only diffusion_flowmol task, which this file does not touch. Four modalities are generated jointly: x (coordinates, continuous flow matching) and a/c/e (atom types, formal charges, bond orders – CTMC discrete flow matching from an all-mask prior).

Wraps modules/models/flowmol_graph3d in the duck-typed Task contract (docs/adding_new_models.md §2.1) and adapts the platform’s graph3d PyG Batch (bond_collate: raw) to / from FlowMol’s batched DGL graph.

Three atom-type widths, all real, all easy to confuse (this cost the first FlowMol port two test attempts):

width	value	where
`len(atom_vocab)`	10	stored data, and the `sample()` one-hot
`n_atom_types`	11	output heads (+1 fake-atom column)
`a` input token embedding	12	`n_atom_types` + 1 CTMC mask token

kekulize: true is mandatory in every config that reaches this task. The released FlowMol3 weights have explicit_aromaticity: False: their to_edge_logits is 4 logits wide and token_embeddings.e is 4 classes plus one mask row. Canonical bond class 4 (AROMATIC) therefore has no input row and no output logit – a class-4 label is a hard index crash, not a distribution wrinkle. Because kekulize_bonds returns bonds unchanged when RDKit cannot kekulize a molecule (graph3d_dataset.py:213-216), which does happen on GEOM-Drugs, Graph3DToDGLAdapter.forward() asserts the invariant rather than trusting it.

Attributes¶

`CANONICAL_FEAT_ORDER`
`Chem`
`ModelTaskFactory`
`N_BOND_TYPES_KEKULIZED`
`logger`

Classes¶

`FlowMolGraph3DTask`	FlowMol3 wrapped in the platform's duck-typed Task contract (§2.1).
`FlowMolGraph3DTaskFactory`	Factory instantiated by `cli/train.py`.
`Graph3DToDGLAdapter`	`graph3d` PyG `Batch` -> batched, fully-connected DGL graph.

Module Contents¶

class MolecularDiffusion.modules.tasks.diffusion_flowmol_graph3d.FlowMolGraph3DTask(atom_vocab: list, n_atom_types: int, n_charge_classes: int, charge_offset: int, n_bond_types: int, fake_atom_p: float, fake_atom_std: float, distort_p: float, distort_t: float, prior_std: float, time_scaled_loss: bool, total_loss_weights: dict, default_n_timesteps: int, stochasticity: float, high_confidence_threshold: float, sdf_output_path: str | None, interpolant_scheduler_config: dict, vector_field_config: dict, n_atoms_hist: dict, task_type: str = 'diffusion_flowmol_graph3d', condition_names: list = [], context_mask_rate: float = 0.0, mask_value: float = 0.0, normalize_condition: str | None = None, adapter_conditions: list | None = None, use_adapter_module: bool = False)¶

Bases: torch.nn.Module

FlowMol3 wrapped in the platform’s duck-typed Task contract (§2.1).

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

forward(batch: Any) → tuple[torch.Tensor, dict]¶: One training step: interpolate, distort, denoise, weight the losses.

predict_and_target(batch: Any)¶: Pure-generative stub: the loss is the only scalar to report.

preprocess(train_set=None, valid_set=None, test_set=None)¶

Build self.property_norms for CFG conditioning (train-side only).

Called generically by cli/train.py if this attribute exists. Does NOT touch node_dist_model/n_node_dist – those come from dataset_stats/graph3d_stats at __init__ time via FlowMolGraph3DTaskFactory, a separate mechanism. Deliberately skips DistributionProperty/prop_dist_model (out of scope – generation always takes an explicit target_value).

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

Unconditional sampling from the all-mask prior.

Returns the platform’s (one_hot, charges, coords, node_mask). one_hot is ``len(atom_vocab)`` wide (10) – the fake-atom column is consumed internally and never surfaces. charges carries signed formal charges (MiDi’s and FlowMol’s precedent for that slot). The generated bond table has no channel in the tuple, so it is stashed on self.last_bond_types and, when sdf_output_path is set, written as an .sdf sidecar beside the platform’s .xyz.

Because a node the model assigns to the fake class is physically deleted (upstream molecule_builder.py:226-231), the produced atom count is <= the requested one. node_mask carries the truth; a mild size shortfall is upstream behaviour, not a port defect.

sample_guidance_conitional(target_function=None, target_value=None, negative_target_value=None, nodesxsample=None, cfg_scale: float = 1, cfg_scale_schedule: str | None = None, guidance_ver: str = 'cfg', n_frames: int = 0, num_steps: int | None = None, **kwargs: Any)¶: Classifier-free-guidance generation. Matches the call signature GenerativeFactory.conditional_generation() hardcodes for task_type == “cfg” (runmodes/generate/tasks_generate.py), and returns (one_hot, charges, x, node_mask) like sample() – “EDM compatibility”.

atom_vocab¶

canonical_feat_order = ['x', 'a', 'c', 'e']¶

charge_offset¶

condition = []¶

context_mask_rate = 0.0¶

property device: torch.device¶

distort_p¶

distort_t¶

fake_atom_index¶

fake_atom_p¶

fm_num_timesteps¶

high_confidence_threshold¶

interpolant_scheduler¶

last_bond_types: list | None = None¶

loss_cat¶

loss_x¶

mask_value = 0.0¶

property model: FlowMolGraph3DTask¶

n_adapter_context¶

n_atom_types¶

n_bond_types¶

n_charge_classes¶

n_concat_context¶

property n_node_dist: dict¶: {n_atoms: count}, used by GenerativeFactory to clamp sizes.

n_real_atom_types¶

node_dist_model¶

normalize_condition = None¶

prior_std¶

prop_dist_model = None¶

property_norms = None¶

sdf_output_path¶

stochasticity¶

task_type = 'diffusion_flowmol_graph3d'¶

time_scaled_loss¶

to_dgl¶

total_loss_weights¶

vector_field¶

class MolecularDiffusion.modules.tasks.diffusion_flowmol_graph3d.FlowMolGraph3DTaskFactory(task_type: str = 'diffusion_flowmol_graph3d', interpolant_scheduler_config: dict | None = None, vector_field_config: dict | None = None, n_charge_classes: int = 6, charge_offset: int = 2, n_bond_types: int = N_BOND_TYPES_KEKULIZED, fake_atom_p: float = 0.3, fake_atom_std: float = 1.0, distort_p: float = 0.7, distort_t: float = 0.25, prior_std: float = 1.0, time_scaled_loss: bool = True, total_loss_weights: dict | None = None, default_n_timesteps: int = 250, stochasticity: float = 30.0, high_confidence_threshold: float = 0.9, sdf_output_path: str | None = None, dataset_stats: dict | None = None, atom_vocab: list | None = None, train_set: torch.utils.data.Dataset | None = None, **kwargs: Any)¶

Factory instantiated by cli/train.py.

Declaring train_set is what makes the declarative injection seam (cli/train.py:624-637) hand over the dataset, which carries graph3d_stats and therefore the molecule-size histogram FlowMol needs for sampling.

sdf_output_path is declared generation-time (docs §2.5b): the task is rebuilt from the checkpoint’s training-time config, where it is null, so without this declaration the generate config’s value never arrives and the bond sidecar – the whole 2D half of the model – is silently dropped.

build() → FlowMolGraph3DTask¶

atom_vocab¶

charge_offset = 2¶

dataset_stats¶

default_n_timesteps = 250¶

distort_p = 0.7¶

distort_t = 0.25¶

fake_atom_p = 0.3¶

fake_atom_std = 1.0¶

generation_time_keys = ('sdf_output_path',)¶

high_confidence_threshold = 0.9¶

interpolant_scheduler_config¶

kwargs¶

n_bond_types = 4¶

n_charge_classes = 6¶

prior_std = 1.0¶

sdf_output_path = None¶

stochasticity = 30.0¶

task: FlowMolGraph3DTask | None = None¶

task_type = 'diffusion_flowmol_graph3d'¶

time_scaled_loss = True¶

total_loss_weights¶

train_set = None¶

vector_field_config¶

class MolecularDiffusion.modules.tasks.diffusion_flowmol_graph3d.Graph3DToDGLAdapter(n_atom_types: int, n_charge_classes: int, charge_offset: int, n_bond_types: int, fake_atom_p: float, fake_atom_std: float)¶

Bases: torch.nn.Module

graph3d PyG Batch -> batched, fully-connected DGL graph.

Per molecule, in upstream’s order (data_processing/dataset.py:99-149):

inject fake atoms first, so the fully-connected edge set covers them;
re-remove the centre of mass (upstream re-centres after the injection);
build the edge set with build_edge_idxs – never re-derived, because get_upper_edge_mask infers the upper/lower split from that ordering alone and silently returns a wrong mask under any other order;
densify the stored upper-triangular bonds into an (n, n) integer adjacency and read the labels back along triu_indices. This is the step that materializes bond class 0 (“no bond”), which is never stored;
mirror the labels onto both edge directions.

forward(batch: Any, *, use_fake_atoms: bool, condition: torch.Tensor | None = None) → dgl.DGLGraph¶

charge_offset¶

fake_atom_index¶

fake_atom_p¶

fake_atom_std¶

n_atom_types¶

n_bond_types¶

n_charge_classes¶

MolecularDiffusion.modules.tasks.diffusion_flowmol_graph3d.CANONICAL_FEAT_ORDER = ['x', 'a', 'c', 'e']¶

MolecularDiffusion.modules.tasks.diffusion_flowmol_graph3d.Chem = None¶

MolecularDiffusion.modules.tasks.diffusion_flowmol_graph3d.ModelTaskFactory¶

MolecularDiffusion.modules.tasks.diffusion_flowmol_graph3d.N_BOND_TYPES_KEKULIZED = 4¶

MolecularDiffusion.modules.tasks.diffusion_flowmol_graph3d.logger¶