Skip to content

Commit 4fe5f53

Browse files
murrayrmmurrayrm
committed
implement uniform legend processing + unit test/doc updates
1 parent fc09a85 commit 4fe5f53

10 files changed

Lines changed: 478 additions & 175 deletions

File tree

‎control/ctrlplot.py‎

Lines changed: 52 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,8 @@
3030
#
3131
# # Figure out the shape of the plot and find/create axes
3232
# fig, ax_array = _process_ax_keyword(ax, (nrows, ncols), rcParams)
33+
# legend_loc, legend_map, show_legend = _process_legend_keywords(
34+
# kwargs, (nrows, ncols), 'center right')
3335
#
3436
# # Customize axes (curvilinear grids, shared axes, etc)
3537
#
@@ -51,14 +53,17 @@
5153
# ax_array[i, j].set_ylabel("y label")
5254
#
5355
# # Create legends
54-
# legend_map = _process_legend_keywords(kwargs)
55-
# for i, j in itertools.product(range(nrows), range(ncols)):
56-
# if legend_map[i, j] is not None:
57-
# lines = ax_array[i, j].get_lines()
58-
# labels = _make_legend_labels(lines)
59-
# if len(labels) > 1:
60-
# legend_array[i, j] = ax.legend(
61-
# lines, labels, loc=legend_map[i, j])
56+
# if show_legend != False:
57+
# legend_array = np.full(ax_array.shape, None, dtype=object)
58+
# for i, j in itertools.product(range(nrows), range(ncols)):
59+
# if legend_map[i, j] is not None:
60+
# lines = ax_array[i, j].get_lines()
61+
# labels = _make_legend_labels(lines)
62+
# if len(labels) > 1:
63+
# legend_array[i, j] = ax.legend(
64+
# lines, labels, loc=legend_map[i, j])
65+
# else:
66+
# legend_array = None
6267
#
6368
# # Update the plot title (only if ax was not given)
6469
# sysnames = [response.sysname for response in data]
@@ -131,7 +136,7 @@ class ControlPlot(object):
131136
figure : :class:`matplotlib:Figure`
132137
Figure on which the Axes are drawn.
133138
legend : :class:`matplotlib:.legend.Legend` (instance or ndarray)
134-
Legend object(s) for the plat. If more than one legend is
139+
Legend object(s) for the plot. If more than one legend is
135140
included, this will be an array with each entry being either None
136141
(for no legend) or a legend object.
137142
@@ -436,8 +441,46 @@ def _get_line_labels(ax, use_color=True):
436441
return lines, [label for label, color in labels_colors]
437442

438443

444+
def _process_legend_keywords(
445+
kwargs, shape=None, default_loc='center right'):
446+
legend_loc = kwargs.pop('legend_loc', None)
447+
if shape is None and 'legend_map' in kwargs:
448+
raise TypeError("unexpected keyword argument 'legend_map'")
449+
else:
450+
legend_map = kwargs.pop('legend_map', None)
451+
show_legend = kwargs.pop('show_legend', None)
452+
453+
# If legend_loc or legend_map were given, always show the legend
454+
if legend_loc is False or legend_map is False:
455+
if show_legend is True:
456+
warnings.warn(
457+
"show_legend ignored; legend_loc or legend_map was given")
458+
show_legend = False
459+
legend_loc = legend_map = None
460+
elif legend_loc is not None or legend_map is not None:
461+
if show_legend is False:
462+
warnings.warn(
463+
"show_legend ignored; legend_loc or legend_map was given")
464+
show_legend = True
465+
466+
if legend_loc is None:
467+
legend_loc = default_loc
468+
elif not isinstance(legend_loc, (int, str)):
469+
raise ValueError("legend_loc must be string or int")
470+
471+
# Make sure the legend map is the right size
472+
if legend_map is not None:
473+
legend_map = np.atleast_2d(legend_map)
474+
if legend_map.shape != shape:
475+
raise ValueError("legend_map shape just match axes shape")
476+
477+
return legend_loc, legend_map, show_legend
478+
479+
439480
# Utility function to make legend labels
440481
def _make_legend_labels(labels, ignore_common=False):
482+
if len(labels) == 1:
483+
return labels
441484

