Create MMD.py

MMD.py

@@ -0,0 +1,168 @@
+import torch
+
+
+def flexible_kernel(X, Y, X_org, Y_org, sigma, sigma0=0.1, epsilon=1e-08):
+    """Flexible kernel calculation as in MMDu."""
+    Dxy = Pdist2(X, Y)
+    Dxy_org = Pdist2(X_org, Y_org)
+    L = 1
+    Kxy = (1 - epsilon) * torch.exp(
+        -((Dxy / sigma0) ** L) - Dxy_org / sigma
+    ) + epsilon * torch.exp(-Dxy_org / sigma)
+    return Kxy
+
+
+def MMD_Diff_Var(Kyy, Kzz, Kxy, Kxz, epsilon=1e-08):
+    """Compute the variance of the difference statistic MMDXY - MMDXZ."""
+    """Referenced from: https://github.com/eugenium/MMD/blob/master/mmd.py"""
+    m = Kxy.shape[0]
+    n = Kyy.shape[0]
+    r = Kzz.shape[0]
+
+    # Remove diagonal elements
+    Kyynd = Kyy - torch.diag(torch.diag(Kyy))
+    Kzznd = Kzz - torch.diag(torch.diag(Kzz))
+
+    u_yy = torch.sum(Kyynd) * (1.0 / (n * (n - 1)))
+    u_zz = torch.sum(Kzznd) * (1.0 / (r * (r - 1)))
+    u_xy = torch.sum(Kxy) / (m * n)
+    u_xz = torch.sum(Kxz) / (m * r)
+
+    t1 = (1.0 / n**3) * torch.sum(Kyynd.T @ Kyynd) - u_yy**2
+    t2 = (1.0 / (n**2 * m)) * torch.sum(Kxy.T @ Kxy) - u_xy**2
+    t3 = (1.0 / (n * m**2)) * torch.sum(Kxy @ Kxy.T) - u_xy**2
+    t4 = (1.0 / r**3) * torch.sum(Kzznd.T @ Kzznd) - u_zz**2
+    t5 = (1.0 / (r * m**2)) * torch.sum(Kxz @ Kxz.T) - u_xz**2
+    t6 = (1.0 / (r**2 * m)) * torch.sum(Kxz.T @ Kxz) - u_xz**2
+    t7 = (1.0 / (n**2 * m)) * torch.sum(Kyynd @ Kxy.T) - u_yy * u_xy
+    t8 = (1.0 / (n * m * r)) * torch.sum(Kxy.T @ Kxz) - u_xz * u_xy
+    t9 = (1.0 / (r**2 * m)) * torch.sum(Kzznd @ Kxz.T) - u_zz * u_xz
+
+    if type(epsilon) == torch.Tensor:
+        epsilon_tensor = epsilon.clone().detach()
+    else:
+        epsilon_tensor = torch.tensor(epsilon, device=Kyy.device)
+    zeta1 = torch.max(t1 + t2 + t3 + t4 + t5 + t6 - 2 * (t7 + t8 + t9), epsilon_tensor)
+    zeta2 = torch.max(
+        (1 / m / (m - 1)) * torch.sum((Kyynd - Kzznd - Kxy.T - Kxy + Kxz + Kxz.T) ** 2)
+        - (u_yy - 2 * u_xy - (u_zz - 2 * u_xz)) ** 2,
+        epsilon_tensor,
+    )
+
+    data = {
+        "t1": t1.item(),
+        "t2": t2.item(),
+        "t3": t3.item(),
+        "t4": t4.item(),
+        "t5": t5.item(),
+        "t6": t6.item(),
+        "t7": t7.item(),
+        "t8": t8.item(),
+        "t9": t9.item(),
+        "zeta1": zeta1.item(),
+        "zeta2": zeta2.item(),
+    }
+
+    Var = (4 * (m - 2) / (m * (m - 1))) * zeta1
+    Var_z2 = Var + (2.0 / (m * (m - 1))) * zeta2
+
+    return Var, Var_z2, data
+
+
+def Pdist2(x, y):
+    """compute the paired distance between x and y."""
+    x_norm = (x**2).sum(1).view(-1, 1)
+    if y is not None:
+        y_norm = (y**2).sum(1).view(1, -1)
+    else:
+        y = x
+        y_norm = x_norm.view(1, -1)
+    Pdist = x_norm + y_norm - 2.0 * torch.mm(x, torch.transpose(y, 0, 1))
+    Pdist[Pdist < 0] = 0
+    return Pdist
+
+
+def MMD_batch2(
+    Fea,
+    len_s,
+    Fea_org,
+    sigma,
+    sigma0=0.1,
+    epsilon=10 ** (-10),
+    is_var_computed=True,
+    use_1sample_U=True,
+    coeff_xy=2,
+):
+    X = Fea[0:len_s, :]
+    Y = Fea[len_s:, :]
+    L = 1  # generalized Gaussian (if L>1)
+
+    nx = X.shape[0]
+    ny = Y.shape[0]
+    Dxx = Pdist2(X, X)
+    Dyy = torch.zeros(Fea.shape[0] - len_s, 1).to(Dxx.device)
+    # Dyy = Pdist2(Y, Y)
+    Dxy = Pdist2(X, Y).transpose(0, 1)
+    Kx = torch.exp(-Dxx / sigma0)
+    Ky = torch.exp(-Dyy / sigma0)
+    Kxy = torch.exp(-Dxy / sigma0)
+
+    nx = Kx.shape[0]
+
+    is_unbiased = False
+    xx = torch.div((torch.sum(Kx)), (nx * nx))
+    yy = Ky.reshape(-1)
+    xy = torch.div(torch.sum(Kxy, dim=1), (nx))
+
+    mmd2 = xx - 2 * xy + yy
+    return mmd2
+
+
+# MMD for three samples
+def MMD_3_Sample_Test(
+    ref_fea,
+    fea_y,
+    fea_z,
+    ref_fea_org,
+    fea_y_org,
+    fea_z_org,
+    sigma,
+    sigma0,
+    epsilon,
+    alpha,
+):
+    """Run three-sample test (TST) using deep kernel kernel."""
+    X = ref_fea.clone().detach()
+    Y = fea_y.clone().detach()
+    Z = fea_z.clone().detach()
+    X_org = ref_fea_org.clone().detach()
+    Y_org = fea_y_org.clone().detach()
+    Z_org = fea_z_org.clone().detach()
+
+    Kyy = flexible_kernel(Y, Y, Y_org, Y_org, sigma, sigma0, epsilon)
+    Kzz = flexible_kernel(Z, Z, Z_org, Z_org, sigma, sigma0, epsilon)
+    Kxy = flexible_kernel(X, Y, X_org, Y_org, sigma, sigma0, epsilon)
+    Kxz = flexible_kernel(X, Z, X_org, Z_org, sigma, sigma0, epsilon)
+
+    Kyynd = Kyy - torch.diag(torch.diag(Kyy))
+    Kzznd = Kzz - torch.diag(torch.diag(Kzz))
+
+    Diff_Var, _, _ = MMD_Diff_Var(Kyy, Kzz, Kxy, Kxz, epsilon)
+
+    u_yy = torch.sum(Kyynd) / (Y.shape[0] * (Y.shape[0] - 1))
+    u_zz = torch.sum(Kzznd) / (Z.shape[0] * (Z.shape[0] - 1))
+    u_xy = torch.sum(Kxy) / (X.shape[0] * Y.shape[0])
+    u_xz = torch.sum(Kxz) / (X.shape[0] * Z.shape[0])
+
+    t = u_yy - 2 * u_xy - (u_zz - 2 * u_xz)
+    if Diff_Var.item() <= 0:
+        Diff_Var = torch.max(epsilon, torch.tensor(1e-08))
+    p_value = torch.distributions.Normal(0, 1).cdf(-t / torch.sqrt((Diff_Var)))
+    t = t / torch.sqrt(Diff_Var)
+
+    if p_value > alpha:
+        h = 0
+    else:
+        h = 1
+
+    return h, p_value.item(), t.item()