MolecularDiffusion.modules.tasks.diffusion_etflow¶

ET-Flow conformer-generation task (task_type: diffusion_etflow).

Wraps modules/models/etflow in the platform’s duck-typed Task contract (docs/adding_new_models.md 2.1) and adapts a graph3d PyG Batch (bond_collate: raw) into the flat tensors ET-Flow’s vector field consumes.

This model does not invent molecules. You hand it a molecule you already have – its atoms and its bonds – and it returns 3D conformers of exactly that molecule. Nothing about the composition is generated, so:

ETFlowTask.sample() raises. There is no unconditional mode: without a molecular graph there is no bond Laplacian, hence no harmonic prior, no node features and nothing to place. Generation goes through the shared ConformerFactory, pointed at by configs/interference/gen_conformer.yaml – the same route DiTMC and NExT-Mol take.
node_dist_model / n_node_dist are deliberately absent. They exist to let GenerativeFactory choose a molecule size; ET-Flow’s size is dictated by the input molecule.

BOND MAPPING (canonical -> ET-Flow). ET-Flow has no bond vocabulary. Every canonical class 1=SINGLE 2=DOUBLE 3=TRIPLE 4=AROMATIC maps to the same model-side value – edge_type = 1, “these two atoms are bonded” (models/utils.py:54-56, reached because upstream passes edge_attr=None on both its training and inference paths). Class 0 is never an edge, which is the platform’s storage rule – a pass-through. Radius-graph edges take edge_type = 0 and are regenerated from the current coordinates at every step. This is not a lossy adaptation on our side: it is upstream’s shipped behaviour, and bond order still reaches the model through node_attr (aromatic flag, hybridization, degree, implicit valence, H count).

Do not set ``kekulize: true``: it would zero the aromatic column and shift hybridization for every ring atom, changing the conditioning distribution the published weights were trained on.

Attributes¶

`ModelTaskFactory`
`TIME_HIGH`
`TIME_LOW`
`logger`

Classes¶

`ETFlowTask`	ET-Flow in the platform's Task contract.
`ETFlowTaskFactory`	Factory instantiated by `cli/train.py`.
`Graph3DBatchToETFlowAdapter`	One `graph3d` PyG `Batch` (or a plain item list) -> flat tensors.

Module Contents¶

class MolecularDiffusion.modules.tasks.diffusion_etflow.ETFlowTask(model_kwargs: dict, sigma: float = 0.1, prior_type: str = 'harmonic', harmonic_alpha: float = 1.0, parity_switch: str | None = 'post_hoc', sample_time_dist: str = 'uniform', max_num_neighbors: int = 32, atom_vocab: list | None = None, task_type: str = 'diffusion_etflow')¶

Bases: torch.nn.Module

ET-Flow in the platform’s Task contract.

self.network is named to match the released checkpoints exactly, so scripts/convert_checkpoint.py is an identity remap that can assert a strict bijection instead of guessing at a rename table.

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

forward(batch: Any)¶

generate_conformers(items: list, *, n_timesteps: int = 50, s_churn: float = 1.0, t_min: float = TIME_LOW, t_max: float = TIME_HIGH, std: float = 1.0, sampler_type: str = 'ode')¶

Euler-integrate the learned field; one output graph per input item.

With sampler_type: ode (upstream’s shipped setting) the t_min/t_max window is inert – the churn branch below is only reachable with sampler_type: stochastic.

predict_and_target(batch: Any, all_loss=None, metric=None)¶

sample(*args, **kwargs)¶

adapter¶

atom_vocab¶

cutoff¶

property device: torch.device¶

harmonic_sampler¶

max_num_neighbors = 32¶

property model: ETFlowTask¶

network¶

parity_switch = 'post_hoc'¶

prior_type = 'harmonic'¶

sample_time_dist = 'uniform'¶

sigma = 0.1¶

task_type = 'diffusion_etflow'¶

class MolecularDiffusion.modules.tasks.diffusion_etflow.ETFlowTaskFactory(task_type: str = 'diffusion_etflow', model: dict | None = None, flow: dict | None = None, atom_vocab: list | None = None, **kwargs: Any)¶

Factory instantiated by cli/train.py.

train_set is deliberately not declared: ET-Flow needs no dataset statistics at construction (no marginals, no valency table, no size histogram), so the declarative injection seam simply does not fire.

build() → ETFlowTask¶

atom_vocab¶

flow_kwargs¶

kwargs¶

model_kwargs¶

task: ETFlowTask | None = None¶

task_type = 'diffusion_etflow'¶

class MolecularDiffusion.modules.tasks.diffusion_etflow.Graph3DBatchToETFlowAdapter¶

Bases: torch.nn.Module

One graph3d PyG Batch (or a plain item list) -> flat tensors.

Two honesty caveats, carried from the integration plan:

Chirality is recovered from geometry, not read from a SMILES. build_rdkit_mol(..., coords=...) runs AssignStereochemistryFrom3D (data/component/graph3d_dataset.py:184-186), so the chiral tags come from the input conformer. Upstream reads them off the GEOM mol. This is the first thing to check if a converted pretrained checkpoint underperforms.
Achiral molecules yield chiral_index of shape (1, 0), and every consumer downstream must treat that as a no-op rather than an error.

build(items: list, device: Any = None) → dict¶: Flat, concatenated tensors in item-major order.

forward(batch: Any) → dict¶

cache¶

MolecularDiffusion.modules.tasks.diffusion_etflow.ModelTaskFactory¶

MolecularDiffusion.modules.tasks.diffusion_etflow.TIME_HIGH = 0.9999¶

MolecularDiffusion.modules.tasks.diffusion_etflow.TIME_LOW = 0.0001¶

MolecularDiffusion.modules.tasks.diffusion_etflow.logger¶