Structure confidence heads

Directed Chain-Pair TM Scores

HardRaw PyTorchAbout 30 min3 public examples plus private tests

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

Your task

For each ordered chain pair, score points in the column chain using candidate alignment frames in the row chain. Select the best row under the declared weighting rule and return its TM score and support flag in chain_labels order.

Function signature

directed_chain_pair_tm_scores(
    pae_probability: torch.Tensor,
    activations: dict,
    masks: dict,
    structure: dict,
    tm_length: str = 'mask',
    count_eps: float = 1e-08,
    tm_interface_length: str = 'global',
    tm_weighted_selection: bool = True,
) -> tuple[torch.Tensor, torch.Tensor]

Worked example

Two chains can have different directional evidence

Use two valid tokens in chains [10,20], chain_labels=[10,20], unit token weights, all pairs allowed, tm_length="mask", tm_interface_length="global", and count_eps=0.01. Centers $[0,d_0]$ give TM kernels [1,1/2], with

d_0=1.24\sqrt[3]{19-15}-1.8.

Give diagonal pairs probabilities [1/2,1/2], pair (0,1) probabilities [1/4,3/4], and (1,0) probabilities [3/4,1/4]. The directed chain matrix is

\frac{1}{1.01}\begin{bmatrix}0.75&0.625\\0.875&0.75\end{bmatrix}\approx\begin{bmatrix}0.742574&0.618812\\0.866337&0.742574\end{bmatrix}.

All entries have support. In bidirectional mode, each off-diagonal entry considers both row choices and selects $0.875/1.01\approx0.866337$ ; diagonal entries remain unchanged.

Parameters

Parameter	Shape and meaning
`pae_probability`torch.Tensor	$[B,N,N,L]$ Supplied PAE probabilities. B is batch size, N token count, and L error bins.
`activations`dict	Supplied intermediate tensors.
`masks`dict	Supplied masks. Each field below defines what True includes or permits.
`structure`dict	Fields supplied together in this dictionary.
`tm_length`strDefault `'mask'`	"mask" uses the valid-token count; "weight_floor" uses the sum of valid token weights rounded down for the TM length scale.
`count_eps`floatDefault `1e-08`	Positive offset in the declared TM/confidence weighted-count denominators, added after summing the relevant support.
`tm_interface_length`strDefault `'global'`	"global" uses the overall TM length scale for interfaces; "chain_pair" derives a separate scale from the two participating chains.
`tm_weighted_selection`boolDefault `True`	Use residue weights when selecting the best aligned TM row.

Return

tuple[torch.Tensor, torch.Tensor]

Return these tensors in the stated order.

Required behavior

pae_probability is $[B,N,N,L]$ , with nonnegative bins summing to 1 on valid token pairs; invalid pair values are ignored. pae_centers supplies L nonnegative aligned-error bin centers. weights is a nonnegative importance weight per token; zero invalid token weights.
For a length $n$ , use $d_0=1.24(\max(n,19)-15)^{1/3}-1.8$ and $K_l=d_0^2/(d_0^2+center_l^2)$ . A pair term is the expected K under that pair’s PAE probabilities.
tm_length="mask"(default) uses valid-token count; "weight_floor" uses floor(sum of valid weights). tm_pair_mask additionally restricts scored ordered pairs.
For row i, average pair terms with weights $w_iw_j$ over allowed j, dividing by their sum plus count_eps (default $10^{-8}$ ). Empty rows score zero. tm_weighted_selection=True (default) selects the row maximizing row_score × w_i; False selects the maximum row_score. Return the UNWEIGHTED score of the selected row; ties choose the first row.
The validity flag is True iff at least one allowed pair has positive weights at both endpoints. No support returns zero scores and False, without inventing a positive denominator or fallback.
tm_interface_length="global"(default) uses the global d0 above. "chain_pair" instead uses the sum of the VALID token counts of the endpoint chains (only one chain count for a same-chain pair), before the same d0 clamp. This local length uses token counts even when tm_length="weight_floor".
Return a matrix in the supplied chain_labels order. Each(left,right) entry restricts row endpoints to the left chain and column endpoints to the right chain. Same-chain diagonal entries are allowed. Apply the same row-selection rule and report support independently per matrix entry.
The output axis convention is frame chain then point chain. Transpose this matrix when comparing with the pinned ESMFold2 per-chain reporting convention.

Why this matters

Different heads answer different questions: pLDDT describes local accuracy, a distogram describes distance, and PAE describes error after alignment in a particular residue’s frame. Their probabilities feed separate summaries; they are not interchangeable confidence scales.

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_chain_pair_tm_scores(
    pae_probability: torch.Tensor,  # (B,N,N,L)
    activations: dict,
    masks: dict,
    structure: dict,
    tm_length: str = 'mask',
    count_eps: float = 1e-08,
    tm_interface_length: str = 'global',
    tm_weighted_selection: bool = True,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Return tuple[torch.Tensor, torch.Tensor]:
    
    chain_pair_tm: (B, C, C).
    chain_pair_valid: (B, C, C)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement directed_chain_pair_tm_scores")

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.

Related exercises

Sources and further reading

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