Protein sequence and inverse folding

Protein backbone edge geometry and 25 RBF blocks

MediumRaw PyTorchAbout 20 min1 public example plus private tests

Free account with email, Google or GitHub. Your code runs in an isolated PyTorch sandbox.

Your task

Describe each supplied residue-neighbor edge using distances between backbone atoms and a constructed virtual Cβ. Each distances[b,i,k,p] is the stabilized distance for atom-pair type p from residue i to its k-th neighbor; each such distance expands into the supplied radial channels. Return virtual Cβ positions, distances, radial features, and effective edge validity.

Function signature

protein_backbone_edge_features(
    backbone: torch.Tensor,
    neighbors: torch.Tensor,
    node_mask: torch.Tensor,
    edge_mask: torch.Tensor,
    centers: torch.Tensor,
    widths: torch.Tensor,
    eps: float = 1e-06,
) -> dict[str, torch.Tensor]

Worked example

Construct Cβ, then encode a CA distance

For one valid residue let N=[-1,0,0], CA=[0,0,0], C=[0,1,0], and O=[0,2,0]. Then b=CA−N=[1,0,0], c=C−CA=[0,1,0], and b×c=[0,0,1], so virtual Cβ is [0.56802827,-0.54067466,-0.58273431]. For a valid neighbor whose CA is [3,0,0], the first distance channel (CA,CA), using illustrative eps=7, is $\sqrt{9+7}=4$ . With one radial center 3 and width 1, its first radial value is $\exp(-(4-3)^2)=e^{-1}\approx0.367879$ . Other pair channels follow the ordered list below.

Parameters

Parameter	Shape and meaning
`backbone`torch.Tensor	$[B, L, 4, 3]$ In $N,\mathrm{CA},C,O$ order.
`neighbors`torch.Tensor	$[B, L, K]$ Integer [ $B,L,K$ ], in-range indices.
`node_mask`torch.Tensor	$[B, L]$ True includes a residue and permits its virtual Cβ.
`edge_mask`torch.Tensor	$[B, L, K]$ True permits the supplied neighbor slot; both endpoint residues must also be valid.
`centers`torch.Tensor	$[R]$ Supplied centers of the radial basis channels.
`widths`torch.Tensor	$[R]$ Supplied positive widths of the radial basis channels.
`eps`floatDefault `1e-06`	Positive value added to squared atom-pair distance before sqrt; the radial basis uses the resulting distance.

Return

dict[str, torch.Tensor]

Return exactly the named fields applicable to this operation.

Required behavior

backbone[ $B,L$ , $4$ , $3$ ] has atom order $N,\mathrm{CA},C,O$ ; $B,L \ge 1$ . neighbors[ $B,L,K$ ] contains source residue $j$ for target $i$ , all indices in $[0,L)$ , $K\ge 0$ . node_mask[ $B,L$ ] and edge_mask[ $B,L,K$ ] are bool, $\mathrm{True}=\mathrm{valid}/\mathrm{enabled}$ .
Floating centers $[R]$ and widths $[R]$ have $R\ge 1$ and $\mathrm{widths}>0$ ; all floating inputs share dtype/device and are finite. Config $\varepsilon>0$ defaults to $10^{-6}$ and enters squared distance before sqrt; no other keys.
For each valid residue, $b=\mathrm{CA}-N$ , $c=C-\mathrm{CA}$ , $a=b$ cross $c$ , $\mathrm{virtual\_CB}=\mathrm{CA}-0.58273431\cdot a+0.56802827\cdot b-0.54067466\cdot c$ . Invalid residues return $\mathrm{virtual\_CB}=0$ .
For each neighbor j of residue i, use the following exact ordered atom pairs (left from i, right from j): (CA,CA), (N,N), (C,C), (O,O), (CB,CB), (CA,N), (CA,C), (CA,O), (CA,CB), (N,C), (N,O), (N,CB), (CB,C), (CB,O), (O,C), (N,CA), (C,CA), (O,CA), (CB,CA), (C,N), (O,N), (CB,N), (C,CB), (O,CB), (C,O).
Each $d=\sqrt{\sum_{\mathrm{xyz}}((\mathrm{atom\_i}-\mathrm{atom\_j})^{2})+\varepsilon}$ . Self edges and duplicate slots are allowed. Effective edge_mask is supplied edge_mask AND both endpoint node masks.
For pair $p$ and basis $r$ , $\mathrm{RBF}[p,r]=\operatorname{exp}(-((d[p]-\mathrm{centers}[r])/\mathrm{widths}[r])^{2})$ : no factor $1/2$ and no cosine cutoff. Return distances[ $B,L,K$ , $25$ ], rbf[ $B,L,K$ , $25\cdot R$ ] flattened pair-major then center-minor, virtual_cb[ $B,L$ , $3$ ], and effective edge_mask.
Both floating edge outputs are exactly zero on invalid slots. $K=0$ returns empty edge tensors, while virtual_CB still exists. Centers/widths are supplied rather than silently fixed to the source defaults ( $2$ . $.22$ and $20/R$ ).
Translation/proper-rotation covariance of CB and invariance of distances/RBFs are required on valid data; mirror invariance is not promised because the virtual $-C-\beta$ cross product is oriented.

Why this matters

A residue’s backbone determines a virtual C-beta position. Distances among five atom types then describe how neighboring residues are arranged, while remaining unchanged when the whole protein moves rigidly.

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 protein_backbone_edge_features(
    backbone: torch.Tensor,  # (B, L, 4, 3)
    neighbors: torch.Tensor,  # (B, L, K)
    node_mask: torch.Tensor,  # (B, L)
    edge_mask: torch.Tensor,  # (B, L, K)
    centers: torch.Tensor,  # (R,)
    widths: torch.Tensor,  # (R,)
    eps: float = 1e-06,
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    virtual_cb: (B, L, 3).
    distances: (B, L, K, 25).
    rbf: (B, L, K, 25* R).
    edge_mask: (B, L, K)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement protein_backbone_edge_features")

Ready to implement it?

3 progressive hints, a solution walkthrough, and grading that checks values, shapes and gradients are waiting in the workspace.

Write and run your solution Sign in

Free account with email, Google or GitHub. Your code runs in an isolated PyTorch sandbox.

Sources and further reading

A primitive exercise, not an end-to-end reproduction of the named models. No affiliation or endorsement is implied.