442485
# Look for a common prefix (up to a space)
443486
common_prefix = commonprefix(labels)
@@ -474,7 +517,6 @@ def _update_plot_title(
474517
rcParams = config._get_param('ctrlplot', 'rcParams', kwargs, pop=True)
475518
if rcParams is None:
476519
rcParams = _ctrlplot_rcParams
477-
print(f"{rcParams['figure.titlesize']=}")
478520

479521
if use_existing:
480522
# Get the current title, if it exists

‎control/descfcn.py‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -423,6 +423,12 @@ def describing_function_plot(
423423
point_label : str, optional
424424
Formatting string used to label intersection points on the Nyquist
425425
plot. Defaults to "%5.2g @ %-5.2g". Set to `None` to omit labels.
426+
ax : matplotlib.axes.Axes, optional
427+
The matplotlib axes to draw the figure on. If not specified and
428+
the current figure has a single axes, that axes is used.
429+
Otherwise, a new figure is created.
430+
title : str, optional
431+
Set the title of the plot. Defaults to plot type and system name(s).
426432
427433
Returns
428434
-------
@@ -475,6 +481,11 @@ def describing_function_plot(
475481
else:
476482
raise TypeError("1, 3, or 4 position arguments required")
477483

484+
# Don't allow legend keyword arguments
485+
for kw in ['legend_loc', 'legend_map', 'show_legend']:
486+
if kw in kwargs:
487+
raise TypeError(f"unexpected keyword argument '{kw}'")
488+
478489
# Create a list of lines for the output
479490
lines = np.empty(2, dtype=object)
480491

‎control/freqplot.py‎

Lines changed: 71 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,8 @@
2121
from .bdalg import feedback
2222
from .ctrlplot import ControlPlot, _add_arrows_to_line2D, _ctrlplot_rcParams, \
2323
_find_axes_center, _get_line_labels, _make_legend_labels, \
24-
_process_ax_keyword, _process_line_labels, _update_plot_title
24+
_process_ax_keyword, _process_legend_keywords, _process_line_labels, \
25+
_update_plot_title
2526
from .ctrlutil import unwrap
2627
from .exception import ControlMIMONotImplemented
2728
from .frdata import FrequencyResponseData
@@ -87,8 +88,7 @@ def bode_plot(
8788
plot=None, plot_magnitude=True, plot_phase=None,
8889
overlay_outputs=None, overlay_inputs=None, phase_label=None,
8990
magnitude_label=None, label=None, display_margins=None,
90-
margins_method='best', legend_map=None, legend_loc=None,
91-
sharex=None, sharey=None, title=None, **kwargs):
91+
margins_method='best', title=None, sharex=None, sharey=None, **kwargs):
9292
"""Bode plot for a system.
9393
9494
Plot the magnitude and phase of the frequency response over a
@@ -142,25 +142,33 @@ def bode_plot(
142142
143143
Other Parameters
144144
----------------
145-
ax : array of Axes
146-
The matplotlib Axes to draw the figure on. If not specified, the
147-
Axes for the current figure are used or, if there is no current
148-
figure with the correct number and shape of Axes, a new figure is
145+
ax : array of matplotlib.axes.Axes, optional
146+
The matplotlib axes to draw the figure on. If not specified, the
147+
axes for the current figure are used or, if there is no current
148+
figure with the correct number and shape of axes, a new figure is
149149
created. The shape of the array must match the shape of the
150150
plotted data.
151-
grid : bool
151+
grid : bool, optional
152152
If True, plot grid lines on gain and phase plots. Default is set by
153153
`config.defaults['freqplot.grid']`.
154-
initial_phase : float
154+
initial_phase : float, optional
155155
Set the reference phase to use for the lowest frequency. If set, the
156156
initial phase of the Bode plot will be set to the value closest to the
157157
value specified. Units are in either degrees or radians, depending on
158158
the `deg` parameter. Default is -180 if wrap_phase is False, 0 if
159159
wrap_phase is True.
160-
label : str or array-like of str
160+
label : str or array_like of str, optional
161161
If present, replace automatically generated label(s) with the given
162162
label(s). If sysdata is a list, strings should be specified for each
163163
system. If MIMO, strings required for each system, output, and input.
164+
legend_map : array of str, optional
165+
Location of the legend for multi-axes plots. Specifies an array
166+
of legend location strings matching the shape of the subplots, with
167+
each entry being either None (for no legend) or a legend location
168+
string (see :func:`~matplotlib.pyplot.legend`).
169+
legend_loc : int or str, optional
170+
Include a legend in the given location. Default is 'center right',
171+
with no legend for a single response. Use False to supress legend.
164172
margins_method : str, optional
165173
Method to use in computing margins (see :func:`stability_margins`).
166174
omega_limits : array_like of two values
@@ -179,6 +187,10 @@ def bode_plot(
179187
rcParams : dict
180188
Override the default parameters used for generating plots.
181189
Default is set by config.default['freqplot.rcParams'].
190+
show_legend : bool, optional
191+
Force legend to be shown if ``True`` or hidden if ``False``. If
192+
``None``, then show legend when there is more than one line on an
193+
axis or ``legend_loc`` or ``legend_map`` has been specified.
182194
title : str, optional
183195
Set the title of the plot. Defaults to plot type and system name(s).
184196
wrap_phase : bool or float
@@ -478,8 +490,10 @@ def bode_plot(
478490
if kw not in kwargs or kwargs[kw] is None:
479491
kwargs[kw] = config.defaults['freqplot.' + kw]
480492

481-
fig, ax_array = _process_ax_keyword(ax, (
482-
nrows, ncols), squeeze=False, rcParams=rcParams, clear_text=True)
493+
fig, ax_array = _process_ax_keyword(
494+
ax, (nrows, ncols), squeeze=False, rcParams=rcParams, clear_text=True)
495+
legend_loc, legend_map, show_legend = _process_legend_keywords(
496+
kwargs, (nrows,ncols), 'center right')
483497

484498
# Get the values for sharing axes limits
485499
share_magnitude = kwargs.pop('share_magnitude', None)
@@ -989,21 +1003,15 @@ def gen_zero_centered_series(val_min, val_max, period):
9891003
# different response (system).
9901004
#
9911005

992-
# Figure out where to put legends
993-
if legend_map is None:
994-
legend_map = np.full(ax_array.shape, None, dtype=object)
995-
if legend_loc == None:
996-
legend_loc = 'center right'
997-
998-
# TODO: add in additional processing later
999-
1000-
# Put legend in the upper right
1001-
legend_map[0, -1] = legend_loc
1002-
10031006
# Create axis legends
1004-
legend_array = np.full(ax_array.shape, None, dtype=object)
1005-
for i in range(nrows):
1006-
for j in range(ncols):
1007+
if show_legend != False:
1008+
# Figure out where to put legends
1009+
if legend_map is None:
1010+
legend_map = np.full(ax_array.shape, None, dtype=object)
1011+
legend_map[0, -1] = legend_loc
1012+
1013+
legend_array = np.full(ax_array.shape, None, dtype=object)
1014+
for i, j in itertools.product(range(nrows), range(ncols)):
10071015
if legend_map[i, j] is None:
10081016
continue
10091017
ax = ax_array[i, j]
@@ -1016,10 +1024,13 @@ def gen_zero_centered_series(val_min, val_max, period):
10161024
ignore_common=line_labels is not None)
10171025

10181026
# Generate the label, if needed
1019-
if len(labels) > 1:
1027+
if show_legend == True or len(labels) > 1:
10201028
with plt.rc_context(rcParams):
1029+
print(f"{lines=}, {labels=}")
10211030
legend_array[i, j] = ax.legend(
10221031
lines, labels, loc=legend_map[i, j])
1032+
else:
1033+
legend_array = None
10231034

10241035
#
10251036
# Legacy return pocessing
@@ -1476,7 +1487,7 @@ def nyquist_response(
14761487

14771488
def nyquist_plot(
14781489
data, omega=None, plot=None, label_freq=0, color=None, label=None,
1479-
return_contour=None, title=None, legend_loc='upper right', ax=None,
1490+
return_contour=None, title=None, ax=None,
14801491
unit_circle=False, mt_circles=None, ms_circles=None, **kwargs):
14811492
"""Nyquist plot for a system.
14821493
@@ -1550,6 +1561,10 @@ def nyquist_plot(
15501561
8 and can be set using config.defaults['nyquist.arrow_size'].
15511562
arrow_style : matplotlib.patches.ArrowStyle, optional
15521563
Define style used for Nyquist curve arrows (overrides `arrow_size`).
1564+
ax : matplotlib.axes.Axes, optional
1565+
The matplotlib axes to draw the figure on. If not specified and
1566+
the current figure has a single axes, that axes is used.
1567+
Otherwise, a new figure is created.
15531568
encirclement_threshold : float, optional
15541569
Define the threshold for generating a warning if the number of net
15551570
encirclements is a non-integer value. Default value is 0.05 and can
@@ -1564,13 +1579,16 @@ def nyquist_plot(
15641579
Amount to indent the Nyquist contour around poles on or near the
15651580
imaginary axis. Portions of the Nyquist plot corresponding to indented
15661581
portions of the contour are plotted using a different line style.
1567-
label : str or array-like of str
1582+
label : str or array_like of str, optional
15681583
If present, replace automatically generated label(s) with the given
15691584
label(s). If sysdata is a list, strings should be specified for each
15701585
system.
15711586
label_freq : int, optiona
15721587
Label every nth frequency on the plot. If not specified, no labels
15731588
are generated.
1589+
legend_loc : int or str, optional
1590+
Include a legend in the given location. Default is 'center right',
1591+
with no legend for a single response. Use False to supress legend.
15741592
max_curve_magnitude : float, optional
15751593
Restrict the maximum magnitude of the Nyquist plot to this value.
15761594
Portions of the Nyquist plot whose magnitude is restricted are
@@ -1609,6 +1627,10 @@ def nyquist_plot(
16091627
return_contour : bool, optional
16101628
(legacy) If 'True', return the encirclement count and Nyquist
16111629
contour used to generate the Nyquist plot.
1630+
show_legend : bool, optional
1631+
Force legend to be shown if ``True`` or hidden if ``False``. If
1632+
``None``, then show legend when there is more than one line on the
1633+
plot or ``legend_loc`` has been specified.
16121634
start_marker : str, optional
16131635
Matplotlib marker to use to mark the starting point of the Nyquist
16141636
plot. Defaults value is 'o' and can be set using
@@ -1775,6 +1797,8 @@ def _parse_linestyle(style_name, allow_false=False):
17751797

17761798
fig, ax = _process_ax_keyword(
17771799
ax_user, shape=(1, 1), squeeze=True, rcParams=rcParams)
1800+
legend_loc, _, show_legend = _process_legend_keywords(
1801+
kwargs, None, 'upper right')
17781802

17791803
# Create a list of lines for the output
17801804
out = np.empty(len(nyquist_responses), dtype=object)
@@ -1950,11 +1974,11 @@ def _parse_linestyle(style_name, allow_false=False):
19501974
lines, labels = _get_line_labels(ax)
19511975

19521976
# Add legend if there is more than one system plotted
1953-
if len(labels) > 1:
1977+
if show_legend == True or (show_legend != False and len(labels) > 1):
19541978
with plt.rc_context(rcParams):
19551979
legend = ax.legend(lines, labels, loc=legend_loc)
19561980
else:
1957-
legend=None
1981+
legend = None
19581982

19591983
# Add the title
19601984
if ax_user is None:
@@ -2193,7 +2217,7 @@ def singular_values_response(
21932217

21942218
def singular_values_plot(
21952219
data, omega=None, *fmt, plot=None, omega_limits=None, omega_num=None,
2196-
ax=None, label=None, title=None, legend_loc='center right', **kwargs):
2220+
ax=None, label=None, title=None, **kwargs):
21972221
"""Plot the singular values for a system.
21982222
21992223
Plot the singular values as a function of frequency for a system or
@@ -2237,16 +2261,20 @@ def singular_values_plot(
22372261
22382262
Other Parameters
22392263
----------------
2264+
ax : matplotlib.axes.Axes, optional
2265+
The matplotlib axes to draw the figure on. If not specified and
2266+
the current figure has a single axes, that axes is used.
2267+
Otherwise, a new figure is created.
22402268
grid : bool
22412269
If True, plot grid lines on gain and phase plots. Default is set by
22422270
`config.defaults['freqplot.grid']`.
2243-
label : str or array-like of str
2271+
label : str or array_like of str, optional
22442272
If present, replace automatically generated label(s) with the given
22452273
label(s). If sysdata is a list, strings should be specified for each
22462274
system.
2247-
legend_loc : str, optional
2248-
For plots with multiple lines, a legend will be included in the
2249-
given location. Default is 'center right'. Use False to supress.
2275+
legend_loc : int or str, optional
2276+
Include a legend in the given location. Default is 'center right',
2277+
with no legend for a single response. Use False to supress legend.
22502278
omega_limits : array_like of two values
22512279
Set limits for plotted frequency range. If Hz=True the limits are
22522280
in Hz otherwise in rad/s. Specifying ``omega`` as a list of two
@@ -2262,6 +2290,10 @@ def singular_values_plot(
22622290
rcParams : dict
22632291
Override the default parameters used for generating plots.
22642292
Default is set up config.default['freqplot.rcParams'].
2293+
show_legend : bool, optional
2294+
Force legend to be shown if ``True`` or hidden if ``False``. If
2295+
``None``, then show legend when there is more than one line on an
2296+
axis or ``legend_loc`` or ``legend_map`` has been specified.
22652297
title : str, optional
22662298
Set the title of the plot. Defaults to plot type and system name(s).
22672299
@@ -2343,6 +2375,8 @@ def singular_values_plot(
23432375
fig, ax_sigma = _process_ax_keyword(
23442376
ax, shape=(1, 1), squeeze=True, rcParams=rcParams)
23452377
ax_sigma.set_label('control-sigma') # TODO: deprecate?
2378+
legend_loc, _, show_legend = _process_legend_keywords(
2379+
kwargs, None, 'center right')
23462380

23472381
# Handle color cycle manually as all singular values
23482382
# of the same systems are expected to be of the same color
@@ -2413,7 +2447,7 @@ def singular_values_plot(
24132447
lines, labels = _get_line_labels(ax_sigma)
24142448

24152449
# Add legend if there is more than one system plotted
2416-
if len(labels) > 1 and legend_loc is not False:
2450+
if show_legend == True or (show_legend != False and len(labels) > 1):
24172451
with plt.rc_context(rcParams):
24182452
legend = ax_sigma.legend(lines, labels, loc=legend_loc)
24192453
else:

0 commit comments

Comments
 (0)