diffusion_utils.py

import torch
import numpy as np
from scipy.stats import betaprime
import numpy as np
from xgboost import XGBRegressor
import warnings
warnings.filterwarnings('ignore')
from catboost import CatBoostClassifier

randn_like=torch.randn_like

SIGMA_MIN=0.002
SIGMA_MAX=80
rho=7
S_churn= 1
S_min=0
S_max=float('inf')
S_noise=1

def impute_mask(net, x, mask, num_samples, dim, num_steps = 50, device = 'cuda:0'):
    step_indices = torch.arange(num_steps, dtype=torch.float32, device=device)
    x_t = torch.randn([num_samples, dim], device=device)

    sigma_min = max(SIGMA_MIN, net.sigma_min)
    sigma_max = min(SIGMA_MAX, net.sigma_max)

    t_steps = (sigma_max ** (1 / rho) + step_indices / (num_steps - 1) * (
                sigma_min ** (1 / rho) - sigma_max ** (1 / rho))) ** rho
    t_steps = torch.cat([net.round_sigma(t_steps), torch.zeros_like(t_steps[:1])])

    mask = mask.to(torch.int).to(device)
    x_t = x_t.to(torch.float32) * t_steps[0]

    N = 10
    with torch.no_grad():

        for i, (t_cur, t_next) in enumerate(zip(t_steps[:-1], t_steps[1:])):
            if i < num_steps - 1:
         
                for j in range(N):
                    n_curr = torch.randn_like(x_t).to(device) * t_cur
                    n_prev = torch.randn_like(x_t).to(device) * t_next

                    x_known_t_prev = x + n_prev
                    x_unknown_t_prev = sample_step(net, num_steps, i, t_cur, t_next, x_t)

                    x_t_prev = x_known_t_prev * (1-mask) + x_unknown_t_prev * mask

                    n = torch.randn_like(x_t) * (t_cur.pow(2) - t_next.pow(2)).sqrt()

                    if j == N - 1:
                        x_t = x_t_prev                                                # turn to x_{t-1}
                    else:
                        x_t = x_t_prev + n                                            # new x_t

    return x_t

def sample(net, num_samples, dim, num_steps = 50, device = 'cuda:0'):
    latents = torch.randn([num_samples, dim], device=device)

    step_indices = torch.arange(num_steps, dtype=torch.float32, device=latents.device)

    sigma_min = max(SIGMA_MIN, net.sigma_min)
    sigma_max = min(SIGMA_MAX, net.sigma_max)

    t_steps = (sigma_max ** (1 / rho) + step_indices / (num_steps - 1) * (
                sigma_min ** (1 / rho) - sigma_max ** (1 / rho))) ** rho
    t_steps = torch.cat([net.round_sigma(t_steps), torch.zeros_like(t_steps[:1])])

    x_next = latents.to(torch.float32) * t_steps[0]

    with torch.no_grad():
        for i, (t_cur, t_next) in enumerate(zip(t_steps[:-1], t_steps[1:])):
            x_next = sample_step(net, num_steps, i, t_cur, t_next, x_next)

    return x_next

def sample_step(net, num_steps, i, t_cur, t_next, x_next):

    x_cur = x_next
    # Increase noise temporarily.
    gamma = min(S_churn / num_steps, np.sqrt(2) - 1) if S_min <= t_cur <= S_max else 0
    t_hat = net.round_sigma(t_cur + gamma * t_cur) 
    x_hat = x_cur + (t_hat ** 2 - t_cur ** 2).sqrt() * S_noise * randn_like(x_cur)
    # Euler step.

    denoised = net(x_hat, t_hat).to(torch.float32)
    d_cur = (x_hat - denoised) / t_hat
    x_next = x_hat + (t_next - t_hat) * d_cur

    # Apply 2nd order correction.
    if i < num_steps - 1:
        denoised = net(x_next, t_next).to(torch.float32)
        d_prime = (x_next - denoised) / t_next
        x_next = x_hat + (t_next - t_hat) * (0.5 * d_cur + 0.5 * d_prime)

    return x_next

class VPLoss:
    def __init__(self, beta_d=19.9, beta_min=0.1, epsilon_t=1e-5):
        self.beta_d = beta_d
        self.beta_min = beta_min
        self.epsilon_t = epsilon_t

    def __call__(self, denosie_fn, data, labels, augment_pipe=None):
        rnd_uniform = torch.rand([data.shape[0], 1, 1, 1], device=data.device)
        sigma = self.sigma(1 + rnd_uniform * (self.epsilon_t - 1))
        weight = 1 / sigma ** 2
        y, augment_labels = augment_pipe(data) if augment_pipe is not None else (data, None)
        n = torch.randn_like(y) * sigma
        D_yn = denosie_fn(y + n, sigma, labels, augment_labels=augment_labels)
        loss = weight * ((D_yn - y) ** 2)
        return loss

    def sigma(self, t):
        t = torch.as_tensor(t)
        return ((0.5 * self.beta_d * (t ** 2) + self.beta_min * t).exp() - 1).sqrt()



