Skip to content

Commit 7f92ce0

Browse files
murrayrmmurrayrm
committed
add trace_labels option to time_response_plot()
1 parent 29054ec commit 7f92ce0

4 files changed

Lines changed: 76 additions & 20 deletions

File tree

‎control/tests/timeplot_test.py‎

Lines changed: 27 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -321,6 +321,22 @@ def test_linestyles():
321321
assert line.get_color() == 'k'
322322
assert line.get_linestyle() == '--'
323323

324+
out = ct.step_response(sys_mimo).plot(
325+
plot_inputs='overlay', overlay_signals=True, overlay_traces=True,
326+
output_props=[{'color': c} for c in ['blue', 'orange']],
327+
input_props=[{'color': c} for c in ['red', 'green']],
328+
trace_props=[{'linestyle': s} for s in ['-', '--']])
329+
return None
330+
# assert out.shape == (1, 1) # TODO: fix
331+
assert out[0,0].get_color() == 'blue' and lines[0].get_linestyle() == '-'
332+
assert out[0,1].get_color() == 'blue' and lines[0].get_linestyle() == '--'
333+
assert out[1,0].get_color() == 'orange' and lines[0].get_linestyle() == '-'
334+
assert out[1,1].get_color() == 'orange' and lines[0].get_linestyle() == '--'
335+
assert out[2,0].get_color() == 'red' and lines[0].get_linestyle() == '-'
336+
assert out[2,1].get_color() == 'red' and lines[0].get_linestyle() == '--'
337+
assert out[3,0].get_color() == 'green' and lines[0].get_linestyle() == '-'
338+
assert out[3,1].get_color() == 'green' and lines[0].get_linestyle() == '--'
339+
324340

325341
def test_relabel():
326342
sys1 = ct.rss(2, inputs='u', outputs='y')
@@ -358,6 +374,11 @@ def test_errors():
358374
with pytest.raises(ValueError, match="unrecognized value"):
359375
stepresp.plot(plot_inputs='unknown')
360376

377+
for kw in ['input_props', 'output_props', 'trace_props']:
378+
propkw = {kw: {'color': 'green'}}
379+
with pytest.warns(UserWarning, match="ignored since fmt string"):
380+
out = stepresp.plot('k-', **propkw)
381+
assert out[0, 0][0].get_color() == 'k'
361382

362383
if __name__ == "__main__":
363384
#
@@ -450,8 +471,10 @@ def test_errors():
450471
"[transpose]")
451472
plt.savefig('timeplot-mimo_ioresp-mt_tr.png')
452473

453-
# Reset line styles
454474
plt.figure()
455-
resp1.plot('g-')
456-
resp2.plot('r--')
457-
# plt.savefig('timeplot-mimo_step-linestyle.png')
475+
out = ct.step_response(sys_mimo).plot(
476+
plot_inputs='overlay', overlay_signals=True, overlay_traces=True,
477+
output_props=[{'color': c} for c in ['blue', 'orange']],
478+
input_props=[{'color': c} for c in ['red', 'green']],
479+
trace_props=[{'linestyle': s} for s in ['-', '--']])
480+
plt.savefig('timeplot-mimo_step-linestyle.png')

‎control/timeplot.py‎

Lines changed: 36 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
import matplotlib as mpl
1313
import matplotlib.pyplot as plt
1414
from os.path import commonprefix
15+
from warnings import warn
1516

1617
from . import config
1718

