csgc.py

# -*- coding: utf-8 -*-
"""csgc.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/10-72bciFo8Tk_f5GXrpjDrJIt8Sjy602
"""

# -*- coding: utf-8 -*-
import torch
import math
import numpy as np
import warnings
from itertools import permutations
from sklearn.utils.extmath import randomized_svd
from sklearn.cluster import KMeans
from scipy.optimize import fsolve
import psutil

def gen_adj_sbm(K,z,device="cpu"):
  """
  INPUT
  K = connection probability matrix (tensor)
  z = labels (tensor)

  OUTPUT
  A:  adjacency matrix (tensor)
  P:  probability matrix (tensor)
  """
  K = K.to(device)
  z = z.to(device)
  k = K.shape[0]
  n = len(z)
  Z = torch.zeros((n,k)).to(device)
  Z[range(n),z.long()-1] = 1
  P = Z @ K @ Z.T
  P.fill_diagonal_(0)
  A = torch.bernoulli(torch.triu(P,diagonal=1))
  A = (A+A.T).to(torch.uint8)
  return A,P

def gen_adj_dcsbm(K,Theta,z,device="cpu"):
  """
  INPUT
  K = connection probability matrix (tensor)
  Theta = degree parameter vector (tensor)
  z = community labels (tensor)

  OUTPUT
  A:  adjacency matrix (tensor)
  P:  probability matrix (tensor)
  """
  K = K.to(device)
  Theta = Theta.to(device)
  z = z.to(device)
  k = K.shape[0]
  n = len(z)
  Z = torch.zeros((n,k)).to(device)
  Z[range(n),z.long()-1] = 1
  P = torch.diag(Theta) @ Z @ K @ Z.T @ torch.diag(Theta)
  P.fill_diagonal_(0)
  A = torch.poisson(torch.triu(P,diagonal=1))
  A = (A+A.T).to(torch.uint8)
  return A,P

def sbm_mle(A,z,device="cpu"): # A,z as tensor
  A = A.to(device)
  z = z.to(device)
  k = len(torch.unique(z))
  n = len(z)
  K = torch.tensor([[torch.sum(A[z==i,:][:,z==j]) / (torch.sum(z==i)*torch.sum(z==j))
           for i in range(1,k+1)] for j in range(1,k+1)]).to(device)
  diagK = torch.tensor([torch.sum(A[z==i,:][:,z==i]) / (torch.sum(z==i)*(torch.sum(z==i)-1))
            if torch.sum(z==i)!=1 else 0 for i in range(1,k+1)]).to(device)
  K[range(k),range(k)] = diagK
  Z = torch.zeros((n,k)).to(device)
  Z[range(n),z.long()-1] = 1
  P = Z @ K @ Z.T
  return K,P # K,P as tensor

def dcsbm_mle(A,z,device="cpu"): # A,z as tensor
  A = A.to(device)
  z = z.to("cpu")
  k = len(torch.unique(z))
  n = len(z)
  ind = [(z==i).nonzero(as_tuple=True)[0] for i in range(1,k+1)]
  K = torch.tensor([[torch.sum(A[z==i,:][:,z==j]) / (torch.sum(z==i)*torch.sum(z==j))
           for i in range(1,k+1)] for j in range(1,k+1)]).to("cpu")
  deg = torch.sum(A,dim=0)
  Theta0 = torch.zeros(n).to(device)
  for i in range(1,k+1):
    if torch.sum(deg[z==i]==0):
      Theta0[z==i] = 0
    else:
      Theta0[z==i] = deg[z==i]/torch.sum(deg[z==i])
  Theta = np.zeros(n)
  def f(x):
    Thetasq_r = np.sum(np.square(x))
    return (k_r-kappa_r*x)*(1-Thetasq_r) - x*m_rr*Thetasq_r+m_rr*np.square(x)
  for r in range(1,k+1):
    init = Theta0[z==r].to("cpu").numpy()
    k_r = deg[z==r].to("cpu").numpy()
    kappa_r = np.sum(k_r)
    m_rr = np.diag(K)[r-1]
    root = fsolve(f,init)
    Theta[z==r] = root
    K[r-1,r-1] = K[r-1,r-1]/(1-np.sum(np.square(root)))
  Theta = torch.Tensor(Theta).to(device)
  Z = torch.zeros((n,k)).to(device)
  Z[range(n),z.long()-1] = 1
  K = K.to(device)
  P = torch.diag(Theta) @ Z @ K @ Z.T @ torch.diag(Theta)
  P.fill_diagonal_(0)
  return K,Theta,P # K,Theta,P as tensor

