allow sisotool to receive kvect as a singleton rather than always an array

sawyerbfuller · sawyerbfuller · commit 385405086bc0 · 2022-12-07T11:20:53.000-08:00
diff --git a/control/sisotool.py b/control/sisotool.py
@@ -9,6 +9,7 @@
 from .bdalg import append, connect
 from .iosys import tf2io, ss2io, summing_junction, interconnect
 from control.statesp import _convert_to_statespace, StateSpace
+import numpy as np
 import matplotlib.pyplot as plt
 import warnings
 
@@ -101,6 +102,9 @@ def sisotool(sys, kvect=None, xlim_rlocus=None, ylim_rlocus=None,
         'margins': margins_bode
     }
 
+    # make sure kvect is an array
+    if kvect is not None and ~hasattr(kvect, '__len__'):
+        kvect = np.atleast_1d(kvect)
     # First time call to setup the bode and step response plots
     _SisotoolUpdate(sys, fig,
         1 if kvect is None else kvect[0], bode_plot_params)
diff --git a/control/tests/sisotool_test.py b/control/tests/sisotool_test.py
@@ -136,6 +136,16 @@ def test_sisotool_tvect(self, tsys):
                            bode_plot_params=dict(), tvect=tvect)
         assert_array_almost_equal(tvect, ax_step.lines[0].get_data()[0])
 
+    @pytest.mark.skipif(plt.get_current_fig_manager().toolbar is None,
+                        reason="Requires the zoom toolbar")
+    def test_sisotool_kvect(self, tsys):
+        # test supply kvect
+        kvect = np.linspace(0, 1, 10)
+        # check if it can receive an array
+        sisotool(tsys, kvect=kvect)
+        # check if it can receive a singleton
+        sisotool(tsys, kvect=1)
+
 
     def test_sisotool_mimo(self,  sys222, sys221):
         # a 2x2 should not raise an error:
