MolecularDiffusion.modules.tasks.diffusion_difflinker

DiffLinker integration with the MolecularDiffusion data pipeline.

DiffLinker (https://github.com/igashov/DiffLinker) is a fragment/linker mask-conditioned E(n)-equivariant diffusion model: given a set of fixed “fragment” atoms, it diffuses only a “linker” subset of atoms that connects them. See docs/model_integrations/difflinker/INTEGRATION_PLAN.md for the full integration plan (data adapters, task-contract mapping, scope) this module implements.

Ported model code lives under MolecularDiffusion.modules.models.difflinker (edm.py, egnn.py, noise.py, linker_size.py).

Attributes

DIFFLINKER_ROW_DATA_COLUMNS

Classes

DiffLinkerTask	Plain nn.Module task wrapper (TABASCO-style, no Task/
DiffLinkerTaskFactory	Factory matching train.py's task_module.build() /
PointCloudToDiffLinkerBatch	Converts a MolCraftDiffusion PointCloud batch dict into DiffLinker's

Module Contents

class MolecularDiffusion.modules.tasks.diffusion_difflinker.DiffLinkerTask(in_node_nf: int, n_dims: int, hidden_nf: int, activation: str, tanh: bool, n_layers: int, attention: bool, norm_constant: float, inv_sublayers: int, sin_embedding: bool, normalization_factor: float, aggregation_method: str, model: str, normalization: str | None, condition_time: bool, anchors_context: bool, diffusion_steps: int, diffusion_noise_schedule: str, diffusion_noise_precision: float, diffusion_loss_type: str, normalize_factors: tuple, center_of_mass: str, atom_vocab: list | None = None)

Bases: torch.nn.Module

Plain nn.Module task wrapper (TABASCO-style, no Task/ core.Configurable base needed – see docs/adding_new_models.md §2.6) implementing the §2.1 contract around DiffLinker’s EDM.

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

forward(batch: dict)

get_extra_state() → dict

predict_and_target(batch: dict)

sample(nodesxsample: torch.Tensor | None = None, batch_size: int | None = None, num_steps: int | None = None, batch: dict | None = None, condition_tensor: torch.Tensor | None = None, condition_mode: str | None = None, outpaint_cfgs: dict | None = None, use_noised_conditioning: bool = False, n_frames: int = 0, n_retrys: int = 0, t_retry: int | None = None, context: torch.Tensor | None = None, **kwargs)

Generate molecules by diffusing a linker between fixed fragment atoms. Dispatched via the generic “outpaint” GenerativeFactory path (runmodes/generate/tasks_generate.py::structural_guidance) – see INTEGRATION_PLAN.md’s Task-contract mapping section for the exact kwarg contract this method mirrors.

Several kwargs accepted here are explicitly out of scope this pass (inert, not implemented with real behavior): batch, num_steps, condition_mode, outpaint_cfgs, use_noised_conditioning, n_frames (no trajectory export), n_retrys/t_retry (no bond-distance retry loop), context (property-conditioning value – DiffLinker has no prop_dist_model, so this is always None here).

set_extra_state(state: dict) → None

anchors_context

property atom_count_histogram: dict | None

atom_vocab = None

center_of_mass

property device: torch.device

edm

in_node_nf

loss_type

property model

n_dims

property n_node_dist: dict | None

ndim_extra = 0

property node_dist_model: MolecularDiffusion.modules.models.difflinker.linker_size.DistributionNodes | None

norm_values

prop_dist_model = None

class MolecularDiffusion.modules.tasks.diffusion_difflinker.DiffLinkerTaskFactory(task_type: str, in_node_nf: int, n_dims: int = 3, hidden_nf: int = 128, activation: str = 'silu', tanh: bool = False, n_layers: int = 6, attention: bool = False, norm_constant: float = 1e-06, inv_sublayers: int = 2, sin_embedding: bool = False, normalization_factor: float = 100, aggregation_method: str = 'sum', model: str = 'egnn_dynamics', normalization: str | None = 'batch_norm', condition_time: bool = True, anchors_context: bool = True, diffusion_steps: int = 500, diffusion_noise_schedule: str = 'polynomial_2', diffusion_noise_precision: float = 1e-05, diffusion_loss_type: str = 'l2', normalize_factors: tuple = (1, 4, 10), center_of_mass: str = 'fragments', atom_vocab: list | None = None, **kwargs)

Factory matching train.py’s task_module.build() / task_module.task instantiation pattern (see diffusion_tabasco.py::ModelTaskFactory for the precedent).

build() → DiffLinkerTask

activation = 'silu'

aggregation_method = 'sum'

anchors_context = True

atom_vocab

attention = False

center_of_mass = 'fragments'

condition_time = True

diffusion_loss_type = 'l2'

diffusion_noise_precision = 1e-05

diffusion_noise_schedule = 'polynomial_2'

diffusion_steps = 500

hidden_nf = 128

in_node_nf

inv_sublayers = 2

kwargs

model = 'egnn_dynamics'

n_dims = 3

n_layers = 6

norm_constant = 1e-06

normalization = 'batch_norm'

normalization_factor = 100

normalize_factors = (1, 4, 10)

sin_embedding = False

tanh = False

task_type

class MolecularDiffusion.modules.tasks.diffusion_difflinker.PointCloudToDiffLinkerBatch(atom_vocab: list)

Converts a MolCraftDiffusion PointCloud batch dict into DiffLinker’s native per-atom field layout.

Two data paths share this one adapter (see INTEGRATION_PLAN.md’s Data adapters section):

• Path A (real ZINC data, converted via convert_zinc_pt_to_asedb.py + use_row_data_features: true): linker_mask/anchors are sliced straight out of node_feature’s trailing two columns.

• Path B (synthetic fallback, any ordinary single-molecule dataset): when those trailing columns aren’t present, a random contiguous slice of each molecule is fabricated as the “linker”, with the fragment atoms bordering the cut tagged as anchors.

atom_vocab

n_vocab

MolecularDiffusion.modules.tasks.diffusion_difflinker.DIFFLINKER_ROW_DATA_COLUMNS = ('linker_mask', 'anchors')