(self, batch, randomize=True)
| 157 | |
| 158 | @torch.no_grad() |
| 159 | def generate(self, batch, randomize=True): |
| 160 | protein = batch['graph'] |
| 161 | if randomize: |
| 162 | protein = rotamer.randomize(protein) |
| 163 | |
| 164 | schedule = self.schedule_1pi_periodic.reverse_t_schedule.to(self.device) |
| 165 | for chi_id in tqdm(range(self.NUM_CHI_ANGLES), desc="Autoregressive generation"): |
| 166 | for j in range(len(schedule) - 1): |
| 167 | t = schedule[j] |
| 168 | dt = schedule[j] - schedule[j + 1] if j + 1 < len(schedule) else 1 |
| 169 | chis = rotamer.get_chis(protein, protein.node_position) # [num_residue, 4] |
| 170 | |
| 171 | # Predict score |
| 172 | sigma = self.schedule_1pi_periodic.t_to_sigma(t).repeat(protein.batch_size) |
| 173 | chi_protein = rotamer.remove_by_chi(protein, chi_id) |
| 174 | pred_score, _ = self.predict({ |
| 175 | "graph": chi_protein, |
| 176 | "sigma": sigma, |
| 177 | "chi_id": chi_id |
| 178 | }) |
| 179 | |
| 180 | # Step backward |
| 181 | chis = self.schedule_1pi_periodic.step(chis, pred_score, t, dt, chi_protein.chi_1pi_periodic_mask) |
| 182 | chis = self.schedule_2pi_periodic.step(chis, pred_score, t, dt, chi_protein.chi_2pi_periodic_mask) |
| 183 | protein = rotamer.set_chis(protein, chis) |
| 184 | return batch |
| 185 | |
| 186 | def get_metric(self, pred_protein, true_protein, metric): |
| 187 | # assert pred_pos.shape == true_pos.shape |
no test coverage detected