Skip to content

Commit 6b32ab1

Browse files
feat(letsplot): implement root-locus-basic (#5121)
## Implementation: `root-locus-basic` - letsplot Implements the **letsplot** version of `root-locus-basic`. **File:** `plots/root-locus-basic/implementations/letsplot.py` **Parent Issue:** #4414 --- :robot: *[impl-generate workflow](https://github.com/MarkusNeusinger/pyplots/actions/runs/23362982902)* --------- Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
1 parent 9ff7398 commit 6b32ab1

2 files changed

Lines changed: 562 additions & 0 deletions

File tree

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

Comments
 (0)