fix shape of output for time_reesponse_plot()

murrayrm · murrayrm · commit 1f44fd600ac1 · 2023-07-09T21:37:20.000-07:00
diff --git a/control/tests/timeplot_test.py b/control/tests/timeplot_test.py
@@ -135,21 +135,23 @@ def test_response_plots(
 
     out = response.plot(**kwargs)
 
+    # Make sure all of the outputs are of the right type
+    nlines_plotted = 0
+    for ax_lines in np.nditer(out, flags=["refs_ok"]):
+        for line in ax_lines.item():
+            assert isinstance(line, mpl.lines.Line2D)
+            nlines_plotted += 1
+
     # Make sure number of plots is correct
     if pltinp is None:
         if fcn in [ct.forced_response, ct.input_output_response]:
             pltinp = True
         else:
             pltinp = False
     ntraces = max(1, response.ntraces)
-    nlines = (response.ninputs if pltinp else 0) * ntraces + \
+    nlines_expected = (response.ninputs if pltinp else 0) * ntraces + \
         (response.noutputs if pltout else 0) * ntraces
-    assert out.size == nlines
-
-    # Make sure all of the outputs are of the right type
-    for ax_lines in np.nditer(out, flags=["refs_ok"]):
-        for line in ax_lines.item():
-            assert isinstance(line, mpl.lines.Line2D)
+    assert nlines_plotted == nlines_expected
 
     # Save the old axes to compare later
     old_axes = plt.gcf().get_axes()
@@ -326,17 +328,17 @@ def test_linestyles():
         output_props=[{'color': c} for c in ['blue', 'orange']],
         input_props=[{'color': c} for c in ['red', 'green']],
         trace_props=[{'linestyle': s} for s in ['-', '--']])
-    return None
-    # assert out.shape == (1, 1)                # TODO: fix
-    assert out[0,0].get_color() == 'blue' and lines[0].get_linestyle() == '-'
-    assert out[0,1].get_color() == 'blue' and lines[0].get_linestyle() == '--'
-    assert out[1,0].get_color() == 'orange' and lines[0].get_linestyle() == '-'
-    assert out[1,1].get_color() == 'orange' and lines[0].get_linestyle() == '--'
-    assert out[2,0].get_color() == 'red' and lines[0].get_linestyle() == '-'
-    assert out[2,1].get_color() == 'red' and lines[0].get_linestyle() == '--'
-    assert out[3,0].get_color() == 'green' and lines[0].get_linestyle() == '-'
-    assert out[3,1].get_color() == 'green' and lines[0].get_linestyle() == '--'
 
+    assert out.shape == (1, 1)
+    lines = out[0, 0]
+    assert lines[0].get_color() == 'blue' and lines[0].get_linestyle() == '-'
+    assert lines[1].get_color() == 'orange' and lines[1].get_linestyle() == '-'
+    assert lines[2].get_color() == 'red' and lines[2].get_linestyle() == '-'
+    assert lines[3].get_color() == 'green' and lines[3].get_linestyle() == '-'
+    assert lines[4].get_color() == 'blue' and lines[4].get_linestyle() == '--'
+    assert lines[5].get_color() == 'orange' and lines[5].get_linestyle() == '--'
+    assert lines[6].get_color() == 'red' and lines[6].get_linestyle() == '--'
+    assert lines[7].get_color() == 'green' and lines[7].get_linestyle() == '--'
 
 def test_relabel():
     sys1 = ct.rss(2, inputs='u', outputs='y')
