Class SO2VESchedule

diffpack/schedule.py:155–294

@R.register('SO2VESchedule')
class SO2VESchedule(SO2Schedule, core.Configurable):
    def __init__(self, pi_periodic=False, cache_folder=None, sigma_min=0.01 * np.pi, sigma_max=np.pi, annealed_temp=3, mode="sde"):
        """
        Args:
            sigma_min (float): minimum standard deviation
            sigma_max (float): maximum standard deviation
        """
        PI = 1/2 * np.pi if pi_periodic else np.pi  # TODO: remove ambiguity
        super().__init__(PI, cache_folder)
        self.sigma_min = sigma_min
        self.sigma_max = sigma_max
        self.sigma_min_log = np.log(sigma_min)
        self.sigma_max_log = np.log(sigma_max)
        self.annealed_temp = annealed_temp
        self.mode = mode

    def t_to_sigma(self, t):
        """Transfer timesteps to standard deviation.

        Args:
            t (Tensor): timesteps ranging from 0 to 1

        Returns:
            Tensor: standard deviation ranging from sigma_min to sigma_max
        """
        return torch.exp(self.sigma_min_log + (self.sigma_max_log - self.sigma_min_log) * t)

    @torch.no_grad()
    def add_noise(self, x, t, x_mask=None):
        """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)`
        """
        sigmas = self.t_to_sigma(t)
        noise = torch.randn_like(x) * sigmas.unsqueeze(-1)
        score = torch.tensor(
            self.score(noise.cpu().numpy(), sigmas.cpu().numpy()), device=x.device, dtype=x.dtype
        )

        if x_mask is not None:
            noise *= x_mask
            score *= x_mask

        x = x + noise
        return x, score

    @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)`