Method step

diffpack/schedule.py:205–241

Denoise 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)` t (Tensor): Timesteps of shape :math:`(num_res)` dt (float): Step size of

(self, x, x_score, t, dt, x_mask=None)

    @torch.no_grad()
    def step(self, x, x_score, t, dt, x_mask=None):
        """Denoise 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)`
            t (Tensor): Timesteps of shape :math:`(num_res)`
            dt (float): Step size of shape :math:`(num_res)`
            x_mask (Tensor): Mask of shape :math:`(num_res, 4)`

        Returns:
            Tensor: Denoised torsion angles of shape :math:`(num_res, 4)`
        """
        sigma = self.t_to_sigma(t)  # (num_res)
        g = sigma * np.sqrt(2 * np.log(self.sigma_max / self.sigma_min))  # (num_res)

        # Temperature Coefficient
        alpha = 1 - (sigma / np.exp(self.sigma_max_log)) ** 2 if self.annealed_temp else None
        annealed_weight = self.annealed_temp / (alpha + (1 - alpha) * self.annealed_temp) if self.annealed_temp else 1

        # Noise
        x_prev = x.clone()
        if self.mode == "ode":
            drift = (0.5 * g ** 2 * dt * (x_score * annealed_weight))
            x_prev += drift
        elif self.mode == "sde":
            noise = torch.normal(mean=0, std=1, size=x_score.shape, device=x_score.device)
            drift = g ** 2 * dt * (x_score * annealed_weight)
            diffusion = g * torch.sqrt(dt) * noise
            x_prev += (drift + diffusion)
        else:
            raise NotImplementedError

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

        return x_prev

    @torch.no_grad()
    def step_correct(self, x, x_score, x_batch, x_mask=None, snr=0.16):