diff --git a/control/tests/timeplot_test.py b/control/tests/timeplot_test.py index ea0a290c9..50dc778d6 100644 --- a/control/tests/timeplot_test.py +++ b/control/tests/timeplot_test.py @@ -321,6 +321,34 @@ def test_combine_time_responses(): ct.combine_time_responses([resp1, resp]) +@pytest.mark.usefixtures('mplcleanup') +def test_time_response_plot_response_list(): + sys = ct.ss(ct.tf([1, 2], [3, 4, 5])) + timepts = np.linspace(0, 10) + resp1 = ct.input_output_response(sys, timepts, np.sin(timepts)) + resp2 = ct.input_output_response(sys, timepts, np.cos(timepts)) + + cplt = ct.time_response_plot( + [resp1, resp2], trace_labels=["sine input", "cosine input"]) + assert cplt.lines.shape == (2, 2) + assert cplt.axes[0, 0].get_title() == "sine input" + assert cplt.axes[0, 1].get_title() == "cosine input" + assert all(len(lines.item()) == 1 for lines in np.nditer( + cplt.lines, flags=["refs_ok"])) + plt.close() + + cplt = ct.time_response_plot([resp1, resp2], plot_inputs=True) + assert cplt.lines.shape == (2, 2) + assert all(len(lines.item()) == 1 for lines in np.nditer( + cplt.lines, flags=["refs_ok"])) + plt.close() + + cplt = ct.time_response_plot([resp1, resp2], plot_inputs='overlay') + assert cplt.lines.shape == (1, 2) + assert all(len(lines.item()) == 2 for lines in np.nditer( + cplt.lines, flags=["refs_ok"])) + + @pytest.mark.parametrize("resp_fcn", [ ct.step_response, ct.initial_response, ct.impulse_response, ct.forced_response, ct.input_output_response]) diff --git a/control/timeplot.py b/control/timeplot.py index 545618f75..1b226ab67 100644 --- a/control/timeplot.py +++ b/control/timeplot.py @@ -52,8 +52,9 @@ def time_response_plot( Parameters ---------- - data : `TimeResponseData` - Data to be plotted. + data : `TimeResponseData` or list of `TimeResponseData` + Data to be plotted. Lists of responses are combined into a + multi-trace response before plotting. plot_inputs : bool or str, optional Sets how and where to plot the inputs: * False: don't plot the inputs @@ -179,6 +180,9 @@ def time_response_plot( # # Process keywords and set defaults # + if isinstance(data, (list, tuple)): + data = combine_time_responses(data, trace_labels=trace_labels) + # Set up defaults ax_user = ax sharex = config._get_param('timeplot', 'sharex', kwargs, pop=True)