diff --git a/CASE_R_package.Rproj b/CASE_R_package.Rproj
index 21a4da0..1bb667d 100644
--- a/CASE_R_package.Rproj
+++ b/CASE_R_package.Rproj
@@ -1,4 +1,5 @@
 Version: 1.0
+ProjectId: 2241ace3-378b-4e48-bb0a-ef8f1b97ac60
 
 RestoreWorkspace: Default
 SaveWorkspace: Default
diff --git a/DESCRIPTION b/DESCRIPTION
index a479035..3405390 100644
--- a/DESCRIPTION
+++ b/DESCRIPTION
@@ -1,7 +1,7 @@
 Type: Package
 Package: CASE
 Title: Cell-type-specific And Shared EQTL fine-mapping
-Version: 0.3.0
+Version: 0.3.1
 Authors@R: c(
     person("Chen", "Lin", , "c.lin@yale.edu", role = c("aut", "cre"),
            comment = c(ORCID = "0000-0001-9821-2578")),
@@ -18,7 +18,7 @@ Imports:
     mvtnorm,
     stats
 Encoding: UTF-8
-RoxygenNote: 7.3.1
+RoxygenNote: 7.3.2
 Suggests: 
     knitr,
     rmarkdown,
diff --git a/R/CASE.R b/R/CASE.R
index 6cb96e2..7b9e197 100644
--- a/R/CASE.R
+++ b/R/CASE.R
@@ -81,14 +81,30 @@ CASE <- function(Z = NULL, R, hatB = NULL, hatS = NULL, N, V = NULL, cs = TRUE,
     hatS = hatBS$hatS
     
     # V = estimate_null_correlation_simple(mash_set_data(do.call(rbind, raw.data$hatB), do.call(rbind, raw.data$hatS)))
-
-    m1 <- CASE_train(hatB = hatB, hatS = hatS, V = V, R = R, N = N, Z = NULL, verbose = verbose, ...)
+    sump = 2 - 2 * pnorm(abs(hatB / hatS))
+    sumfdr = apply(sump, 2, function(x) p.adjust(x, method = "fdr"))
+    C = ncol(sumfdr)
+    
+    # no strong signals in all cell types, or low SNP numbers 
+    if (sum(sumfdr <= 0.2) == 0){
+      if (verbose){
+        cat("No FDR-significant variants in the inputs.", "\n")
+      }
+      m1 = list(pi = 1, n.iter = 0)
+    } else if (nrow(sump) <= 20){
+      if (verbose){
+        cat("Too few SNPs in the data (<=20).", "\n")
+      }
+      m1 = list(pi = 1, n.iter = 0)
+    } else{
+      m1 <- CASE_train(hatB = hatB, hatS = hatS, V = V, R = R, N = N, Z = NULL, verbose = verbose, ...)
+    }
     
     res <- CASE_test(hatB = hatB, hatS = hatS, R = R, N = N, CASE_training = m1, Z = NULL, verbose = verbose, ...)
     t2 = Sys.time()
     res$time = difftime(t2, t1, units = "secs")
     if (cs){
-      res$sets <- get_credible_sets(res$pip, R = R, verbose = verbose)
+      res$sets <- get_credible_sets(pips = res$pip, R = R, verbose = verbose)
     }
     
     return(res)
diff --git a/R/CASE_models.R b/R/CASE_models.R
index 1cdb8be..fb22f16 100644
--- a/R/CASE_models.R
+++ b/R/CASE_models.R
@@ -80,10 +80,17 @@ CASE_train <- function(Z = NULL, R, hatB = NULL, hatS = NULL, N, V = NULL, verbo
   ## Train with only marginally significant SNPs
   M0 = nrow(R)
   ME_p <- 2 - 2 * pnorm(abs(hatB / hatS), 0, 1)
-  significant_thres = ifelse("significant_thres" %in% names(args), args$significant_thres, 1e-1)
-  idx <- which(ME_p <= significant_thres, arr.ind = TRUE)[, 1] %>% unique
+  significant_thres = ifelse("significant_thres" %in% names(args), 
+                             args$significant_thres, 
+                             ifelse(sqrt(6) / sqrt(diag(N)) >= 0.05, sqrt(6) / sqrt(diag(N)), 0.05))
+  if (C == 1){
+    idx <- which(ME_p <= significant_thres)
+  } else{
+    idx <- which(t(apply(ME_p, 1, function(x) x <= significant_thres)), arr.ind = TRUE)[, 1] %>% unique
+  }
+  
   if (length(idx) == 1){
-    idx = c(idx - 1, idx, idx + 1)
+    idx = idx + (-2):2
     idx = idx[idx >= 1 & idx <= nrow(hatB)]
   }
   hatB = hatB[idx, ]
@@ -96,7 +103,7 @@ CASE_train <- function(Z = NULL, R, hatB = NULL, hatS = NULL, N, V = NULL, verbo
     if (verbose){
       cat("No marginally significant variants in the inputs.", "\n")
     }
-    return(list(pi = 1, U = list(matrix(0, C, C)), V = V, n.iter = 0))
+    return(list(pi = 1, n.iter = 0))
   }
   
   J <- 0
@@ -310,10 +317,6 @@ CASE_test <- function(Z = NULL, R, hatB = NULL, hatS = NULL, N, CASE_training, v
     cat("Start Posterior Analysis.", "\n")
   }
   
-  U = CASE_training$U
-  V = CASE_training$V
-  pi = CASE_training$pi
-  
   if (is.null(Z)){
     Z = hatB / hatS
   }
@@ -322,15 +325,19 @@ CASE_test <- function(Z = NULL, R, hatB = NULL, hatS = NULL, N, CASE_training, v
   hatBS = transform_Z(Z, N)
   hatB = hatBS$hatB
   hatS = hatBS$hatS
-
+  
   C <- ncol(hatB)
   M <- nrow(R)
   
+  pi = CASE_training$pi
   if (length(pi) <= 1){
     post_mean = pip = matrix(0, M, C)
-    return(list(pi = pi, U = U, V = V, pip = pip, post_mean = post_mean))
+    return(list(pip = pip, post_mean = post_mean))
   }
   
+  U = CASE_training$U
+  V = CASE_training$V
+
   L = length(U)
   ## MC step
   gBc = gB_coef(U, V)
diff --git a/data/example_data.rda b/data/example_data.rda
index 5ce2d3d..fb2ab7c 100644
Binary files a/data/example_data.rda and b/data/example_data.rda differ