def spectral_sbm(A,k,device="cpu"): # A as tensor
  A = A.to(device)
  A = A+torch.mean(torch.sum(A.float(),1))/A.shape[0]
  V = torch.tensor(randomized_svd(A.to("cpu").numpy(), n_components=k, random_state=None)[2])
  labels = KMeans(n_clusters=k,n_init="auto").fit(V.T).labels_
  return torch.tensor(labels+1).to(torch.uint8) # labels as tensor

def spectral_dcsbm(A,k,device="cpu"): # A as tensor
  A = A.to(device)
  A = A+torch.mean(torch.sum(A.float(),1))/A.shape[0]
  V = torch.tensor(randomized_svd(A.to("cpu").numpy(), n_components=k, random_state=None)[2])
  vr = torch.sqrt(torch.sum(torch.square(V.T),1))
  reg_V = torch.diag(1/vr) @ V.T
  labels = KMeans(n_clusters=k,n_init="auto").fit(reg_V).labels_
  return torch.tensor(labels+1).to(torch.uint8) # labels as tensor

def csgc(tsA,tsP,var_structure="bernoulli",device="cpu"): # tsA,tsP as tensor
  if device=="cpu":
    if psutil.virtual_memory()[0]*0.7/(torch.numel(tsP)*4)/11 < 1:
      raise ValueError("Matrix size too large!")
  if device=="cuda:0":
    if torch.cuda.get_device_properties(0).total_memory*0.7/(torch.numel(tsP)*4)/11 < 1:
       raise ValueError("Matrix size too large!")
  tsA = tsA.to(device)
  tsP = tsP.to(device)
  def csgc_V(Yh):
    # Yh: Hadamard
    Yh_2 = Yh*Yh
    Yh_3 = Yh_2*Yh
    Yh_4s = torch.sum(Yh_3*Yh,(-2,-1))
    # Yh: matrix multiplication
    Yh_2p = Yh @ Yh
    Yh_3p = Yh_2p @ Yh
    Yh_4p = Yh_3p @ Yh
    Yhp_2_1s = torch.sum(Yh_2@Yh,(-2,-1))
    Yhp_2_2 = Yh_2 @ Yh_2
    Yhp_2_2s = torch.sum(Yhp_2_2,(-2,-1))
    Yhp_2_2p = Yh_2 @ Yh_2p
    Yhp_3_1s = torch.sum(Yh_3@Yh,(-2,-1))
    # Yh: row sum
    Yhr = torch.sum(Yh,dim=-1)
    Yhrrs = torch.sum(torch.sum(Yh_2,dim=-1)*Yhr*Yhr,-1)
    # Yh: diagonal
    Yhd_3p = torch.diagonal(Yh_3p,dim1=-2,dim2=-1)
    Yhd_3p_1s = torch.sum(torch.diag_embed(Yhd_3p)@Yh,(-2,-1))
    # V for 10 subgraphs
    V_twostar = torch.sum(torch.triu(Yh_2p,diagonal=1),(-2,-1))
    V_triangle = torch.diagonal(Yh_3p,dim1=-2,dim2=-1).sum(-1)/6
    V_fourcycle = (torch.diagonal(Yh_4p,dim1=-2,dim2=-1).sum(-1)-4*torch.sum(torch.triu(Yhp_2_2,diagonal=1),(-2,-1))-Yh_4s)/8
    V_threepath = torch.sum(torch.triu(Yh_3p,diagonal=1),(-2,-1))+torch.sum(torch.triu(Yh_3,diagonal=1),(-2,-1))-Yhp_2_1s
    V_threestar = (torch.sum(Yhr.pow(3),-1)-3*Yhp_2_1s+2*torch.sum(Yh_3,(-2,-1)))/6
    V_triangleapp = (Yhd_3p_1s-2*torch.diagonal(Yh_2p@Yh_2,dim1=-2,dim2=-1).sum(-1))/2
    V_twotriangle = (torch.sum(Yh_2p*Yh_2p*Yh,(-2,-1))-torch.diagonal(Yhp_2_2@Yh,dim1=-2,dim2=-1).sum(-1))/4
    V_fivecycle = (torch.diagonal(Yh_4p@Yh,dim1=-2,dim2=-1).sum(-1)-5*torch.sum(torch.diag_embed(Yhd_3p)@Yh_2,(-2,-1))+ \
            5*torch.diagonal(Yh_3@Yh_2p,dim1=-2,dim2=-1).sum(-1))/10
    V_fourpath = (2*torch.sum(torch.triu(Yh_4p,diagonal=1),(-2,-1))-2*torch.sum(Yhp_2_2p,(-2,-1))+2*Yhp_3_1s+3*Yhp_2_2s+ \
            3*torch.diagonal(Yhp_2_2p,dim1=-2,dim2=-1).sum(-1)-2*Yh_4s-Yhrrs-2*Yhd_3p_1s)/2
    V_fourstar = (torch.sum(Yhr.pow(4),-1)-6*Yhrrs+3*Yhp_2_2s+8*Yhp_3_1s-6*Yh_4s)/24
    V = torch.stack((V_twostar,V_triangle,V_fourcycle,V_threepath,V_threestar,
              V_triangleapp,V_twotriangle,V_fivecycle,V_fourpath,V_fourstar)).t()
    return V
  def csgc_sd(vr):
    # vr: Hadamard
    vr_2 = vr*vr
    vr_3 = vr_2*vr
    vr_4s = torch.sum(vr_3*vr,(-2,-1))
    # vr: matrix multiplication
    vr_2p = vr @ vr
    vr_3p = vr_2p @ vr
    vr_4p = vr_3p @ vr
    vrp_2_1s = torch.sum(vr_2@vr,(-2,-1))
    vrp_2_2 = vr_2 @ vr_2
    vrp_2_2s = torch.sum(vrp_2_2,(-2,-1))
    vrp_2_2p = vr_2 @ vr_2p
    vrp_3_1s = torch.sum(vr_3@vr,(-2,-1))
    # vr: row sum
    vrr = torch.sum(vr,dim=-1)
    vrrrs = torch.sum(torch.sum(vr_2,dim=-1)*vrr*vrr,-1)
    # vr: diagonal
    vrd_3p = torch.diagonal(vr_3p,dim1=-2,dim2=-1)
    vrd_3p_1s = torch.sum(torch.diag_embed(vrd_3p)@vr,(-2,-1))
    # vr for 10 subgraphs
    vr_twostar = torch.sum(torch.triu(vr_2p,diagonal=1),(-2,-1))
    vr_triangle = torch.diagonal(vr_3p,dim1=-2,dim2=-1).sum(-1)/6
    vr_fourcycle = (torch.diagonal(vr_4p,dim1=-2,dim2=-1).sum(-1)-4*torch.sum(torch.triu(vrp_2_2,diagonal=1),(-2,-1))-vr_4s)/8
    vr_threepath = torch.sum(torch.triu(vr_3p,diagonal=1),(-2,-1))+torch.sum(torch.triu(vr_3,diagonal=1),(-2,-1))-vrp_2_1s
    vr_threestar = (torch.sum(vrr.pow(3),-1)-3*vrp_2_1s+2*torch.sum(vr_3,(-2,-1)))/6
    vr_triangleapp = (vrd_3p_1s-2*torch.diagonal(vr_2p@vr_2,dim1=-2,dim2=-1).sum(-1))/2
    vr_twotriangle = (torch.sum(vr_2p*vr_2p*vr,(-2,-1))-torch.diagonal(vrp_2_2@vr,dim1=-2,dim2=-1).sum(-1))/4
    vr_fivecycle = (torch.diagonal(vr_4p@vr,dim1=-2,dim2=-1).sum(-1)-5*torch.sum(torch.diag_embed(vrd_3p)@vr_2,(-2,-1))+ \
            5*torch.diagonal(vr_3@vr_2p,dim1=-2,dim2=-1).sum(-1))/10
    vr_fourpath = (2*torch.sum(torch.triu(vr_4p,diagonal=1),(-2,-1))-2*torch.sum(vrp_2_2p,(-2,-1))+2*vrp_3_1s+3*vrp_2_2s+ \
            3*torch.diagonal(vrp_2_2p,dim1=-2,dim2=-1).sum(-1)-2*vr_4s-vrrrs-2*vrd_3p_1s)/2
    vr_fourstar = (torch.sum(vrr.pow(4),-1)-6*vrrrs+3*vrp_2_2s+8*vrp_3_1s-6*vr_4s)/24
    sd = torch.sqrt(torch.stack((vr_twostar,vr_triangle,vr_fourcycle,vr_threepath,vr_threestar,
                    vr_triangleapp,vr_twotriangle,vr_fivecycle,vr_fourpath,vr_fourstar)).t())
    return sd
  if var_structure == "bernoulli":
    if torch.sum(torch.sum(tsA>1,(-2,-1))/(tsA.size(dim=-1)*tsA.size(dim=-2))>0.01)>0:
      warnings.warn("Results may be inaccurate: too many multiedges and selfedges.")
    tsA = tsA - torch.diag_embed(torch.diagonal(tsA,dim1=-2,dim2=-1))  # remove diagonal for tsA
    tsP = tsP - torch.diag_embed(torch.diagonal(tsP,dim1=-2,dim2=-1))  # remove diagonal for tsP
    tsP = torch.where(tsP<1,tsP,1.)  # elements in tsP capped as 1
    return csgc_V(Yh=torch.where(tsP<1,tsA-tsP,0.))/csgc_sd(vr=tsP*(1-tsP)) # csgc as tensor
  elif var_structure == "poisson":
    return csgc_V(Yh=tsA-tsP)/csgc_sd(vr=tsP) # csgc as tensor
  else:
    raise ValueError("variance structure can only be bernoulli or poisson!")

