Directed molecular bonds and reverse-edge bookkeeping

Your task

Function signature

directed_bond_graph(
    atom_features: torch.Tensor,
    bond_features: torch.Tensor,
    bonds: torch.Tensor,
    atom_mask: torch.Tensor,
    bond_mask: torch.Tensor,
) -> dict[str, torch.Tensor]

Worked example

Chain and isolated atom

import torch

atom_features = torch.tensor([[0.6163984432976174, -0.11737156230437856], [-0.871515752829823, 0.22737251091226712], [-0.43380893696960837, -0.5594381581483507], [0.16133873905171972, 0.33521053319906396]], dtype=torch.float64, requires_grad=True)
bond_features = torch.tensor([[-0.2877030313877437, -0.16133740997286983], [-0.23865414504605093, 0.07281459402452221]], dtype=torch.float64, requires_grad=True)
bonds = torch.tensor([[0, 1], [1, 2]], dtype=torch.int64)
atom_mask = torch.tensor([True, True, True, True], dtype=torch.bool)
bond_mask = torch.tensor([True, True], dtype=torch.bool)

result = directed_bond_graph(atom_features=atom_features, bond_features=bond_features, bonds=bonds, atom_mask=atom_mask, bond_mask=bond_mask)

Expected return

{
    'edge_index': torch.tensor([[0, 1, 1, 2], [1, 0, 2, 1]], dtype=torch.int64),
    'reverse_index': torch.tensor([1, 0, 3, 2], dtype=torch.int64),
    'edge_mask': torch.tensor([True, True, True, True], dtype=torch.bool),
    'edge_features': torch.tensor([[-0.2877030313877437, -0.16133740997286983], [-0.2877030313877437, -0.16133740997286983], [-0.23865414504605093, 0.07281459402452221], [-0.23865414504605093, 0.07281459402452221]], dtype=torch.float64),
    'edge_inputs': torch.tensor([[0.6163984432976174, -0.11737156230437856, -0.2877030313877437, -0.16133740997286983], [-0.871515752829823, 0.22737251091226712, -0.2877030313877437, -0.16133740997286983], [-0.871515752829823, 0.22737251091226712, -0.23865414504605093, 0.07281459402452221], [-0.43380893696960837, -0.5594381581483507, -0.23865414504605093, 0.07281459402452221]], dtype=torch.float64),
    'in_degree': torch.tensor([1, 2, 1, 0], dtype=torch.int64),
}

Parameters

Return

dict[str, torch.Tensor]

Required behavior

• Construct a fixed-capacity directed graph from prepared descriptors and an undirected bond list. Inputs are atom_features X [N,A], bond_features F [M,B] (\mathrm{same} \mathrm{CPU} \mathrm{float32}/\mathrm{float64}, \mathrm{finite}), bonds [M,2] int64, atom_mask [N] bool, and bond_mask [M] bool.
• A,B \ge 1; N,M may be zero. Every bond with bond_mask[m]=True must list two distinct indices u_m,v_m in [0,N); otherwise raise ValueError. False bond-mask rows may contain arbitrary integer endpoints. An atom mask of True means real/valid.
• Effective bond validity is q_{m} = \mathrm{bond\_mask}[m] AND atom_mask [u_{m}] AND atom_mask [v_{m}]. Thus a listed bond touching a padded atom is ignored.
• Parallel bond slots are permitted and preserve separate identities; they are a graph-indexing stress test, not a claim that chemical bond orders are represented as parallel edges. No sorting or deduplication occurs. Retain D=2M slots in input order: edge 2 m is u_{m} \to v_{m} and edge 2m+1 is v_{m} \to u_{m}.
• Return edge_index [2,D], whose first row is source and second row destination; invalid columns are both zero. Return reverse_index [D] with r(2m)=2m+1 and r(2m+1)=2m, including invalid slots. Return edge_mask [D] with each q_m repeated twice.
• An invalid index zero is only a sentinel: when N=0 it does not denote a real atom.
• For each valid directed edge e arising from m, \mathrm{edge\_features}[e]=F[m] and \mathrm{edge\_inputs}[e]=\operatorname{concat}(X[\mathrm{source}(e)], F[m]) along the last axis. Both floating rows are exactly zero for invalid edges.
• Return \mathrm{in\_degree}[v] = \mathrm{number} of valid directed edges whose destination is v, as int64 [N]. Isolated and padded atoms have zero degree. When M=0, outputs have their stated empty edge dimensions; N=0 also permits M>0 provided every bond_mask entry is False.
• Empty/all-masked floating outputs retain graph-connected zero derivatives to both descriptor tensors. All finite masked feature payloads are ignored. There are no configuration keys. Chemical parsing, aromaticity/valence rules, and categorical descriptor encoding are not performed here; their prepared outputs are supplied inputs.

Why this matters

Starter code

You write the body of this function in the BioTorch workspace. Your code runs in an isolated PyTorch sandbox; Run checks the public examples and Submit grades the full test suite, including gradients.

import torch


def directed_bond_graph(
    atom_features: torch.Tensor,  # (N, A)
    bond_features: torch.Tensor,  # (M, B)
    bonds: torch.Tensor,  # (M, 2)
    atom_mask: torch.Tensor,  # (N,)
    bond_mask: torch.Tensor,  # (M,)
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    edge_index: (2, 2M).
    reverse_index: (2M).
    edge_mask: (2M).
    edge_features: (2M, B).
    edge_inputs: (2M, A+B).
    in_degree: (N,)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement directed_bond_graph")

Sources and further reading

• doi.org
• chemprop/chemprop · molecule.py
• chemprop/chemprop · mixins.py