AlphaJudge PAE handling for AF3x outputs (cross-linker tokens) #30

I noticed an issue when running AlphaJudge (1.4.1; the code is identical in 1.4.2) on AF3x predictions. AF3x models each cross-link as a cross-linker ligand, so the pae matrix in confidences.json has more tokens than there are protein residues (e.g. 257 tokens for 227 residues). In AF3Parser._normalize_pae_af3, when the token matrix does not match the residue count, the code falls back to filling each chain-pair block with the minimum PAE of that block. As a result the per-residue PAE is lost and every interface residue pair gets the best value in the interface, which inflates ipSAE, LIS and pDockQ2 and lowers the average interface PAE. The contact_probs matrix does not have this problem, because it is already aligned via _align_token_pair_matrix_to_residues.

My local fix is to call the same _align_token_pair_matrix_to_residues function for the PAE when the shapes differ, before the chain-pair minimum fallback. This keeps only the residue tokens, matched by chain and residue number. Outputs where tokens and residues already match (standard AF3) are unaffected. I tested it on an AF3 run and on a copy of it with cross-linker tokens and atoms added. With the patch, all AlphaJudge scores for the modified copy were identical to the original. Without it, the average interface PAE dropped from about 28–29 Å to 19–27 Å and pDockQ2 increased. The 16-line patch is attached.

A related point: iptm_ptm and confidence_score are global in AF3x and include the cross-linker chains, so they are not directly comparable to plain AF3 runs. The per-interface iptm from chain_pair_iptm is fine.

alphajudge_af3x_pae.patch