Skip to content

Commit 75d6e5e

Browse files
feat(bokeh): implement heatmap-correlation
- Add heatmap-correlation specification - Implement correlation matrix heatmap using bokeh rect glyphs - Support diverging color scheme centered at 0 - Add text annotations for correlation values - Include colorbar legend Note: PNG export requires selenium and webdriver. HTML output works without additional dependencies.
1 parent 989c85d commit 75d6e5e

2 files changed

Lines changed: 298 additions & 0 deletions

File tree

Lines changed: 226 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,226 @@
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

specs/heatmap-correlation.md

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,72 @@
1+
# heatmap-correlation
2+
3+
## Title
4+
Correlation Matrix Heatmap
5+
6+
## Description
7+
Create a heatmap visualization showing pairwise correlations between multiple variables. Display correlation coefficients as colors with values annotated in cells.
8+
9+
## Required Parameters
10+
- `data` (pd.DataFrame): Input data with numeric columns for correlation calculation
11+
- `columns` (List[str], optional): Specific columns to include in correlation. If None, use all numeric columns
12+
13+
## Optional Parameters
14+
- `method` (str): Correlation method ('pearson', 'spearman', 'kendall'). Default: 'pearson'
15+
- `annot_format` (str): Format string for annotations. Default: '.2f'
16+
- `cmap` (str): Colormap name. Default: 'RdBu_r' (diverging blue-white-red)
17+
- `vmin` (float): Minimum value for color scale. Default: -1
18+
- `vmax` (float): Maximum value for color scale. Default: 1
19+
- `center` (float): Center value for diverging colormap. Default: 0
20+
- `square` (bool): Make cells square. Default: True
21+
- `linewidths` (float): Width of lines between cells. Default: 0.5
22+
- `linecolor` (str): Color of lines between cells. Default: 'white'
23+
- `cbar_label` (str): Label for colorbar. Default: 'Correlation'
24+
- `title` (str): Title for the plot. Default: None
25+
- `figsize` (Tuple[int, int]): Figure size in inches. Default: (10, 8)
26+
27+
## Returns
28+
Figure object specific to the plotting library
29+
30+
## Data Characteristics
31+
- Input should contain multiple numeric columns
32+
- Columns may have different scales (correlation normalizes)
33+
- Missing values handled via correlation calculation
34+
35+
## Visual Requirements
36+
- Diverging color scheme centered at 0
37+
- Values displayed in each cell
38+
- Clear cell boundaries
39+
- Color bar showing scale
40+
- Variable names as tick labels on both axes
41+
- Square cells for better readability
42+
43+
## Example Use Case
44+
```python
45+
# Create sample data with correlated features
46+
import numpy as np
47+
import pandas as pd
48+
49+
np.random.seed(42)
50+
n = 100
51+
52+
data = pd.DataFrame({
53+
'temperature': np.random.normal(20, 5, n),
54+
'humidity': np.random.normal(60, 10, n),
55+
'pressure': np.random.normal(1013, 20, n),
56+
'wind_speed': np.random.normal(10, 3, n)
57+
})
58+
59+
# Add correlations
60+
data['humidity'] = 100 - 2 * data['temperature'] + np.random.normal(0, 5, n)
61+
data['wind_speed'] = 0.5 * data['temperature'] + np.random.normal(0, 2, n)
62+
63+
# Generate heatmap
64+
fig = create_plot(data)
65+
```
66+
67+
## Implementation Notes
68+
- Calculate correlation matrix from input data
69+
- Handle both full dataset or subset of columns
70+
- Ensure color scale is symmetric around 0
71+
- Annotations should be readable (appropriate font size)
72+
- Consider matrix size for layout (4-10 variables optimal)

0 commit comments

Comments
 (0)