Method step_correct

diffpack/schedule.py:244–279

Correct step for the input tensor (torsion angles). Args: x (Tensor): Torsion angles of shape :math:`(num_res, 4)` x_score (Tensor): Score of shape :math:`(num_res, 4)` x_batch (Tensor): Batch of shape :math:`(num_res)` x_mask (Tensor): Mask o

(self, x, x_score, x_batch, x_mask=None, snr=0.16)

    @torch.no_grad()
    def step_correct(self, x, x_score, x_batch, x_mask=None, snr=0.16):
        """Correct step for the input tensor (torsion angles).

        Args:
            x (Tensor): Torsion angles of shape :math:`(num_res, 4)`
            x_score (Tensor): Score of shape :math:`(num_res, 4)`
            x_batch (Tensor): Batch of shape :math:`(num_res)`
            x_mask (Tensor): Mask of shape :math:`(num_res, 4)`
            snr (float): Signal to noise ratio

        Returns:
            Tensor: Corrected torsion angles of shape :math:`(num_res, 4)`
        """
        x_batch = x_batch.reshape(-1, 4)

        # Calculate Score Norm
        x_score_2 = x_score ** 2
        score_norm = torch.sqrt(scatter_add(x_score_2[x_mask], x_batch[x_mask], dim=0)).mean()

        # Calculate Noise Norm
        noise = torch.randn_like(x_score)
        noise_2 = noise ** 2
        noise_norm = torch.sqrt(scatter_add(noise_2[x_mask], x_batch[x_mask], dim=0)).mean()

        # Step Size
        step_size = (snr * noise_norm / score_norm) ** 2 * 2

        # Correct Step
        x_prev = x.clone()
        x_prev += step_size * x_score
        x_prev += ((step_size * 2) ** 0.5) * noise

        if x_mask is not None:
            x_prev[~x_mask] = x[~x_mask]

        return x_prev

    def sample_train_t(self, shape):
        """Sample timesteps from uniform distribution.