(self, batch, randomize=True)
| 261 | |
| 262 | @torch.no_grad() |
| 263 | def generate(self, batch, randomize=True): |
| 264 | protein = batch['graph'] |
| 265 | if randomize: |
| 266 | protein = rotamer.randomize(protein) |
| 267 | |
| 268 | best_protein = protein.clone() |
| 269 | best_rmsd = torch.zeros(protein.num_residue, device=self.device) + 1e6 |
| 270 | for _ in tqdm(range(self.num_sample), desc="Confidence sampling"): |
| 271 | batch = super().generate(batch, randomize=True) # TODO: do we need to randomize? |
| 272 | protein = batch['graph'] |
| 273 | rmsd = self.predict_rmsd(batch) |
| 274 | residue_update_mask = rmsd < best_rmsd # [num_residue] |
| 275 | atom_update_mask = residue_update_mask[protein.atom2residue] # [num_atom] |
| 276 | best_protein.node_position[atom_update_mask] = protein.node_position[atom_update_mask] |
| 277 | best_rmsd[residue_update_mask] = rmsd[residue_update_mask] |
| 278 | |
| 279 | best_batch = { |
| 280 | "graph": best_protein, |
| 281 | "rmsd": best_rmsd |
| 282 | } |
| 283 | return best_batch |
nothing calls this directly
no test coverage detected