Skip to content

Commit a7634ac

Browse files
feat(altair): implement root-locus-basic (#5134)
## Implementation: `root-locus-basic` - altair Implements the **altair** version of `root-locus-basic`. **File:** `plots/root-locus-basic/implementations/altair.py` **Parent Issue:** #4414 --- :robot: *[impl-generate workflow](https://github.com/MarkusNeusinger/pyplots/actions/runs/23362981953)* --------- 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 afdf33e commit a7634ac

2 files changed

Lines changed: 537 additions & 0 deletions

File tree

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

Comments
 (0)