Skip to content

Commit 416dff8

Browse files
murrayrmmurrayrm
committed
updated rlocus/pzmap docs, docstrings, tests
1 parent 8776875 commit 416dff8

11 files changed

Lines changed: 264 additions & 62 deletions

‎control/pzmap.py‎

Lines changed: 92 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -8,16 +8,6 @@
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-
#
2111

2212
import numpy as np
2313
from numpy import real, imag, linspace, exp, cos, sin, sqrt
@@ -34,7 +24,7 @@
3424
from .freqplot import _freqplot_defaults, _get_line_labels
3525
from . import config
3626

37-
__all__ = ['pole_zero_map', 'pole_zero_plot', 'pzmap']
27+
__all__ = ['pole_zero_map', 'pole_zero_plot', 'pzmap', 'PoleZeroData']
3828

3929

4030
# Define default parameter values for this module
@@ -50,22 +40,63 @@
5040
# Classes for keeping track of pzmap plots
5141
#
5242
# The PoleZeroData class keeps track of the information that is on a
53-
# pole-zero plot.
43+
# pole/zero plot.
5444
#
5545
# In addition to the locations of poles and zeros, you can also save a set
5646
# of gains and loci for use in generating a root locus plot. The gain
5747
# variable is a 1D array consisting of a list of increasing gains. The
5848
# loci variable is a 2D array indexed by [gain_idx, root_idx] that can be
5949
# plotted using the `pole_zero_plot` function.
6050
#
61-
# The PoleZeroList class is used to return a list of pole-zero plots. It
51+
# The PoleZeroList class is used to return a list of pole/zero plots. It
6252
# is a lightweight wrapper on the built-in list class that includes a
6353
# `plot` method, allowing plotting a set of root locus diagrams.
6454
#
6555
class PoleZeroData:
56+
"""Pole/zero data object.
57+
58+
This class is used as the return type for computing pole/zero responses
59+
and root locus diagrams. It contains information on the location of
60+
system poles and zeros, as well as the gains and loci for root locus
61+
diagrams.
62+
63+
Attributes
64+
----------
65+
poles : ndarray
66+
1D array of system poles.
67+
zeros : ndarray
68+
1D array of system zeros.
69+
gains : ndarray, optional
70+
1D array of gains for root locus plots.
71+
loci : ndarray, optiona
72+
2D array of poles, with each row corresponding to a gain.
73+
sysname : str, optional
74+
System name.
75+
sys : StateSpace or TransferFunction
76+
System corresponding to the data.
77+
78+
"""
6679
def __init__(
6780
self, poles, zeros, gains=None, loci=None, dt=None, sysname=None,
6881
sys=None):
82+
"""Create a pole/zero map object.
83+
84+
Parameters
85+
----------
86+
poles : ndarray
87+
1D array of system poles.
88+
zeros : ndarray
89+
1D array of system zeros.
90+
gains : ndarray, optional
91+
1D array of gains for root locus plots.
92+
loci : ndarray, optiona
93+
2D array of poles, with each row corresponding to a gain.
94+
sysname : str, optional
95+
System name.
96+
sys : StateSpace or TransferFunction
97+
System corresponding to the data.
98+
99+
"""
69100
self.poles = poles
70101
self.zeros = zeros
71102
self.gains = gains
@@ -79,17 +110,51 @@ def __iter__(self):
79110
return iter((self.poles, self.zeros))
80111

81112
def plot(self, *args, **kwargs):
113+
"""Plot the pole/zero data.
114+
115+
See :func:`~control.pole_zero_plot` for description of arguments
116+
and keywords.
117+
118+
"""
119+
# If this is a root locus plot, use rlocus defaults for grid
120+
if self.loci is not None:
121+
from .rlocus import _rlocus_defaults
122+
kwargs = kwargs.copy()
123+
kwargs['grid'] = config._get_param(
124+
'rlocus', 'grid', kwargs.get('grid', None), _rlocus_defaults)
125+
82126
return pole_zero_plot(self, *args, **kwargs)
83127

84128

