(self, batch, all_loss=None, metric=None)
| 114 | return batch |
| 115 | |
| 116 | def predict(self, batch, all_loss=None, metric=None): |
| 117 | protein = batch['graph'] |
| 118 | chi_id = batch['chi_id'] |
| 119 | sigma = batch['sigma'] # [num_graph] |
| 120 | if self.graph_construction_model: |
| 121 | protein = self.graph_construction_model(protein) |
| 122 | |
| 123 | # Model forward |
| 124 | node_sigma = sigma[protein.atom2graph] # [num_node] |
| 125 | node_feature = self.sigma_embedding_list[chi_id](protein.node_feature.float(), node_sigma) |
| 126 | node_feature = self.model_list[chi_id](protein, node_feature, all_loss=all_loss, metric=metric)["node_feature"] |
| 127 | residue_feature = scatter_mean(node_feature, protein.atom2residue, dim=0, dim_size=protein.num_residue) |
| 128 | pred = self.torsion_mlp_list[chi_id](residue_feature) |
| 129 | |
| 130 | # Scaled by norm |
| 131 | torsion_sigma = sigma[protein.residue2graph].unsqueeze(-1).expand(-1, self.NUM_CHI_ANGLES) # [num_residue, 4] |
| 132 | score_norm_1pi = torch.tensor(self.schedule_1pi_periodic.score_norm(torsion_sigma), device=self.device) |
| 133 | score_norm_2pi = torch.tensor(self.schedule_2pi_periodic.score_norm(torsion_sigma), device=self.device) |
| 134 | score_norm = torch.where(protein.chi_1pi_periodic_mask, score_norm_1pi, score_norm_2pi) |
| 135 | pred_score = pred * score_norm.sqrt() |
| 136 | |
| 137 | # Mask out non-related chis |
| 138 | pred_score = pred_score * protein.chi_mask.to(pred_score.dtype) |
| 139 | |
| 140 | return pred_score, score_norm |
| 141 | |
| 142 | def target(self, batch): |
| 143 | protein = batch["graph"] |
no test coverage detected