Skip to content

Commit 57dac87

Browse files
murrayrmmurrayrm
committed
initial refactoring of root_locus_{map,plot}
1 parent 1c4322c commit 57dac87

7 files changed

Lines changed: 308 additions & 481 deletions

File tree

‎control/grid.py‎

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -99,15 +99,19 @@ def sgrid():
9999
ax.axis[:].major_ticklabels.set_visible(visible)
100100
ax.axis[:].major_ticks.set_visible(False)
101101
ax.axis[:].invert_ticklabel_direction()
102+
ax.axis[:].major_ticklabels.set_color('gray')
102103

103104
ax.axis["wnxneg"] = axis = ax.new_floating_axis(0, 180)
104105
axis.set_ticklabel_direction("-")
105106
axis.label.set_visible(False)
107+
106108
ax.axis["wnxpos"] = axis = ax.new_floating_axis(0, 0)
107109
axis.label.set_visible(False)
110+
108111
ax.axis["wnypos"] = axis = ax.new_floating_axis(0, 90)
109112
axis.label.set_visible(False)
110-
axis.set_axis_direction("left")
113+
axis.set_axis_direction("right")
114+
111115
ax.axis["wnyneg"] = axis = ax.new_floating_axis(0, 270)
112116
axis.label.set_visible(False)
113117
axis.set_axis_direction("left")
@@ -149,9 +153,10 @@ def _final_setup(ax):
149153
plt.axis('equal')
150154

151155

152-
def nogrid(dt=None):
156+
def nogrid(dt=None, ax=None):
153157
fig = plt.gcf()
154-
ax = plt.axes()
158+
if ax is None:
159+
ax = fig.gca()
155160

156161
# Draw the unit circle for discrete time systems
157162
if isdtime(dt=dt, strict=True):

‎control/pzmap.py‎

Lines changed: 166 additions & 57 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,16 @@
88
# storing and plotting pole/zero and root locus diagrams. (The actual
99
# computation of root locus diagrams is in rlocus.py.)
1010
#
11+
# TODO (Sep 2023):
12+
# * Test out ability to set line styles
13+
# - Make compatible with other plotting (and refactor?)
14+
# - Allow line fmt to be overwritten (including color=CN for different
15+
# colors for each segment?)
16+
# * Add ability to set style of root locus click point
17+
# - Sort out where default parameter values should live (pzmap vs rlocus)
18+
# * Decide whether click functionality should be in rlocus.py
19+
# * Add back print_gain option to sisotool (and any other options)
20+
#
1121

1222
import numpy as np
1323
from numpy import real, imag, linspace, exp, cos, sin, sqrt
@@ -24,32 +34,44 @@
2434
from .freqplot import _freqplot_defaults, _get_line_labels
2535
from . import config
2636

27-
__all__ = ['pole_zero_map', 'root_locus_map', 'pole_zero_plot', 'pzmap']
37+
__all__ = ['pole_zero_map', 'pole_zero_plot', 'pzmap']
2838

2939

3040
# Define default parameter values for this module
3141
_pzmap_defaults = {
32-
'pzmap.grid': False, # Plot omega-damping grid
33-
'pzmap.marker_size': 6, # Size of the markers
34-
'pzmap.marker_width': 1.5, # Width of the markers
42+
'pzmap.grid': False, # Plot omega-damping grid
43+
'pzmap.marker_size': 6, # Size of the markers
44+
'pzmap.marker_width': 1.5, # Width of the markers
45+
'pzmap.expansion_factor': 2, # Amount to scale plots beyond features
3546
}
3647

