55Python: 3.10+
66"""
77
8+ from typing import TYPE_CHECKING , Optional
9+
810import altair as alt
9- import pandas as pd
1011import numpy as np
11- from typing import TYPE_CHECKING , Optional
12+ import pandas as pd
13+
1214
1315if TYPE_CHECKING :
1416 from altair import Chart
@@ -21,10 +23,10 @@ def create_plot(
2123 title : Optional [str ] = None ,
2224 xlabel : Optional [str ] = None ,
2325 ylabel : Optional [str ] = None ,
24- color_scheme : str = ' set2' ,
26+ color_scheme : str = " set2" ,
2527 width : int = 600 ,
2628 height : int = 400 ,
27- ** kwargs
29+ ** kwargs ,
2830) -> Chart :
2931 """
3032 Create a basic box plot showing statistical distribution of multiple groups using altair.
@@ -66,92 +68,72 @@ def create_plot(
6668 raise KeyError (f"Column '{ col } ' not found. Available columns: { available } " )
6769
6870 # Create the box plot using Altair's mark_boxplot
69- base = alt .Chart (data ).mark_boxplot (
70- extent = 1.5 , # 1.5 * IQR for whiskers
71- outliers = True ,
72- size = 40 ,
73- opacity = 0.7
74- ).encode (
75- x = alt .X (
76- f'{ groups } :N' ,
77- title = xlabel or groups ,
78- axis = alt .Axis (
79- labelAngle = 0 if data [groups ].nunique () <= 5 else - 45 ,
80- labelLimit = 200
81- )
82- ),
83- y = alt .Y (
84- f'{ values } :Q' ,
85- title = ylabel or values ,
86- scale = alt .Scale (zero = False )
87- ),
88- color = alt .Color (
89- f'{ groups } :N' ,
90- scale = alt .Scale (scheme = color_scheme ),
91- legend = None # Hide legend as it's redundant with x-axis
92- ),
93- tooltip = [
94- alt .Tooltip (f'{ groups } :N' , title = 'Group' ),
95- alt .Tooltip (f'count({ values } ):Q' , title = 'Count' ),
96- alt .Tooltip (f'min({ values } ):Q' , title = 'Min' , format = '.2f' ),
97- alt .Tooltip (f'q1({ values } ):Q' , title = 'Q1' , format = '.2f' ),
98- alt .Tooltip (f'median({ values } ):Q' , title = 'Median' , format = '.2f' ),
99- alt .Tooltip (f'q3({ values } ):Q' , title = 'Q3' , format = '.2f' ),
100- alt .Tooltip (f'max({ values } ):Q' , title = 'Max' , format = '.2f' )
101- ]
71+ base = (
72+ alt .Chart (data )
73+ .mark_boxplot (
74+ extent = 1.5 , # 1.5 * IQR for whiskers
75+ outliers = True ,
76+ size = 40 ,
77+ opacity = 0.7 ,
78+ )
79+ .encode (
80+ x = alt .X (
81+ f"{ groups } :N" ,
82+ title = xlabel or groups ,
83+ axis = alt .Axis (labelAngle = 0 if data [groups ].nunique () <= 5 else - 45 , labelLimit = 200 ),
84+ ),
85+ y = alt .Y (f"{ values } :Q" , title = ylabel or values , scale = alt .Scale (zero = False )),
86+ color = alt .Color (
87+ f"{ groups } :N" ,
88+ scale = alt .Scale (scheme = color_scheme ),
89+ legend = None , # Hide legend as it's redundant with x-axis
90+ ),
91+ tooltip = [
92+ alt .Tooltip (f"{ groups } :N" , title = "Group" ),
93+ alt .Tooltip (f"count({ values } ):Q" , title = "Count" ),
94+ alt .Tooltip (f"min({ values } ):Q" , title = "Min" , format = ".2f" ),
95+ alt .Tooltip (f"q1({ values } ):Q" , title = "Q1" , format = ".2f" ),
96+ alt .Tooltip (f"median({ values } ):Q" , title = "Median" , format = ".2f" ),
97+ alt .Tooltip (f"q3({ values } ):Q" , title = "Q3" , format = ".2f" ),
98+ alt .Tooltip (f"max({ values } ):Q" , title = "Max" , format = ".2f" ),
99+ ],
100+ )
102101 )
103102
104103 # Add sample size annotations
105- text = alt .Chart (data ).mark_text (
106- align = 'center' ,
107- baseline = 'top' ,
108- dy = 10 ,
109- fontSize = 10 ,
110- opacity = 0.7
111- ).encode (
112- x = alt .X (f'{ groups } :N' ),
113- y = alt .Y (f'min({ values } ):Q' ),
114- text = alt .Text ('count():Q' , format = 'd' )
115- ).transform_aggregate (
116- count = 'count()' ,
117- groupby = [groups ]
104+ text = (
105+ alt .Chart (data )
106+ .mark_text (align = "center" , baseline = "top" , dy = 10 , fontSize = 10 , opacity = 0.7 )
107+ .encode (x = alt .X (f"{ groups } :N" ), y = alt .Y (f"min({ values } ):Q" ), text = alt .Text ("count():Q" , format = "d" ))
108+ .transform_aggregate (count = "count()" , groupby = [groups ])
118109 )
119110
120111 # Combine box plot with annotations
121- chart = (base + text ).properties (
122- width = width ,
123- height = height ,
124- title = alt .TitleParams (
125- text = title or 'Box Plot Distribution' ,
126- fontSize = 16 ,
127- anchor = 'middle'
112+ chart = (
113+ (base + text )
114+ .properties (
115+ width = width ,
116+ height = height ,
117+ title = alt .TitleParams (text = title or "Box Plot Distribution" , fontSize = 16 , anchor = "middle" ),
118+ )
119+ .configure_view (strokeWidth = 0 )
120+ .configure_axis (grid = True , gridOpacity = 0.3 , gridDash = [3 , 3 ], domainWidth = 1 , tickWidth = 1 )
121+ .configure_boxplot (
122+ median = {"color" : "red" , "strokeWidth" : 2 },
123+ box = {"strokeWidth" : 1.5 },
124+ outliers = {"fill" : "red" , "fillOpacity" : 0.5 , "size" : 50 },
128125 )
129- ).configure_view (
130- strokeWidth = 0
131- ).configure_axis (
132- grid = True ,
133- gridOpacity = 0.3 ,
134- gridDash = [3 , 3 ],
135- domainWidth = 1 ,
136- tickWidth = 1
137- ).configure_boxplot (
138- median = dict (color = 'red' , strokeWidth = 2 ),
139- box = dict (strokeWidth = 1.5 ),
140- outliers = dict (fill = 'red' , fillOpacity = 0.5 , size = 50 )
141126 )
142127
143128 return chart
144129
145130
146- if __name__ == ' __main__' :
131+ if __name__ == " __main__" :
147132 # Sample data for testing with different distributions per group
148133 np .random .seed (42 ) # For reproducibility
149134
150135 # Generate sample data with 4 groups
151- data_dict = {
152- 'Group' : [],
153- 'Value' : []
154- }
136+ data_dict = {"Group" : [], "Value" : []}
155137
156138 # Group A: Normal distribution, mean=50, std=10
157139 group_a_data = np .random .normal (50 , 10 , 40 )
@@ -173,28 +155,29 @@ def create_plot(
173155
174156 # Combine all data
175157 for group , values in zip (
176- ['Group A' , 'Group B' , 'Group C' , 'Group D' ],
177- [group_a_data , group_b_data , group_c_data , group_d_data ]
158+ ["Group A" , "Group B" , "Group C" , "Group D" ],
159+ [group_a_data , group_b_data , group_c_data , group_d_data ],
160+ strict = False ,
178161 ):
179- data_dict [' Group' ].extend ([group ] * len (values ))
180- data_dict [' Value' ].extend (values )
162+ data_dict [" Group" ].extend ([group ] * len (values ))
163+ data_dict [" Value" ].extend (values )
181164
182165 data = pd .DataFrame (data_dict )
183166
184167 # Create plot
185168 chart = create_plot (
186169 data ,
187- values = ' Value' ,
188- groups = ' Group' ,
189- title = ' Statistical Distribution Comparison Across Groups' ,
190- ylabel = ' Measurement Value' ,
191- xlabel = ' Categories'
170+ values = " Value" ,
171+ groups = " Group" ,
172+ title = " Statistical Distribution Comparison Across Groups" ,
173+ ylabel = " Measurement Value" ,
174+ xlabel = " Categories" ,
192175 )
193176
194177 # Save for inspection
195- chart .save (' plot.html' )
178+ chart .save (" plot.html" )
196179 print ("Interactive plot saved to plot.html" )
197180
198181 # Also save as PNG
199- chart .save (' plot.png' , scale_factor = 2.0 )
200- print ("Static plot saved to plot.png" )
182+ chart .save (" plot.png" , scale_factor = 2.0 )
183+ print ("Static plot saved to plot.png" )
0 commit comments