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

Functions

dense_edge_mask(→ torch.Tensor)

Build the flat (B*N*N, 1) edge mask for the dense path.

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, pocket_conditioned: bool = False, context_node_nf: int | None = None, graph_type: str = 'FC', upstream_int8_edge_mask: bool = False)

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.

outpaint_cfgs is read for two keys: anchor_indices – a list of 1-indexed positions into the reference fragment file, matching upstream’s –anchors CLI convention (others/difflinker/generate.py:129-132). Required when this checkpoint’s center_of_mass == “anchors” (see INTEGRATION_PLAN.md Revision 5, Point 1/2); unused when center_of_mass == “fragments”. max_generate_attempts – a bounded full-restart retry count around self.edm.sample_chain(…), catching FoundNaNException (INTEGRATION_PLAN.md Revision 6, Point 1). Mirrors upstream’s own generate.py retry loop (fixed budget, full restart, no partial/ t_retry-style resume – EDM.sample_chain draws fresh noise internally every call). Defaults to 1 (today’s single-shot behavior, unchanged for any config that doesn’t set it). This is independent of the still-inert n_retrys/t_retry kwargs below, which the shared structural_guidance dispatch always forces to 0/None – do not confuse the two.

When self.pocket_conditioned (the three genuinely pocket- conditioned Pockets checkpoints, INTEGRATION_PLAN.md Revision 9), outpaint_cfgs additionally requires pocket_db – a path to a single-row-per-pocket ASE db (see convert_pocket_pt_to_asedb.py) – and accepts pocket_index (0-indexed, default 0). The pocket is inserted as a fixed, non-diffused block between the fragment and linker regions; nodesxsample keeps its existing meaning (fragment + linker atom count, NOT including the pocket).

When self.size_predictor is set (an optional LinkerSizePredictor, INTEGRATION_PLAN.md Revision 8), the incoming nodesxsample is overridden with a fragment-geometry-conditioned prediction instead of whatever GenerativeFactory sampled from node_dist_model.

Several other kwargs accepted here are explicitly out of scope this pass (inert, not implemented with real behavior): batch, num_steps, condition_mode, 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

context_node_nf = None

property device: torch.device

edm

graph_type = 'FC'

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

pocket_conditioned = False

prop_dist_model = None

size_predictor: MolecularDiffusion.modules.models.difflinker.size_gnn.LinkerSizePredictor | None = None

upstream_int8_edge_mask = False

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, pocket_conditioned: bool = False, context_node_nf: int | None = None, graph_type: str = 'FC', upstream_int8_edge_mask: bool = False, size_gnn_checkpoint: str | 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

context_node_nf = None

diffusion_loss_type = 'l2'

diffusion_noise_precision = 1e-05

diffusion_noise_schedule = 'polynomial_2'

diffusion_steps = 500

graph_type = 'FC'

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)

pocket_conditioned = False

sin_embedding = False

size_gnn_checkpoint = None

tanh = False

task_type

upstream_int8_edge_mask = False

class MolecularDiffusion.modules.tasks.diffusion_difflinker.PointCloudToDiffLinkerBatch(atom_vocab: list, *, upstream_int8_edge_mask: bool = False)

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

upstream_int8_edge_mask = False

MolecularDiffusion.modules.tasks.diffusion_difflinker.dense_edge_mask(atom_mask: torch.Tensor, *, upstream_int8: bool = False) → torch.Tensor

Build the flat (B*N*N, 1) edge mask for the dense path.

upstream_int8 reproduces a real integer-overflow bug in DiffLinker’s own src/datasets.py:366-369, which every released non-pocket weight was trained under:

edge_mask = atom_mask[:, None, :] * atom_mask[:, :, None]
diag_mask = ~torch.eye(n, dtype=const.TORCH_INT)   # TORCH_INT = int8
edge_mask *= diag_mask

~ on an int8 tensor is bitwise NOT, not logical NOT, so diag_mask is -1 off-diagonal and -2 on-diagonal rather than True/False. The resulting mask therefore negates every real message and passes self-loops through at -2. That sign convention is baked into the released weights: fed the arithmetically-correct +1/0 mask they diverge to NaN on the first reverse step, and fed this one they generate cleanly. Weights trained by this platform learnt the correct mask, so this defaults to False and only the converted upstream checkpoints turn it on.

The pocket checkpoints are unaffected either way: DynamicsWithPockets builds its own distance-cutoff edges and is called with edge_mask=None.

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