Skip to content

Commit 9f99219

Browse files
committed
unit tests + bug fixes
1 parent dacf17c commit 9f99219

3 files changed

Lines changed: 196 additions & 78 deletions

File tree

control/tests/timeplot_test.py

Lines changed: 150 additions & 56 deletions
Original file line numberDiff line numberDiff line change
@@ -5,53 +5,134 @@
55
import control as ct
66
import matplotlib as mpl
77
import matplotlib.pyplot as plt
8-
9-
# Step responses
10-
@pytest.mark.parametrize("nin, nout", [(1, 1), (1, 2), (2, 1), (2, 2), (2, 3)])
8+
import numpy as np
9+
10+
# Detailed test of (almost) all functionality
11+
# (uncomment rows for developmental testing, but otherwise takes too long)
12+
@pytest.mark.parametrize(
13+
"sys", [
14+
# ct.rss(1, 1, 1, strictly_proper=True, name="rss"),
15+
ct.nlsys(
16+
lambda t, x, u, params: -x + u, None,
17+
inputs=1, outputs=1, states=1, name="nlsys"),
18+
# ct.rss(2, 1, 2, strictly_proper=True, name="rss"),
19+
ct.rss(2, 2, 1, strictly_proper=True, name="rss"),
20+
# ct.drss(2, 2, 2, name="drss"),
21+
# ct.rss(2, 2, 3, strictly_proper=True, name="rss"),
22+
])
1123
@pytest.mark.parametrize("transpose", [True, False])
1224
@pytest.mark.parametrize("plot_inputs", [None, True, False, 'overlay'])
13-
def test_simple_response(nout, nin, transpose, plot_inputs):
14-
sys = ct.rss(4, nout, nin)
15-
stepresp = ct.step_response(sys)
16-
stepresp.plot(plot_inputs=plot_inputs, transpose=transpose)
25+
@pytest.mark.parametrize("plot_outputs", [True, False])
26+
@pytest.mark.parametrize("combine_signals", [True, False])
27+
@pytest.mark.parametrize("combine_traces", [True, False])
28+
@pytest.mark.parametrize("second_system", [False, True])
29+
@pytest.mark.parametrize("fcn", [
30+
ct.step_response, ct.impulse_response, ct.initial_response,
31+
ct.forced_response, ct.input_output_response])
32+
def test_response_plots(
33+
fcn, sys, plot_inputs, plot_outputs, combine_signals, combine_traces,
34+
transpose, second_system, clear=True):
35+
# Figure out the time range to use and check some special cases
36+
if not isinstance(sys, ct.lti.LTI):
37+
if fcn == ct.impulse_response:
38+
pytest.skip("impulse response not implemented for nlsys")
39+
40+
# Nonlinear systems require explicit time limits
41+
T = 10
42+
timepts = np.linspace(0, T)
43+
else:
44+
# Linear systems figure things out on their own
45+
T = None
46+
timepts = np.linspace(0, 10) # for input_output_response
47+
48+
# Save up the keyword arguments
49+
kwargs = dict(
50+
plot_inputs=plot_inputs, plot_outputs=plot_outputs, transpose=transpose,
51+
combine_signals=combine_signals, combine_traces=combine_traces)
52+
53+
# Create the response
54+
if fcn is ct.input_output_response and \
55+
not isinstance(sys, ct.NonlinearIOSystem):
56+
# Skip transfer functions and other non-state space systems
57+
return None
58+
if fcn in [ct.input_output_response, ct.forced_response]:
59+
U = np.zeros((sys.ninputs, timepts.size))
60+
for i in range(sys.ninputs):
61+
U[i] = np.cos(timepts * i + i)
62+
args = [timepts, U]
63+
64+
elif fcn == ct.initial_response:
65+
args = [T, np.ones(sys.nstates)] # T, X0
66+
67+
elif not isinstance(sys, ct.lti.LTI):
68+
args = [T] # nonlinear systems require final time
69+
70+
else: # step, initial, impulse responses
71+
args = []
72+
73+
# Create a new figure (in case previous one is of the same size) and plot
74+
if not clear:
75+
plt.figure()
76+
response = fcn(sys, *args)
77+
78+
# Look for cases where there are no data to plot
79+
if not plot_outputs and (
80+
plot_inputs is False or response.ninputs == 0 or
81+
plot_inputs is None and response.plot_inputs is False):
82+
with pytest.raises(ValueError, match=".* no data to plot"):
83+
out = response.plot(**kwargs)
84+
return None
85+
elif not plot_outputs and plot_inputs == 'overlay':
86+
with pytest.raises(ValueError, match="can't overlay inputs"):
87+
out = response.plot(**kwargs)
88+
return None
89+
elif plot_inputs in [True, 'overlay'] and response.ninputs == 0:
90+
with pytest.raises(ValueError, match=".* but no inputs"):
91+
out = response.plot(**kwargs)
92+
return None
93+
94+
out = response.plot(**kwargs)
95+
96+
# TODO: add some basic checks here
1797

