diff --git a/DESCRIPTION b/DESCRIPTION
index 97d2989..3d3cea6 100644
--- a/DESCRIPTION
+++ b/DESCRIPTION
@@ -1,6 +1,6 @@
 Package: dlptools
 Title: DLPTools: Handling DLP+ Data
-Version: 1.1.0
+Version: 1.1.1
 Authors@R: c(
     person("Ben", "Furman", email = "bfurman@bccrc.ca", role = c("aut", "cre")),
     person("Aparicio Lab Past and Present", role = "aut")
diff --git a/R/heatmap_plot.R b/R/heatmap_plot.R
index 25a77c2..bd0447c 100644
--- a/R/heatmap_plot.R
+++ b/R/heatmap_plot.R
@@ -135,6 +135,7 @@ generate_state_hm <- function(
 #'
 #' @param total_hm ComplexHeatmap::Heatmap of the combined tree and states.
 #' @param file_name optional string of where to save a png image of the heatmap.
+#' @param png_height optional height of png file
 #' @return ComplexHeatmap::draw or nothing if a file is written.
 #' @export
 output_hm_image <- function(
@@ -389,7 +390,6 @@ build_left_annot <- function(
       col = anno_cols_list
     )
   }
-
   return(left_annot)
 }
 
@@ -514,7 +514,12 @@ make_clone_palette <- function(clone_ids) {
 #' in the heatmap. Will default to the name of the column being plotted.
 #' @param legend_11plus bool. Default FALSE. For HMMCopy state values,
 #' state 11 is really 11+, so we replace 11s with 11+ for the plot in the
-#' legend.
+#' #' legend.
+#' @param scale_y bool. Default FALSE. Scale the y size of the saved image accoring
+#' to the number of cells in the heatmap
+#' @param cell_height, int. Default = 10 Heigh in pixels per cell if scale_y is True
+#' @param default_png_height, int: Default = 1600. If scale_y not used, height
+#' of final png
 #' @export
 plot_state_hm <- function(
     states_df,
@@ -522,7 +527,7 @@ plot_state_hm <- function(
     phylogeny = NULL,
     file_name = NULL,
     anno_df = NULL, # optional, can also specify columns in the dataframe
-    anno_colors_list = list(), # for custom colors of annotations
+    anno_colors_list = list(), # for custom colors of annotations. Can also specify custom breaks in the cmap
     anno_columns = NULL,
     clone_column = NULL,
     clones_df = NULL, # optional, can also specify columns in the dataframe
@@ -537,6 +542,9 @@ plot_state_hm <- function(
     hm_legend_title = NULL,
     legend_11plus = FALSE,
     color_tree_clones = FALSE,
+    scale_y = FALSE,
+    cell_height = 10, 
+    default_png_height = 1600,
     ...) {
   check_args()
 
@@ -634,9 +642,19 @@ plot_state_hm <- function(
     legend_11plus = legend_11plus
   )
 
+  # If scale_y is true, scale image height to number of cells
+  if (scale_y) {
+    message("Scaling Image Size by number of cells")
+    ncells <- length(unique(states_df$cell_id))
+    png_height <- 200 + ncells*cell_height
+  } else {
+    png_height <- default_png_height
+  }
+
   output_hm_image(
     hm_p,
     file_name = file_name,
+    png_height = png_height,
     ...
   )
 }