Get per-residue energy during scoring

[ben-jsmith]

Thanks for the great work creating the model!

Would it be possible to get the score/energy per position (or ideally for each of the 20 amino acids at each position)? Not sure if this is information that can be extracted.

[richardshuai]

Hi Ben! Yes it’s possible, just pushed a new commit to main with an option to save out local conditional energies for score.yaml and score_ensemble.yaml by setting save_local_conditionals to true. These will be of shape [N, 32] where N is the length of the protein and 32 will be our vocabulary size. You can interpret this as the conditional energy for each possible amino acid at position i, assuming all other positions are fixed (note that lower energy = more favorable).

You can see what tokens correspond to what indices using import caliby.data.const as const; print(const.AF3_ENCODING.token_to_idx) . Note that we don’t use the non-protein tokens in the vocabulary, so you should just look at the 21 tokens given by const.AF3_ENCODING.protein_tokens . Let me know if that doesn’t work