1898
# Add additional data (and provide infon in the title)
19-
newsys = ct.rss(4, nout, nin)
20-
out = ct.step_response(newsys, stepresp.time[-1]).plot(
21-
plot_inputs=plot_inputs, transpose=transpose)
99+
if second_system:
100+
newsys = ct.rss(
101+
sys.nstates, sys.noutputs, sys.ninputs, strictly_proper=True)
102+
if fcn not in [ct.initial_response, ct.forced_response,
103+
ct.input_output_response] and \
104+
isinstance(sys, ct.lti.LTI):
105+
# Reuse the previously computed time to make plots look nicer
106+
fcn(newsys, *args, T=response.time[-1]).plot(**kwargs)
107+
else:
108+
# Compute and plot new response (time is one of the arguments)
109+
fcn(newsys, *args).plot(**kwargs)
110+
111+
# TODO: add some basic checks here
22112

23113
# Update the title so we can see what is going on
24114
fig = out[0, 0][0].axes.figure
25115
fig.suptitle(
26-
fig._suptitle._text + f" [{nout}x{nin}, {plot_inputs=}, {transpose=}]",
116+
fig._suptitle._text +
117+
f" [{sys.noutputs}x{sys.ninputs}, cs={combine_signals}, "
118+
f"ct={combine_traces}, pi={plot_inputs}, tr={transpose}]",
27119
fontsize='small')
28120

29-
@pytest.mark.parametrize("transpose", [True, False])
30-
def test_combine_signals(transpose):
31-
sys = ct.rss(4, 2, 3)
32-
stepresp = ct.step_response(sys)
33-
stepresp.plot(
34-
combine_signals=True, transpose=transpose,
35-
title=f"Step response: combine_signals = True; transpose={transpose}")
36-
37-
38-
@pytest.mark.parametrize("transpose", [True, False])
39-
def test_combine_traces(transpose):
40-
sys = ct.rss(4, 2, 3)
41-
stepresp = ct.step_response(sys)
42-
stepresp.plot(
43-
combine_traces=True, transpose=transpose,
44-
title=f"Step response: combine_traces = True; transpose={transpose}")
121+
# Get rid of the figure to free up memory
122+
if clear:
123+
plt.clf()
45124

46125

47-
@pytest.mark.parametrize("transpose", [True, False])
48-
def test_combine_signals_traces(transpose):
49-
sys = ct.rss(4, 5, 3)
50-
stepresp = ct.step_response(sys)
51-
stepresp.plot(
52-
combine_signals=True, combine_traces=True, transpose=transpose,
53-
title=f"Step response: combine_signals/traces = True;" +
54-
f"transpose={transpose}")
126+
def test_legend_map():
127+
sys_mimo = ct.tf2ss(
128+
[[[1], [0.1]], [[0.2], [1]]],
129+
[[[1, 0.6, 1], [1, 1, 1]], [[1, 0.4, 1], [1, 2, 1]]], name="MIMO")
130+
response = ct.step_response(sys_mimo)
131+
response.plot(
132+
legend_map=np.array([['center', 'upper right'],
133+
[None, 'center right']]),
134+
plot_inputs=True, combine_signals=True, transpose=True,
135+
title='MIMO step response with custom legend placement')
55136

56137

57138
def test_errors():
@@ -81,26 +162,39 @@ def test_errors():
81162
# Start by clearing existing figures
82163
plt.close('all')
83164

