Class ICPlan

transport/path.py:18–136

Linear Coupling Plan

#################### Coupling Plans ####################

class ICPlan:
    """Linear Coupling Plan"""
    def __init__(self, sigma=0.0):
        self.sigma = sigma

    def compute_alpha_t(self, t):
        """Compute the data coefficient along the path"""
        return t, 1
    
    def compute_sigma_t(self, t):
        """Compute the noise coefficient along the path"""
        return 1 - t, -1
    
    def compute_d_alpha_alpha_ratio_t(self, t):
        """Compute the ratio between d_alpha and alpha"""
        return 1 / t

    def compute_drift(self, x, t):
        """We always output sde according to score parametrization; """
        t = expand_t_like_x(t, x)
        alpha_ratio = self.compute_d_alpha_alpha_ratio_t(t)
        sigma_t, d_sigma_t = self.compute_sigma_t(t)
        drift = alpha_ratio * x
        diffusion = alpha_ratio * (sigma_t ** 2) - sigma_t * d_sigma_t # beta_t in Table 2, from VE-SDE.

        return -drift, diffusion

    def compute_diffusion(self, x, t, form="constant", norm=1.0): #compute w_t in Eq 4.
        """Compute the diffusion term of the SDE
        Args:
          x: [batch_dim, ...], data point
          t: [batch_dim,], time vector
          form: str, form of the diffusion term
          norm: float, norm of the diffusion term
        """
        t = expand_t_like_x(t, x)
        choices = {
            "constant": norm,
            "SBDM": norm * self.compute_drift(x, t)[1], #follow the calculation of w_t for SBDM
            "sigma": norm * self.compute_sigma_t(t)[0], # This suggests the choice wt = σt in (4) to cancel this singularity (see Appendix A.3)
            "linear": norm * (1 - t), #Table 2, seems never used 
            "decreasing": 0.25 * (norm * th.cos(np.pi * t) + 1) ** 2, #Table2, seems never used 
            "inccreasing-decreasing": norm * th.sin(np.pi * t) ** 2, #seems never used 
        }

        try:
            diffusion = choices[form]
        except KeyError:
            raise NotImplementedError(f"Diffusion form {form} not implemented")
        
        return diffusion

    def get_score_from_velocity(self, velocity, x, t): #Eq 9 in SiT paper by simple algebra
        """Wrapper function: transfrom velocity prediction model to score
        Args:
            velocity: [batch_dim, ...] shaped tensor; velocity model output
            x: [batch_dim, ...] shaped tensor; x_t data point
            t: [batch_dim,] time tensor