37-
48+
#
3849
# Classes for keeping track of pzmap plots
39-
class RootLocusList(list):
40-
def plot(self, *args, **kwargs):
41-
return pole_zero_plot(self, *args, **kwargs)
42-
43-
44-
class RootLocusData:
50+
#
51+
# The PoleZeroData class keeps track of the information that is on a
52+
# pole-zero plot.
53+
#
54+
# In addition to the locations of poles and zeros, you can also save a set
55+
# of gains and loci for use in generating a root locus plot. The gain
56+
# variable is a 1D array consisting of a list of increasing gains. The
57+
# loci variable is a 2D array indexed by [gain_idx, root_idx] that can be
58+
# plotted using the `pole_zero_plot` function.
59+
#
60+
# The PoleZeroList class is used to return a list of pole-zero plots. It
61+
# is a lightweight wrapper on the built-in list class that includes a
62+
# `plot` method, allowing plotting a set of root locus diagrams.
63+
#
64+
class PoleZeroData:
4565
def __init__(
46-
self, poles, zeros, gains=None, loci=None, dt=None, sysname=None):
66+
self, poles, zeros, gains=None, loci=None, dt=None, sysname=None,
67+
sys=None):
4768
self.poles = poles
4869
self.zeros = zeros
4970
self.gains = gains
5071
self.loci = loci
5172
self.dt = dt
5273
self.sysname = sysname
74+
self.sys = sys
5375

5476
# Implement functions to allow legacy assignment to tuple
5577
def __iter__(self):
@@ -59,54 +81,25 @@ def plot(self, *args, **kwargs):
5981
return pole_zero_plot(self, *args, **kwargs)
6082

6183

84+
class PoleZeroList(list):
85+
def plot(self, *args, **kwargs):
86+
return pole_zero_plot(self, *args, **kwargs)
87+
88+
6289
# Pole/zero map
6390
def pole_zero_map(sysdata):
91+
# TODO: add docstring (from old pzmap?)
6492
# Convert the first argument to a list
6593
syslist = sysdata if isinstance(sysdata, (list, tuple)) else [sysdata]
6694