diff --git a/control/timeplot.py b/control/timeplot.py
@@ -48,7 +48,6 @@ def time_response_plot(
         data, *fmt, ax=None, plot_inputs=None, plot_outputs=True,
         transpose=False, overlay_traces=False, overlay_signals=False,
         legend_map=None, legend_loc=None, add_initial_zero=True,
-        input_props=None, output_props=None, trace_props=None,
         trace_labels=None, title=None, relabel=True, **kwargs):
     """Plot the time response of an input/output system.
 
@@ -162,19 +161,19 @@ def time_response_plot(
     time_label = config._get_param(
         'timeplot', 'time_label', kwargs, _timeplot_defaults, pop=True)
 
-    if input_props and len(fmt) > 0:
+    if kwargs.get('input_props', None) and len(fmt) > 0:
         warn("input_props ignored since fmt string was present")
     input_props = config._get_param(
         'timeplot', 'input_props', kwargs, _timeplot_defaults, pop=True)
     iprop_len = len(input_props)
 
-    if output_props and len(fmt) > 0:
+    if kwargs.get('output_props', None) and len(fmt) > 0:
         warn("output_props ignored since fmt string was present")
     output_props = config._get_param(
         'timeplot', 'output_props', kwargs, _timeplot_defaults, pop=True)
     oprop_len = len(output_props)
 
-    if trace_props and len(fmt) > 0:
+    if kwargs.get('trace_props', None) and len(fmt) > 0:
         warn("trace_props ignored since fmt string was present")
     trace_props = config._get_param(
         'timeplot', 'trace_props', kwargs, _timeplot_defaults, pop=True)
@@ -306,35 +305,33 @@ def time_response_plot(
     # Map inputs/outputs and traces to axes
     #
     # This set of code takes care of all of the various options for how to
-    # plot the data.  The arrays ax_outputs and ax_inputs are used to map
+    # plot the data.  The arrays output_map and input_map are used to map
     # the different signals that are plotted onto the axes created above.
     # This code is complicated because it has to handle lots of different
     # variations.
     #
 
     # Create the map from trace, signal to axes, accounting for overlay_*
-    ax_outputs = np.empty((noutputs, ntraces), dtype=object)
-    ax_inputs = np.empty((ninputs, ntraces), dtype=object)
+    output_map = np.empty((noutputs, ntraces), dtype=tuple)
+    input_map = np.empty((ninputs, ntraces), dtype=tuple)
 
     for i in range(noutputs):
         for j in range(ntraces):
             signal_index = i if not overlay_signals else 0
             trace_index = j if not overlay_traces else 0
             if transpose:
-                ax_outputs[i, j] = \
-                    ax_array[trace_index, signal_index + ninput_axes]
+                output_map[i, j] = (trace_index, signal_index + ninput_axes)
             else:
-                ax_outputs[i, j] = ax_array[signal_index, trace_index]
+                output_map[i, j] = (signal_index, trace_index)
 
     for i in range(ninputs):
         for j in range(ntraces):
             signal_index = noutput_axes + (i if not overlay_signals else 0)
             trace_index = j if not overlay_traces else 0
             if transpose:
-                ax_inputs[i, j] = \
-                    ax_array[trace_index, signal_index - noutput_axes]
+                input_map[i, j] = (trace_index, signal_index - noutput_axes)
             else:
-                ax_inputs[i, j] = ax_array[signal_index, trace_index]
+                input_map[i, j] = (signal_index, trace_index)
 
     #
     # Plot the data
@@ -361,7 +358,10 @@ def time_response_plot(
         inputs = data.u.reshape(data.ninputs, ntraces, -1)
 
     # Create a list of lines for the output
-    out = np.empty((noutputs + ninputs, ntraces), dtype=object)
+    out = np.empty((nrows, ncols), dtype=object)
+    for i in range(nrows):
+        for j in range(ncols):
+            out[i, j] = []      # unique list in each element
 
     # Utility function for creating line label
     def _make_line_label(signal_index, signal_labels, trace_index):
@@ -402,7 +402,7 @@ def _make_line_label(signal_index, signal_labels, trace_index):
             else:
                 line_props = kwargs
 
-            out[i, trace] = ax_outputs[i, trace].plot(
+            out[output_map[i, trace]] += ax_array[output_map[i, trace]].plot(
                 data.time, outputs[i][trace], *fmt, label=label, **line_props)
 
         # Plot the input
@@ -425,7 +425,7 @@ def _make_line_label(signal_index, signal_labels, trace_index):
             else:
                 line_props = kwargs
 
-            out[noutputs + i, trace] = ax_inputs[i, trace].plot(
+            out[input_map[i, trace]] += ax_array[input_map[i, trace]].plot(
                 x, y, *fmt, label=label, **line_props)
 
     # Stop here if the user wants to control everything
@@ -457,23 +457,23 @@ def _make_line_label(signal_index, signal_labels, trace_index):
         if overlay_signals and plot_inputs:
             label = overlaid_title if overlaid else "Inputs"
             for trace in range(ntraces):
-                ax_inputs[0, trace].set_ylabel(label)
+                ax_array[input_map[0, trace]].set_ylabel(label)
         else:
             for i in range(ninputs):
                 label = overlaid_title if overlaid else data.input_labels[i]
                 for trace in range(ntraces):
-                    ax_inputs[i, trace].set_ylabel(label)
+                    ax_array[input_map[i, trace]].set_ylabel(label)
 
         # Label the outputs
         if overlay_signals and plot_outputs:
             label = overlaid_title if overlaid else "Outputs"
             for trace in range(ntraces):
-                ax_outputs[0, trace].set_ylabel(label)
+                ax_array[output_map[0, trace]].set_ylabel(label)
         else:
             for i in range(noutputs):
                 label = overlaid_title if overlaid else data.output_labels[i]
                 for trace in range(ntraces):
-                    ax_outputs[i, trace].set_ylabel(label)
+                    ax_array[output_map[i, trace]].set_ylabel(label)
 
         # Set the trace titles, if needed
         if ntraces > 1 and not overlay_traces:
@@ -507,20 +507,20 @@ def _make_line_label(signal_index, signal_labels, trace_index):
 
         # Label the outputs
         if overlay_signals and plot_outputs:
-            ax_outputs[0, 0].set_ylabel("Outputs")
+            ax_array[output_map[0, 0]].set_ylabel("Outputs")
         else:
             for i in range(noutputs):
-                ax_outputs[i, 0].set_ylabel(
+                ax_array[output_map[i, 0]].set_ylabel(
                     overlaid_title if overlaid else data.output_labels[i])
 
         # Label the inputs
         if overlay_signals and plot_inputs:
             label = overlaid_title if overlaid else "Inputs"
-            ax_inputs[0, 0].set_ylabel(label)
+            ax_array[input_map[0, 0]].set_ylabel(label)
         else:
             for i in range(ninputs):
                 label = overlaid_title if overlaid else data.input_labels[i]
-                ax_inputs[i, 0].set_ylabel(label)
+                ax_array[input_map[i, 0]].set_ylabel(label)
 
     #
     # Create legends
