Interface TM Score

Your task

Function signature

interface_tm_score(
    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,
) -> dict[str, torch.Tensor]

Worked example

Two directed interface scores

import torch

pae_probability = torch.tensor([[[[0.5, 0.5], [0.25, 0.75]], [[0.75, 0.25], [0.5, 0.5]]]], dtype=torch.float64, requires_grad=True)
activations = {
    'pae_centers': torch.tensor([0.0, 0.1683773044405672], dtype=torch.float64, requires_grad=True),
    'weights': torch.tensor([[1.0, 1.0]], dtype=torch.float64, requires_grad=True),
}
masks = {
    'token_mask': torch.tensor([[True, True]], dtype=torch.bool),
    'tm_pair_mask': torch.tensor([[[True, True], [True, True]]], dtype=torch.bool),
}
structure = {
    'chain_ids': torch.tensor([[0, 1]], dtype=torch.int64),
}
tm_length = 'mask'
count_eps = 1e-08
tm_interface_length = 'global'
tm_weighted_selection = True

result = interface_tm_score(pae_probability=pae_probability, activations=activations, masks=masks, structure=structure, tm_length=tm_length, count_eps=count_eps, tm_interface_length=tm_interface_length, tm_weighted_selection=tm_weighted_selection)

Expected return

{
    'iptm_per_row': torch.tensor([[0.6249999937500001, 0.8749999912500002]], dtype=torch.float64),
    'iptm': torch.tensor([0.8749999912500002], dtype=torch.float64),
    'iptm_valid': torch.tensor([True], dtype=torch.bool),
}

Parameters

Return

dict[str, torch.Tensor]

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".
• Only different-chain pairs may contribute to ipTM; use chain_ids to exclude same-chain pairs.
• Use CPU float32/float64 tensors with one shared floating dtype/device. Preserve strided layouts and do not mutate supplied tensors.
• Use raw tensor operations, not a ready-made implementation of the requested primitive. Supplied parameters and random variates must not be reinitialized.

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 interface_tm_score(
    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,
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    iptm_per_row: (B, N).
    iptm: (B,).
    iptm_valid: (B,)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement interface_tm_score")

Related exercises

• Builds onPredicted TM Score
• Combined exerciseLearned structural confidence heads and model-aware ranking

Sources and further reading

• www.nature.com
• www.nature.com
• google-deepmind/alphafold · modules.py
• google-deepmind/alphafold · confidence.py
• facebookresearch/esm · esmfold.py
• facebookresearch/esm · esmfold.py
• google-deepmind/alphafold3 · confidence_head.py
• google-deepmind/alphafold3 · confidences.py
• google-deepmind/alphafold3 · confidences.py
• google-deepmind/alphafold3 · confidences.py
• huggingface/transformers · modeling_esmfold2.py
• Biohub/esm · vqvae.py
• Biohub/esm · vqvae.py
• RosettaCommons/RFdiffusion · AuxiliaryPredictor.py
• RosettaCommons/RFdiffusion · RoseTTAFoldModel.py