python-control
diff --git a/‎control/tests/timeplot_test.py
Lines changed: 33 additions & 0 deletions b/‎control/tests/timeplot_test.py
Lines changed: 33 additions & 0 deletions
diff --git a/‎control/timeplot.py
Lines changed: 20 additions & 2 deletions b/‎control/timeplot.py
Lines changed: 20 additions & 2 deletions
diff --git a/‎examples/cds110_lti-systems.ipynb
Lines changed: 29 additions & 30 deletions b/‎examples/cds110_lti-systems.ipynb
Lines changed: 29 additions & 30 deletions
diff --git a/‎examples/springmass-coupled.png
57 KB b/‎examples/springmass-coupled.png
57 KB
@@ -588,6 +588,39 @@ def test_errors():
                        match="(has no property|unexpected keyword)"):
         stepresp.plot(unknown=None)
 
+
+def test_legend_customization():
+    sys = ct.rss(4, 2, 1, name='sys')
+    timepts = np.linspace(0, 10)
+    U = np.sin(timepts)
+    resp = ct.input_output_response(sys, timepts, U)
+
+    # Generic input/output plot
+    out = resp.plot(overlay_signals=True)
+    axs = ct.get_plot_axes(out)
+    assert axs[0, 0].get_legend()._loc == 7                 # center right
+    assert len(axs[0, 0].get_legend().get_texts()) == 2
+    assert axs[1, 0].get_legend() == None
+    plt.close()
+
+    # Hide legend
+    out = resp.plot(overlay_signals=True, show_legend=False)
+    axs = ct.get_plot_axes(out)
+    assert axs[0, 0].get_legend() == None
+    assert axs[1, 0].get_legend() == None
+    plt.close()
+
+    # Put legend in both axes
+    out = resp.plot(
+        overlay_signals=True, legend_map=[['center left'], ['center right']])
+    axs = ct.get_plot_axes(out)
+    assert axs[0, 0].get_legend()._loc == 6                 # center left
+    assert len(axs[0, 0].get_legend().get_texts()) == 2
+    assert axs[1, 0].get_legend()._loc == 7                 # center right
+    assert len(axs[1, 0].get_legend().get_texts()) == 1
+    plt.close()
+
+
 if __name__ == "__main__":
     #
     # Interactive mode: generate plots for manual viewing
 
@@ -50,7 +50,8 @@ 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, label=None,
-        trace_labels=None, title=None, relabel=True, **kwargs):
+        trace_labels=None, title=None, relabel=True, show_legend=None,
+        **kwargs):
     """Plot the time response of an input/output system.
 
     This function creates a standard set of plots for the input/output
@@ -131,6 +132,10 @@ def time_response_plot(
     relabel : bool, optional
         By default, existing figures and axes are relabeled when new data
         are added.  If set to `False`, just plot new data on existing axes.
+    show_legend : bool, optional
+        Force legend to be shown if ``True`` or hidden if ``False``.  If
+        ``None``, then show legend when there is more than one line on an
+        axis or ``legend_loc`` or ``legend_map`` have been specified.
     time_label : str, optional
         Label to use for the time axis.
     trace_props : array of dicts
@@ -565,6 +570,9 @@ def _make_line_label(signal_index, signal_labels, trace_index):
         legend_map = np.full(ax_array.shape, None, dtype=object)
         if legend_loc == None:
             legend_loc = 'center right'
+        else:
+            show_legend = True if show_legend is None else show_legend
+
         if transpose:
             if (overlay_signals or plot_inputs == 'overlay') and overlay_traces:
                 # Put a legend in each plot for inputs and outputs
@@ -611,6 +619,14 @@ def _make_line_label(signal_index, signal_labels, trace_index):
             else:
                 # Put legend in the upper right
                 legend_map[0, -1] = legend_loc
+    else:
+        # Make sure the legend map is the right size
+        legend_map = np.atleast_2d(legend_map)
+        if legend_map.shape != ax_array.shape:
+            raise ValueError("legend_map shape just match axes shape")
+
+        # Turn legend on unless overridden by user
+        show_legend = True if show_legend is None else show_legend
 
     # Create axis legends
     for i in range(nrows):
@@ -621,7 +637,9 @@ def _make_line_label(signal_index, signal_labels, trace_index):
                 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:
+            if show_legend != False and \
+               (len(labels) > 1 or show_legend) and \
+               legend_map[i, j] != None:
                 with plt.rc_context(rcParams):
                     ax.legend(labels, loc=legend_map[i, j])