unify frequency response processing + unit tests

murrayrm · murrayrm · commit f367cf657c29 · 2021-01-17T08:18:51.000-08:00
diff --git a/control/frdata.py b/control/frdata.py
@@ -50,7 +50,7 @@
 from numpy import angle, array, empty, ones, \
     real, imag, absolute, eye, linalg, where, dot, sort
 from scipy.interpolate import splprep, splev
-from .lti import LTI
+from .lti import LTI, _process_frequency_response
 from . import config
 
 __all__ = ['FrequencyResponseData', 'FRD', 'frd']
@@ -391,14 +391,10 @@ def eval(self, omega, squeeze=None):
                     for k, w in enumerate(omega_array):
                         frraw = splev(w, self.ifunc[i, j], der=0)
                         out[i, j, k] = frraw[0] + 1.0j * frraw[1]
-        if not hasattr(omega, '__len__'):
-            # omega is a scalar, squeeze down array along last dim
-            out = np.squeeze(out, axis=2)
-        if squeeze and self.issiso():
-            out = out[0][0]
-        return out
-
-    def __call__(self, s, squeeze=True):
+
+        return _process_frequency_response(self, omega, out, squeeze=squeeze)
+
+    def __call__(self, s, squeeze=None):
         """Evaluate system's transfer function at complex frequencies.
         
         Returns the complex frequency response `sys(s)` of system `sys` with
diff --git a/control/lti.py b/control/lti.py
@@ -15,6 +15,7 @@
 import numpy as np
 from numpy import absolute, real, angle, abs
 from warnings import warn
+from . import config
 
 __all__ = ['issiso', 'timebase', 'common_timebase', 'timebaseEqual',
            'isdtime', 'isctime', 'pole', 'zero', 'damp', 'evalfr',
@@ -596,3 +597,18 @@ def dcgain(sys):
         at the origin
     """
     return sys.dcgain()
+
+
+# Process frequency responses in a uniform way
+def _process_frequency_response(sys, omega, out, squeeze=None):
+    if not hasattr(omega, '__len__'):
+        # received a scalar x, squeeze down the array along last dim
+        out = np.squeeze(out, axis=2)
+
+    # Get rid of unneeded dimensions
+    if squeeze is None:
+        squeeze = config.defaults['control.squeeze']
+    if squeeze and sys.issiso():
+        return out[0][0]
+    else:
+        return out
diff --git a/control/statesp.py b/control/statesp.py
@@ -62,7 +62,7 @@
 from scipy.signal import cont2discrete
 from scipy.signal import StateSpace as signalStateSpace
 from warnings import warn
-from .lti import LTI, common_timebase, isdtime
+from .lti import LTI, common_timebase, isdtime, _process_frequency_response
 from . import config
 from copy import deepcopy
 
@@ -646,9 +646,9 @@ def __call__(self, x, squeeze=None):
         Returns the complex frequency response `sys(x)` where `x` is `s` for
         continuous-time systems and `z` for discrete-time systems.
         
-        In general the system may be multiple input, multiple output (MIMO), where
-        `m = self.inputs` number of inputs and `p = self.outputs` number of
-        outputs.
+        In general the system may be multiple input, multiple output
+        (MIMO), where `m = self.inputs` number of inputs and `p =
+        self.outputs` number of outputs.
 
         To evaluate at a frequency omega in radians per second, enter
         ``x = omega * 1j``, for continuous-time systems, or
@@ -671,19 +671,9 @@ def __call__(self, x, squeeze=None):
             only if system is SISO and ``squeeze=True``.
 
         """
-        # Set value of squeeze argument if not set
-        if squeeze is None:
-            squeeze = config.defaults['control.squeeze']
-
         # Use Slycot if available
         out = self.horner(x)
-        if not hasattr(x, '__len__'):
-            # received a scalar x, squeeze down the array along last dim
-            out = np.squeeze(out, axis=2)
-        if squeeze and self.issiso():
-            return out[0][0]
-        else:
-            return out
+        return _process_frequency_response(self, x, out, squeeze=squeeze)
 
     def slycot_laub(self, x):
         """Evaluate system's transfer function at complex frequency
diff --git a/control/tests/lti_test.py b/control/tests/lti_test.py
@@ -2,12 +2,14 @@
 
 import numpy as np
 import pytest