84-
print ("Simple step responses")
85-
for size in [(1, 1), (1, 2), (2, 1), (2, 2), (2, 3)]:
86-
for transpose in [False, True]:
87-
for plot_inputs in [None, True, False, 'overlay']:
88-
plt.figure()
89-
test_simple_response(
90-
*size, transpose=transpose, plot_inputs=plot_inputs)
165+
# Define a set of systems to test
166+
sys_siso = ct.tf2ss([1], [1, 2, 1], name="SISO")
167+
sys_mimo = ct.tf2ss(
168+
[[[1], [0.1]], [[0.2], [1]]],
169+
[[[1, 0.6, 1], [1, 1, 1]], [[1, 0.4, 1], [1, 2, 1]]], name="MIMO")
170+
171+
# Define and run a selected set of interesting tests
172+
# def test_response_plots(
173+
# fcn, sys, plot_inputs, plot_outputs, combine_signals,
174+
# combine_traces, transpose, second_system, clear=True):
175+
N, T, F = None, True, False
176+
test_cases = [
177+
# response fcn system in out cs ct tr ss
178+
(ct.step_response, sys_siso, N, T, F, F, F, F), # 1
179+
(ct.step_response, sys_siso, T, F, F, F, F, F), # 2
180+
(ct.step_response, sys_siso, T, T, F, F, F, T), # 3
181+
(ct.step_response, sys_siso, 'overlay', T, F, F, F, T), # 4
182+
(ct.step_response, sys_mimo, F, T, F, F, F, F), # 5
183+
(ct.step_response, sys_mimo, T, T, F, F, F, F), # 6
184+
(ct.step_response, sys_mimo, 'overlay', T, F, F, F, F), # 7
185+
(ct.step_response, sys_mimo, T, T, T, F, F, F), # 8
186+
(ct.step_response, sys_mimo, T, T, T, T, F, F), # 9
187+
(ct.step_response, sys_mimo, T, T, F, F, T, F), # 10
188+
(ct.step_response, sys_mimo, T, T, T, F, T, F), # 11
189+
(ct.step_response, sys_mimo, 'overlay', T, T, F, T, F), # 12
190+
(ct.forced_response, sys_mimo, N, T, T, F, T, F), # 13
191+
(ct.forced_response, sys_mimo, 'overlay', T, F, F, F, F), # 14
192+
]
193+
for args in test_cases:
194+
test_response_plots(*args, clear=F)
91195

92-
print ("Combine signals")
93-
for transpose in [False, True]:
94-
plt.figure()
95-
test_combine_signals(transpose)
96-
97-
print ("Combine traces")
98-
for transpose in [False, True]:
99-
plt.figure()
100-
test_combine_traces(transpose)
101-
102-
print ("Combine signals and traces")
103-
for transpose in [False, True]:
104-
plt.figure()
105-
test_combine_signals_traces(transpose)
196+
#
197+
# Run a few more special cases to show off capabilities
198+
#
106199

200+
test_legend_map() # show ability to set legend location

control/timeplot.py

Lines changed: 28 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@
1818
# [i] Multi-trace graphs using different line styles
1919
# [i] Plotting function return Line2D elements
2020
# [i] Axis labels/legends based on what is plotted (siso, mimo, multi-trace)
21-
# [ ] Ability to select (index) output and/or trace (and time?)
21+
# [x] Ability to select (index) output and/or trace (and time?)
2222
# [i] Legends should not contain redundant information (nor appear redundantly)
2323

