R/hetset_methods.R

In ZytoHMGU/hetset: Identification of Heterogeneous Subsets in Data

# hetset-methods ========================

estimate_densities <- function(H){
  if (class(H)!="SummarizedExperiment"){
    stop(" object of class SE required, try hetset() to convert object /n")
  }
  if(is.null(H@metadata$slf)){
    H@metadata$slf <- NA
  }
  n_dim <- sum(!is.na(H@metadata$slf))
  S_assigned <- (sum(H$prt=="A",na.rm = T)>2) & (sum(H$prt=="B",na.rm = T)>2)
  NA.flag <- F
  if (sum(is.na(H$prt))>0){
    H$prt <- factor(H$prt,levels = c("A","B","NA"))
    H$prt[is.na(H$prt)] <- "NA"
    NA.flag <- T
  }
  if (n_dim < 1){
    print(" warning: no features selected ")
    H@metadata$prm.full$mean <- rowMeans(assays(H)[[1]],na.rm = T)
    H@metadata$prm.full$cov <- cov(t(assays(H)[[1]]),use = "complete.obs")
    # estimate parameters for subpopulations --------------------------------
    if (S_assigned){
      H@metadata$prm.A$mean <- rowMeans(assays(H[,H$prt=="A"])[[1]],na.rm = T)
      H@metadata$prm.A$cov <- cov(t(assays(H[,H$prt=="A"])[[1]]),
                                  use = "complete.obs")
      H@metadata$prm.B$mean <- rowMeans(assays(H[,H$prt=="B"])[[1]],na.rm = T)
      H@metadata$prm.B$cov <- cov(t(assays(H[,H$prt=="B"])[[1]]),
                                  use = "complete.obs")
      H@metadata$prp.A <- mean(H$prt=="A",na.rm = T)
    }
  } else {
    H@metadata$prm.full$mean <- rowMeans(assays(H[H@metadata$slf,])[[1]],
                                         na.rm = T)
    H@metadata$prm.full$cov <- cov(t(assays(H[H@metadata$slf,])[[1]]),
                                   use = "complete.obs")
    # estimate parameters for subpopulations --------------------------------
    if (S_assigned){
      levels(H$prt) <- c("A","B","NA")
      H@metadata$prm.A$mean <- rowMeans(assays(H[H@metadata$slf,
                                                 H$prt=="A"])[[1]],na.rm = T)
      H@metadata$prm.A$cov <- cov(t(assays(H[H@metadata$slf,H$prt=="A"])[[1]]),
                                  use = "complete.obs")
      H@metadata$prm.B$mean <- rowMeans(assays(H[H@metadata$slf,
                                                 H$prt=="B"])[[1]],na.rm = T)
      H@metadata$prm.B$cov <- cov(t(assays(H[H@metadata$slf,H$prt=="B"])[[1]]),
                                  use = "complete.obs")
      H@metadata$prp.A <- mean(H$prt=="A",na.rm = T)
      H$prt[H$prt == "NA"] <- NA
      H$prt <- droplevels(H$prt)
    }
  }
  if (NA.flag){
    H$prt[H$prt=="NA"] <- NA
  }
  if (S_assigned){ # prevent errors caused by atomary subdistributions
    n_features <- length(H@metadata$slf)
    if (n_features == 1){
      H@metadata$prm.A$cov <- (max(sqrt(H@metadata$prm.A$cov),
                                   sqrt(H@metadata$prm.full$cov)/1000))^2
      H@metadata$prm.B$cov <- (max(sqrt(H@metadata$prm.B$cov),
                                   sqrt(H@metadata$prm.full$cov)/1000))^2
    } else {
      for (i in 1:n_features){
        H@metadata$prm.A$cov[i,i] <- (max(sqrt(H@metadata$prm.A$cov[i,i]),
                                     sqrt(H@metadata$prm.full$cov[i,i])/1000))^2
        H@metadata$prm.B$cov[i,i] <- (max(sqrt(H@metadata$prm.B$cov[i,i]),
                                     sqrt(H@metadata$prm.full$cov[i,i])/1000))^2
      }
    }
  }
  H
}

reassign_samples <- function(H){
  if(is.null(H@metadata$slf)){
    H@metadata$slf <- NA
  }
  if (sum(!is.na(H@metadata$slf)) < 2){
    d_a <- dnorm(x = t(assays(H[H@metadata$slf,])[[1]]),
                 mean = H@metadata$prm.A$mean,
                 sd = sqrt(H@metadata$prm.A$cov))
    d_b <- dnorm(x = t(assays(H[H@metadata$slf,])[[1]]),
                 mean = H@metadata$prm.B$mean,
                 sd = sqrt(H@metadata$prm.B$cov))
  } else {
    d_a <- mvtnorm::dmvnorm(x = t(assays(H[H@metadata$slf,])[[1]]),
                            mean = H@metadata$prm.A$mean,
                            sigma = H@metadata$prm.A$cov)
    d_b <- mvtnorm::dmvnorm(x = t(assays(H[H@metadata$slf,])[[1]]),
                            mean = H@metadata$prm.B$mean,
                            sigma = H@metadata$prm.B$cov)
  }
  # avoid numeric instability due to atomic distributions
  prob_b <- d_b/(d_a+d_b)
  H$prt <- factor(c("A","B")[1 + rbinom(n = length(prob_b),size = 1,
                                        prob = prob_b)],levels = c("A","B"))
  H
}

calculate_dist <- function(H){
  P <- H@metadata
  n_dim <- sum(!is.na(P$slf))
  if (n_dim==1){
    H@metadata$sqHell <- 1 -
      (P$prm.A$cov^(1/4)*P$prm.B$cov^(1/4)) /
      (((P$prm.A$cov+P$prm.B$cov)/2)^(1/2)) *
      exp(-1/4*(P$prm.A$mean-P$prm.B$mean)^2 /
            (P$prm.A$cov+P$prm.B$cov))
  } else {
    H@metadata$sqHell <- 1 - (det(P$prm.A$cov)^(1/4)*det(P$prm.B$cov)^(1/4)) /
      (det((P$prm.A$cov+P$prm.B$cov)/2)^(1/2)) *
      exp(-1/8*t(P$prm.A$mean-P$prm.B$mean) %*%
            matlib::inv((P$prm.A$cov+P$prm.B$cov)/2) %*%
            (P$prm.A$mean-P$prm.B$mean))
  }
  H
}

censor_data <- function(H,k = 4){
  M <- assays(H)[[1]]
  W <- apply(X = M, MARGIN=1, FUN = censor_feature, k = k)
  assays(H)[[1]] <- t(W)
  H
}

censor_feature <- function(x,k = 4){
  m <- mean(x,na.rm = T)
  s <- sd(x,na.rm = T)
  x[((x-m)/s) < (-k)] <- m-k*s
  x[((x-m)/s) > (k)] <- m+k*s
  x
}