Add noise to the input tensor (torsion angles). Args: x (Tensor): Torsion angles of shape :math:`(num_res, 4)` x_mask (Tensor): Mask of shape :math:`(num_res, 4)` t (Tensor): Timesteps of shape :math:`(num_res)`
(self, x, t, x_mask=None)
| 181 | |
| 182 | @torch.no_grad() |
| 183 | def add_noise(self, x, t, x_mask=None): |
| 184 | """Add noise to the input tensor (torsion angles). |
| 185 | |
| 186 | Args: |
| 187 | x (Tensor): Torsion angles of shape :math:`(num_res, 4)` |
| 188 | x_mask (Tensor): Mask of shape :math:`(num_res, 4)` |
| 189 | t (Tensor): Timesteps of shape :math:`(num_res)` |
| 190 | """ |
| 191 | sigmas = self.t_to_sigma(t) |
| 192 | noise = torch.randn_like(x) * sigmas.unsqueeze(-1) |
| 193 | score = torch.tensor( |
| 194 | self.score(noise.cpu().numpy(), sigmas.cpu().numpy()), device=x.device, dtype=x.dtype |
| 195 | ) |
| 196 | |
| 197 | if x_mask is not None: |
| 198 | noise *= x_mask |
| 199 | score *= x_mask |
| 200 | |
| 201 | x = x + noise |
| 202 | return x, score |
| 203 | |
| 204 | @torch.no_grad() |
| 205 | def step(self, x, x_score, t, dt, x_mask=None): |
nothing calls this directly
no test coverage detected