|
1 | | -""" pyplots.ai |
| 1 | +""" anyplot.ai |
2 | 2 | sankey-basic: Basic Sankey Diagram |
3 | | -Library: altair 6.0.0 | Python 3.13.11 |
4 | | -Quality: 91/100 | Created: 2025-12-23 |
| 3 | +Library: altair 6.1.0 | Python 3.13.13 |
| 4 | +Quality: 82/100 | Updated: 2026-04-30 |
5 | 5 | """ |
6 | 6 |
|
| 7 | +import os |
| 8 | + |
7 | 9 | import altair as alt |
8 | 10 | import pandas as pd |
9 | 11 |
|
10 | 12 |
|
| 13 | +# Theme tokens |
| 14 | +THEME = os.getenv("ANYPLOT_THEME", "light") |
| 15 | +PAGE_BG = "#FAF8F1" if THEME == "light" else "#1A1A17" |
| 16 | +ELEVATED_BG = "#FFFDF6" if THEME == "light" else "#242420" |
| 17 | +INK = "#1A1A17" if THEME == "light" else "#F0EFE8" |
| 18 | +INK_SOFT = "#4A4A44" if THEME == "light" else "#B8B7B0" |
| 19 | + |
11 | 20 | # Data - Energy flow from sources to sectors |
12 | 21 | flows = [ |
13 | 22 | {"source": "Coal", "target": "Residential", "value": 20}, |
|
26 | 35 |
|
27 | 36 | df = pd.DataFrame(flows) |
28 | 37 |
|
29 | | -# Target output: 4800x2700 px (16:9 aspect ratio) with scale_factor=3.0 |
30 | | -# Internal canvas: 1600x900 pixels |
| 38 | +# Canvas dimensions: 1600x900 internal → 4800x2700 px at scale_factor=3.0 |
31 | 39 | width = 1600 |
32 | 40 | height = 900 |
33 | 41 | node_width = 80 |
|
37 | 45 | sources = df["source"].unique().tolist() |
38 | 46 | targets = df["target"].unique().tolist() |
39 | 47 |
|
40 | | -# Calculate totals for positioning |
41 | 48 | source_totals = df.groupby("source")["value"].sum().to_dict() |
42 | 49 | target_totals = df.groupby("target")["value"].sum().to_dict() |
43 | 50 | total_flow = df["value"].sum() |
44 | 51 |
|
45 | | -# Available height for nodes - reserve space for title (top) and margins |
46 | 52 | top_margin = 100 |
47 | 53 | bottom_margin = 60 |
48 | 54 | available_height = height - top_margin - bottom_margin |
|
71 | 77 | target_positions[tgt] = {"y": current_y, "height": node_height} |
72 | 78 | current_y += node_height + node_padding |
73 | 79 |
|
74 | | -# Color palettes - Python Blue (#306998) as primary, Yellow (#FFD43B) for accent |
75 | | -source_colors = {"Coal": "#306998", "Gas": "#4A8BC6", "Nuclear": "#2D5986", "Renewable": "#FFD43B"} |
| 80 | +# Okabe-Ito palette for source colors — distinct, colorblind-safe |
| 81 | +source_colors = { |
| 82 | + "Coal": "#009E73", # Okabe-Ito #1 (brand green) |
| 83 | + "Gas": "#D55E00", # Okabe-Ito #2 (vermillion) |
| 84 | + "Nuclear": "#0072B2", # Okabe-Ito #3 (blue) |
| 85 | + "Renewable": "#CC79A7", # Okabe-Ito #4 (reddish purple) |
| 86 | +} |
76 | 87 |
|
77 | | -target_colors = {"Residential": "#4ECDC4", "Commercial": "#95E1D3", "Industrial": "#FF6B6B", "Transport": "#FFA07A"} |
| 88 | +# Target node colors — muted, distinct from source palette |
| 89 | +target_colors = {"Residential": "#7EC8C8", "Commercial": "#A8D8A8", "Industrial": "#E8C07A", "Transport": "#C8A8E8"} |
78 | 90 |
|
79 | | -# Create node rectangles data |
| 91 | +# Build node rectangles data |
80 | 92 | nodes_data = [] |
81 | | - |
82 | 93 | for src in sources: |
83 | 94 | pos = source_positions[src] |
84 | 95 | nodes_data.append( |
|
115 | 126 |
|
116 | 127 | nodes_df = pd.DataFrame(nodes_data) |
117 | 128 |
|
118 | | -# Create flow paths using polygons |
119 | | -# Track current position within each node for stacking flows |
| 129 | +# Generate smoothstep S-curve polygon points for each flow band |
120 | 130 | source_y_offsets = {src: source_positions[src]["y"] for src in sources} |
121 | 131 | target_y_offsets = {tgt: target_positions[tgt]["y"] for tgt in targets} |
122 | 132 |
|
123 | | -# Generate polygon points for each flow (closed path) |
124 | 133 | all_flow_data = [] |
125 | 134 | num_curve_points = 40 |
126 | 135 |
|
|
129 | 138 | tgt = row["target"] |
130 | 139 | val = row["value"] |
131 | 140 |
|
132 | | - # Flow height proportional to value within each node |
133 | 141 | src_height = (val / source_totals[src]) * source_positions[src]["height"] |
134 | 142 | tgt_height = (val / target_totals[tgt]) * target_positions[tgt]["height"] |
135 | 143 |
|
136 | | - # Start and end Y positions for this flow band |
137 | 144 | src_y_top = source_y_offsets[src] |
138 | 145 | src_y_bottom = src_y_top + src_height |
139 | 146 | tgt_y_top = target_y_offsets[tgt] |
140 | 147 | tgt_y_bottom = tgt_y_top + tgt_height |
141 | 148 |
|
142 | | - # Update offsets for stacking next flow from same source/target |
143 | 149 | source_y_offsets[src] += src_height |
144 | 150 | target_y_offsets[tgt] += tgt_height |
145 | 151 |
|
146 | 152 | x_start = node_width |
147 | 153 | x_end = width - node_width |
148 | 154 |
|
149 | | - # Generate top curve points (left to right) using smoothstep interpolation |
150 | 155 | top_points = [] |
151 | 156 | for i in range(num_curve_points): |
152 | 157 | t = i / (num_curve_points - 1) |
153 | 158 | x = x_start + t * (x_end - x_start) |
154 | | - # Smoothstep creates smooth S-curve for natural flow appearance |
155 | 159 | bezier_t = t * t * (3 - 2 * t) |
156 | 160 | y = src_y_top + bezier_t * (tgt_y_top - src_y_top) |
157 | 161 | top_points.append((x, y)) |
158 | 162 |
|
159 | | - # Generate bottom curve points (right to left to close the polygon) |
160 | 163 | bottom_points = [] |
161 | 164 | for i in range(num_curve_points - 1, -1, -1): |
162 | 165 | t = i / (num_curve_points - 1) |
|
165 | 168 | y = src_y_bottom + bezier_t * (tgt_y_bottom - src_y_bottom) |
166 | 169 | bottom_points.append((x, y)) |
167 | 170 |
|
168 | | - # Combine top + bottom into closed polygon for filled area rendering |
169 | 171 | all_points = top_points + bottom_points |
170 | 172 | for pt_idx, (x, y) in enumerate(all_points): |
171 | 173 | all_flow_data.append( |
|
174 | 176 |
|
175 | 177 | flows_df = pd.DataFrame(all_flow_data) |
176 | 178 |
|
177 | | -# Create flow polygons using mark_line with filled=True |
| 179 | +# Flow polygons colored by source |
178 | 180 | links_chart = ( |
179 | 181 | alt.Chart(flows_df) |
180 | 182 | .mark_line(filled=True, opacity=0.55, strokeWidth=0) |
|
184 | 186 | color=alt.Color( |
185 | 187 | "source:N", |
186 | 188 | scale=alt.Scale(domain=list(source_colors.keys()), range=list(source_colors.values())), |
187 | | - legend=alt.Legend( |
188 | | - title="Energy Source", |
189 | | - titleFontSize=18, |
190 | | - labelFontSize=16, |
191 | | - orient="bottom-right", |
192 | | - titleColor="#333333", |
193 | | - labelColor="#333333", |
194 | | - ), |
| 189 | + legend=alt.Legend(title="Energy Source", titleFontSize=18, labelFontSize=16, orient="bottom-right"), |
195 | 190 | ), |
196 | 191 | detail="flow_id:N", |
197 | 192 | order="order:Q", |
198 | 193 | ) |
199 | 194 | ) |
200 | 195 |
|
201 | | -# Create node rectangles |
| 196 | +# Node rectangles |
202 | 197 | nodes_chart = ( |
203 | 198 | alt.Chart(nodes_df) |
204 | | - .mark_rect(stroke="#333333", strokeWidth=2) |
| 199 | + .mark_rect(stroke=INK_SOFT, strokeWidth=2) |
205 | 200 | .encode( |
206 | 201 | x=alt.X("x:Q", scale=alt.Scale(domain=[0, width])), |
207 | 202 | y=alt.Y("y:Q", scale=alt.Scale(domain=[0, height])), |
|
212 | 207 | ) |
213 | 208 | ) |
214 | 209 |
|
215 | | -# Create source labels (right-aligned to the left of nodes) |
| 210 | +# Source labels (right of left nodes) |
216 | 211 | source_labels_df = nodes_df[nodes_df["side"] == "source"] |
217 | 212 | source_labels = ( |
218 | 213 | alt.Chart(source_labels_df) |
219 | | - .mark_text(fontSize=20, fontWeight="bold", color="#333333", align="left", baseline="middle") |
| 214 | + .mark_text(fontSize=20, fontWeight="bold", align="left", baseline="middle") |
220 | 215 | .encode( |
221 | 216 | x=alt.X("label_x:Q", scale=alt.Scale(domain=[0, width])), |
222 | 217 | y=alt.Y("label_y:Q", scale=alt.Scale(domain=[0, height])), |
223 | 218 | text="name:N", |
| 219 | + color=alt.value(INK), |
224 | 220 | ) |
225 | 221 | ) |
226 | 222 |
|
227 | | -# Create target labels (left-aligned to the right of nodes) |
| 223 | +# Target labels (left of right nodes) |
228 | 224 | target_labels_df = nodes_df[nodes_df["side"] == "target"] |
229 | 225 | target_labels = ( |
230 | 226 | alt.Chart(target_labels_df) |
231 | | - .mark_text(fontSize=20, fontWeight="bold", color="#333333", align="right", baseline="middle") |
| 227 | + .mark_text(fontSize=20, fontWeight="bold", align="right", baseline="middle") |
232 | 228 | .encode( |
233 | 229 | x=alt.X("label_x:Q", scale=alt.Scale(domain=[0, width])), |
234 | 230 | y=alt.Y("label_y:Q", scale=alt.Scale(domain=[0, height])), |
235 | 231 | text="name:N", |
| 232 | + color=alt.value(INK), |
236 | 233 | ) |
237 | 234 | ) |
238 | 235 |
|
239 | | -# Combine all layers |
| 236 | +# Compose all layers with theme-adaptive chrome |
240 | 237 | chart = ( |
241 | 238 | alt.layer(links_chart, nodes_chart, source_labels, target_labels) |
242 | 239 | .properties( |
243 | 240 | width=width, |
244 | 241 | height=height, |
| 242 | + background=PAGE_BG, |
245 | 243 | title=alt.Title( |
246 | | - text="sankey-basic · altair · pyplots.ai", |
| 244 | + text="sankey-basic · altair · anyplot.ai", |
247 | 245 | subtitle="Energy Flow from Sources to Sectors", |
248 | 246 | fontSize=28, |
249 | 247 | subtitleFontSize=20, |
250 | 248 | anchor="middle", |
251 | | - color="#333333", |
252 | | - subtitleColor="#666666", |
| 249 | + color=INK, |
| 250 | + subtitleColor=INK_SOFT, |
253 | 251 | ), |
254 | | - autosize=alt.AutoSizeParams(type="fit", contains="padding"), |
255 | 252 | ) |
256 | | - .configure_view(strokeWidth=0) |
257 | | - .configure_legend(padding=15, cornerRadius=5, fillColor="#FFFFFF", strokeColor="#DDDDDD") |
| 253 | + .configure_view(strokeWidth=0, fill=PAGE_BG) |
| 254 | + .configure_legend( |
| 255 | + padding=15, cornerRadius=5, fillColor=ELEVATED_BG, strokeColor=INK_SOFT, labelColor=INK_SOFT, titleColor=INK |
| 256 | + ) |
258 | 257 | ) |
259 | 258 |
|
260 | | -# Save as PNG (4800x2700 px with scale_factor=3.0) and HTML |
261 | | -chart.save("plot.png", scale_factor=3.0) |
262 | | -chart.save("plot.html") |
| 259 | +# Save outputs (PNG at 4800×2700 px + HTML for interactivity) |
| 260 | +chart.save(f"plot-{THEME}.png", scale_factor=3.0) |
| 261 | +chart.save(f"plot-{THEME}.html") |
0 commit comments