def csgc_greedy(A,z0,var_structure="bernoulli",device="cpu"): # A,z0 as tensor
  A = A.to(device)
  z0 = z0.to(device)
  k = len(torch.unique(z0))
  n = A.shape[0]
  stats0 = csgc(A,sbm_mle(A,z0,device)[1],var_structure,device) # initial csgc statistics
  chisq0 = torch.sum(torch.square(stats0)) # initial chi-square value
  zadj = z0
  chisqadj = chisq0
  chisq_diff = .01
  chisq_log = zadj_log = torch.tensor(()).to(device) # records for chi-square values and labels for each iteration
  if device == "cuda:0": # The maximum number of P matrices allowed for each calculation of csgc function (GPU and CPU)
    maxsize = math.floor(torch.cuda.get_device_properties(0).total_memory*0.7/(n*n*4)/11)
  if device == "cpu":
    maxsize = math.floor(psutil.virtual_memory()[0]*0.7/(n*n*4)/11)
  batchnum = math.ceil(n*(k-1)/maxsize) # Number of batches needed for each iteration below
  print("Number of batches per iteration: ", batchnum)
  while chisq_diff>0:
    zhat_mat = torch.zeros((n*(k-1),n)).to(device) # list all possibilities with only one location changed for each iteration
    all_k = torch.arange(1,k+1).to(device)
    for iter in list(range(1,n*(k-1)+1)):
      r = math.ceil(iter/(k-1))
      s = iter-(r-1)*(k-1)
      zhat = torch.clone(zadj)
      zhat[r-1] = all_k[all_k!=zhat[r-1]][s-1]
      zhat_mat[iter-1] = zhat
    zhat_mat_batch = torch.split(zhat_mat,maxsize) # split all possibilities of labels into batches
    chisq_all = []
    for batch in range(batchnum):
      tsP = torch.stack([sbm_mle(A,item,device)[1] for item in zhat_mat_batch[batch]]) # 3d tensor: last 2d records corresponding P matrix for all possible labels per batch
      chisq_all.append(torch.sum(torch.square(csgc(A,tsP,var_structure,device)),1)) # calculate corresponding chi-square values per batch
      del tsP # delete records to save memory
    chisq_all = torch.cat(chisq_all) # combine all batches of chi-square values for one iteration
    chisq_diff = torch.max(chisqadj-chisq_all) # record the least chi-square value, and calculate the difference with last iteration
    update_ind = torch.argmin(chisq_all) # record the change of location corresponding to the least chi-square value
    chisqadj = torch.min(chisq_all) # update chi-square values
    print("chisq stats: ", chisqadj.item(), ", change location: ", math.ceil((update_ind+1)/(k-1)))
    zadj = zhat_mat[update_ind].to(torch.uint8) # update labels
    # print("adjusted labels: ", zadj)
    chisq_log = torch.cat((chisq_log,torch.tensor([chisqadj]).to(device))) # add to records for chi-square values
    zadj_log = torch.cat((zadj_log,zadj)) # add to records for labels
  zadj_log = zadj_log.reshape(len(chisq_log),n)
  chisq_out = chisq_log[len(chisq_log)-2] # the least chi-square value is the second last in records
  z_out = zadj_log[len(chisq_log)-2].to(torch.uint8) # final labels are the second last row in records
  stats_out = csgc(A,sbm_mle(A,z_out,device)[1],var_structure,device) # calculate the corresponding csgc statistics
  return chisq_out,z_out,stats_out #chisq,z,stats_out as tensor

