Heatmap: Taken create_heatmap() from Jupyter notebook created by Martin Fajcik

analyze_winscore.py

@@ -207,4 +207,113 @@ def create_scatter_plot_with_curve_with_variances_named(category, variance_acros
     
     return p
 
+def create_heatmap(data_matrix, original_scores, selected_rows=None, hide_scores_tasks=[], width=800, height=1400):
+    plot_width = width
+    plot_height = height
+    n_rows, n_cols = data_matrix.shape
+
+    # Clean column names (remove 'benczechmark_' prefix)
+    # data_matrix.columns = data_matrix.columns.str.replace('benczechmark_', '', regex=False)
+    # original_scores.columns = original_scores.columns.str.replace('benczechmark_', '', regex=False)
+
+    if selected_rows is not None:
+        # Select only the specified rows (models)
+        data_matrix = data_matrix[selected_rows]
+        original_scores = original_scores[selected_rows]
+
+    # Set up the figure with tasks as x-axis and models as y-axis
+    p = figure(
+        width=plot_width, height=plot_height, 
+        x_range=list(data_matrix.index), y_range=list(data_matrix.columns), 
+        toolbar_location="below", tools="pan,wheel_zoom,box_zoom,reset,save", 
+        x_axis_label="Model", y_axis_label="Category"
+    )
+
+    # Create the color mapper for the heatmap
+    color_mapper = LinearColorMapper(palette='Viridis256', low=0, high=1)  # Light for low values, dark for high
+
+    # Flatten the matrix for Bokeh plotting
+    heatmap_data = {
+        'x': [],
+        'y': [],
+        'colors': [],
+        'model_names': [],  # Updated: Reflects model names now
+        'scores': [],
+    }
+    label_data = {
+        'x': [],
+        'y': [],
+        'value': [],
+        'text_color': [],  # New field for label text colors
+    }
+    
+    # Iterate through the data_matrix to populate heatmap and label data
+    for row_idx, (model_name, task_scores) in enumerate(data_matrix.iterrows()):
+        for col_idx, score in enumerate(task_scores):
+            heatmap_data['x'].append(model_name)  # Model goes to x-axis
+            heatmap_data['y'].append(data_matrix.columns[col_idx])  # Task goes to y-axis
+            heatmap_data['colors'].append(score)
+            heatmap_data['model_names'].append(model_name)  # Model names added to hover info
+
+            # Get the original score
+            original_score = original_scores.loc[model_name, data_matrix.columns[col_idx]]
+            plot_score = data_matrix.loc[model_name, data_matrix.columns[col_idx]]
+            heatmap_data['scores'].append(original_score)
+            task_name = data_matrix.columns[col_idx]
+
+            if task_name not in hide_scores_tasks:
+                label_data['x'].append(model_name)
+                label_data['y'].append(task_name)
+                label_data['value'].append(round(original_score))  # Round the score
+
+                # Determine text color based on score
+                if plot_score <= 0.6:  # Threshold for light/dark text
+                    label_data['text_color'].append('white')  # Light color for lower scores
+                else:
+                    label_data['text_color'].append('black')  # Dark color for higher scores
+
+    heatmap_source = ColumnDataSource(heatmap_data)
+    label_source = ColumnDataSource(label_data)
+
+    # Create the heatmap
+    p.rect(x='x', y='y', width=1, height=1, source=heatmap_source,
+           line_color=None, fill_color={'field': 'colors', 'transform': color_mapper})
+
+    # Add color bar
+    # Add color bar with custom ticks
+    color_bar = ColorBar(
+        color_mapper=color_mapper, 
+        width=8, location=(0, 0),
+        ticker=FixedTicker(ticks=[0, 0.2, 0.4, 0.6, 0.8, 1]),  # Fixed ticks at 0, 20, 40, 60, 80, 100
+        major_label_overrides={0: '0', 0.2: '20', 0.4: '40', 0.6: '60', 0.8: '80', 1: '100'}  # Custom labels for ticks
+    )
+    #p.add_layout(color_bar, 'right')
+
+    # Add HoverTool for interactivity
+    hover = HoverTool()
+    hover.tooltips = [("Model", "@x"), ("Task", "@y"), ("DS", "@scores")]  # Updated tooltip
+    p.add_tools(hover)
+
+    # Add labels with dynamic text color
+    labels = LabelSet(x='x', y='y', text='value', source=label_source,
+                      text_color='text_color', text_align='center', text_baseline='middle')
+    p.add_layout(labels)
+
+    # Customize the plot appearance
+    p.xgrid.grid_line_color = None
+    p.ygrid.grid_line_color = None
+    p.xaxis.major_label_orientation = "vertical"
+    p.yaxis.major_label_text_font_size = "13pt"
+    p.xaxis.major_label_text_font_size = "13pt"
+
+        # Set the axis label font size
+    p.xaxis.axis_label_text_font_size = "18pt"  # Set font size for x-axis label
+    p.yaxis.axis_label_text_font_size = "18pt"  # Set font size for y-axis label
+    p.xaxis.axis_label_text_font_style = "normal"  # Set x-axis label to normal
+    p.yaxis.axis_label_text_font_style = "normal"  # Set y-axis label to normal
+    
+    #p.yaxis.visible = False  # Hide the y-axis labels
+
+    return p
+
 # EOF