fix missing plot title in bode_plot() with display_margins

murrayrm · murrayrm · commit 0b122038a2ca · 2025-02-02T23:12:28.000-08:00
diff --git a/control/freqplot.py b/control/freqplot.py
@@ -134,7 +134,8 @@ def bode_plot(
         If True, draw gain and phase margin lines on the magnitude and phase
         graphs and display the margins at the top of the graph.  If set to
         'overlay', the values for the gain and phase margin are placed on
-        the graph.  Setting display_margins turns off the axes grid.
+        the graph.  Setting `display_margins` turns off the axes grid, unless
+        `grid` is explicitly set to True.
     **kwargs : `matplotlib.pyplot.plot` keyword properties, optional
         Additional keywords passed to `matplotlib` to specify line properties.
 
@@ -276,6 +277,24 @@ def bode_plot(
     # Make a copy of the kwargs dictionary since we will modify it
     kwargs = dict(kwargs)
 
+    # Legacy keywords for margins
+    display_margins = config._process_legacy_keyword(
+        kwargs, 'margins', 'display_margins', display_margins)
+    if kwargs.pop('margin_info', False):
+        warnings.warn(
+            "keyword 'margin_info' is deprecated; "
+            "use 'display_margins='overlay'")
+        if display_margins is False:
+            raise ValueError(
+                "conflicting_keywords: `display_margins` and `margin_info`")
+
+    # Turn off grid if display margins, unless explicitly overridden
+    if display_margins and 'grid' not in kwargs:
+        kwargs['grid'] = False
+
+    margins_method = config._process_legacy_keyword(
+        kwargs, 'method', 'margins_method', margins_method)
+
     # Get values for params (and pop from list to allow keyword use in plot)
     dB = config._get_param(
         'freqplot', 'dB', kwargs, _freqplot_defaults, pop=True)
@@ -316,19 +335,6 @@ def bode_plot(
                 "sharex cannot be present with share_frequency")
         kwargs['share_frequency'] = sharex
 
-    # Legacy keywords for margins
-    display_margins = config._process_legacy_keyword(
-        kwargs, 'margins', 'display_margins', display_margins)
-    if kwargs.pop('margin_info', False):
-        warnings.warn(
-            "keyword 'margin_info' is deprecated; "
-            "use 'display_margins='overlay'")
-        if display_margins is False:
-            raise ValueError(
-                "conflicting_keywords: `display_margins` and `margin_info`")
-    margins_method = config._process_legacy_keyword(
-        kwargs, 'method', 'margins_method', margins_method)
-
     if not isinstance(data, (list, tuple)):
         data = [data]
 
@@ -727,7 +733,7 @@ def _make_line_label(response, output_index, input_index):
                         label='_nyq_mag_' + sysname)
 
                 # Add a grid to the plot
-                ax_mag.grid(grid and not display_margins, which='both')
+                ax_mag.grid(grid, which='both')
 
             # Phase
             if plot_phase:
@@ -742,7 +748,7 @@ def _make_line_label(response, output_index, input_index):
                         label='_nyq_phase_' + sysname)
 
                 # Add a grid to the plot
-                ax_phase.grid(grid and not display_margins, which='both')
+                ax_phase.grid(grid, which='both')
 
         #
         # Display gain and phase margins (SISO only)
@@ -753,6 +759,10 @@ def _make_line_label(response, output_index, input_index):
                 raise NotImplementedError(
                     "margins are not available for MIMO systems")
 
+            if display_margins == 'overlay' and len(data) > 1:
+                raise NotImplementedError(
+                    f"{display_margins=} not supported for multi-trace plots")
+
             # Compute stability margins for the system
             margins = stability_margins(response, method=margins_method)
             gm, pm, Wcg, Wcp = (margins[i] for i in [0, 1, 3, 4])
@@ -844,12 +854,12 @@ def _make_line_label(response, output_index, input_index):
 
             else:
                 # Put the title underneath the suptitle (one line per system)
-                ax = ax_mag if ax_mag else ax_phase
-                axes_title = ax.get_title()
+                ax_ = ax_mag if ax_mag else ax_phase
+                axes_title = ax_.get_title()
                 if axes_title is not None and axes_title != "":
                     axes_title += "\n"
                 with plt.rc_context(rcParams):