+from .conftest import editsdefaults
 
+import control as ct
 from control import c2d, tf, tf2ss, NonlinearIOSystem
 from control.lti import (LTI, common_timebase, damp, dcgain, isctime, isdtime,
                          issiso, pole, timebaseEqual, zero)
 from control.tests.conftest import slycotonly
-
+from control.exception import slycot_check
 
 class TestLTI:
 
@@ -153,3 +155,47 @@ def test_isdtime(self, objfun, arg, dt, ref, strictref):
             strictref = not strictref
         assert isctime(obj) == ref
         assert isctime(obj, strict=True) == strictref
+
+    @pytest.mark.usefixtures("editsdefaults")
+    @pytest.mark.parametrize("fcn", [ct.ss, ct.tf, ct.frd])
+    @pytest.mark.parametrize("nstate, nout, ninp, squeeze, shape", [
+        [1, 1, 1, None, (8,)],                          # SISO
+        [2, 1, 1, True, (8,)],
+        [3, 1, 1, False, (1, 1, 8)],
+        [1, 2, 1, None, (2, 1, 8)],                     # SIMO
+        [2, 2, 1, True, (2, 1, 8)],
+        [3, 2, 1, False, (2, 1, 8)],
+        [1, 1, 2, None, (1, 2, 8)],                     # MISO
+        [2, 1, 2, True, (1, 2, 8)],
+        [3, 1, 2, False, (1, 2, 8)],
+        [1, 2, 2, None, (2, 2, 8)],                     # MIMO
+        [2, 2, 2, True, (2, 2, 8)],
+        [3, 2, 2, False, (2, 2, 8)]
+    ])
+    def test_squeeze(self, fcn, nstate, nout, ninp, squeeze, shape):
+        # Compute the length of the frequency array
+        omega = np.logspace(-2, 2, 8)
+
+        # Create the system to be tested
+        if fcn == ct.frd:
+            sys = fcn(ct.rss(nstate, nout, ninp), omega)
+        elif fcn == ct.tf and (nout > 1 or ninp > 1) and not slycot_check():
+            pytest.skip("Conversion of MIMO systems to transfer functions "
+                        "requires slycot.")
+        else:
+            sys = fcn(ct.rss(nstate, nout, ninp))
+
+        # Pass squeeze argument and make sure the shape is correct
+        mag, phase, _ = sys.frequency_response(omega, squeeze=squeeze)
+        assert mag.shape == shape
+        assert phase.shape == shape
+        assert sys(omega * 1j, squeeze=squeeze).shape == shape
+        assert ct.evalfr(sys, omega * 1j, squeeze=squeeze).shape == shape
+
+        # Changing config.default to False should return 3D frequency response
+        ct.config.set_defaults('control', squeeze=False)
+        mag, phase, _ = sys.frequency_response(omega)
+        assert mag.shape == (sys.outputs, sys.inputs, 8)
+        assert phase.shape == (sys.outputs, sys.inputs, 8)
+        assert sys(omega * 1j).shape == (sys.outputs, sys.inputs, 8)
+        assert ct.evalfr(sys, omega * 1j).shape == (sys.outputs, sys.inputs, 8)
diff --git a/control/xferfcn.py b/control/xferfcn.py
@@ -63,7 +63,7 @@
 from warnings import warn
 from itertools import chain
 from re import sub
-from .lti import LTI, common_timebase, isdtime
+from .lti import LTI, common_timebase, isdtime, _process_frequency_response
 from . import config
 
 __all__ = ['TransferFunction', 'tf', 'ss2tf', 'tfdata']
@@ -265,19 +265,8 @@ def __call__(self, x, squeeze=None):
             only if system is SISO and ``squeeze=True``.
 
         """
-        # Set value of squeeze argument if not set
-        if squeeze is None:
-            squeeze = config.defaults['control.squeeze']
-
         out = self.horner(x)
-        if not hasattr(x, '__len__'):
-            # received a scalar x, squeeze down the array along last dim
-            out = np.squeeze(out, axis=2)
-        if squeeze and self.issiso():
-            # return a scalar/1d array of outputs
-            return out[0][0]
-        else:
-            return out
+        return _process_frequency_response(self, x, out, squeeze=squeeze)
 
     def horner(self, x):
         """Evaluate system's transfer function at complex frequency
