add option to show heatmap text

Author: dengzq1234

treeprofiler/layouts/staple_layouts.py

@@ -264,7 +264,7 @@ class LayoutHeatmap(TreeLayout):
     def __init__(self, name=None, column=0, width=70, height=None, 
             padding_x=1, padding_y=0, prop=None, internal_rep=None,
             value_color=None, value_range=[], color_range=None, minval=0, maxval=None, 
-            absence_color="#EBEBEB",
+            absence_color="#EBEBEB", show_text=False,
             legend=True):
 
         super().__init__(name)
@@ -284,6 +284,7 @@ def __init__(self, name=None, column=0, width=70, height=None,
         self.height = height
         self.padding_x = padding_x
         self.padding_y = padding_y
+        self.show_text = show_text
 
     def set_tree_style(self, tree, tree_style):
         super().set_tree_style(tree, tree_style)
@@ -305,9 +306,10 @@ def set_tree_style(self, tree, tree_style):
                                     )
     def set_node_style(self, node):
         heatmap_num = node.props.get(self.prop)
-
+        
         if heatmap_num is not None and heatmap_num != 'NaN':
             heatmap_num = float(heatmap_num)
+            text = "%.2f" % (heatmap_num) if self.show_text else ''
             if node.is_leaf:
                 # heatmap
                 tooltip = ""
@@ -318,13 +320,14 @@ def set_node_style(self, node):
 
                 gradient_color = self.value_color.get(heatmap_num)
                 if gradient_color:
-                    identF = RectFace(width=self.width, height=self.height, 
+                    identF = RectFace(width=self.width, height=self.height, text=text,
                     color=gradient_color, padding_x=self.padding_x, padding_y=self.padding_y, tooltip=tooltip)
                     node.add_face(identF, column = self.column,  position = 'aligned')
 
         elif node.is_leaf and node.props.get(self.internal_prop) is not None:
             heatmap_num = node.props.get(self.internal_prop)
             heatmap_num = float(heatmap_num)
+            text = "%.2f" % (heatmap_num) if self.show_text else ''
             # heatmap
             tooltip = ""
             if node.name:
@@ -335,13 +338,14 @@ def set_node_style(self, node):
             gradient_color = self.value_color.get(heatmap_num)
 
             if gradient_color:
-                identF = RectFace(width=self.width, height=self.height, 
+                identF = RectFace(width=self.width, height=self.height, text=text,
                 color=gradient_color, padding_x=self.padding_x, padding_y=self.padding_y, tooltip=tooltip)
                 node.add_face(identF, column = self.column,  position = 'aligned', collapsed_only=True)
 
         elif node.props.get(self.internal_prop) is not None:
             heatmap_num = node.props.get(self.internal_prop)
             heatmap_num = float(heatmap_num)
+            text = "%.2f" % (heatmap_num) if self.show_text else ''
             # heatmap
             tooltip = ""
             if node.name:
@@ -352,7 +356,7 @@ def set_node_style(self, node):
             gradient_color = self.value_color.get(heatmap_num)
 
             if gradient_color:
-                identF = RectFace(width=self.width, height=self.height,
+                identF = RectFace(width=self.width, height=self.height, text=text,
                 color=gradient_color, padding_x=self.padding_x, padding_y=self.padding_y, tooltip=tooltip)
                 node.add_face(identF, column = self.column,  position = 'aligned', collapsed_only=True)
 

treeprofiler/tree_plot.py

@@ -261,6 +261,11 @@ def poplulate_plot_args(plot_args_p):
         nargs='+',
         required=False,
         help="<prop1> <prop2> names of numerical properties which need to be read as heatmap-layout")
+    group.add_argument('--show-heatmap-text',
+        default=False,
+        action='store_true',
+        help="Show gradient in heatmap-layout")
+    
     group.add_argument('--barplot-layout',
         nargs='+',
         required=False,
@@ -514,7 +519,7 @@ def run(args):
             heatmap_layouts, level = get_heatmap_layouts(tree, 
             args.heatmap_layout, level, column_width=args.column_width, 
             padding_x=args.padding_x, padding_y=args.padding_y, 
-            internal_rep=internal_num_rep, color_config=color_config, norm_method='min-max')
+            internal_rep=internal_num_rep, color_config=color_config, norm_method='min-max', show_text=args.show_heatmap_text)
             layouts.extend(heatmap_layouts)
             for prop in args.heatmap_layout:
                 visualized_props.append(prop)
@@ -524,7 +529,7 @@ def run(args):
             heatmap_mean_layouts, level = get_heatmap_layouts(tree, 
             args.heatmap_mean_layout, level, column_width=args.column_width, 
             padding_x=args.padding_x, padding_y=args.padding_y, 
-            internal_rep=internal_num_rep, color_config=color_config, norm_method='mean')
+            internal_rep=internal_num_rep, color_config=color_config, norm_method='mean', show_text=args.show_heatmap_text)
             layouts.extend(heatmap_mean_layouts)
             for prop in args.heatmap_mean_layout:
                 visualized_props.append(prop)
@@ -534,7 +539,7 @@ def run(args):
             heatmap_zscore_layouts, level = get_heatmap_layouts(tree, 
             args.heatmap_zscore_layout, level, column_width=args.column_width, 
             padding_x=args.padding_x, padding_y=args.padding_y, 
-            internal_rep=internal_num_rep, color_config=color_config, norm_method='zscore')
+            internal_rep=internal_num_rep, color_config=color_config, norm_method='zscore', show_text=args.show_heatmap_text)
             layouts.extend(heatmap_zscore_layouts)
             for prop in args.heatmap_zscore_layout:
                 visualized_props.append(prop)
@@ -1856,7 +1861,7 @@ def process_prop_values(tree, prop):
 
     return layouts, level, prop_color_dict
 
-def get_heatmap_layouts(tree, props, level, column_width=70, padding_x=1, padding_y=0, internal_rep='avg', color_config=None, norm_method='min-max', global_scaling=True):
+def get_heatmap_layouts(tree, props, level, column_width=70, padding_x=1, padding_y=0, internal_rep='avg', color_config=None, norm_method='min-max', show_text=False, global_scaling=True):
     # Helper functions for normalization
     def min_max_normalize(value, minval, maxval):
         return 0 if maxval - minval == 0 else (value - minval) / (maxval - minval)
@@ -1958,15 +1963,15 @@ def parse_color_config(prop, color_config, minval, maxval):
                 index = np.abs(index_values - normalized_value).argmin() + 1
                 if value not in value2color:  # Ensure color is not overwritten
                     value2color[value] = gradientscolor.get(index, "")
-
+        
         # Add layout for the current property
         layout = staple_layouts.LayoutHeatmap(
             name=f'Heatmap_{prop}_{norm_method}', column=level,
             width=column_width, padding_x=padding_x, padding_y=padding_y,
             internal_rep=internal_rep, prop=prop,
             maxval=maxval, minval=minval, value_color=value2color,
             value_range=[minval, maxval], color_range=gradientscolor,
-            absence_color=nan_color
+            absence_color=nan_color, show_text=show_text,
         )
         layouts.append(layout)
         level += 1