|
| 1 | +""" pyplots.ai |
| 2 | +root-locus-basic: Root Locus Plot for Control Systems |
| 3 | +Library: letsplot 4.9.0 | Python 3.14.3 |
| 4 | +Quality: 90/100 | Created: 2026-03-20 |
| 5 | +""" |
| 6 | + |
| 7 | +import numpy as np |
| 8 | +import pandas as pd |
| 9 | +from lets_plot import * # noqa: F403 |
| 10 | +from lets_plot.export import ggsave as export_ggsave |
| 11 | + |
| 12 | + |
| 13 | +LetsPlot.setup_html() # noqa: F405 |
| 14 | + |
| 15 | +# Data - Root locus for G(s) = (s + 3) / [s(s + 1)(s + 2)(s + 4)] |
| 16 | +# Open-loop poles: 0, -1, -2, -4 | Open-loop zero: -3 |
| 17 | +num_coeffs = np.array([1, 3]) # s + 3 |
| 18 | +den_coeffs = np.polymul(np.polymul([1, 0], [1, 1]), np.polymul([1, 2], [1, 4])) # s(s+1)(s+2)(s+4) |
| 19 | + |
| 20 | +open_loop_poles = np.array([0.0, -1.0, -2.0, -4.0]) |
| 21 | +open_loop_zeros = np.array([-3.0]) |
| 22 | + |
| 23 | +# Compute closed-loop poles for varying gain K |
| 24 | +k_values = np.concatenate( |
| 25 | + [np.linspace(0, 0.5, 200), np.linspace(0.5, 5, 400), np.linspace(5, 30, 400), np.linspace(30, 120, 400)] |
| 26 | +) |
| 27 | + |
| 28 | +all_real = [] |
| 29 | +all_imag = [] |
| 30 | +all_gain = [] |
| 31 | +all_branch = [] |
| 32 | + |
| 33 | +for k in k_values: |
| 34 | + char_poly = np.polyadd(den_coeffs, k * np.array([0, 0, 0, 1, 3])) |
| 35 | + roots = np.roots(char_poly) |
| 36 | + roots = np.sort_complex(roots) |
| 37 | + for b, root in enumerate(roots): |
| 38 | + all_real.append(root.real) |
| 39 | + all_imag.append(root.imag) |
| 40 | + all_gain.append(k) |
| 41 | + all_branch.append(f"Branch {b + 1}") |
| 42 | + |
| 43 | +df = pd.DataFrame({"real": all_real, "imaginary": all_imag, "gain": all_gain, "branch": all_branch}) |
| 44 | + |
| 45 | +# Open-loop poles and zeros for markers |
| 46 | +poles_df = pd.DataFrame( |
| 47 | + { |
| 48 | + "real": open_loop_poles, |
| 49 | + "imaginary": [0.0] * len(open_loop_poles), |
| 50 | + "type": ["Open-loop pole"] * len(open_loop_poles), |
| 51 | + } |
| 52 | +) |
| 53 | +zeros_df = pd.DataFrame( |
| 54 | + { |
| 55 | + "real": open_loop_zeros, |
| 56 | + "imaginary": [0.0] * len(open_loop_zeros), |
| 57 | + "type": ["Open-loop zero"] * len(open_loop_zeros), |
| 58 | + } |
| 59 | +) |
| 60 | + |
| 61 | +# Find imaginary axis crossings (where real part ≈ 0 and imag ≠ 0) |
| 62 | +crossing_mask = (np.abs(df["real"]) < 0.08) & (np.abs(df["imaginary"]) > 0.3) |
| 63 | +crossings = df[crossing_mask].copy() |
| 64 | +if len(crossings) > 0: |
| 65 | + crossings = crossings.sort_values("imaginary") |
| 66 | + pos_cross = crossings[crossings["imaginary"] > 0].head(1) |
| 67 | + neg_cross = crossings[crossings["imaginary"] < 0].head(1) |
| 68 | + crossing_pts = pd.concat([pos_cross, neg_cross]) |
| 69 | +else: |
| 70 | + crossing_pts = pd.DataFrame(columns=df.columns) |
| 71 | + |
| 72 | +# Direction arrows - sample points at specific gain values for each branch |
| 73 | +arrow_gains = [5, 15, 50] |
| 74 | +arrow_rows = [] |
| 75 | +for ag in arrow_gains: |
| 76 | + idx = np.argmin(np.abs(k_values - ag)) |
| 77 | + subset = df[(df["gain"] >= k_values[max(0, idx - 1)]) & (df["gain"] <= k_values[min(len(k_values) - 1, idx + 1)])] |
| 78 | + for _, row in subset.drop_duplicates(subset="branch").iterrows(): |
| 79 | + k_next = k_values[min(len(k_values) - 1, idx + 5)] |
| 80 | + next_pts = df[(np.abs(df["gain"] - k_next) < 1.0) & (df["branch"] == row["branch"])] |
| 81 | + if len(next_pts) > 0: |
| 82 | + npt = next_pts.iloc[0] |
| 83 | + dx = npt["real"] - row["real"] |
| 84 | + dy = npt["imaginary"] - row["imaginary"] |
| 85 | + mag = np.sqrt(dx**2 + dy**2) |
| 86 | + if mag > 0.01: |
| 87 | + scale = 0.25 / mag |
| 88 | + arrow_rows.append( |
| 89 | + { |
| 90 | + "x": row["real"], |
| 91 | + "y": row["imaginary"], |
| 92 | + "xend": row["real"] + dx * scale, |
| 93 | + "yend": row["imaginary"] + dy * scale, |
| 94 | + } |
| 95 | + ) |
| 96 | + |
| 97 | +arrows_df = pd.DataFrame(arrow_rows) if arrow_rows else pd.DataFrame(columns=["x", "y", "xend", "yend"]) |
| 98 | + |
| 99 | +# Damping ratio lines (constant zeta) |
| 100 | +zeta_values = [0.2, 0.4, 0.6, 0.8] |
| 101 | +zeta_lines = [] |
| 102 | +for zeta in zeta_values: |
| 103 | + theta = np.arccos(zeta) |
| 104 | + r_max = 5.0 |
| 105 | + zeta_lines.append( |
| 106 | + {"x": 0, "y": 0, "xend": -r_max * np.cos(theta), "yend": r_max * np.sin(theta), "label": f"ζ={zeta}"} |
| 107 | + ) |
| 108 | + zeta_lines.append( |
| 109 | + {"x": 0, "y": 0, "xend": -r_max * np.cos(theta), "yend": -r_max * np.sin(theta), "label": f"ζ={zeta}"} |
| 110 | + ) |
| 111 | +zeta_df = pd.DataFrame(zeta_lines) |
| 112 | + |
| 113 | +# Zeta labels (upper half only, positioned along lines with offset to avoid crowding) |
| 114 | +zeta_label_df = pd.DataFrame( |
| 115 | + [ |
| 116 | + { |
| 117 | + "x": -r_max * (0.4 + i * 0.08) * np.cos(np.arccos(z)), |
| 118 | + "y": r_max * (0.4 + i * 0.08) * np.sin(np.arccos(z)), |
| 119 | + "label": f"ζ={z}", |
| 120 | + } |
| 121 | + for i, z in enumerate(zeta_values) |
| 122 | + ] |
| 123 | +) |
| 124 | + |
| 125 | +# Natural frequency circles (constant ωn) |
| 126 | +wn_values = [1, 2, 3, 4] |
| 127 | +wn_rows = [] |
| 128 | +for wn in wn_values: |
| 129 | + theta = np.linspace(np.pi / 2, 3 * np.pi / 2, 100) |
| 130 | + wn_rows.extend([{"real": wn * np.cos(t), "imaginary": wn * np.sin(t), "wn": f"ωn={wn}"} for t in theta]) |
| 131 | +wn_df = pd.DataFrame(wn_rows) |
| 132 | + |
| 133 | +# Branch colors - colorblind-safe palette (no red-green pair) |
| 134 | +branch_colors = ["#306998", "#E69F00", "#CC79A7", "#56B4E9"] |
| 135 | + |
| 136 | +# Plot |
| 137 | +plot = ( |
| 138 | + ggplot() # noqa: F405 |
| 139 | + # Natural frequency arcs |
| 140 | + + geom_path( # noqa: F405 |
| 141 | + aes(x="real", y="imaginary", group="wn"), # noqa: F405 |
| 142 | + data=wn_df, |
| 143 | + color="#E0E0E0", |
| 144 | + size=0.4, |
| 145 | + linetype="dashed", |
| 146 | + tooltips="none", |
| 147 | + ) |
| 148 | + # Damping ratio lines |
| 149 | + + geom_segment( # noqa: F405 |
| 150 | + aes(x="x", y="y", xend="xend", yend="yend"), # noqa: F405 |
| 151 | + data=zeta_df, |
| 152 | + color="#E0E0E0", |
| 153 | + size=0.4, |
| 154 | + linetype="dashed", |
| 155 | + tooltips="none", |
| 156 | + ) |
| 157 | + # Damping ratio labels |
| 158 | + + geom_text( # noqa: F405 |
| 159 | + aes(x="x", y="y", label="label"), # noqa: F405 |
| 160 | + data=zeta_label_df, |
| 161 | + size=11, |
| 162 | + color="#A0A0A0", |
| 163 | + inherit_aes=False, |
| 164 | + family="monospace", |
| 165 | + ) |
| 166 | + # Stable region shading (left half-plane) |
| 167 | + + geom_rect( # noqa: F405 |
| 168 | + aes(xmin="xmin", ymin="ymin", xmax="xmax", ymax="ymax"), # noqa: F405 |
| 169 | + data=pd.DataFrame({"xmin": [-5.0], "xmax": [0], "ymin": [-4], "ymax": [4]}), |
| 170 | + fill="#E8F4E8", |
| 171 | + alpha=0.3, |
| 172 | + inherit_aes=False, |
| 173 | + tooltips="none", |
| 174 | + ) |
| 175 | + # Imaginary axis (stability boundary) |
| 176 | + + geom_vline(xintercept=0, color="#B0B0B0", size=0.7, linetype="solid") # noqa: F405 |
| 177 | + + geom_hline(yintercept=0, color="#CCCCCC", size=0.4) # noqa: F405 |
| 178 | + # Stability boundary label |
| 179 | + + geom_text( # noqa: F405 |
| 180 | + aes(x="x", y="y", label="label"), # noqa: F405 |
| 181 | + data=pd.DataFrame({"x": [0.15], "y": [3.7], "label": ["jω"]}), |
| 182 | + size=14, |
| 183 | + color="#888888", |
| 184 | + inherit_aes=False, |
| 185 | + hjust=0, |
| 186 | + family="serif", |
| 187 | + ) |
| 188 | + + geom_text( # noqa: F405 |
| 189 | + aes(x="x", y="y", label="label"), # noqa: F405 |
| 190 | + data=pd.DataFrame({"x": [1.1], "y": [-0.25], "label": ["σ"]}), |
| 191 | + size=14, |
| 192 | + color="#888888", |
| 193 | + inherit_aes=False, |
| 194 | + hjust=0, |
| 195 | + family="serif", |
| 196 | + ) |
| 197 | + # Region labels for storytelling |
| 198 | + + geom_text( # noqa: F405 |
| 199 | + aes(x="x", y="y", label="label"), # noqa: F405 |
| 200 | + data=pd.DataFrame({"x": [-4.3], "y": [2.8], "label": ["STABLE"]}), |
| 201 | + size=16, |
| 202 | + color="#4CAF50", |
| 203 | + alpha=0.4, |
| 204 | + inherit_aes=False, |
| 205 | + family="monospace", |
| 206 | + fontface="bold", |
| 207 | + ) |
| 208 | + + geom_text( # noqa: F405 |
| 209 | + aes(x="x", y="y", label="label"), # noqa: F405 |
| 210 | + data=pd.DataFrame({"x": [0.55], "y": [2.8], "label": ["UNSTABLE"]}), |
| 211 | + size=12, |
| 212 | + color="#D55E00", |
| 213 | + alpha=0.35, |
| 214 | + inherit_aes=False, |
| 215 | + family="monospace", |
| 216 | + fontface="bold", |
| 217 | + ) |
| 218 | + # Root locus branches with interactive tooltips |
| 219 | + + geom_path( # noqa: F405 |
| 220 | + aes(x="real", y="imaginary", color="branch"), # noqa: F405 |
| 221 | + data=df, |
| 222 | + size=1.8, |
| 223 | + alpha=0.9, |
| 224 | + tooltips=layer_tooltips() # noqa: F405 |
| 225 | + .format("gain", ".2f") |
| 226 | + .format("real", ".3f") |
| 227 | + .format("imaginary", ".3f") |
| 228 | + .line("Branch: @branch") |
| 229 | + .line("Gain K: @gain") |
| 230 | + .line("Re: @real") |
| 231 | + .line("Im: @imaginary"), |
| 232 | + ) |
| 233 | + + scale_color_manual( # noqa: F405 |
| 234 | + values=branch_colors, name="Locus Branch" |
| 235 | + ) |
| 236 | + # Direction arrows (larger, more visible) |
| 237 | + + geom_segment( # noqa: F405 |
| 238 | + aes(x="x", y="y", xend="xend", yend="yend"), # noqa: F405 |
| 239 | + data=arrows_df, |
| 240 | + color="#333333", |
| 241 | + size=1.2, |
| 242 | + arrow=arrow(length=14, ends="last", type="open"), # noqa: F405 |
| 243 | + inherit_aes=False, |
| 244 | + tooltips="none", |
| 245 | + ) |
| 246 | + # Open-loop poles (× markers) |
| 247 | + + geom_point( # noqa: F405 |
| 248 | + aes(x="real", y="imaginary"), # noqa: F405 |
| 249 | + data=poles_df, |
| 250 | + shape=4, # cross/x marker |
| 251 | + size=8, |
| 252 | + color="#222222", |
| 253 | + stroke=2.5, |
| 254 | + inherit_aes=False, |
| 255 | + tooltips=layer_tooltips().line("Open-loop pole").line("s = @real"), # noqa: F405 |
| 256 | + ) |
| 257 | + # Open-loop zeros (○ markers) |
| 258 | + + geom_point( # noqa: F405 |
| 259 | + aes(x="real", y="imaginary"), # noqa: F405 |
| 260 | + data=zeros_df, |
| 261 | + shape=1, # open circle |
| 262 | + size=8, |
| 263 | + color="#222222", |
| 264 | + stroke=2.5, |
| 265 | + inherit_aes=False, |
| 266 | + tooltips=layer_tooltips().line("Open-loop zero").line("s = @real"), # noqa: F405 |
| 267 | + ) |
| 268 | + # Imaginary axis crossings |
| 269 | + + geom_point( # noqa: F405 |
| 270 | + aes(x="real", y="imaginary"), # noqa: F405 |
| 271 | + data=crossing_pts, |
| 272 | + shape=18, # diamond |
| 273 | + size=8, |
| 274 | + color="#D55E00", |
| 275 | + inherit_aes=False, |
| 276 | + tooltips=layer_tooltips().line("Stability boundary crossing").line("K ≈ @gain").line("Im: @imaginary"), # noqa: F405 |
| 277 | + ) |
| 278 | + # Crossing gain annotation |
| 279 | + + geom_text( # noqa: F405 |
| 280 | + aes(x="x", y="y", label="label"), # noqa: F405 |
| 281 | + data=pd.DataFrame( |
| 282 | + { |
| 283 | + "x": [0.25], |
| 284 | + "y": [crossing_pts["imaginary"].max() if len(crossing_pts) > 0 else 2.0], |
| 285 | + "label": [f"K ≈ {crossing_pts['gain'].iloc[0]:.1f}" if len(crossing_pts) > 0 else ""], |
| 286 | + } |
| 287 | + ), |
| 288 | + size=13, |
| 289 | + color="#D55E00", |
| 290 | + inherit_aes=False, |
| 291 | + hjust=0, |
| 292 | + family="monospace", |
| 293 | + fontface="bold", |
| 294 | + ) |
| 295 | + # Labels and styling |
| 296 | + + labs( # noqa: F405 |
| 297 | + x="Real Axis (σ)", |
| 298 | + y="Imaginary Axis (jω)", |
| 299 | + title="root-locus-basic · letsplot · pyplots.ai", |
| 300 | + caption="G(s) = (s + 3) / [s(s + 1)(s + 2)(s + 4)] · × = poles · ○ = zeros · Stable region shaded", |
| 301 | + ) |
| 302 | + + coord_fixed(ratio=1, xlim=[-5.0, 1.5], ylim=[-4, 4]) # noqa: F405 |
| 303 | + + ggsize(1600, 900) # noqa: F405 |
| 304 | + + theme_minimal() # noqa: F405 |
| 305 | + + theme( # noqa: F405 |
| 306 | + axis_text=element_text(size=16, color="#666666", family="monospace"), # noqa: F405 |
| 307 | + axis_title=element_text(size=20, color="#333333", face="bold"), # noqa: F405 |
| 308 | + plot_title=element_text(size=26, color="#1A1A1A", face="bold"), # noqa: F405 |
| 309 | + plot_caption=element_text(size=14, color="#777777", face="italic"), # noqa: F405 |
| 310 | + legend_text=element_text(size=14), # noqa: F405 |
| 311 | + legend_title=element_text(size=16, face="bold"), # noqa: F405 |
| 312 | + legend_position="right", |
| 313 | + panel_grid_major=element_blank(), # noqa: F405 |
| 314 | + panel_grid_minor=element_blank(), # noqa: F405 |
| 315 | + plot_background=element_rect(fill="#FCFCFC", color="#FCFCFC"), # noqa: F405 |
| 316 | + panel_background=element_rect(fill="#FCFCFC", color="#FCFCFC"), # noqa: F405 |
| 317 | + axis_ticks=element_line(color="#CCCCCC", size=0.3), # noqa: F405 |
| 318 | + axis_line=element_line(color="#BBBBBB", size=0.5), # noqa: F405 |
| 319 | + plot_margin=[35, 45, 25, 25], |
| 320 | + ) |
| 321 | +) |
| 322 | + |
| 323 | +# Save |
| 324 | +export_ggsave(plot, filename="plot.png", path=".", scale=3) |
| 325 | +export_ggsave(plot, filename="plot.html", path=".") |
0 commit comments