Add noise to protein and update protein Args: batch (dict): batch t (torch.Tensor): [num_graph] random number in [0, 1] Returns: batch (dict): dict with the following attributes: Protein: chi_1pi_periodic_mask
(self, batch, t, chi_id=None)
| 80 | return all_loss, metric |
| 81 | |
| 82 | def add_noise(self, batch, t, chi_id=None): |
| 83 | """Add noise to protein and update protein |
| 84 | |
| 85 | Args: |
| 86 | batch (dict): batch |
| 87 | t (torch.Tensor): [num_graph] random number in [0, 1] |
| 88 | |
| 89 | Returns: |
| 90 | batch (dict): dict with the following attributes: |
| 91 | Protein: |
| 92 | chi_1pi_periodic_mask (torch.Tensor): [num_residue, 4] bool |
| 93 | chi_2pi_periodic_mask (torch.Tensor): [num_residue, 4] bool |
| 94 | chi_mask (torch.Tensor): [num_residue, 4] bool |
| 95 | chi_id (int): chi angle to be trained |
| 96 | sigma (torch.Tensor): [num_graph] sigma |
| 97 | score (torch.Tensor): [num_residue, 4] score |
| 98 | """ |
| 99 | protein = batch['graph'] |
| 100 | if chi_id is not None: |
| 101 | protein = rotamer.remove_by_chi(protein, chi_id) |
| 102 | chis = rotamer.get_chis(protein, protein.node_position) # [num_residue, 4] |
| 103 | |
| 104 | # Add noise to chis |
| 105 | chis, score_1pi = self.schedule_1pi_periodic.add_noise(chis, t, protein.chi_1pi_periodic_mask) |
| 106 | chis, score_2pi = self.schedule_2pi_periodic.add_noise(chis, t, protein.chi_2pi_periodic_mask) |
| 107 | score = torch.where(protein.chi_1pi_periodic_mask, score_1pi, score_2pi) |
| 108 | protein = rotamer.set_chis(protein, chis) # TODO：maybe have bug |
| 109 | |
| 110 | batch['protein'] = protein |
| 111 | batch['chi_id'] = chi_id |
| 112 | batch['sigma'] = self.schedule_1pi_periodic.t_to_sigma(t) |
| 113 | batch['score'] = score |
| 114 | return batch |
| 115 | |
| 116 | def predict(self, batch, all_loss=None, metric=None): |
| 117 | protein = batch['graph'] |