MolecularDiffusion.modules.tasks.diffusion_flowmol¶

Coordinate-only FlowMol flow-matching task integration.

Wraps the SE(3)-equivariant GVP endpoint vector field (modules/models/flowmol) in the duck-typed Task contract (docs/adding_new_models.md §2.1) and adapts the platform’s PointCloud dict batch to / from FlowMol’s batched-DGL-graph container.

Coordinate-only scope (per the approved INTEGRATION_PLAN): the generated modalities are x (coords), a (atom types), c (formal charge). Bonds (the e modality) are NOT generated – edges are a geometry-only message-passing convenience. Since the PointCloud pipeline carries no formal-charge channel, c is trained against a constant neutral target (QM9-style neutral molecules) and is not surfaced in the sample() return.

Classes¶

`DGLToPointCloudAdapter`	Integrated DGL graph -> padded PointCloud sample tuple pieces.
`FlowMolFlowMatchingTask`
`FlowMolTaskFactory`	Factory instantiated by `cli/train.py` (declares `train_set` so the
`PointCloudToDGLAdapter`	PointCloud dict batch -> batched, fully-connected DGL graph.

Module Contents¶

class MolecularDiffusion.modules.tasks.diffusion_flowmol.DGLToPointCloudAdapter¶

Bases: torch.nn.Module

Integrated DGL graph -> padded PointCloud sample tuple pieces.

forward(g: dgl.DGLGraph, n_atom_types: int) → Dict[str, torch.Tensor]¶

class MolecularDiffusion.modules.tasks.diffusion_flowmol.FlowMolFlowMatchingTask(n_atom_types: int, n_charges: int, neutral_charge_index: int, interpolant_scheduler_config: dict, vector_field_config: dict, prior_std: float, time_scaled_loss: bool, total_loss_weights: dict, default_n_timesteps: int, dataset_stats: dict, atom_vocab: list | None = None)¶

Bases: torch.nn.Module

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

forward(batch: Dict[str, torch.Tensor])¶

predict_and_target(batch: Dict[str, torch.Tensor])¶

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

atom_vocab = None¶

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

property device¶

fm_num_timesteps¶

interpolant_scheduler¶

loss_a¶

loss_c¶

loss_x¶

property model¶

n_atom_types¶

n_charges¶

property n_node_dist¶

neutral_charge_index¶

property node_dist_model¶

prior_std¶

prop_dist_model = None¶

task_type = 'diffusion_flowmol'¶

time_scaled_loss¶

to_dgl¶

to_pc¶

total_loss_weights¶

vector_field¶

class MolecularDiffusion.modules.tasks.diffusion_flowmol.FlowMolTaskFactory(task_type: str, interpolant_scheduler_config: dict, vector_field_config: dict, n_charges: int = 6, neutral_charge_index: int = 0, prior_std: float = 1.0, time_scaled_loss: bool = True, total_loss_weights: dict | None = None, default_n_timesteps: int = 250, dataset_stats: dict | None = None, atom_vocab: list | None = None, train_set: torch.utils.data.Dataset | None = None, **kwargs)¶

Factory instantiated by cli/train.py (declares train_set so the declarative seam injects the dataset for atom-count histogram stats).

build()¶

compute_dataset_stats(dataset)¶: Build the atom-count histogram from the training set (like TABASCO’s compute_dataset_stats).

atom_vocab¶

dataset_stats¶

default_n_timesteps = 250¶

interpolant_scheduler_config¶

kwargs¶

n_charges = 6¶

neutral_charge_index = 0¶

prior_std = 1.0¶

task_type¶

time_scaled_loss = True¶

total_loss_weights¶

train_set = None¶

vector_field_config¶

class MolecularDiffusion.modules.tasks.diffusion_flowmol.PointCloudToDGLAdapter(n_atom_types: int, n_charges: int, neutral_charge_index: int)¶

Bases: torch.nn.Module

PointCloud dict batch -> batched, fully-connected DGL graph.

Padding is stripped via node_mask; per-molecule COM is removed from coordinates so x_1_true lives in the same zero-COM subspace as the Gaussian prior. No edata is set – bonds are dropped.

forward(batch: Dict[str, torch.Tensor]) → dgl.DGLGraph¶

n_atom_types¶

n_charges¶

neutral_charge_index¶