1+ """
2+ heatmap-correlation: Correlation Matrix Heatmap
3+ Library: bokeh
4+ """
5+
6+ import numpy as np
7+ import pandas as pd
8+ from bokeh .plotting import figure
9+ from bokeh .models import ColumnDataSource , LinearColorMapper , ColorBar , BasicTicker
10+ from bokeh .models import Label
11+ from bokeh .io import export_png
12+ from bokeh .transform import transform
13+ from bokeh .palettes import RdBu11
14+ from typing import List , Optional , Tuple , TYPE_CHECKING
15+
16+ if TYPE_CHECKING :
17+ from bokeh .plotting import figure as FigureType
18+
19+
20+ def create_plot (
21+ data : pd .DataFrame ,
22+ columns : Optional [List [str ]] = None ,
23+ method : str = 'pearson' ,
24+ annot_format : str = '.2f' ,
25+ cmap : str = 'RdBu_r' ,
26+ vmin : float = - 1 ,
27+ vmax : float = 1 ,
28+ center : float = 0 ,
29+ square : bool = True ,
30+ linewidths : float = 0.5 ,
31+ linecolor : str = 'white' ,
32+ cbar_label : str = 'Correlation' ,
33+ title : Optional [str ] = None ,
34+ figsize : Tuple [int , int ] = (10 , 8 ),
35+ ** kwargs
36+ ) -> 'FigureType' :
37+ """
38+ Create a heatmap visualization showing pairwise correlations between multiple variables.
39+
40+ Args:
41+ data: Input DataFrame
42+ columns: Specific columns to include in correlation. If None, use all numeric columns
43+ method: Correlation method ('pearson', 'spearman', 'kendall'). Default: 'pearson'
44+ annot_format: Format string for annotations. Default: '.2f'
45+ cmap: Colormap name. Default: 'RdBu_r' (diverging blue-white-red)
46+ vmin: Minimum value for color scale. Default: -1
47+ vmax: Maximum value for color scale. Default: 1
48+ center: Center value for diverging colormap. Default: 0
49+ square: Make cells square. Default: True
50+ linewidths: Width of lines between cells. Default: 0.5
51+ linecolor: Color of lines between cells. Default: 'white'
52+ cbar_label: Label for colorbar. Default: 'Correlation'
53+ title: Title for the plot. Default: None
54+ figsize: Figure size in inches. Default: (10, 8)
55+ **kwargs: Additional parameters
56+
57+ Returns:
58+ bokeh Figure object
59+
60+ Raises:
61+ ValueError: If data is empty
62+ KeyError: If specified columns not found
63+
64+ Example:
65+ >>> data = pd.DataFrame({'A': [1, 2, 3], 'B': [2, 4, 6]})
66+ >>> fig = create_plot(data)
67+ """
68+ # Input validation
69+ if data .empty :
70+ raise ValueError ("Data cannot be empty" )
71+
72+ # Select columns for correlation
73+ if columns is not None :
74+ for col in columns :
75+ if col not in data .columns :
76+ available = ", " .join (data .columns )
77+ raise KeyError (f"Column '{ col } ' not found. Available: { available } " )
78+ corr_data = data [columns ]
79+ else :
80+ # Use all numeric columns
81+ corr_data = data .select_dtypes (include = [np .number ])
82+
83+ # Calculate correlation matrix
84+ corr_matrix = corr_data .corr (method = method )
85+
86+ # Prepare data for bokeh
87+ var_names = list (corr_matrix .columns )
88+ n_vars = len (var_names )
89+
90+ # Create x and y coordinates for rectangles
91+ x_coords = []
92+ y_coords = []
93+ colors = []
94+ values = []
95+
96+ for i , row_name in enumerate (var_names ):
97+ for j , col_name in enumerate (var_names ):
98+ x_coords .append (col_name )
99+ y_coords .append (row_name )
100+ value = corr_matrix .iloc [i , j ]
101+ values .append (value )
102+
103+ # Create ColumnDataSource
104+ source = ColumnDataSource (data = {
105+ 'x' : x_coords ,
106+ 'y' : y_coords ,
107+ 'value' : values ,
108+ 'formatted' : [annot_format .format (v ) for v in values ]
109+ })
110+
111+ # Convert figsize from inches to pixels (assuming 100 dpi for display)
112+ width_px = int (figsize [0 ] * 160 ) # 16:9 aspect ratio adjustment
113+ height_px = int (figsize [1 ] * 100 )
114+
115+ # Create figure
116+ p = figure (
117+ width = width_px ,
118+ height = height_px ,
119+ title = title ,
120+ x_range = var_names ,
121+ y_range = list (reversed (var_names )), # Reverse to match traditional matrix display
122+ toolbar_location = "right" ,
123+ tools = "hover,save,pan,box_zoom,reset,wheel_zoom"
124+ )
125+
126+ # Create color mapper (using RdBu reversed palette)
127+ mapper = LinearColorMapper (
128+ palette = list (reversed (RdBu11 )),
129+ low = vmin ,
130+ high = vmax
131+ )
132+
133+ # Add rectangles for heatmap
134+ p .rect (
135+ x = 'x' ,
136+ y = 'y' ,
137+ width = 1 ,
138+ height = 1 ,
139+ source = source ,
140+ fill_color = transform ('value' , mapper ),
141+ line_color = linecolor ,
142+ line_width = linewidths
143+ )
144+
145+ # Add text annotations
146+ from bokeh .models import Text
147+ text_glyph = Text (
148+ x = 'x' ,
149+ y = 'y' ,
150+ text = 'formatted' ,
151+ text_align = 'center' ,
152+ text_baseline = 'middle' ,
153+ text_font_size = '10pt' ,
154+ text_color = 'black'
155+ )
156+ p .add_glyph (source , text_glyph )
157+
158+ # Add color bar
159+ color_bar = ColorBar (
160+ color_mapper = mapper ,
161+ ticker = BasicTicker (),
162+ label_standoff = 12 ,
163+ border_line_color = None ,
164+ location = (0 , 0 ),
165+ title = cbar_label ,
166+ title_text_font_size = '10pt'
167+ )
168+ p .add_layout (color_bar , 'right' )
169+
170+ # Style the plot
171+ p .grid .visible = False
172+ p .axis .axis_line_color = None
173+ p .axis .major_tick_line_color = None
174+ p .axis .minor_tick_line_color = None
175+ p .xaxis .major_label_orientation = np .pi / 4
176+ p .xaxis .axis_label = None
177+ p .yaxis .axis_label = None
178+
179+ # Configure hover tool
180+ p .hover .tooltips = [
181+ ('Variables' , '@x, @y' ),
182+ ('Correlation' , '@value{0.000}' )
183+ ]
184+
185+ return p
186+
187+
188+ if __name__ == '__main__' :
189+ # Sample data for testing
190+ np .random .seed (42 )
191+ n = 100
192+
193+ data = pd .DataFrame ({
194+ 'temperature' : np .random .normal (20 , 5 , n ),
195+ 'humidity' : np .random .normal (60 , 10 , n ),
196+ 'pressure' : np .random .normal (1013 , 20 , n ),
197+ 'wind_speed' : np .random .normal (10 , 3 , n )
198+ })
199+
200+ # Add correlations
201+ data ['humidity' ] = 100 - 2 * data ['temperature' ] + np .random .normal (0 , 5 , n )
202+ data ['wind_speed' ] = 0.5 * data ['temperature' ] + np .random .normal (0 , 2 , n )
203+
204+ # Create plot
205+ fig = create_plot (
206+ data ,
207+ title = 'Correlation Matrix Heatmap'
208+ )
209+
210+ # Save - ALWAYS use 'plot.png'!
211+ # Note: bokeh PNG export requires selenium and a webdriver
212+ # In CI/production environments, these should be pre-installed
213+ from bokeh .io import output_file , save
214+
215+ # For development/testing: save HTML first
216+ output_file ('plot.html' )
217+ save (fig )
218+
219+ try :
220+ # Try PNG export if dependencies are available
221+ export_png (fig , filename = 'plot.png' )
222+ print ("Plot saved to plot.png" )
223+ except (ImportError , RuntimeError ) as e :
224+ print ("Note: PNG export requires selenium and a webdriver." )
225+ print ("HTML version saved to plot.html" )
226+ # In CI, this will be handled by the workflow
0 commit comments