diff --git a/R/contrasts.R b/R/contrasts.R
index 996ebc0..99e264d 100644
--- a/R/contrasts.R
+++ b/R/contrasts.R
@@ -204,6 +204,7 @@ estimate_abundances <- function(ccm, newdata, min_log_abund=-5, cell_group="cell
     
     # vcov_type <- grep('vcov', names(attributes(vhat_coef)), value=TRUE)
     v_hat <- ccm@vhat
+    v_hat[is.na(v_hat)] <- 0
     v_hat_method <- ccm@vhat_method
     
     se_fit = sqrt(Matrix::diag(as.matrix(X %*% v_hat %*% Matrix::t(X))))
diff --git a/R/deg.R b/R/deg.R
index 05dd161..cbc4724 100755
--- a/R/deg.R
+++ b/R/deg.R
@@ -78,6 +78,10 @@ pseudobulk_ccs_for_states <- function(ccs,
       gene_agg_fun = gene_agg_fun
     )
 
+    agg_expr_mat = agg_expr_mat[,Matrix::colSums(agg_expr_mat) > 0]
+    agg_coldata = agg_coldata %>%
+      filter(pseudobulk_id %in% colnames(agg_expr_mat))
+
     agg_rowdata = agg_rowdata[row.names(agg_expr_mat),]
   }
 
diff --git a/R/edges.R b/R/edges.R
index 006275b..eff922b 100644
--- a/R/edges.R
+++ b/R/edges.R
@@ -121,7 +121,31 @@ get_paga_graph <- function(cds, reduction_method = "UMAP", partition_q_value=0.0
 #'
 #' @export
 initial_pcor_graph = function(ccs) {
-  paga_edges = get_paga_graph(ccs@cds) %>% igraph::as_data_frame() %>% as_tibble()
+  
+  cds = ccs@cds
+  paga_graph = get_paga_graph(cds) 
+  
+  # if the ccs isnt cluster based -- contract the cluster-based paga_graph 
+  # to be cell_group based
+  if (ccs@info$cell_group != "cell_state") {
+    
+    colData(cds)$cell_state = monocle3:::clusters(cds)
+    
+    cs_ccs = new_cell_count_set(cds, 
+                             cell_group = "cell_state", 
+                             sample_group = "embryo")
+    paga_graph = paga_graph %>% 
+      igraph::as_data_frame() %>% 
+      filter(from %in% rownames(cs_ccs), to %in% rownames(cs_ccs)) %>% 
+      igraph::graph_from_data_frame()
+    
+    paga_graph = contract_state_graph(cs_ccs, paga_graph, group_nodes_by = ccs@info$cell_group)
+    
+  }
+  
+  paga_edges = paga_graph %>% igraph::as_data_frame() %>% as_tibble()
+  
+  
 
   # filter out values that aren't in the cds anymore
   cell_groups = unique(colData(ccs@cds)[[ccs@info$cell_group]])
diff --git a/R/plotting.R b/R/plotting.R
index d55f6cc..31088a7 100755
--- a/R/plotting.R
+++ b/R/plotting.R
@@ -37,6 +37,8 @@ plot_abundance = function(ccs,
 
   plot_df = ccs@metadata[["cell_group_assignments"]] %>% dplyr::select(cell_group)
   plot_df$cell = row.names(plot_df)
+  
+  plot_df = plot_df[rownames(reducedDim(ccs@cds, type="UMAP")),]
 
   plot_df$umap2D_1 <- reducedDim(ccs@cds, type="UMAP")[plot_df$cell,x]
   plot_df$umap2D_2 <- reducedDim(ccs@cds, type="UMAP")[plot_df$cell,y]
diff --git a/R/utils.R b/R/utils.R
index a718ddb..4938cf3 100755
--- a/R/utils.R
+++ b/R/utils.R
@@ -64,7 +64,7 @@ get_distances <- function(ccs, method="euclidean", matrix=T) {
 #' @param group_cells_by
 #' @noRd
 #'
-aggregated_expr_data <- function(cds, group_cells_by = "cell_type_broad"){
+aggregated_expr_data <- function(cds, group_cells_by = "cell_type_broad", gene_group_df=NULL){
 
   cds = cds[, !is.na(colData(cds)$timepoint)]
   cds = cds[, !is.na(colData(cds)[[group_cells_by]])]
@@ -77,6 +77,7 @@ aggregated_expr_data <- function(cds, group_cells_by = "cell_type_broad"){
   cell_group_df$cell_group <- as.character(cell_group_df$cell_group)
   cluster_binary_exprs = as.matrix(aggregate_gene_expression(cds,
                                                              cell_group_df = cell_group_df,
+                                                             gene_group_df = gene_group_df,
                                                              norm_method = "binary",
                                                              scale_agg_values=FALSE))
 
@@ -87,8 +88,10 @@ aggregated_expr_data <- function(cds, group_cells_by = "cell_type_broad"){
 
   cluster_mean_exprs = as.matrix(aggregate_gene_expression(cds,
                                                            cell_group_df = cell_group_df,
+                                                           gene_group_df= gene_group_df,
                                                            norm_method = "size_only",
                                                            scale_agg_values=FALSE))
+  rownames(cluster_mean_exprs) = rownames(cluster_binary_exprs)
 
   cluster_expr_table = tibble::rownames_to_column(as.data.frame(cluster_mean_exprs))
   cluster_expr_table = tidyr::gather(cluster_expr_table, "cell_group",