85129
class PoleZeroList(list):
130+
"""List of PoleZeroData objects."""
86131
def plot(self, *args, **kwargs):
132+
"""Plot pole/zero data.
133+
134+
See :func:`~control.pole_zero_plot` for description of arguments
135+
and keywords.
136+
137+
"""
87138
return pole_zero_plot(self, *args, **kwargs)
88139

89140

90141
# Pole/zero map
91142
def pole_zero_map(sysdata):
92-
# TODO: add docstring (from old pzmap?)
143+
"""Compute the pole/zero map for an LTI system.
144+
145+
Parameters
146+
----------
147+
sys : LTI system (StateSpace or TransferFunction)
148+
Linear system for which poles and zeros are computed.
149+
150+
Returns
151+
-------
152+
pzmap_data : PoleZeroMap
153+
Pole/zero map containing the poles and zeros of the system. Use
154+
`pzmap_data.plot()` or `pole_zero_plot(pzmap_data)` to plot the
155+
pole/zero map.
156+
157+
"""
93158
# Convert the first argument to a list
94159
syslist = sysdata if isinstance(sysdata, (list, tuple)) else [sysdata]
95160

@@ -113,7 +178,6 @@ def pole_zero_plot(
113178
marker_size=None, marker_width=None, legend_loc='upper right',
114179
xlim=None, ylim=None, interactive=None, ax=None, scaling=None,
115180
initial_gain=None, **kwargs):
116-
# TODO: update docstring (see other response/plot functions for style)
117181
"""Plot a pole/zero map for a linear system.
118182
119183
If the system data include root loci, a root locus diagram for the
@@ -137,7 +201,7 @@ def pole_zero_plot(
137201
(legacy) If ``True`` a graph is generated with Matplotlib,
138202
otherwise the poles and zeros are only computed and returned.
139203
If this argument is present, the legacy value of poles and
140-
zero is returned.
204+
zeros is returned.
141205
142206
Returns
143207
-------
@@ -287,6 +351,7 @@ def pole_zero_plot(
287351

288352
# Plot the responses (and keep track of axes limits)
289353
xlim, ylim = ax.get_xlim(), ax.get_ylim()
354+
loci_count = 0
290355
for idx, response in enumerate(pzmap_responses):
291356
poles = response.poles
292357
zeros = response.zeros
@@ -331,9 +396,17 @@ def pole_zero_plot(
331396

332397
# TODO: add arrows to root loci (reuse Nyquist arrow code?)
333398

334-
# Set up the limits for the plot
335-
ax.set_xlim(xlim if xlim_user is None else xlim_user)
336-
ax.set_ylim(ylim if ylim_user is None else ylim_user)
399+
# Set the axis limits to something reasonable
400+
if any([response.loci is not None for response in pzmap_responses]):
401+
# Set up the limits for the plot using information from loci
402+
ax.set_xlim(xlim if xlim_user is None else xlim_user)
403+
ax.set_ylim(ylim if ylim_user is None else ylim_user)
404+
else:
405+
# No root loci => only set axis limits if users specified them
406+
if xlim_user is not None:
407+
ax.set_xlim(xlim_user)
408+
if ylim_user is not None:
409+
ax.set_ylim(ylim_user)
337410

338411
# List of systems that are included in this plot
339412
lines, labels = _get_line_labels(ax)
@@ -409,7 +482,6 @@ def _click_dispatcher(event):
409482

410483

411484
# Utility function to find gain corresponding to a click event
412-
# TODO: project onto the root locus plot (here or above?)
413485
def _find_root_locus_gain(event, sys, ax):
414486
# Get the current axis limits to set various thresholds
415487
xlim, ylim = ax.get_xlim(), ax.get_ylim()
@@ -469,7 +541,6 @@ def _mark_root_locus_gain(ax, sys, K):
469541

470542

471543
# Return a string identifying a clicked point
472-
# TODO: project onto the root locus plot (here or above?)
473544
def _create_root_locus_label(sys, K, s):
474545
# Figure out the damping ratio
475546
if isdtime(sys, strict=True):
@@ -482,7 +553,6 @@ def _create_root_locus_label(sys, K, s):
482553

483554

484555
# Utility function to compute limits for root loci
485-
# TODO: (note that sys is now available => code here may not be needed)
486556
def _compute_root_locus_limits(response):
487557
loci = response.loci
488558

‎control/rlocus.py‎

Lines changed: 57 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -25,23 +25,46 @@
2525
from .xferfcn import _convert_to_transfer_function
2626
from .exception import ControlMIMONotImplemented
2727
from . import config
28+
from .lti import LTI
2829
import warnings
2930

3031
__all__ = ['root_locus_map', 'root_locus_plot', 'root_locus', 'rlocus']
3132

3233
# Default values for module parameters
33-
# TODO: merge these with pzmap parameters (?)
3434
_rlocus_defaults = {
3535
'rlocus.grid': True,
36-
'rlocus.plotstr': 'C0', # default color cycle [TODO: not used?]
37-
'rlocus.print_gain': True,
38-
'rlocus.plot': True
3936
}
4037

4138

4239
# Root locus map
43-
# TODO: add docstring
4440
def root_locus_map(sysdata, gains=None):
41+
"""Compute the root locus map for an LTI system.
42+
43+
Calculate the root locus by finding the roots of 1 + k * G(s) where G
44+
is a linear system with transfer function num(s)/den(s) and each k is
45+
an element of kvect.
46+
47+
Parameters
48+
----------
49+
sys : LTI system or list of LTI systems
50+
Linear input/output systems (SISO only, for now).
51+
kvect : array_like, optional
52+
Gains to use in computing plot of closed-loop poles.
53+
54+
Returns
55+
-------
56+
rldata : PoleZeroData or list of PoleZeroData
57+
Root locus data object(s) corresponding to the . The loci of
58+
the root locus diagram are available in the array
59+
`rldata.loci`, indexed by the gain index and the locus index,
60+
and the gains are in the array `rldata.gains`.
61+
62+
Notes
63+
-----
64+
For backward compatibility, the `rldata` return object can be
65+
assigned to the tuple `roots, gains`.
66+
67+
"""
4568
from .pzmap import PoleZeroData, PoleZeroList
4669

4770
# Convert the first argument to a list
@@ -75,6 +98,7 @@ def root_locus_map(sysdata, gains=None):
7598

7699
def root_locus_plot(
77100
sysdata, kvect=None, grid=None, plot=None, **kwargs):
101+
78102
"""Root locus plot.
79103
80104
Calculate the root locus by finding the roots of 1 + k * G(s) where G
@@ -83,7 +107,7 @@ def root_locus_plot(
83107
84108
Parameters
85109
----------
86-
sys : LTI object
110+
sysdata : PoleZeroMap or LTI object or list
87111
Linear input/output systems (SISO only, for now).
88112
kvect : array_like, optional
89113
Gains to use in computing plot of closed-loop poles.
@@ -93,16 +117,9 @@ def root_locus_plot(
93117
ylim : tuple or list, optional
94118
Set limits of y axis, normally with tuple
95119
(see :doc:`matplotlib:api/axes_api`).
96-
plotstr : :func:`matplotlib.pyplot.plot` format string, optional
97-
plotting style specification
98-
TODO: check
99-
plot : boolean, optional
100-
If True (default), plot root locus diagram.
101-
TODO: legacy
102-
print_gain : bool
103-
If True (default), report mouse clicks when close to the root locus
104-
branches, calculate gain, damping and print.
105-
TODO: update
120+
plot : bool, optional
121+
(legacy) If given, `root_locus_plot` returns the legacy return values
122+
of roots and gains. If False, just return the values with no plot.
106123
grid : bool or str, optional
107124
If `True` plot omega-damping grid, if `False` show imaginary axis
108125
for continuous time systems, unit circle for discrete time systems.
@@ -111,19 +128,29 @@ def root_locus_plot(
111128
ax : :class:`matplotlib.axes.Axes`
112129
Axes on which to create root locus plot
113130
initial_gain : float, optional
114-
Specify the initial gain to use when marking current gain. [TODO: update]
131+
Mark the point on the root locus diagram corresponding to the
132+
given gain.
115133
116-
Returns (TODO: update)
134+
Returns
117135
-------
118-
roots : ndarray
119-
Closed-loop root locations, arranged in which each row corresponds
120-
to a gain in gains.
121-
gains : ndarray
122-
Gains used. Same as kvect keyword argument if provided.
136+
lines : List of Line2D
137+
Array of Line2D objects for each set of markers in the plot. The
138+
shape of the array is given by (nsys, 2) where nsys is the number
139+
of systems or Nyquist responses passed to the function. The second
140+
index specifies the pzmap object type:
141+
142+
* lines[idx, 0]: poles
143+
* lines[idx, 1]: zeros
144+
145+
roots, gains : ndarray
146+
(legacy) If the `plot` keyword is given, returns the
147+
closed-loop root locations, arranged such that each row
148+
corresponds to a gain in gains, and the array of gains (ame as
149+
kvect keyword argument if provided).
123150
124151
Notes
125152
-----
126-
The root_locus function calls matplotlib.pyplot.axis('equal'), which
153+
The root_locus_plot function calls matplotlib.pyplot.axis('equal'), which
127154
means that trying to reset the axis limits may not behave as expected.
128155
To change the axis limits, use matplotlib.pyplot.gca().axis('auto') and
129156
then set the axis limits to the desired values.
@@ -132,10 +159,14 @@ def root_locus_plot(
132159
from .pzmap import pole_zero_plot
133160

134161
# Set default parameters
135-
# TODO: move this to pole_zero_plot()
136162
grid = config._get_param('rlocus', 'grid', grid, _rlocus_defaults)
137163

138-
responses = root_locus_map(sysdata, gains=kvect)
164+
if isinstance(sysdata, list) and all(
165+
[isinstance(sys, LTI) for sys in sysdata]) or \
166+
isinstance(sysdata, LTI):
167+
responses = root_locus_map(sysdata, gains=kvect)
168+
else:
169+
responses = sysdata
139170

140171
#
141172
# Process `plot` keyword

‎control/tests/kwargs_test.py‎

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -176,6 +176,8 @@ def test_matplotlib_kwargs(function, nsysargs, moreargs, kwargs, mplcleanup):
176176
(control.frequency_response, control.bode, True),
177177
(control.frequency_response, control.bode_plot, True),
178178
(control.nyquist_response, control.nyquist_plot, False),
179+
(control.pole_zero_map, control.pole_zero_plot, False),
180+
(control.root_locus_map, control.root_locus_plot, False),
179181
])
180182
def test_response_plot_kwargs(data_fcn, plot_fcn, mimo):
181183
# Create a system for testing
@@ -194,16 +196,18 @@ def test_response_plot_kwargs(data_fcn, plot_fcn, mimo):
194196
plot_fcn(response)
195197

196198
# Now add an unrecognized keyword and make sure there is an error
197-
with pytest.raises(AttributeError,
198-
match="(has no property|unexpected keyword)"):
199+
with pytest.raises(
200+
(AttributeError, TypeError),
201+
match="(has no property|unexpected keyword|unrecognized keyword)"):
199202
plot_fcn(response, unknown=None)
200203

201204
# Call the plotting function via the response and make sure it works
202205
response.plot()
203206

204207
# Now add an unrecognized keyword and make sure there is an error
205-
with pytest.raises(AttributeError,
206-
match="(has no property|unexpected keyword)"):
208+
with pytest.raises(
209+
(AttributeError, TypeError),
210+
match="(has no property|unexpected keyword|unrecognized keyword)"):
207211
response.plot(unknown=None)
208212

209213
#
@@ -277,6 +281,7 @@ def test_response_plot_kwargs(data_fcn, plot_fcn, mimo):
277281
'flatsys.LinearFlatSystem.__init__': test_unrecognized_kwargs,
278282
'NonlinearIOSystem.linearize': test_unrecognized_kwargs,
279283
'NyquistResponseData.plot': test_response_plot_kwargs,
284+
'PoleZeroData.plot': test_response_plot_kwargs,
280285
'InterconnectedSystem.__init__':
281286
interconnect_test.test_interconnect_exceptions,
282287
'StateSpace.__init__':

0 commit comments

Comments
 (0)