def wasserstein_uniform(p): # p as numpy/tensor
  u = np.arange(0,1+1/len(p),1/len(p))
  s = np.sort(np.concatenate((p,u)))
  au = np.diff(s)*np.convolve(s,np.ones(2),'valid')/2
  aph = np.full(len(au), np.nan)
  aph[np.nonzero(np.in1d(s,np.sort(np.append(p,0))))[0]] = u
  idx = np.where(~np.isnan(aph),np.arange(np.isnan(aph).size),0)
  np.maximum.accumulate(idx,out=idx)
  aph = aph[idx]
  ap = np.diff(s)*aph
  area = 2*np.sum(np.abs(au-ap))
  return area # area as numpy

def ccr(label_true, label_est): # label_true/label_est as tensor
    k = len(np.unique(label_true))
    n = len(label_true)
    all_perm = list(permutations(range(1, k+1)))
    ccr_all = np.zeros(len(all_perm))
    label_est = np.array(label_est)
    for i, perm in enumerate(all_perm):
        label_est_mapped = np.array([perm[val-1] for val in label_est])
        ccr_all[i] = np.sum(label_est_mapped == np.array(label_true))
    ccr = np.max(ccr_all)/n
    return ccr # ccr as value



##################################################################
# Example for csgc package using GPU
if __name__ == '__main__':
  torch.manual_seed(123)
  # Generate adjacency matrix A and probability matrix P from
  # parameter matrix K and label vector.
  k = 4; n = 200
  z = torch.repeat_interleave(torch.arange(1,k+1),int(n/k)).to(torch.uint8)
  K = torch.tensor([[.8,.5,.1,.1],
            [.5,.1,.1,.1],
            [.1,.1,.8,.5],
            [.1,.1,.5,.1]])
  A,P = gen_adj_sbm(K,z,device="cuda:0")

  # Calculate the maximum likelihood estimated K and P from A
  # and true labels z.
  Khat,Phat = sbm_mle(A,z,device="cuda:0")

  # We can also get the estimated labels given A and number of
  # communities k using spectral clustering method, and compare the
  # estimated labels with true labels.
  zhat = spectral_sbm(A,4,device="cuda:0")
  print("Correct classification rate of spectral clustering: ", ccr(z,zhat), ".\n")

  # If we use the estimated labels to further get the MLE for P,
  # we can calculate the centred subgraph count statistics and check
  # if these statistics deviate from 0 or not.
  Khat2,Phat2 = sbm_mle(A,zhat,device="cuda:0")
  print("A matrix, P matrix:\n", csgc(A,Phat2,device="cuda:0").cpu().numpy(),)

  # We allow both A and P to be 3-dim tensor
  print("A matrix, P tensor:\n", csgc(A,torch.stack((Phat,Phat2)),device="cuda:0"), ".")
  print("A tensor, P matrix:\n", csgc(torch.stack((A,A)),Phat2,device="cuda:0"), ".")
  print("A tensor, P tensor:\n", csgc(torch.stack((A,A)),torch.stack((Phat,Phat2)),device="cuda:0"), ".\n")

  # Sum of squares of these statistics should approximately follow a
  # chi-square distribution if SBM fits the network data. Suppose we
  # run the above simulation multiple times, we can get a batch of
  # p-values, which are supposed to follow a uniform distribution.
  # We can calculate the Wasserstein distance between these p-values
  # and a uniform distribution on [0,1].
  import torch.distributions as dist
  p = torch.zeros(100)
  for i in range(100):
    A1 = gen_adj_sbm(K,z,device="cuda:0")[0]
    t = csgc(A1,sbm_mle(A1,z,device="cuda:0")[1],device="cuda:0")
    p[i] = 1-dist.Chi2(10).cdf(torch.sum(t**2).to("cpu"))
  print("Wasserstein distance between p-values and U[0,1]: ", wasserstein_uniform(p), ".\n")

  # If these csgc statistics deviate from 0, we can use our csgc
  # greedy algorithm to adjust label estimations, and see if correct
  # classification rate increases.
  chisqout,zout,statsout = csgc_greedy(A,zhat,device="cuda:0")
  print("The Chisq value after greedy algorithm: ", chisqout, ".")
  print("The adjusted label vector after greedy algorithm: ", zout, ".")
  print("The csgc statistics after greedy algorithm: ", statsout, ".")
  print("Correct classification rate after greedy algorithm: ", ccr(z,zout.to("cpu")), ".\n")

  # For DCSBM, we can do similar things:
  # Generate A and P.
  Theta = torch.FloatTensor(n).uniform_(0.2,1)
  dA,dP = gen_adj_dcsbm(K,Theta,z,device="cuda:0")
  # Calculate numerical MLE for P (no closed form)
  dPhat = dcsbm_mle(dA,z,device="cuda:0")[2]
  # get spectral clustering estimated labels, and calculate the
  # corrected classification rate.
  dzhat = spectral_dcsbm(dA,4,device="cuda:0")
  print("Correct classification rate of spectral clustering for DCSBM: ", ccr(z,dzhat), ".")