-                    ax.set_title(
+                    ax_.set_title(
                         axes_title + f"{sysname}: "
                         "Gm = %.2f %s(at %.2f %s), "
                         "Pm = %.2f %s (at %.2f %s)" %
diff --git a/control/tests/freqplot_test.py b/control/tests/freqplot_test.py
@@ -1,6 +1,7 @@
 # freqplot_test.py - test out frequency response plots
 # RMM, 23 Jun 2023
 
+import re
 import matplotlib as mpl
 import matplotlib.pyplot as plt
 import numpy as np
@@ -597,6 +598,12 @@ def test_suptitle():
             TypeError, match="unexpected keyword|no property"):
         cplt.set_plot_title("New title", unknown=None)
 
+    # Make sure title is still there if we display margins underneath
+    sys = ct.rss(2, 1, 1, name='sys')
+    cplt = ct.bode_plot(sys, display_margins=True)
+    assert re.match(r"^Bode plot for sys$", cplt.figure._suptitle._text)
+    assert re.match(r"^sys: Gm = .*, Pm = .*$", cplt.axes[0, 0].get_title())
+
 
 @pytest.mark.parametrize("plt_fcn", [ct.bode_plot, ct.singular_values_plot])
 def test_freqplot_errors(plt_fcn):
@@ -617,6 +624,7 @@ def test_freqplot_errors(plt_fcn):
     with pytest.raises(ValueError, match="invalid limits"):
         plt_fcn(response, omega_limits=[1e2, 1e-2])
 
+
 def test_freqresplist_unknown_kw():
     sys1 = ct.rss(2, 1, 1)
     sys2 = ct.rss(2, 1, 1)
@@ -626,6 +634,52 @@ def test_freqresplist_unknown_kw():
     with pytest.raises(AttributeError, match="unexpected keyword"):
         resp.plot(unknown=True)
 
+@pytest.mark.parametrize("nsys, display_margins, gridkw, match", [
+    (1, True, {}, None),
+    (1, False, {}, None),
+    (1, False, {}, None),
+    (1, True, {'grid': True}, None),
+    (1, 'overlay', {}, None),
+    (1, 'overlay', {'grid': True}, None),
+    (1, 'overlay', {'grid': False}, None),
+    (2, True, {}, None),
+    (2, 'overlay', {}, "not supported for multi-trace plots"),
+    (2, True, {'grid': 'overlay'}, None),
+    (3, True, {'grid': True}, None),
+])
+def test_display_margins(nsys, display_margins, gridkw, match):
+    sys1 = ct.tf([10], [1, 1, 1, 1], name='sys1')
+    sys2 = ct.tf([20], [2, 2, 2, 1], name='sys2')
+    sys3 = ct.tf([30], [2, 3, 3, 1], name='sys3')
+
+    sysdata = [sys1, sys2, sys3][0:nsys]
+
+    plt.figure()
+    if match is None:
+        cplt = ct.bode_plot(sysdata, display_margins=display_margins, **gridkw)
+    else:
+        with pytest.raises(NotImplementedError, match=match):
+            ct.bode_plot(sysdata, display_margins=display_margins, **gridkw)
+        return
+
+    cplt.set_plot_title(
+        cplt.figure._suptitle._text + f" [d_m={display_margins}, {gridkw=}")
+
+    # Make sure the grid is there if it should be
+    if gridkw.get('grid') or not display_margins:
+        assert all(
+            [line.get_visible() for line in cplt.axes[0, 0].get_xgridlines()])
+    else:
+        assert not any(
+            [line.get_visible() for line in cplt.axes[0, 0].get_xgridlines()])
+
+    # Make sure margins are displayed
+    if display_margins == True:
+        ax_title = cplt.axes[0, 0].get_title()
+        assert len(ax_title.split('\n')) == nsys
+    elif display_margins == 'overlay':
+        assert cplt.axes[0, 0].get_title() == ''
+
 
 if __name__ == "__main__":
     #
@@ -680,3 +734,6 @@ def test_freqresplist_unknown_kw():
     # of them for use in the documentation).
     #
     test_mixed_systypes()
+    test_display_margins(2, True, {})
+    test_display_margins(2, 'overlay', {})
+    test_display_margins(2, True, {'grid': True})
