MolecularDiffusion.modules.models.syncogen.api.ops.graph_ops¶

Pure tensor operations for graph representation conversions.

All conversion functions: - Accept (N, …) for single graphs or (B, N, …) for batched - Return tensors of matching batch structure - Are stateless and side-effect free

Attributes¶

`MASK_CHANNEL`
`NO_EDGE_CHANNEL`

Functions¶

`apply_edge_givens`(→ torch.Tensor)	Apply edge invariants: diagonals are NO-EDGE.
`bb_indices_to_onehot`(→ torch.Tensor)	Convert BB indices to one-hot.
`bb_onehot_to_indices`(→ torch.Tensor)	Convert BB one-hot to indices.
`compute_compatibility_masks`(bb_onehot, rxn_onehot, ...)	Compute compatibility masks for nodes and edges given selected reactions and BBs.
`create_masked_graph`(→ Tuple[torch.Tensor, ...)	Create masked graph one-hots and a node mask for padding.
`rxn_indices_to_onehot`(→ torch.Tensor)	Convert reaction flat indices to one-hot.
`rxn_onehot_to_indices`(→ torch.Tensor)	Convert reaction one-hot to flat indices.
`rxn_onehot_to_tuple`(→ Union[torch.Tensor, ...)	Convert reaction one-hot to tuple format.
`rxn_tuple_to_onehot`(→ torch.Tensor)	Convert reaction tuples to one-hot matrices.

Module Contents¶

MolecularDiffusion.modules.models.syncogen.api.ops.graph_ops.apply_edge_givens(rxn_onehot: torch.Tensor, node_mask: torch.Tensor) → torch.Tensor¶

Apply edge invariants: diagonals are NO-EDGE.

Parameters:

rxn_onehot – (N, N, D) or (B, N, N, D) reaction one-hot tensor
node_mask – (N,) or (B, N) boolean/float mask where 1 = valid node (unused, kept for API)

Returns:

rxn_onehot with diagonals set to NO-EDGE (modified in-place and returned)

MolecularDiffusion.modules.models.syncogen.api.ops.graph_ops.bb_indices_to_onehot(bb_indices: torch.Tensor, vocab_size: int, pad: int = 1) → torch.Tensor¶

Convert BB indices to one-hot.

Parameters:

bb_indices – (N,) or (B, N) tensor of indices
vocab_size – Number of building blocks in vocabulary
pad – Number of padding dimensions (e.g., 1 for MASK token)

Returns:

(N, D) or (B, N, D) one-hot tensor where D = vocab_size + pad

MolecularDiffusion.modules.models.syncogen.api.ops.graph_ops.bb_onehot_to_indices(bb_onehot: torch.Tensor) → torch.Tensor¶

Convert BB one-hot to indices.

Parameters:: bb_onehot – (N, D) or (B, N, D)
Returns:: (N,) or (B, N) tensor of indices

MolecularDiffusion.modules.models.syncogen.api.ops.graph_ops.compute_compatibility_masks(bb_onehot: torch.Tensor, rxn_onehot: torch.Tensor, compatibility: torch.Tensor, no_edge_channel: int = -2)¶

Compute compatibility masks for nodes and edges given selected reactions and BBs.

Inputs are one-hot tensors; outputs are boolean masks of the same shapes.

Parameters:

bb_onehot – (N, D_bb) or (B, N, D_bb)
rxn_onehot – (N, N, D_rxn) or (B, N, N, D_rxn)
compatibility – (N_bbs, N_rxns, N_centers) tensor with values in {0,1,2,3} 0 = incompatible 1 = compatible as reactant 1 2 = compatible as reactant 2 3 = compatible as both
no_edge_channel – index of the NO-EDGE channel in rxn_onehot

Returns:

bool tensor same shape as bb_onehot compatibility_mask_rxn: bool tensor same shape as rxn_onehot

Return type:

compatibility_mask_bb

MolecularDiffusion.modules.models.syncogen.api.ops.graph_ops.create_masked_graph(max_nodes: int, vocab_num_bbs: int, vocab_num_rxns: int, vocab_num_centers: int, bb_pad: int = 1, rxn_pad: int = 2, batch_size: int | None = None, n_nodes: int | List[int] | None = None, device: torch.device | None = None) → Tuple[torch.Tensor, torch.Tensor, torch.Tensor]¶

Create masked graph one-hots and a node mask for padding.

Convention: - ALL BBs (valid and padding): one-hot to MASK channel - ALL edges (valid, padding, and diagonals): one-hot to MASK channel

Diagonals get set to NO-EDGE after the first sampling step via apply_edge_givens.

Parameters:

max_nodes – Maximum number of nodes (padding length)
vocab_num_bbs – Building block vocabulary size
vocab_num_rxns – Reaction vocabulary size
vocab_num_centers – Number of reaction centers
bb_pad – Extra dims for BB features (e.g., MASK token)
rxn_pad – Extra dims for reaction features (e.g., NO-EDGE, MASK)
batch_size – Optional batch size for batched initialization
n_nodes – Actual number of nodes per sample (int for single, list for batched)
device – Target device for tensors

Returns:

(N, D_bb) or (B, N, D_bb) with MASK tokens for ALL nodes rxn_onehot: (N, N, D_rxn) or (B, N, N, D_rxn) with MASK for ALL edges (including padding) node_mask: (N,) or (B, N) boolean mask indicating valid nodes

Return type:

bb_onehot

MolecularDiffusion.modules.models.syncogen.api.ops.graph_ops.rxn_indices_to_onehot(rxn_indices: torch.Tensor, vocab_num_rxns: int, vocab_num_centers: int, rxn_pad: int = 2) → torch.Tensor¶

Convert reaction flat indices to one-hot.

Parameters:

rxn_indices – (N, N) or (B, N, N) tensor of flat indices
vocab_num_rxns – Number of reactions in vocabulary
vocab_num_centers – Number of reaction centers
rxn_pad – Number of padding dimensions (default 2 for NO-EDGE and PAD)

Returns:

(N, N, D) or (B, N, N, D) one-hot tensor

MolecularDiffusion.modules.models.syncogen.api.ops.graph_ops.rxn_onehot_to_indices(rxn_onehot: torch.Tensor) → torch.Tensor¶

Convert reaction one-hot to flat indices.

Parameters:: rxn_onehot – (N, N, D) or (B, N, N, D)
Returns:: (N, N) or (B, N, N) tensor of flat indices

MolecularDiffusion.modules.models.syncogen.api.ops.graph_ops.rxn_onehot_to_tuple(rxn_onehot: torch.Tensor, n_centers: int) → torch.Tensor | List[torch.Tensor]¶

Convert reaction one-hot to tuple format.

Parameters:

rxn_onehot – (N, N, D) or (B, N, N, D)
n_centers – Number of reaction centers

Returns:

(E, 5) tensor with [rxn_id, node1, node2, center1, center2] Batched: List of (E_i, 5) tensors

Return type:

Single

MolecularDiffusion.modules.models.syncogen.api.ops.graph_ops.rxn_tuple_to_onehot(rxn_tuple: torch.Tensor | List[torch.Tensor], n_nodes: int | List[int] | torch.Tensor, vocab_num_rxns: int, vocab_num_centers: int, rxn_pad: int = 2) → torch.Tensor¶

Convert reaction tuples to one-hot matrices.

Parameters:

rxn_tuple – (E, 5) for single or List[(E_i, 5)] for batched
n_nodes – Number of nodes (int for single, list/tensor for batched)
vocab_num_rxns – Number of reactions in vocabulary
vocab_num_centers – Number of reaction centers
rxn_pad – Number of padding dimensions

Returns:

(N, N, D) or (B, N, N, D) one-hot tensor

MolecularDiffusion.modules.models.syncogen.api.ops.graph_ops.MASK_CHANNEL = -1¶

MolecularDiffusion.modules.models.syncogen.api.ops.graph_ops.NO_EDGE_CHANNEL = -2¶