allow label keyword to override generated labels

murrayrm · murrayrm · commit ed8a1c18e112 · 2024-06-29T12:41:18.000-07:00
diff --git a/control/freqplot.py b/control/freqplot.py
@@ -2647,12 +2647,13 @@ def _get_line_labels(ax, use_color=True):
 
 
 # Turn label keyword into array indexed by trace, output, input
-def _process_line_labels(label, nsys, ninputs=0, noutputs=0):
+# TODO: move to ctrlutil.py and update parameter names to reflect general use
+def _process_line_labels(label, ntraces, ninputs=0, noutputs=0):
     if label is None:
         return None
 
     if isinstance(label, str):
-        label = [label]
+        label = [label] * ntraces          # single label for all traces
 
     # Convert to an ndarray, if not done aleady
     try:
@@ -2664,12 +2665,14 @@ def _process_line_labels(label, nsys, ninputs=0, noutputs=0):
     # TODO: allow more sophisticated broadcasting (and error checking)
     try:
         if ninputs > 0 and noutputs > 0:
-            if line_labels.ndim == 1:
-                line_labels = line_labels.reshape(nsys, 1, 1)
-            line_labels = np.broadcast_to(
-                line_labels,(nsys, ninputs, noutputs))
+            if line_labels.ndim == 1 and line_labels.size == ntraces:
+                line_labels = line_labels.reshape(ntraces, 1, 1)
+                line_labels = np.broadcast_to(
+                    line_labels, (ntraces, ninputs, noutputs))
+            else:
+                line_labels = line_labels.reshape(ntraces, ninputs, noutputs)
     except:
-        if line_labels.shape[0] != nsys:
+        if line_labels.shape[0] != ntraces:
             raise ValueError("number of labels must match number of traces")
         else:
             raise ValueError("labels must be given for each input/output pair")
diff --git a/control/tests/timeplot_test.py b/control/tests/timeplot_test.py
@@ -424,6 +424,113 @@ def test_rcParams():
         my_rcParams['ytick.labelsize']
     assert fig._suptitle.get_fontsize() == my_rcParams['figure.titlesize']
 
+
+@pytest.mark.parametrize("resp_fcn", [
+    ct.step_response, ct.initial_response, ct.impulse_response,
+    ct.forced_response, ct.input_output_response])
+@pytest.mark.usefixtures("editsdefaults")
+def test_timeplot_trace_labels(resp_fcn):
+    plt.close('all')
+    sys1 = ct.rss(2, 2, 2, strictly_proper=True, name='sys1')
+    sys2 = ct.rss(2, 2, 2, strictly_proper=True, name='sys2')
+
+    # Figure out the expected shape of the system
+    match resp_fcn:
+        case ct.step_response | ct.impulse_response:
+            shape = (2, 2)
+            kwargs = {}
+        case ct.initial_response:
+            shape = (2, 1)
+            kwargs = {}
+        case ct.forced_response | ct.input_output_response:
+            shape = (4, 1)      # outputs and inputs both plotted
+            T = np.linspace(0, 10)
+            U = [np.sin(T), np.cos(T)]
+            kwargs = {'T': T, 'U': U}
+
+    # Use figure frame for suptitle to speed things up
+    ct.set_defaults('freqplot', suptitle_frame='figure')
+
+    # Make sure default labels are as expected
+    out = resp_fcn([sys1, sys2], **kwargs).plot()
+    axs = ct.get_plot_axes(out)
+    if axs.ndim == 1:
+        legend = axs[0].get_legend().get_texts()
+    else:
+        legend = axs[0, -1].get_legend().get_texts()
+    assert legend[0].get_text() == 'sys1'
+    assert legend[1].get_text() == 'sys2'
+    plt.close()
+
+    # Override labels all at once
+    out = resp_fcn([sys1, sys2], **kwargs).plot(label=['line1', 'line2'])
+    axs = ct.get_plot_axes(out)
+    if axs.ndim == 1:
+        legend = axs[0].get_legend().get_texts()
+    else:
+        legend = axs[0, -1].get_legend().get_texts()
+    assert legend[0].get_text() == 'line1'
+    assert legend[1].get_text() == 'line2'
+    plt.close()
+
+    # Override labels one at a time
+    out = resp_fcn(sys1, **kwargs).plot(label='line1')
+    out = resp_fcn(sys2, **kwargs).plot(label='line2')
+    axs = ct.get_plot_axes(out)
+    if axs.ndim == 1:
+        legend = axs[0].get_legend().get_texts()
+    else:
+        legend = axs[0, -1].get_legend().get_texts()
+    assert legend[0].get_text() == 'line1'
+    assert legend[1].get_text() == 'line2'
+    plt.close()
+
+
+def test_full_label_override():
+    sys1 = ct.rss(2, 2, 2, strictly_proper=True, name='sys1')
+    sys2 = ct.rss(2, 2, 2, strictly_proper=True, name='sys2')
+
+    labels_2d = np.array([
+        ["outsys1u1y1", "outsys1u1y2", "outsys1u2y1", "outsys1u2y2",
+         "outsys2u1y1", "outsys2u1y2", "outsys2u2y1", "outsys2u2y2"],
+        ["inpsys1u1y1", "inpsys1u1y2", "inpsys1u2y1", "inpsys1u2y2",
+        "inpsys2u1y1", "inpsys2u1y2", "inpsys2u2y1", "inpsys2u2y2"]])
+
+
+    labels_4d = np.empty((2, 2, 2, 2), dtype=object)
+    for i, sys in enumerate(['sys1', 'sys2']):
+        for j, trace in enumerate(['u1', 'u2']):
+            for k, out in enumerate(['y1', 'y2']):
+                labels_4d[i, j, k, 0] = "out" + sys + trace + out
+                labels_4d[i, j, k, 1] = "inp" + sys + trace + out
+
+    # Test 4D labels
+    out = ct.step_response([sys1, sys2]).plot(
+        overlay_signals=True, overlay_traces=True, plot_inputs=True,
+        label=labels_4d)
+    axs = ct.get_plot_axes(out)
+    assert axs.shape == (2, 1)
+    legend_text = axs[0, 0].get_legend().get_texts()
+    for i, label in enumerate(labels_2d[0]):
+        assert legend_text[i].get_text() == label
+    legend_text = axs[1, 0].get_legend().get_texts()
+    for i, label in enumerate(labels_2d[1]):
+        assert legend_text[i].get_text() == label
+
+    # Test 2D labels
+    out = ct.step_response([sys1, sys2]).plot(
+        overlay_signals=True, overlay_traces=True, plot_inputs=True,
+        label=labels_2d)
+    axs = ct.get_plot_axes(out)
+    assert axs.shape == (2, 1)
+    legend_text = axs[0, 0].get_legend().get_texts()
+    for i, label in enumerate(labels_2d[0]):
+        assert legend_text[i].get_text() == label
+    legend_text = axs[1, 0].get_legend().get_texts()
+    for i, label in enumerate(labels_2d[1]):
+        assert legend_text[i].get_text() == label
+
+
 def test_relabel():
     sys1 = ct.rss(2, inputs='u', outputs='y')
     sys2 = ct.rss(1, 1, 1)      # uses default i/o labels
