(self, pred_protein, true_protein, metric)
| 184 | return batch |
| 185 | |
| 186 | def get_metric(self, pred_protein, true_protein, metric): |
| 187 | # assert pred_pos.shape == true_pos.shape |
| 188 | pred_pos = pred_protein.node_position |
| 189 | true_pos = true_protein.node_position |
| 190 | protein = true_protein |
| 191 | |
| 192 | pred_pos_per_residue = torch.zeros(protein.num_residue, len(atom_name_vocab), 3, device=protein.device) |
| 193 | true_pos_per_residue = torch.zeros(protein.num_residue, len(atom_name_vocab), 3, device=protein.device) |
| 194 | pred_pos_per_residue[protein.atom2residue, protein.atom_name] = pred_pos |
| 195 | true_pos_per_residue[protein.atom2residue, protein.atom_name] = true_pos |
| 196 | symm_true_pos_per_residue = _get_symm_atoms(true_pos_per_residue, protein.residue_type) |
| 197 | |
| 198 | # Symmetric alignment |
| 199 | rmsd_per_residue = _rmsd_per_residue(pred_pos_per_residue, true_pos_per_residue, protein.sidechain37_mask) |
| 200 | sym_rmsd_per_residue = _rmsd_per_residue(pred_pos_per_residue, symm_true_pos_per_residue, |
| 201 | protein.sidechain37_mask) |
| 202 | sym_replace_mask = rmsd_per_residue > sym_rmsd_per_residue |
| 203 | rmsd_per_residue[sym_replace_mask] = sym_rmsd_per_residue[sym_replace_mask] |
| 204 | true_pos_per_residue[sym_replace_mask] = symm_true_pos_per_residue[sym_replace_mask] |
| 205 | true_pos = true_pos_per_residue[protein.atom2residue, protein.atom_name] |
| 206 | metric["atom_rmsd_per_residue"] = rmsd_per_residue |
| 207 | |
| 208 | pred_chi = rotamer.get_chis(protein, pred_pos) |
| 209 | true_chi = rotamer.get_chis(protein, true_pos) |
| 210 | chi_diff = (pred_chi - true_chi).abs() |
| 211 | chi_ae = torch.minimum(chi_diff, 2 * np.pi - chi_diff) |
| 212 | chi_ae_periodic = torch.minimum(chi_ae, np.pi - chi_ae) |
| 213 | chi_ae[protein.chi_1pi_periodic_mask] = chi_ae_periodic[protein.chi_1pi_periodic_mask] |
| 214 | metric["chi_ae_rad"] = chi_ae[protein.chi_mask] # [num_residue, 4] |
| 215 | metric["chi_ae_deg"] = chi_ae[protein.chi_mask] * 180 / np.pi # [num_residue, 4] |
| 216 | for i in range(self.NUM_CHI_ANGLES): |
| 217 | metric[f"chi_{i}_ae_rad"] = chi_ae[:, i][protein.chi_mask[:, i]] |
| 218 | metric[f"chi_{i}_ae_deg"] = chi_ae[:, i][protein.chi_mask[:, i]] * 180 / np.pi |
| 219 | |
| 220 | return metric |
| 221 | |
| 222 | |
| 223 | @R.register("tasks.ConfidencePrediction") |
no test coverage detected