2424
import numpy as np
@@ -53,7 +53,7 @@
5353
# Plot the input/output response of a system
5454
def ioresp_plot(
5555
data, ax=None, plot_inputs=None, plot_outputs=True, transpose=False,
56-
combine_traces=False, combine_signals=False, legend_spec=None,
56+
combine_traces=False, combine_signals=False, legend_map=None,
5757
legend_loc=None, add_initial_zero=True, title=None, relabel=True,
5858
**kwargs):
5959
"""Plot the time response of an input/output system.
@@ -111,14 +111,14 @@ def ioresp_plot(
111111
are added. If set to `False`, just plot new data on existing axes.
112112
time_label : str, optional
113113
Label to use for the time axis.
114-
legend_spec : array of str, option
114+
legend_map : array of str, option
115115
Location of the legend for multi-trace plots. Specifies an array
116116
of legend location strings matching the shape of the subplots, with
117117
each entry being either None (for no legend) or a legend location
118118
string (see :func:`~matplotlib.pyplot.legend`).
119119
legend_loc : str
120120
Location of the legend within the axes for which it appears. This
121-
value is used if legend_spec is None.
121+
value is used if legend_map is None.
122122
add_initial_zero : bool
123123
Add an initial point of zero at the first time point for all
124124
inputs. This is useful when the initial value of the input is
@@ -216,7 +216,13 @@ def ioresp_plot(
216216
ntraces = max(1, data.ntraces) # treat data.ntraces == 0 as 1 trace
217217
if ninputs == 0 and noutputs == 0:
218218
raise ValueError(
219-
"plot_inputs and plot_outputs both True; no data to plot")
219+
"plot_inputs and plot_outputs both False; no data to plot")
220+
elif plot_inputs == 'overlay' and noutputs == 0:
221+
raise ValueError(
222+
"can't overlay inputs with no outputs")
223+
elif plot_inputs in [True, 'overlay'] and data.ninputs == 0:
224+
raise ValueError(
225+
"input plotting requested but no inputs in time response data")
220226

221227
# Figure how how many rows and columns to use + offsets for inputs/outputs
222228
if plot_inputs == 'overlay' and not combine_signals:
@@ -226,8 +232,8 @@ def ioresp_plot(
226232
elif combine_signals:
227233
nrows = int(plot_outputs) # Start with outputs
228234
nrows += int(plot_inputs == True) # Add plot for inputs if needed
229-
noutput_axes = 1 if plot_outputs else 0
230-
ninput_axes = 1 if plot_inputs else 0
235+
noutput_axes = 1 if plot_outputs and plot_inputs is True else 0
236+
ninput_axes = 1 if plot_inputs is True else 0
231237
else:
232238
nrows = noutputs + ninputs # Plot inputs separately
233239
noutput_axes = noutputs if plot_outputs else 0
@@ -321,7 +327,10 @@ def ioresp_plot(
321327

322328
# Reshape the inputs and outputs for uniform indexing
323329
outputs = data.y.reshape(data.noutputs, ntraces, -1)
324-
inputs = data.u.reshape(data.ninputs, ntraces, -1)
330+
if data.u is None or not plot_inputs:
331+
inputs = None
332+
else:
333+
inputs = data.u.reshape(data.ninputs, ntraces, -1)
325334

326335
# Create a list of lines for the output
327336
out = np.empty((noutputs + ninputs, ntraces), dtype=object)
@@ -430,7 +439,7 @@ def _make_line_label(signal_index, signal_labels, trace_index):
430439
if ntraces > 1 and not combine_traces:
431440
for trace in range(ntraces):
432441
with plt.rc_context(_timeplot_rcParams):
433-
ax_outputs[0, trace].set_title(
442+
ax_array[0, trace].set_title(
434443
f"Trace {trace}" if data.trace_labels is None
435444
else data.trace_labels[trace])
436445

@@ -454,7 +463,7 @@ def _make_line_label(signal_index, signal_labels, trace_index):
454463
#
455464
# Create legends
456465
#
457-
# Legends can be placed manually by passing a legend_spec array that
466+
# Legends can be placed manually by passing a legend_map array that
458467
# matches the shape of the suplots, with each item being a string
459468
# indicating the location of the legend for that axes (or None for no
460469
# legend).
@@ -472,21 +481,23 @@ def _make_line_label(signal_index, signal_labels, trace_index):
472481
#
473482

474483
# Figure out where to put legends
475-
if legend_spec is None:
484+
if legend_map is None:
476485
legend_map = np.full(ax_array.shape, None, dtype=object)
477486
if legend_loc == None:
478487
legend_loc = 'center right'
479488
if transpose:
480489
if (combine_signals or plot_inputs == 'overlay') and combine_traces:
481490
# Put a legend in each plot for inputs and outputs
482-
legend_map[0, ninput_axes] = legend_loc
491+
if plot_outputs is True:
492+
legend_map[0, ninput_axes] = legend_loc
483493
if plot_inputs is True:
484494
legend_map[0, 0] = legend_loc
485495
elif combine_signals:
486496
# Put a legend in rightmost input/output plot
487-
legend_map[0, ninput_axes] = legend_loc
488497
if plot_inputs is True:
489498
legend_map[0, 0] = legend_loc
499+
if plot_outputs is True:
500+
legend_map[0, ninput_axes] = legend_loc
490501
elif plot_inputs == 'overlay':
491502
# Put a legend on the top of each column
492503
for i in range(ntraces):
@@ -500,12 +511,14 @@ def _make_line_label(signal_index, signal_labels, trace_index):
500511
else: # regular layout
501512
if (combine_signals or plot_inputs == 'overlay') and combine_traces:
502513
# Put a legend in each plot for inputs and outputs
503-
legend_map[0, -1] = legend_loc
514+
if plot_outputs is True:
515+
legend_map[0, -1] = legend_loc
504516
if plot_inputs is True:
505517
legend_map[noutput_axes, -1] = legend_loc
506518
elif combine_signals:
507519
# Put a legend in rightmost input/output plot
508-
legend_map[0, -1] = legend_loc
520+
if plot_outputs is True:
521+
legend_map[0, -1] = legend_loc
509522
if plot_inputs is True:
510523
legend_map[noutput_axes, -1] = legend_loc
511524
elif plot_inputs == 'overlay':

0 commit comments

Comments
 (0)