Getting PAE matrix

[gezmi]

Hi,

Nice work! Would it be possible to also return the PAE matrix for each decoy?

Thank you!

[jproney]

Sure! I think the easiest way to do this is with a slight modification of the Colab notebook that is linked in the README. There's a function in that notebook called _get_score under the "Import Libraries" header. If you replace that with this very slightly modified version:

  def _get_score(self):
    score = copy_dict(self.model.aux["log"])

    score["plddt"] = score["plddt"]
    score["pae"] = 31.0 * score["pae"]
    score["rmsd_io"] = score.pop("rmsd",None)

    i_xyz = self.model._inputs["batch"]["all_atom_positions"][:,1]
    o_xyz = np.array(self.model.aux["atom_positions"][:,1])

    # TMscore to input/output
    if hasattr(self,"wt_batch"):
      n_xyz = self.wt_batch["all_atom_positions"][:,1]
      score["tm_i"] = tmscore(n_xyz,i_xyz)["tms"]
      score["tm_o"] = tmscore(n_xyz,o_xyz)["tms"]

    # TMscore between input and output
    score["tm_io"] = tmscore(i_xyz,o_xyz)["tms"]

    # composite score
    score["composite"] = score["ptm"] * score["plddt"] * score["tm_io"]
    score["pae_full"] = np.copy(self.model.aux["pae"])
    return score

then the full PAE matrix will be part of the output. The only modification here is in the second-to-last line. Then, after running the rest of the notebook, the SCORES variable will have pae_full as one of its features, which you can use however you want. Let me know if that helps!