Geometric confidence targets and observation masks

Your task

Function signature

structure_confidence_targets(
    activations: dict,
    masks: dict,
    plddt_bins: int = 50,
    lddt_radius: float = 15.0,
    lddt_eps: float = 1e-10,
) -> dict[str, torch.Tensor]

Worked example

Distance error and aligned error answer different questions

Parameters

Return

dict[str, torch.Tensor]

Required behavior

• Construct geometric confidence targets on B\ge 1 batches and N\ge 0 supplied representative points. Coordinates are in angstroms. R [b,i] is a 3x3 local-to-global rotation whose columns are local basis vectors; t [b,i] is its origin. Frames are supplied, not reconstructed. All boolean masks use \mathrm{true}=\mathrm{valid}.
• Let e [b] be resolution_eligible: a supplied experiment/training-policy gate. o=\mathrm{exists} AND resolved AND e; v=o AND pred_mask; f=\mathrm{frame\_mask} AND e. Frame validity must already incorporate the atom availability needed to construct both predicted and target frames.
• Target coordinates outside o, predicted coordinates outside v, and frames outside f may contain arbitrary nonfinite padding and must be sanitized before arithmetic. Parameters/valid coordinates are finite. Model-specific representative-atom choices (CA, pseudo-\beta, token atoms, all atoms) are made before this function.
• For ordered pairs i,j compute D2_ij=\Vert \mathrm{target\_i}-\mathrm{target\_j}\Vert ^2; predicted distances d_ij=\Vert \mathrm{pred\_i}-\mathrm{pred\_j}\Vert ; target distances t_{\mathrm{ij}}=\sqrt{\mathrm{D2\_ij}}. Distogram mask is o_i AND o_j.
• PDE mask is v_i AND v_j, with squared distance-error \mathrm{E2\_ij}=(d_{\mathrm{ij}}-t_{\mathrm{ij}})^{2}. PAE mask is f_i AND v_j: A2_ij=\Vert \mathrm{Rpred\_i}^{\top}(\mathrm{pred\_j}-\mathrm{tpred\_i})-\mathrm{Rtarget\_i}^{\top}(\mathrm{target\_j}-\mathrm{ttarget\_i})\Vert ^2.
• Axis i chooses the alignment frame and j chooses the point. Return all three squared continuous diagnostics, zero outside their individual masks. PAE generally is not symmetric. Distogram/PDE are symmetric. Self pairs are included for distance/error bin targets.
• Quantize distances/errors using strictly increasing nonnegative angstrom edges: \mathrm{label}(\mathrm{q2},\mathrm{edges})=\sum_{k}[\mathrm{q2}>\mathrm{edges\_k}^{2}]. Equality stays in the LOWER bin (strict '>'); this differs from bucketizers using '\ge '.
• distogram_bins uses D2 with distogram_edges; pde_bins and pae_bins use \mathrm{E2}/\mathrm{A2} with error_edges. Empty edge vectors give one catch-all bin. Every invalid label is -1 and its returned mask is false.
• Construct local hard lDDT with neighbors M_{\mathrm{ij}}=v_{i} AND v_j AND i\ne j AND t_{\mathrm{ij}}<\mathrm{lddt\_radius}. A pair contributes mean_k[|d_{\mathrm{ij}}-t_{\mathrm{ij}}|<\mathrm{lddt\_thresholds\_k}].
• Then \mathrm{local\_lddt\_i}=(\mathrm{lddt\_eps}+\sum_{j} M_{\mathrm{ij}}\cdot \mathrm{pair\_score\_ij})/(\mathrm{lddt\_eps}+\sum_{j} M_{\mathrm{ij}}). Invalid v_i returns zero. A valid point with no neighbors returns one, matching the explicit \epsilon-over-\epsilon convention; it is not evidence of a reliable isolated prediction.
• \mathrm{plddt\_bins}=\operatorname{min}(\operatorname{floor}(\mathrm{local\_lddt}\cdot \mathrm{plddt\_bins}),\mathrm{plddt\_bins}-1), or -1 when invalid. Threshold comparisons are strict; no smooth approximation is used.
• Finally \mathrm{resolved\_mask}=\mathrm{exists} AND e and \mathrm{resolved\_bins}=1 for experimentally resolved,0 for unresolved existing sites,-1 outside that mask. An unresolved but chemically existing atom is a valid NEGATIVE resolution target, while absent chemistry is not a target. resolved must be a subset of exists.
• Do not confuse these masks with the coordinate-supervision masks. Config: \mathrm{plddt\_bins}=50 is a positive integer; \mathrm{lddt\_radius}=15.0 angstroms and \mathrm{lddt\_eps}=10^{-10} are positive. lddt_thresholds is a nonempty positive float vector.
• Invalid configured counts/radius/\epsilon or empty thresholds raise ValueError; other shape/domain violations are outside the contract. Empty N and wholly ineligible/missing structures return empty/zero diagnostics, -1 invalid labels and false masks.
• The diagnostics use exact Euclidean lengths, with zero defined at coincidence; gradient fixtures avoid coincident DISTINCT valid points. This does not produce a differentiable lDDT loss or backpropagate through hard labels.

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 structure_confidence_targets(
    activations: dict,
    masks: dict,
    plddt_bins: int = 50,
    lddt_radius: float = 15.0,
    lddt_eps: float = 1e-10,
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    target_distance_squared: (B, N, N).
    distance_error_squared: (B, N, N).
    aligned_error_squared: (B, N, N).
    distogram_bins: (B, N, N).
    pde_bins: (B, N, N).
    pae_bins: (B, N, N).
    local_lddt: (B, N).
    plddt_bins: (B, N).
    resolved_bins: (B, N).
    distogram_mask: (B, N, N).
    pde_mask: (B, N, N).
    pae_mask: (B, N, N).
    plddt_mask: (B, N).
    resolved_mask: (B, N)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement structure_confidence_targets")

Sources and further reading

• www.nature.com
• www.nature.com
• google-deepmind/alphafold · modules.py
• google-deepmind/alphafold · modules.py
• aqlaboratory/openfold · loss.py
• google-deepmind/alphafold3 · confidence_head.py
• huggingface/transformers · modeling_esmfold2.py