@@ -48,7 +49,7 @@ def time_response_plot(
4849
transpose=False, overlay_traces=False, overlay_signals=False,
4950
legend_map=None, legend_loc=None, add_initial_zero=True,
5051
input_props=None, output_props=None, trace_props=None,
51-
title=None, relabel=True, **kwargs):
52+
trace_labels=None, title=None, relabel=True, **kwargs):
5253
"""Plot the time response of an input/output system.
5354
5455
This function creates a standard set of plots for the input/output
@@ -161,14 +162,20 @@ def time_response_plot(
161162
time_label = config._get_param(
162163
'timeplot', 'time_label', kwargs, _timeplot_defaults, pop=True)
163164

165+
if input_props and len(fmt) > 0:
166+
warn("input_props ignored since fmt string was present")
164167
input_props = config._get_param(
165168
'timeplot', 'input_props', kwargs, _timeplot_defaults, pop=True)
166169
iprop_len = len(input_props)
167170

171+
if output_props and len(fmt) > 0:
172+
warn("output_props ignored since fmt string was present")
168173
output_props = config._get_param(
169174
'timeplot', 'output_props', kwargs, _timeplot_defaults, pop=True)
170175
oprop_len = len(output_props)
171176

177+
if trace_props and len(fmt) > 0:
178+
warn("trace_props ignored since fmt string was present")
172179
trace_props = config._get_param(
173180
'timeplot', 'trace_props', kwargs, _timeplot_defaults, pop=True)
174181
tprop_len = len(trace_props)
@@ -365,10 +372,14 @@ def _make_line_label(signal_index, signal_labels, trace_index):
365372
label += signal_labels[signal_index]
366373

367374
# Add the trace label if this is a multi-trace figure
368-
if overlay_traces and ntraces > 1:
375+
if overlay_traces and ntraces > 1 or trace_labels:
369376
label += ", " if label != "" else ""
370-
label += f"trace {trace_index}" if data.trace_labels is None \
371-
else data.trace_labels[trace_index]
377+
if trace_labels:
378+
label += trace_labels[trace_index]
379+
elif data.trace_labels:
380+
label += data.trace_labels[trace_index]
381+
else:
382+
label += f"trace {trace_index}"
372383

373384
# Add the system name (will strip off later if redundant)
374385
label += ", " if label != "" else ""
@@ -471,18 +482,28 @@ def _make_line_label(signal_index, signal_labels, trace_index):
471482
label = ax_array[trace, 0].get_ylabel()
472483

473484
# Add on the trace title
474-
label = f"Trace {trace}" if data.trace_labels is None \
475-
else data.trace_labels[trace] + "\n" + label
485+
if trace_labels:
486+
label = trace_labels[trace] + "\n" + label
487+
elif data.trace_labels:
488+
label = data.trace_labels[trace] + "\n" + label
489+
else:
490+
label = f"Trace {trace}" + "\n" + label
491+
476492
ax_array[trace, 0].set_ylabel(label)
477493

478494
else: # regular plot (outputs over inputs)
479495
# Set the trace titles, if needed
480496
if ntraces > 1 and not overlay_traces:
481497
for trace in range(ntraces):
498+
if trace_labels:
499+
label = trace_labels[trace]
500+
elif data.trace_labels:
501+
label = data.trace_labels[trace]
502+
else:
503+
label = f"Trace {trace}"
504+
482505
with plt.rc_context(_timeplot_rcParams):
483-
ax_array[0, trace].set_title(
484-
f"Trace {trace}" if data.trace_labels is None
485-
else data.trace_labels[trace])
506+
ax_array[0, trace].set_title(label)
486507

487508
# Label the outputs
488509
if overlay_signals and plot_outputs:
@@ -622,22 +643,22 @@ def _make_line_label(signal_index, signal_labels, trace_index):
622643
if fig is not None and title is not None:
623644
# Get the current title, if it exists
624645
old_title = None if fig._suptitle is None else fig._suptitle._text
646+
new_title = title
625647

626648
if old_title is not None:
627649
# Find the common part of the titles
628-
common_prefix = commonprefix([old_title, title])
650+
common_prefix = commonprefix([old_title, new_title])
629651

630652
# Back up to the last space
631653
last_space = common_prefix.rfind(' ')
632654
if last_space > 0:
633655
common_prefix = common_prefix[:last_space]
634-
title_suffix = title[len(common_prefix):]
656+
common_len = len(common_prefix)
635657

636658
# Add the new part of the title (usually the system name)
637-
separator = ',' if len(common_prefix) > 0 else ';'
638-
new_title = old_title + separator + title_suffix
639-
else:
640-
new_title = title
659+
if old_title[common_len:] != new_title[common_len:]:
660+
separator = ',' if len(common_prefix) > 0 else ';'
661+
new_title = old_title + separator + new_title[common_len:]
641662

642663
# Add the title
643664
with plt.rc_context(_timeplot_rcParams):

‎doc/plotting.rst‎

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -77,7 +77,7 @@ following plot::
7777

7878
Input/output response plots created with either the
7979
:func:`~control.forced_response` or the
80-
:func:`~control.input_output_response` include the input signals by
80+
:func:`~control.input_output_response` functions include the input signals by
8181
default. These can be plotted on separate axes, but also "overlaid" on the
8282
output axes (useful when the input and output signals are being compared to
8383
each other). The following plot shows the use of `plot_inputs='overlay'`
@@ -119,6 +119,18 @@ that are used when combining signals and traces are set by the
119119
`input_props`, `output_props` and `trace_props` parameters for
120120
:func:`~control.time_response_plot`.
121121

122+
Additional customization is possible using the `input_props`,
123+
`output_props`, and `trace_props` keywords to set complementary line colors
124+
and styles for various signals and traces::
125+
126+
out = ct.step_response(sys_mimo).plot(
127+
plot_inputs='overlay', overlay_signals=True, overlay_traces=True,
128+
output_props=[{'color': c} for c in ['blue', 'orange']],
129+
input_props=[{'color': c} for c in ['red', 'green']],
130+
trace_props=[{'linestyle': s} for s in ['-', '--']])
131+
132+
.. image:: timeplot-mimo_step-linestyle.png
133+
122134
Plotting functions
123135
==================
124136

‎doc/timeplot-mimo_step-linestyle.png‎

37.8 KB
Loading

0 commit comments

Comments
 (0)