consolidate time/frequency response processing for easier updates

murrayrm · murrayrm · commit b337c9cd525a · 2021-01-14T21:12:30.000-08:00
diff --git a/control/iosys.py b/control/iosys.py
@@ -37,7 +37,7 @@
 from warnings import warn
 
 from .statesp import StateSpace, tf2ss
-from .timeresp import _check_convert_array
+from .timeresp import _check_convert_array, _process_time_response
 from .lti import isctime, isdtime, common_timebase
 from . import config
 
@@ -1353,7 +1353,7 @@ def __init__(self, io_sys, ss_sys=None):
 
 
 def input_output_response(sys, T, U=0., X0=0, params={}, method='RK45',
-                          return_x=False, squeeze=None):
+                          transpose=False, return_x=False, squeeze=None):
 
     """Compute the output response of a system to a given input.
 
@@ -1398,9 +1398,6 @@ def input_output_response(sys, T, U=0., X0=0, params={}, method='RK45',
     if not isinstance(sys, InputOutputSystem):
         raise TypeError("System of type ", type(sys), " not valid")
 
-    if squeeze is None:
-        squeeze = config.defaults['control.squeeze']
-
     # Compute the time interval and number of steps
     T0, Tf = T[0], T[-1]
     n_steps = len(T)
@@ -1423,12 +1420,8 @@ def input_output_response(sys, T, U=0., X0=0, params={}, method='RK45',
         for i in range(len(T)):
             u = U[i] if len(U.shape) == 1 else U[:, i]
             y[:, i] = sys._out(T[i], [], u)
-        if squeeze and y.shape[0] == 1:
-            y = y[0]
-        if return_x:
-            return T, y, []
-        else:
-            return T, y
+        return _process_time_response(sys, T, y, [], transpose=transpose,
+                                      return_x=return_x, squeeze=squeeze)
 
     # create X0 if not given, test if X0 has correct shape
     X0 = _check_convert_array(X0, [(nstates,), (nstates, 1)],
@@ -1503,14 +1496,8 @@ def ivp_rhs(t, x): return sys._rhs(t, x, u(t))
     else:                       # Neither ctime or dtime??
         raise TypeError("Can't determine system type")
 
-    # Get rid of extra dimensions in the output, of desired
-    if squeeze and y.shape[0] == 1:
-        y = y[0]
-
-    if return_x:
-        return soln.t, y, soln.y
-    else:
-        return soln.t, y
+    return _process_time_response(sys, soln.t, y, soln.y, transpose=transpose,
+                                  return_x=return_x, squeeze=squeeze)
 
 
 def find_eqpt(sys, x0, u0=[], y0=None, t=0, params={},
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
 
@@ -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/timeresp.py b/control/timeresp.py
@@ -277,8 +277,6 @@ def forced_response(sys, T=None, U=0., X0=0., transpose=False,
     if not isinstance(sys, LTI):
         raise TypeError('Parameter ``sys``: must be a ``LTI`` object. '
                         '(For example ``StateSpace`` or ``TransferFunction``)')
-    if squeeze is None:
-        squeeze = config.defaults['control.squeeze']
 
     sys = _convert_to_statespace(sys)
     A, B, C, D = np.asarray(sys.A), np.asarray(sys.B), np.asarray(sys.C), \
@@ -433,7 +431,16 @@ def forced_response(sys, T=None, U=0., X0=0., transpose=False,
         xout = np.transpose(xout)
         yout = np.transpose(yout)
 
+    return _process_time_response(sys, tout, yout, xout, transpose=transpose,
+                                  return_x=return_x, squeeze=squeeze)
+
+
+# Process time responses in a uniform way
+def _process_time_response(sys, tout, yout, xout, transpose=None,
+                           return_x=False, squeeze=None):
     # Get rid of unneeded dimensions
+    if squeeze is None:
+        squeeze = config.defaults['control.squeeze']
     if squeeze and yout.shape[0] == 1:
         yout = yout[0]
 
@@ -443,10 +450,8 @@ def forced_response(sys, T=None, U=0., X0=0., transpose=False,
         yout = np.transpose(yout)
         xout = np.transpose(xout)
 
-    if return_x:
-        return tout, yout, xout
-
-    return tout, yout
+    # Return time, output, and (optionally) state
+    return (tout, yout, xout) if return_x else (tout, yout)
 
 
 def _get_ss_simo(sys, input=None, output=None):
@@ -640,7 +645,7 @@ def step_info(sys, T=None, T_num=None, SettlingTimeThreshold=0.02,
         'SettlingMin': yout[tr_upper_index:].min(),
         'SettlingMax': yout.max(),
         'Overshoot': 100. * (yout.max() - InfValue) / (InfValue - yout[0]),
-        'Undershoot': yout.min(), # not very confident about this
+        'Undershoot': yout.min(),       # not very confident about this
         'Peak': yout[PeakIndex],
         'PeakTime':  T[PeakIndex],
         'SteadyStateValue': InfValue
@@ -728,13 +733,8 @@ def initial_response(sys, T=None, X0=0., input=0, output=None, T_num=None,
         T = _default_time_vector(sys, N=T_num, tfinal=T, is_step=False)
     U = np.zeros_like(T)
 
-    T, yout, _xout = forced_response(sys, T, U, X0, transpose=transpose,
-                                     return_x=True, squeeze=squeeze)
-
-    if return_x:
-        return T, yout, _xout
-
-    return T, yout
+    return forced_response(sys, T, U, X0, transpose=transpose,
+                           return_x=return_x, squeeze=squeeze)
 
 
 def impulse_response(sys, T=None, X0=0., input=None, output=None, T_num=None,
@@ -841,15 +841,11 @@ def impulse_response(sys, T=None, X0=0., input=None, output=None, T_num=None,
         new_X0 = B + X0
     else:
         new_X0 = X0
-        U[0] = 1./sys.dt # unit area impulse
+        U[0] = 1./sys.dt        # unit area impulse
 
-    T, yout, _xout = forced_response(sys, T, U, new_X0, transpose=transpose,
-                                     return_x=True, squeeze=squeeze)
+    return forced_response(sys, T, U, new_X0, transpose=transpose,
+                           return_x=return_x, squeeze=squeeze)
 
-    if return_x:
-        return T, yout, _xout
-
-    return T, yout
 
 # utility function to find time period and time increment using pole locations
 def _ideal_tfinal_and_dt(sys, is_step=True):
@@ -949,7 +945,7 @@ def _ideal_tfinal_and_dt(sys, is_step=True):
 
         if p_int.size > 0:
             tfinal = tfinal * 5
-    else: # cont time
+    else:       # cont time
         sys_ss = _convert_to_statespace(sys)
         # Improve conditioning via balancing and zeroing tiny entries
         # See <w,v> for [[1, 2, 0], [9, 1, 0.01], [1, 2, 10*np.pi]]
@@ -977,7 +973,7 @@ def _ideal_tfinal_and_dt(sys, is_step=True):
         dc = np.zeros_like(p, dtype=float)
         # well-conditioned nonzero poles, np.abs just in case
         ok = np.abs(eig_sens) <= 1/sqrt_eps
-        # the averaged t->inf response of each simple eigval on each i/o channel
+        # averaged t->inf response of each simple eigval on each i/o channel
         # See, A = [[-1, k], [0, -2]], response sizes are k-dependent (that is
         # R/L eigenvector dependent)
         dc[ok] = norm(v[ok, :], axis=1)*norm(w[:, ok], axis=0)*eig_sens[ok]
@@ -1003,7 +999,7 @@ def _ideal_tfinal_and_dt(sys, is_step=True):
         texp_mode = log_decay_percent / np.abs(psub[~iw & ~ints].real)
         tfinal += texp_mode.tolist()
         dt += minimum(texp_mode / 50,
-                    (2 * np.pi / pts_per_cycle / wnsub[~iw & ~ints])).tolist()
+                      (2*np.pi / pts_per_cycle / wnsub[~iw & ~ints])).tolist()
 
         # All integrators?
         if len(tfinal) == 0:
@@ -1014,30 +1010,32 @@ def _ideal_tfinal_and_dt(sys, is_step=True):
 
     return tfinal, dt
 
+
 def _default_time_vector(sys, N=None, tfinal=None, is_step=True):
     """Returns a time vector that has a reasonable number of points.
     if system is discrete-time, N is ignored """
 
     N_max = 5000
-    N_min_ct = 100 # min points for cont time systems
-    N_min_dt = 20 # more common to see just a few samples in discrete-time
+    N_min_ct = 100      # min points for cont time systems
+    N_min_dt = 20       # more common to see fewer samples in discrete-time
 
     ideal_tfinal, ideal_dt = _ideal_tfinal_and_dt(sys, is_step=is_step)
 
     if isdtime(sys, strict=True):
         # only need to use default_tfinal if not given; N is ignored.
         if tfinal is None:
             # for discrete time, change from ideal_tfinal if N too large/small
-            N = int(np.clip(ideal_tfinal/sys.dt, N_min_dt, N_max))# [N_min, N_max]
+            N = int(np.clip(ideal_tfinal/sys.dt, N_min_dt, N_max))
             tfinal = sys.dt * N
         else:
             N = int(tfinal/sys.dt)
-            tfinal = N * sys.dt # make tfinal an integer multiple of sys.dt
+            tfinal = N * sys.dt   # make tfinal an integer multiple of sys.dt
     else:
         if tfinal is None:
             # for continuous time, simulate to ideal_tfinal but limit N
             tfinal = ideal_tfinal
         if N is None:
-            N = int(np.clip(tfinal/ideal_dt, N_min_ct, N_max)) # N<-[N_min, N_max]
+            # N <- [N_min, N_max]
+            N = int(np.clip(tfinal/ideal_dt, N_min_ct, N_max))
 
     return np.linspace(0, tfinal, N, endpoint=False)
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
