diff --git a/control/tests/timeplot_test.py b/control/tests/timeplot_test.py index ea0a290c9..30dcd4a9c 100644 --- a/control/tests/timeplot_test.py +++ b/control/tests/timeplot_test.py @@ -369,6 +369,24 @@ def test_list_responses(resp_fcn): assert cplt.lines[row, col][0].get_color() == 'tab:blue' assert cplt.lines[row, col][1].get_color() == 'tab:orange' + # Make sure the public plotting function also accepts response lists + plt.figure() + cplt = ct.time_response_plot(resp_combined) + assert cplt.lines.shape == shape + for row in range(2): # just look at the outputs + for col in range(shape[1]): + assert cplt.lines[row, col][0].get_color() == 'tab:blue' + assert cplt.lines[row, col][1].get_color() == 'tab:orange' + + # Plain Python lists of time responses should follow the same path + plt.figure() + cplt = ct.time_response_plot([resp1, resp2]) + assert cplt.lines.shape == shape + for row in range(2): # just look at the outputs + for col in range(shape[1]): + assert cplt.lines[row, col][0].get_color() == 'tab:blue' + assert cplt.lines[row, col][1].get_color() == 'tab:orange' + @pytest.mark.slycot @pytest.mark.usefixtures('mplcleanup') diff --git a/control/timeplot.py b/control/timeplot.py index 545618f75..35afccfdf 100644 --- a/control/timeplot.py +++ b/control/timeplot.py @@ -52,7 +52,7 @@ def time_response_plot( Parameters ---------- - data : `TimeResponseData` + data : `TimeResponseData` or list of `TimeResponseData` Data to be plotted. plot_inputs : bool or str, optional Sets how and where to plot the inputs: @@ -176,6 +176,24 @@ def time_response_plot( """ from .ctrlplot import _process_ax_keyword, _process_line_labels + # If given a list of time responses, plot them sequentially on the same + # axes using the same path as TimeResponseList.plot(). + if isinstance(data, list): + if len(data) == 0: + raise ValueError("response list must contain at least one response") + + from .timeresp import TimeResponseList + + list_kwargs = dict( + ax=ax, plot_inputs=plot_inputs, plot_outputs=plot_outputs, + transpose=transpose, overlay_traces=overlay_traces, + overlay_signals=overlay_signals, add_initial_zero=add_initial_zero, + trace_labels=trace_labels, title=title, relabel=relabel) + if label is not None: + list_kwargs['label'] = label + + return TimeResponseList(data).plot(*fmt, **list_kwargs, **kwargs) + # # Process keywords and set defaults #