Sequential directed-bond torsion rotations

Your task

Function signature

directed_bond_torsion_rotation(
    positions: torch.Tensor,
    directed_edges: torch.Tensor,
    rotate_mask: torch.Tensor,
    torsion_angles: torch.Tensor,
    valid_mask: torch.Tensor,
    edge_mask: torch.Tensor,
    eps: float = 1e-08,
) -> dict[str, torch.Tensor]

Worked example

The first rotation changes the next bond axis

Parameters

Return

dict[str, torch.Tensor]

Required behavior

• Inputs positions[B,N,3], directed_edges[B,K,2] int64, rotate_mask[B,K,N] bool and torsion_angles[B,K] in radians. Process bonds j=0..K-1 in the supplied order using the CURRENT updated coordinates.
• For edge(u,v), the inward positive axis is a=\frac{x_u-x_v}{\|x_u-x_v\|} and the pivot is x_v. Apply \operatorname{Exp}(\mathrm{angle\_j}\cdot a) to selected offsets x_n-x_v; leave unselected atoms unchanged. Column action is used: row-coordinate offsets multiply the rotation matrix transpose.
• The cross-product Rodrigues form is x_{\mathrm{new}}=x+\operatorname{sin}(\operatorname{angle})\cdot (a\times d)+2\cdot \operatorname{sin}(\operatorname{angle}/2)^{2}\cdot (a\times (a\times d)), d=x-x_{v}. Zero angle must retain its correct angular derivative, not skip the differentiable computation.
• valid_mask[B,N]=True means observed; edge_mask[B,K]=True enables a bond. An active bond requires in-range, distinct, valid endpoints, rotate_mask[u]=False and rotate_mask[v]=True; violations raise ValueError. Inactive bond indices may be -1 and their angles may be nonfinite; ignore them before gathering or trigonometry.
• Invalid atom coordinates may be nonfinite and are zeroed before use; rotate_mask values on invalid atoms are ignored. Other observed coordinates/angles must be finite. \mathrm{eps}>0, default 10^{-8}, is a bond-length threshold.
• An active but coincident/short bond with \mathrm{length}\le \varepsilon is skipped, its axis is zero and axis_valid=False. This defined fallback replaces the original undefined zero-length normalization. Otherwise axis_valid=True and bond_axes records that step's unit inward axis.
• Return final positions, trajectory[B,K,N,3] after each processed bond, bond_axes[B,K,3] and axis_valid[B,K]. B\ge 1; empty N and K are allowed, but N=0 requires all edges inactive. K=0 returns masked positions and empty history.
• No bond detection, graph search, rigid realignment or random draws are hidden. \varepsilon is the only key; nonpositive \varepsilon raises ValueError.

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_torsion_rotation(
    positions: torch.Tensor,  # (B, N, 3)
    directed_edges: torch.Tensor,  # (B, K, 2)
    rotate_mask: torch.Tensor,  # (B, K, N)
    torsion_angles: torch.Tensor,  # (B, K)
    valid_mask: torch.Tensor,  # (B, N)
    edge_mask: torch.Tensor,  # (B, K)
    eps: float = 1e-08,
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    positions: (B, N, 3).
    trajectory: (B, K, N, 3).
    bond_axes: (B, K, 3).
    axis_valid: (B, K)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement directed_bond_torsion_rotation")

Sources and further reading

• arxiv.org
• gcorso/DiffDock · torsion.py
• gcorso/DiffDock · diffusion_utils.py