6795
responses = []
6896
for idx, sys in enumerate(syslist):
6997
responses.append(
70-
RootLocusData(
98+
PoleZeroData(
7199
sys.poles(), sys.zeros(), dt=sys.dt, sysname=sys.name))
72100

73101
if isinstance(sysdata, (list, tuple)):
74-
return RootLocusList(responses)
75-
else:
76-
return responses[0]
77-
78-
79-
# Root locus map
80-
# TODO: use rlocus.py computation instead
81-
def root_locus_map(sysdata, gains=None):
82-
# Convert the first argument to a list
83-
syslist = sysdata if isinstance(sysdata, (list, tuple)) else [sysdata]
84-
85-
responses = []
86-
for idx, sys in enumerate(syslist):
87-
from .rlocus import _systopoly1d, _default_gains
88-
from .rlocus import _RLFindRoots, _RLSortRoots
89-
90-
if not sys.issiso():
91-
raise ControlMIMONotImplemented(
92-
"sys must be single-input single-output (SISO)")
93-
94-
# Convert numerator and denominator to polynomials if they aren't
95-
nump, denp = _systopoly1d(sys[0, 0])
96-
97-
if gains is None:
98-
kvect, root_array, _, _ = _default_gains(nump, denp, None, None)
99-
else:
100-
kvect = np.atleast_1d(gains)
101-
root_array = _RLFindRoots(nump, denp, kvect)
102-
root_array = _RLSortRoots(root_array)
103-
104-
responses.append(RootLocusData(
105-
sys.poles(), sys.zeros(), kvect, root_array,
106-
dt=sys.dt, sysname=sys.name))
107-
108-
if isinstance(sysdata, (list, tuple)):
109-
return RootLocusList(responses)
102+
return PoleZeroList(responses)
110103
else:
111104
return responses[0]
112105

@@ -117,12 +110,14 @@ def root_locus_map(sysdata, gains=None):
117110
def pole_zero_plot(
118111
data, plot=None, grid=None, title=None, marker_color=None,
119112
marker_size=None, marker_width=None, legend_loc='upper right',
120-
xlim=None, ylim=None, **kwargs):
113+
xlim=None, ylim=None, interactive=False, ax=None,
114+
initial_gain=None, **kwargs):
115+
# TODO: update docstring (see other response/plot functions for style)
121116
"""Plot a pole/zero map for a linear system.
122117
123118
Parameters
124119
----------
125-
sysdata: List of RootLocusData objects or LTI systems
120+
sysdata: List of PoleZeroData objects or LTI systems
126121
List of pole/zero response data objects generated by pzmap_response
127122
or rootlocus_response() that are to be plotted. If a list of systems
128123
is given, the poles and zeros of those systems will be plotted.
@@ -165,6 +160,7 @@ def pole_zero_plot(
165160
freqplot_rcParams = config._get_param(
166161
'freqplot', 'rcParams', kwargs, _freqplot_defaults,
167162
pop=True, last=True)
163+
user_ax = ax
168164

169165
# If argument was a singleton, turn it into a tuple
170166
if not isinstance(data, (list, tuple)):
@@ -175,7 +171,7 @@ def pole_zero_plot(
175171
sys, (StateSpace, TransferFunction)) for sys in data]):
176172
# Get the response, popping off keywords used there
177173
pzmap_responses = pole_zero_map(data)
178-
elif all([isinstance(d, RootLocusData) for d in data]):
174+
elif all([isinstance(d, PoleZeroData) for d in data]):
179175
pzmap_responses = data
180176
else:
181177
raise TypeError("unknown system data type")
@@ -198,8 +194,13 @@ def pole_zero_plot(
198194

199195
# Initialize the figure
200196
# TODO: turn into standard utility function (from plotutil.py?)
201-
fig = plt.gcf()
202-
axs = fig.get_axes()
197+
if user_ax is None:
198+
fig = plt.gcf()
199+
axs = fig.get_axes()
200+
else:
201+
fig = ax.figure
202+
axs = [ax]
203+
203204
if len(axs) > 1:
204205
# Need to generate a new figure
205206
fig, axs = plt.figure(), []
@@ -275,6 +276,10 @@ def pole_zero_plot(
275276
xlim = [min(xlim[0], resp_xlim[0]), max(xlim[1], resp_xlim[1])]
276277
ylim = [min(ylim[0], resp_ylim[0]), max(ylim[1], resp_ylim[1])]
277278

279+
# Plot the initial gain, if given
280+
if initial_gain is not None:
281+
_mark_root_locus_gain(ax, response.sys, initial_gain)
282+
278283
# TODO: add arrows to root loci (reuse Nyquist arrow code?)
279284

280285
# Set up the limits for the plot
@@ -313,8 +318,36 @@ def pole_zero_plot(
313318
# Add the title
314319
if title is None:
315320
title = "Pole/zero plot for " + ", ".join(labels)
316-
with plt.rc_context(freqplot_rcParams):
317-
fig.suptitle(title)
321+
if user_ax is None:
322+
with plt.rc_context(freqplot_rcParams):
323+
fig.suptitle(title)
324+
325+
# Add dispather to handle choosing a point on the diagram
326+
if interactive:
327+
if len(pzmap_responses) > 1:
328+
raise NotImplementedError(
329+
"interactive mode only allowed for single system")
330+
elif pzmap_responses[0].sys == None:
331+
raise SystemError("missing system information")
332+
else:
333+
sys = pzmap_responses[0].sys
334+
335+
# Define function to handle mouse clicks
336+
def _click_dispatcher(event):
337+
# Find the gain corresponding to the clicked point
338+
K, s = _find_root_locus_gain(event, sys, ax)
339+
340+
if K is not None:
341+
# Mark the gain on the root locus diagram
342+
_mark_root_locus_gain(ax, sys, K)
343+
344+
# Display the parameters in the axes title
345+
with plt.rc_context(freqplot_rcParams):
346+
ax.set_title(_create_root_locus_label(sys, K, s))
347+
348+
ax.figure.canvas.draw()
349+
350+
fig.canvas.mpl_connect('button_release_event', _click_dispatcher)
318351

319352
# Legacy processing: return locations of poles and zeros as a tuple
320353
if plot is True:
@@ -326,7 +359,82 @@ def pole_zero_plot(
326359
return out
327360

328361

362+
# Utility function to find gain corresponding to a click event
363+
# TODO: project onto the root locus plot (here or above?)
364+
def _find_root_locus_gain(event, sys, ax):
365+
# Get the current axis limits to set various thresholds
366+
xlim, ylim = ax.get_xlim(), ax.get_ylim()
367+
368+
# Catch type error when event click is in the figure but not in an axis
369+
try:
370+
s = complex(event.xdata, event.ydata)
371+
K = -1. / sys(s)
372+
K_xlim = -1. / sys(
373+
complex(event.xdata + 0.05 * abs(xlim[1] - xlim[0]), event.ydata))
374+
K_ylim = -1. / sys(
375+
complex(event.xdata, event.ydata + 0.05 * abs(ylim[1] - ylim[0])))
376+
377+
except TypeError:
378+
K = float('inf')
379+
K_xlim = float('inf')
380+
K_ylim = float('inf')
381+
382+
#
383+
# Compute tolerances for deciding if we clicked on the root locus
384+
#
385+
# This is a bit of black magic that sets some limits for how close we
386+
# need to be to the root locus in order to consider it a click on the
387+
# actual curve. Otherwise, we will just ignore the click.
388+
389+
x_tolerance = 0.1 * abs((xlim[1] - xlim[0]))
390+
y_tolerance = 0.1 * abs((ylim[1] - ylim[0]))
391+
gain_tolerance = np.mean([x_tolerance, y_tolerance]) * 0.1 + \
392+
0.1 * max([abs(K_ylim.imag/K_ylim.real), abs(K_xlim.imag/K_xlim.real)])
393+
394+
# Decide whether to pay attention to this event
395+
if abs(K.real) > 1e-8 and abs(K.imag / K.real) < gain_tolerance and \
396+
event.inaxes == ax.axes and K.real > 0.:
397+
return K.real, s
398+
399+
else:
400+
return None, s
401+
402+
403+
# Mark points corresponding to a given gain on root locus plot
404+
def _mark_root_locus_gain(ax, sys, K):
405+
from .rlocus import _systopoly1d, _RLFindRoots
406+
407+
# Remove any previous gain points
408+
for line in reversed(ax.lines):
409+
if line.get_label() == '_gain_point':
410+
line.remove()
411+
del line
412+
413+
# Visualise clicked point, displaying all roots
414+
# TODO: allow marker parameters to be set
415+
nump, denp = _systopoly1d(sys)
416+
root_array = _RLFindRoots(nump, denp, K.real)
417+
ax.plot(
418+
[root.real for root in root_array], [root.imag for root in root_array],
419+
marker='s', markersize=6, zorder=20, label='_gain_point', color='k')
420+
421+
422+
# Return a string identifying a clicked point
423+
# TODO: project onto the root locus plot (here or above?)
424+
def _create_root_locus_label(sys, K, s):
425+
# Figure out the damping ratio
426+
if isdtime(sys, strict=True):
427+
zeta = -np.cos(np.angle(np.log(s)))
428+
else:
429+
zeta = -1 * s.real / abs(s)
430+
431+
return "Clicked at: %.4g%+.4gj gain = %.4g damping = %.4g" % \
432+
(s.real, s.imag, K.real, zeta)
433+
434+
329435
# Utility function to compute limits for root loci
436+
# TODO: compare to old code and recapture functionality (especially asymptotes)
437+
# TODO: (note that sys is now available => code here may not be needed)
330438
def _compute_root_locus_limits(loci):
331439
# Go through each locus
332440
xlim, ylim = [0, 0], 0
@@ -345,7 +453,8 @@ def _compute_root_locus_limits(loci):
345453
ylim = max(ylim, np.max(ypeaks))
346454

347455
# Adjust the limits to include some space around features
348-
rho = 1.5
456+
# TODO: use _k_max and project out to max k for all value?
457+
rho = config._get_param('pzmap', 'expansion_factor')
349458
xlim[0] = rho * xlim[0] if xlim[0] < 0 else 0
350459
xlim[1] = rho * xlim[1] if xlim[1] > 0 else 0
351460
ylim = rho * ylim if ylim > 0 else np.max(np.abs(xlim))

0 commit comments

Comments
 (0)