class VELoss:
    def __init__(self, sigma_min=0.02, sigma_max=100, D=128, N=3072, opts=None):
        self.sigma_min = sigma_min
        self.sigma_max = sigma_max
        self.D = D
        self.N = N
        print(f"In VE loss: D:{self.D}, N:{self.N}")

    def __call__(self, denosie_fn, data, labels = None, augment_pipe=None, stf=False, pfgmpp=False, ref_data=None):
        if pfgmpp:

            # N, 
            rnd_uniform = torch.rand(data.shape[0], device=data.device)
            sigma = self.sigma_min * ((self.sigma_max / self.sigma_min) ** rnd_uniform)

            r = sigma.double() * np.sqrt(self.D).astype(np.float64)
            # Sampling form inverse-beta distribution
            samples_norm = np.random.beta(a=self.N / 2., b=self.D / 2.,
                                          size=data.shape[0]).astype(np.double)

            samples_norm = np.clip(samples_norm, 1e-3, 1-1e-3)

            inverse_beta = samples_norm / (1 - samples_norm + 1e-8)
            inverse_beta = torch.from_numpy(inverse_beta).to(data.device).double()
            # Sampling from p_r(R) by change-of-variable
            samples_norm = r * torch.sqrt(inverse_beta + 1e-8)
            samples_norm = samples_norm.view(len(samples_norm), -1)
            # Uniformly sample the angle direction
            gaussian = torch.randn(data.shape[0], self.N).to(samples_norm.device)
            unit_gaussian = gaussian / torch.norm(gaussian, p=2, dim=1, keepdim=True)
            # Construct the perturbation for x
            perturbation_x = unit_gaussian * samples_norm
            perturbation_x = perturbation_x.float()

            sigma = sigma.reshape((len(sigma), 1, 1, 1))
            weight = 1 / sigma ** 2
            y, augment_labels = augment_pipe(data) if augment_pipe is not None else (data, None)
            n = perturbation_x.view_as(y)
            D_yn = denosie_fn(y + n, sigma, labels,  augment_labels=augment_labels)
        else:
            rnd_uniform = torch.rand([data.shape[0], 1, 1, 1], device=data.device)
            sigma = self.sigma_min * ((self.sigma_max / self.sigma_min) ** rnd_uniform)
            weight = 1 / sigma ** 2
            y, augment_labels = augment_pipe(data) if augment_pipe is not None else (data, None)
            n = torch.randn_like(y) * sigma
            D_yn = denosie_fn(y + n, sigma, labels, augment_labels=augment_labels)

        loss = weight * ((D_yn - y) ** 2)
        return loss


class EDMLoss:
    def __init__(self, P_mean=-1.2, P_std=1.2, sigma_data=0.5, hid_dim = 100, gamma=5, opts=None):
        self.P_mean = P_mean
        self.P_std = P_std
        self.sigma_data = sigma_data
        self.hid_dim = hid_dim
        self.gamma = gamma
        self.opts = opts


    def __call__(self, denoise_fn, data, batch_mask=None):
        
        rnd_normal = torch.randn(data.shape[0], device=data.device)
        sigma = (rnd_normal * self.P_std + self.P_mean).exp()

        weight = (sigma ** 2 + self.sigma_data ** 2) / (sigma * self.sigma_data) ** 2

        y = data
        n = torch.randn_like(y) * sigma.unsqueeze(1)
        
        D_yn = denoise_fn(y + n, sigma)

        target = y
        loss = weight.unsqueeze(1) * ((D_yn - target) ** 2)

        return loss
    


def refinement(rec_X, mask,clf='catboost',len_num=0):
    
    X = rec_X.copy()
    n_samples, n_features = X.shape
    
    for col in range(n_features):
        missing_idx = np.where(mask[:, col] == 0)[0]
        observed_idx = np.where(mask[:, col] == 1)[0]

        if len(missing_idx) == 0:
            # print("found non missing column: ",col)
            continue  # Nothing to impute

        X_obs = X[observed_idx]
        X_miss = X[missing_idx]

        X_obs_input = np.delete(X_obs, col, axis=1)
        y_obs = X[observed_idx, col]
        X_miss_input = np.delete(X_miss, col, axis=1)

        unique_vals = sorted(set(y_obs))
        if col>=len_num:
            if clf == 'catboost':

                val_to_label = {val: idx for idx, val in enumerate(unique_vals)}
                label_to_val = {idx: val for val, idx in val_to_label.items()}

                y_obs_mapped = np.array([val_to_label[v] for v in y_obs])
                
                if len(set(y_obs_mapped)) == 1:
                    y_pred = np.array([label_to_val[y_obs_mapped[0]]] * len(missing_idx))
                else:
                    model = CatBoostClassifier(logging_level='Silent')
                    model.fit(X_obs_input, y_obs_mapped,verbose=False)
                    y_pred_labels = model.predict(X_miss_input)
                    y_pred = np.array([label_to_val[l.item() if isinstance(l, np.ndarray) else l] for l in y_pred_labels])
                    X[missing_idx, col] = y_pred        

        else:
            model =  XGBRegressor()
            model.fit(X_obs_input, y_obs,verbose=False)
            y_pred = model.predict(X_miss_input)
            X[missing_idx, col] = y_pred

        

    return X