diff --git a/control/timeplot.py b/control/timeplot.py
@@ -49,7 +49,7 @@
 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,
+        legend_map=None, legend_loc=None, add_initial_zero=True, label=None,
         trace_labels=None, title=None, relabel=True, **kwargs):
     """Plot the time response of an input/output system.
 
@@ -112,6 +112,11 @@ def time_response_plot(
     input_props : array of dicts
         List of line properties to use when plotting combined inputs.  The
         default values are set by config.defaults['timeplot.input_props'].
+    label : str or array_like of str
+        If present, replace automatically generated label(s) with the given
+        label(s).  If more than one line is being generated, an array of
+        labels should be provided with label[trace, :, 0] representing the
+        output labels and label[trace, :, 1] representing the input labels.
     legend_map : array of str, option
         Location of the legend for multi-trace plots.  Specifies an array
         of legend location strings matching the shape of the subplots, with
@@ -152,7 +157,7 @@ def time_response_plot(
        config.defaults[''timeplot.rcParams'].
 
     """
-    from .freqplot import _process_ax_keyword
+    from .freqplot import _process_ax_keyword, _process_line_labels
     from .iosys import InputOutputSystem
     from .timeresp import TimeResponseData
 
@@ -342,6 +347,7 @@ def time_response_plot(
             out[i, j] = []      # unique list in each element
 
     # Utility function for creating line label
+    # TODO: combine with freqplot version?
     def _make_line_label(signal_index, signal_labels, trace_index):
         label = ""              # start with an empty label
 
@@ -375,11 +381,22 @@ def _make_line_label(signal_index, signal_labels, trace_index):
     output_offset = fig._output_offset = getattr(fig, '_output_offset', 0)
     input_offset = fig._input_offset = getattr(fig, '_input_offset', 0)
 
+    #
+    # Plot the lines for the response
+    #
+
+    # Process labels
+    line_labels = _process_line_labels(
+        label, ntraces, max(ninputs, noutputs), 2)
+
     # Go through each trace and each input/output
     for trace in range(ntraces):
         # Plot the output
         for i in range(noutputs):
-            label = _make_line_label(i, data.output_labels, trace)
+            if line_labels is None:
+                label = _make_line_label(i, data.output_labels, trace)
+            else:
+                label = line_labels[trace, i, 0]
 
             # Set up line properties for this output, trace
             if len(fmt) == 0:
@@ -397,7 +414,10 @@ def _make_line_label(signal_index, signal_labels, trace_index):
 
         # Plot the input
         for i in range(ninputs):
-            label = _make_line_label(i, data.input_labels, trace)
+            if line_labels is None:
+                label = _make_line_label(i, data.input_labels, trace)
+            else:
+                label = line_labels[trace, i, 1]
 
             if add_initial_zero and data.ntraces > i \
                and data.trace_types[i] == 'step':
@@ -596,16 +616,15 @@ def _make_line_label(signal_index, signal_labels, trace_index):
     for i in range(nrows):
         for j in range(ncols):
             ax = ax_array[i, j]
-            # Get the labels to use
             labels = [line.get_label() for line in ax.get_lines()]
-            labels = _make_legend_labels(labels, plot_inputs == 'overlay')
+            if line_labels is None:
+                labels = _make_legend_labels(labels, plot_inputs == 'overlay')
 
             # Update the labels to remove common strings
             if len(labels) > 1 and legend_map[i, j] != None:
                 with plt.rc_context(rcParams):
                     ax.legend(labels, loc=legend_map[i, j])
 
-
     #
     # Update the plot title (= figure suptitle)
     #
diff --git a/control/timeresp.py b/control/timeresp.py
@@ -746,8 +746,10 @@ class TimeResponseList(list):
     """
     def plot(self, *args, **kwargs):
         out_full = None
-        for response in self:
-            out = TimeResponseData.plot(response, *args, **kwargs)
+        label = kwargs.pop('label', [None] * len(self))
+        for i, response in enumerate(self):
+            out = TimeResponseData.plot(
+                response, *args, label=label[i], **kwargs)
             if out_full is None:
                 out_full = out
             else:
