|
| 1 | +""" pyplots.ai |
| 2 | +root-locus-basic: Root Locus Plot for Control Systems |
| 3 | +Library: altair 6.0.0 | Python 3.14.3 |
| 4 | +Quality: 86/100 | Created: 2026-03-20 |
| 5 | +""" |
| 6 | + |
| 7 | +import altair as alt |
| 8 | +import numpy as np |
| 9 | +import pandas as pd |
| 10 | + |
| 11 | + |
| 12 | +# Data - Root locus for G(s) = 1 / (s(s+1)(s+2)) |
| 13 | +# Open-loop poles at s = 0, -1, -2; no zeros |
| 14 | +# Characteristic equation: s³ + 3s² + 2s + K = 0 |
| 15 | +den_coeffs = [1.0, 3.0, 2.0, 0.0] |
| 16 | + |
| 17 | +gains = np.concatenate( |
| 18 | + [ |
| 19 | + np.linspace(0.001, 0.5, 150), |
| 20 | + np.linspace(0.5, 2, 200), |
| 21 | + np.linspace(2, 6, 150), |
| 22 | + np.linspace(6, 20, 150), |
| 23 | + np.linspace(20, 80, 100), |
| 24 | + ] |
| 25 | +) |
| 26 | +n_roots = 3 |
| 27 | +all_roots = np.zeros((len(gains), n_roots), dtype=complex) |
| 28 | + |
| 29 | +for i, k in enumerate(gains): |
| 30 | + poly = np.array(den_coeffs, dtype=float) |
| 31 | + poly[-1] += k |
| 32 | + all_roots[i] = np.roots(poly) |
| 33 | + |
| 34 | +# Sort roots into continuous branches via nearest-neighbor matching |
| 35 | +all_roots[0] = np.sort(all_roots[0].real) |
| 36 | +for i in range(1, len(gains)): |
| 37 | + prev, curr = all_roots[i - 1], all_roots[i] |
| 38 | + dist = np.abs(prev[:, None] - curr[None, :]) |
| 39 | + order = np.zeros(n_roots, dtype=int) |
| 40 | + used = set() |
| 41 | + for j in range(n_roots): |
| 42 | + dists = [(dist[j, m], m) for m in range(n_roots) if m not in used] |
| 43 | + _, best = min(dists) |
| 44 | + used.add(best) |
| 45 | + order[j] = best |
| 46 | + all_roots[i] = curr[order] |
| 47 | + |
| 48 | +# Build branch dataframe |
| 49 | +rows = [] |
| 50 | +for b in range(n_roots): |
| 51 | + for i in range(len(gains)): |
| 52 | + rows.append( |
| 53 | + { |
| 54 | + "real": float(all_roots[i, b].real), |
| 55 | + "imaginary": float(all_roots[i, b].imag), |
| 56 | + "gain": float(gains[i]), |
| 57 | + "branch": f"Branch {b + 1}", |
| 58 | + "idx": i, |
| 59 | + } |
| 60 | + ) |
| 61 | +locus_df = pd.DataFrame(rows) |
| 62 | + |
| 63 | +# Open-loop poles |
| 64 | +poles_df = pd.DataFrame( |
| 65 | + {"real": [0.0, -1.0, -2.0], "imaginary": [0.0, 0.0, 0.0], "label": ["Pole (s=0)", "Pole (s=−1)", "Pole (s=−2)"]} |
| 66 | +) |
| 67 | + |
| 68 | +# Imaginary axis crossing: ω = √2, K = 6 |
| 69 | +omega_cross = np.sqrt(2) |
| 70 | +crossing_df = pd.DataFrame( |
| 71 | + {"real": [0.0, 0.0], "imaginary": [omega_cross, -omega_cross], "label": ["jω = j√2 (K=6)", "jω = −j√2 (K=6)"]} |
| 72 | +) |
| 73 | + |
| 74 | +# Breakaway point: d/ds[s(s+1)(s+2)] = 3s²+6s+2 = 0 → s ≈ -0.423 |
| 75 | +breakaway_df = pd.DataFrame({"bx": [(-6 + np.sqrt(12)) / 6], "by": [0.0], "label": ["Breakaway (s ≈ −0.42)"]}) |
| 76 | + |
| 77 | +# Damping ratio guide lines (ζ = 0.2, 0.4, 0.6, 0.8) |
| 78 | +damping_rows = [] |
| 79 | +for zeta in [0.2, 0.4, 0.6, 0.8]: |
| 80 | + angle = np.pi - np.arccos(zeta) |
| 81 | + for side, sign in [("upper", 1), ("lower", -1)]: |
| 82 | + seg = f"ζ={zeta}_{side}" |
| 83 | + damping_rows.append({"gx": 0.0, "gy": 0.0, "seg": seg, "ord": 0}) |
| 84 | + damping_rows.append({"gx": 5.0 * np.cos(angle), "gy": sign * 5.0 * np.sin(angle), "seg": seg, "ord": 1}) |
| 85 | +damping_df = pd.DataFrame(damping_rows) |
| 86 | + |
| 87 | +# Damping ratio labels at end of guide lines |
| 88 | +damping_label_rows = [] |
| 89 | +for zeta in [0.4, 0.8]: |
| 90 | + angle = np.pi - np.arccos(zeta) |
| 91 | + damping_label_rows.append({"lx": 4.6 * np.cos(angle), "ly": 4.6 * np.sin(angle), "label": f"ζ={zeta}"}) |
| 92 | +damping_label_df = pd.DataFrame(damping_label_rows) |
| 93 | + |
| 94 | +# Natural frequency arcs (ωn = 1, 2, 3, 4) in left half-plane |
| 95 | +wn_rows = [] |
| 96 | +for wn in [1.0, 2.0, 3.0, 4.0]: |
| 97 | + theta = np.linspace(np.pi / 2, 3 * np.pi / 2, 60) |
| 98 | + for j, t in enumerate(theta): |
| 99 | + wn_rows.append({"gx": wn * np.cos(t), "gy": wn * np.sin(t), "wn": f"ωn={wn}", "ord": j}) |
| 100 | +wn_df = pd.DataFrame(wn_rows) |
| 101 | + |
| 102 | +# Real axis segments: (-1, 0) and (-∞, -2) |
| 103 | +real_axis_df = pd.DataFrame( |
| 104 | + { |
| 105 | + "rx": [-1.0, 0.0, -5.0, -2.0], |
| 106 | + "ry": [0.0, 0.0, 0.0, 0.0], |
| 107 | + "seg": ["seg1", "seg1", "seg2", "seg2"], |
| 108 | + "ord": [0, 1, 0, 1], |
| 109 | + } |
| 110 | +) |
| 111 | + |
| 112 | +# Arrow direction indicators along complex branches |
| 113 | +arrows = [] |
| 114 | +for b in range(n_roots): |
| 115 | + for idx in [350, 500]: |
| 116 | + if idx + 5 < len(gains): |
| 117 | + r0 = all_roots[idx, b] |
| 118 | + if abs(r0.imag) > 0.3: |
| 119 | + arrows.append({"ax": float(r0.real), "ay": float(r0.imag), "branch": f"Branch {b + 1}"}) |
| 120 | +arrow_df = pd.DataFrame(arrows) if arrows else pd.DataFrame({"ax": [], "ay": [], "branch": []}) |
| 121 | + |
| 122 | +# Equal-scaling axes centered on origin (square canvas, equal domain = equal scaling) |
| 123 | +x_scale = alt.Scale(domain=[-5.0, 5.0], nice=False) |
| 124 | +y_scale = alt.Scale(domain=[-5.0, 5.0], nice=False) |
| 125 | + |
| 126 | +branch_palette = ["#306998", "#e07b39", "#2ca02c"] |
| 127 | +branch_domain = ["Branch 1", "Branch 2", "Branch 3"] |
| 128 | + |
| 129 | +# Layer: Locus branches — FIRST so its axis config takes effect |
| 130 | +locus_layer = ( |
| 131 | + alt.Chart(locus_df) |
| 132 | + .mark_line(strokeWidth=2.8, opacity=0.92) |
| 133 | + .encode( |
| 134 | + x=alt.X( |
| 135 | + "real:Q", |
| 136 | + scale=x_scale, |
| 137 | + title="Real Axis (σ)", |
| 138 | + axis=alt.Axis( |
| 139 | + labelFontSize=16, |
| 140 | + titleFontSize=21, |
| 141 | + titleFontWeight="bold", |
| 142 | + titleColor="#2a2a2a", |
| 143 | + labelColor="#444444", |
| 144 | + grid=False, |
| 145 | + tickCount=6, |
| 146 | + titlePadding=14, |
| 147 | + domainColor="#888888", |
| 148 | + tickColor="#888888", |
| 149 | + ), |
| 150 | + ), |
| 151 | + y=alt.Y( |
| 152 | + "imaginary:Q", |
| 153 | + scale=y_scale, |
| 154 | + title="Imaginary Axis (jω)", |
| 155 | + axis=alt.Axis( |
| 156 | + labelFontSize=16, |
| 157 | + titleFontSize=21, |
| 158 | + titleFontWeight="bold", |
| 159 | + titleColor="#2a2a2a", |
| 160 | + labelColor="#444444", |
| 161 | + grid=False, |
| 162 | + tickCount=6, |
| 163 | + titlePadding=14, |
| 164 | + domainColor="#888888", |
| 165 | + tickColor="#888888", |
| 166 | + ), |
| 167 | + ), |
| 168 | + color=alt.Color( |
| 169 | + "branch:N", |
| 170 | + scale=alt.Scale(domain=branch_domain, range=branch_palette), |
| 171 | + legend=alt.Legend( |
| 172 | + title="Branch", |
| 173 | + titleFontSize=16, |
| 174 | + labelFontSize=14, |
| 175 | + symbolSize=180, |
| 176 | + symbolStrokeWidth=3, |
| 177 | + orient="top-right", |
| 178 | + offset=5, |
| 179 | + ), |
| 180 | + ), |
| 181 | + order="idx:Q", |
| 182 | + tooltip=[ |
| 183 | + alt.Tooltip("branch:N", title="Branch"), |
| 184 | + alt.Tooltip("real:Q", title="σ", format=".3f"), |
| 185 | + alt.Tooltip("imaginary:Q", title="jω", format=".3f"), |
| 186 | + alt.Tooltip("gain:Q", title="Gain K", format=".2f"), |
| 187 | + ], |
| 188 | + ) |
| 189 | +) |
| 190 | + |
| 191 | +# Layer: Damping ratio lines |
| 192 | +damping_layer = ( |
| 193 | + alt.Chart(damping_df) |
| 194 | + .mark_line(strokeWidth=0.8, strokeDash=[6, 4], color="#d0d0d0") |
| 195 | + .encode(x=alt.X("gx:Q", scale=x_scale), y=alt.Y("gy:Q", scale=y_scale), detail="seg:N", order="ord:Q") |
| 196 | +) |
| 197 | + |
| 198 | +# Layer: Damping ratio labels |
| 199 | +damping_label_layer = ( |
| 200 | + alt.Chart(damping_label_df) |
| 201 | + .mark_text(fontSize=12, color="#aaaaaa", fontStyle="italic", align="center") |
| 202 | + .encode(x=alt.X("lx:Q", scale=x_scale), y=alt.Y("ly:Q", scale=y_scale), text="label:N") |
| 203 | +) |
| 204 | + |
| 205 | +# Layer: Natural frequency arcs |
| 206 | +wn_layer = ( |
| 207 | + alt.Chart(wn_df) |
| 208 | + .mark_line(strokeWidth=0.8, strokeDash=[4, 4], color="#d0d0d0") |
| 209 | + .encode(x=alt.X("gx:Q", scale=x_scale), y=alt.Y("gy:Q", scale=y_scale), detail="wn:N", order="ord:Q") |
| 210 | +) |
| 211 | + |
| 212 | +# Layer: Real axis segments |
| 213 | +real_axis_layer = ( |
| 214 | + alt.Chart(real_axis_df) |
| 215 | + .mark_line(strokeWidth=5, color="#306998", opacity=0.25) |
| 216 | + .encode(x=alt.X("rx:Q", scale=x_scale), y=alt.Y("ry:Q", scale=y_scale), detail="seg:N", order="ord:Q") |
| 217 | +) |
| 218 | + |
| 219 | +# Layer: Open-loop poles (× markers) |
| 220 | +poles_layer = ( |
| 221 | + alt.Chart(poles_df) |
| 222 | + .mark_point(shape="cross", size=450, strokeWidth=3.5, color="#d62728", filled=False) |
| 223 | + .encode( |
| 224 | + x=alt.X("real:Q", scale=x_scale), |
| 225 | + y=alt.Y("imaginary:Q", scale=y_scale), |
| 226 | + tooltip=[alt.Tooltip("label:N", title=""), alt.Tooltip("real:Q", title="σ")], |
| 227 | + ) |
| 228 | +) |
| 229 | + |
| 230 | +# Layer: Imaginary axis crossings |
| 231 | +crossing_layer = ( |
| 232 | + alt.Chart(crossing_df) |
| 233 | + .mark_point(shape="diamond", size=400, strokeWidth=2.5, color="#d62728", filled=True) |
| 234 | + .encode( |
| 235 | + x=alt.X("real:Q", scale=x_scale), |
| 236 | + y=alt.Y("imaginary:Q", scale=y_scale), |
| 237 | + tooltip=[alt.Tooltip("label:N", title="Crossing")], |
| 238 | + ) |
| 239 | +) |
| 240 | + |
| 241 | +# Layer: Crossing labels |
| 242 | +crossing_text = ( |
| 243 | + alt.Chart(crossing_df) |
| 244 | + .mark_text(fontSize=17, fontWeight="bold", color="#c5211e", align="left", dx=20, font="sans-serif") |
| 245 | + .encode(x=alt.X("real:Q", scale=x_scale), y=alt.Y("imaginary:Q", scale=y_scale), text="label:N") |
| 246 | +) |
| 247 | + |
| 248 | +# Layer: Breakaway point |
| 249 | +breakaway_layer = ( |
| 250 | + alt.Chart(breakaway_df) |
| 251 | + .mark_point(shape="square", size=220, color="#555555", filled=True, opacity=0.8) |
| 252 | + .encode( |
| 253 | + x=alt.X("bx:Q", scale=x_scale), y=alt.Y("by:Q", scale=y_scale), tooltip=[alt.Tooltip("label:N", title="Point")] |
| 254 | + ) |
| 255 | +) |
| 256 | + |
| 257 | +# Layer: Arrow direction indicators |
| 258 | +arrow_up_df = arrow_df[arrow_df["ay"] > 0] if len(arrow_df) > 0 else arrow_df |
| 259 | +arrow_down_df = arrow_df[arrow_df["ay"] <= 0] if len(arrow_df) > 0 else arrow_df |
| 260 | + |
| 261 | +arrow_up_layer = ( |
| 262 | + alt.Chart(arrow_up_df) |
| 263 | + .mark_point(shape="triangle-up", size=250, filled=True, opacity=0.85) |
| 264 | + .encode( |
| 265 | + x=alt.X("ax:Q", scale=x_scale), |
| 266 | + y=alt.Y("ay:Q", scale=y_scale), |
| 267 | + color=alt.Color("branch:N", scale=alt.Scale(domain=branch_domain, range=branch_palette), legend=None), |
| 268 | + ) |
| 269 | +) |
| 270 | + |
| 271 | +arrow_down_layer = ( |
| 272 | + alt.Chart(arrow_down_df) |
| 273 | + .mark_point(shape="triangle-down", size=250, filled=True, opacity=0.85) |
| 274 | + .encode( |
| 275 | + x=alt.X("ax:Q", scale=x_scale), |
| 276 | + y=alt.Y("ay:Q", scale=y_scale), |
| 277 | + color=alt.Color("branch:N", scale=alt.Scale(domain=branch_domain, range=branch_palette), legend=None), |
| 278 | + ) |
| 279 | +) |
| 280 | + |
| 281 | +# Compose — locus_layer first so its axis config renders |
| 282 | +chart = ( |
| 283 | + ( |
| 284 | + locus_layer |
| 285 | + + damping_layer |
| 286 | + + damping_label_layer |
| 287 | + + wn_layer |
| 288 | + + real_axis_layer |
| 289 | + + poles_layer |
| 290 | + + crossing_layer |
| 291 | + + crossing_text |
| 292 | + + breakaway_layer |
| 293 | + + arrow_up_layer |
| 294 | + + arrow_down_layer |
| 295 | + ) |
| 296 | + .properties( |
| 297 | + width=1200, |
| 298 | + height=1200, |
| 299 | + title=alt.Title( |
| 300 | + "root-locus-basic · altair · pyplots.ai", |
| 301 | + fontSize=28, |
| 302 | + fontWeight="bold", |
| 303 | + color="#1a1a1a", |
| 304 | + subtitle="G(s) = 1 / s(s+1)(s+2) · Closed-Loop Pole Trajectories vs Gain K", |
| 305 | + subtitleFontSize=18, |
| 306 | + subtitleColor="#555555", |
| 307 | + subtitlePadding=10, |
| 308 | + anchor="start", |
| 309 | + offset=10, |
| 310 | + ), |
| 311 | + ) |
| 312 | + .configure_view(strokeWidth=0) |
| 313 | + .interactive() |
| 314 | +) |
| 315 | + |
| 316 | +chart.save("plot.png", scale_factor=3.0) |
| 317 | +chart.save("plot.html") |
0 commit comments