consistent squeezing for state property + legacy interface + doc updates

murrayrm · murrayrm · commit ce5a95c31738 · 2021-08-26T15:30:10.000-07:00
diff --git a/control/tests/trdata_test.py b/control/tests/trdata_test.py
@@ -51,14 +51,17 @@ def test_trdata_shapes(nin, nout, squeeze):
     # Check shape of class properties
     if sys.issiso():
         assert res.outputs.shape == (ntimes,)
+        assert res._legacy_states.shape == (sys.nstates, ntimes)
         assert res.states.shape == (sys.nstates, ntimes)
         assert res.inputs is None
     elif res.squeeze is True:
         assert res.outputs.shape == (ntimes, )
+        assert res._legacy_states.shape == (sys.nstates, ntimes)
         assert res.states.shape == (sys.nstates, ntimes)
         assert res.inputs is None
     else:
         assert res.outputs.shape == (sys.noutputs, ntimes)
+        assert res._legacy_states.shape == (sys.nstates, ntimes)
         assert res.states.shape == (sys.nstates, ntimes)
         assert res.inputs is None
 
@@ -78,21 +81,26 @@ def test_trdata_shapes(nin, nout, squeeze):
         # Check shape of inputs and outputs
         if sys.issiso() and squeeze is not False:
             assert res.outputs.shape == (ntimes, )
+            assert res.states.shape == (sys.nstates, ntimes)
             assert res.inputs.shape == (ntimes, )
         elif res.squeeze is True:
             assert res.outputs.shape == \
                 np.empty((sys.noutputs, sys.ninputs, ntimes)).squeeze().shape
+            assert res.states.shape == \
+                np.empty((sys.nstates, sys.ninputs, ntimes)).squeeze().shape
             assert res.inputs.shape == \
                 np.empty((sys.ninputs, sys.ninputs, ntimes)).squeeze().shape
         else:
             assert res.outputs.shape == (sys.noutputs, sys.ninputs, ntimes)
+            assert res.states.shape == (sys.nstates, sys.ninputs, ntimes)
             assert res.inputs.shape == (sys.ninputs, sys.ninputs, ntimes)
 
-        # Check state space dimensions (not affected by squeeze)
+        # Check legacy state space dimensions (not affected by squeeze)
         if sys.issiso():
-            assert res.states.shape == (sys.nstates, ntimes)
+            assert res._legacy_states.shape == (sys.nstates, ntimes)
         else:
-            assert res.states.shape == (sys.nstates, sys.ninputs, ntimes)
+            assert res._legacy_states.shape == \
+                (sys.nstates, sys.ninputs, ntimes)
 
     #
     # Forced response
@@ -107,14 +115,18 @@ def test_trdata_shapes(nin, nout, squeeze):
 
     if sys.issiso() and squeeze is not False:
         assert res.outputs.shape == (ntimes,)
+        assert res.states.shape == (sys.nstates, ntimes)
         assert res.inputs.shape == (ntimes,)
     elif squeeze is True:
         assert res.outputs.shape == \
             np.empty((sys.noutputs, 1, ntimes)).squeeze().shape
+        assert res.states.shape == \
+            np.empty((sys.nstates, 1, ntimes)).squeeze().shape
         assert res.inputs.shape == \
             np.empty((sys.ninputs, 1, ntimes)).squeeze().shape
     else:                       # MIMO or squeeze is False
         assert res.outputs.shape == (sys.noutputs, ntimes)
+        assert res.states.shape == (sys.nstates, ntimes)
         assert res.inputs.shape == (sys.ninputs, ntimes)
 
     # Check state space dimensions (not affected by squeeze)
@@ -141,16 +153,28 @@ def test_response_copy():
     # Squeeze
     response_siso_as_mimo = response_siso(squeeze=False)
     assert response_siso_as_mimo.outputs.shape == (1, 1, siso_ntimes)
-    assert response_siso_as_mimo.states.shape == (4, siso_ntimes)
+    assert response_siso_as_mimo.states.shape == (4, 1, siso_ntimes)
+    assert response_siso_as_mimo._legacy_states.shape == (4, siso_ntimes)
 
     response_mimo_squeezed = response_mimo(squeeze=True)
     assert response_mimo_squeezed.outputs.shape == (2, mimo_ntimes)
-    assert response_mimo_squeezed.states.shape == (4, 1, mimo_ntimes)
+    assert response_mimo_squeezed.states.shape == (4, mimo_ntimes)
+    assert response_mimo_squeezed._legacy_states.shape == (4, 1, mimo_ntimes)
 
     # Squeeze and transpose
     response_mimo_sqtr = response_mimo(squeeze=True, transpose=True)
     assert response_mimo_sqtr.outputs.shape == (mimo_ntimes, 2)
-    assert response_mimo_sqtr.states.shape == (mimo_ntimes, 4, 1)
+    assert response_mimo_sqtr.states.shape == (mimo_ntimes, 4)
+    assert response_mimo_sqtr._legacy_states.shape == (mimo_ntimes, 4, 1)
+
+    # Return_x
+    t, y = response_mimo
+    t, y = response_mimo()
+    t, y, x = response_mimo(return_x=True)
+    with pytest.raises(ValueError, match="too many"):
+        t, y = response_mimo(return_x=True)
+    with pytest.raises(ValueError, match="not enough"):
+        t, y, x = response_mimo
 
     # Unknown keyword
     with pytest.raises(ValueError, match="unknown"):
diff --git a/control/timeresp.py b/control/timeresp.py
@@ -88,7 +88,7 @@
 
 
 class TimeResponseData():
-    """Class for returning time responses.
+    """A class for returning time responses.
 
     This class maintains and manipulates the data corresponding to the
     temporal response of an input/output system.  It is used as the return
@@ -140,20 +140,18 @@ class TimeResponseData():
         performs squeeze processing.
 
     squeeze : bool, optional
-        By default, if a system is single-input, single-output (SISO) then
-        the inputs and outputs are returned as a 1D array (indexed by time)
-        and if a system is multi-input or multi-output, then the inputs are
-        returned as a 2D array (indexed by input and time) and the outputs
-        are returned as either a 2D array (indexed by output and time) or a
-        3D array (indexed by output, trace, and time).  If ``squeeze=True``,
-        access to the output response will remove single-dimensional entries
-        from the shape of the inputs and outputs even if the system is not
-        SISO. If ``squeeze=False``, the input is returned as a 2D or 3D
-        array (indexed by the input [if multi-input], trace [if
-        multi-trace] and time) and the output as a 2D or 3D array (indexed
-        by the output, trace [if multi-trace], and time) even if the system
-        is SISO. The default value can be set using
-        config.defaults['control.squeeze_time_response'].
+        By default, if a system is single-input, single-output (SISO)
+        then the outputs (and inputs) are returned as a 1D array
+        (indexed by time) and if a system is multi-input or
+        multi-output, then the outputs are returned as a 2D array
+        (indexed by output and time) or a 3D array (indexed by output,
+        trace, and time).  If ``squeeze=True``, access to the output
+        response will remove single-dimensional entries from the shape
+        of the inputs and outputs even if the system is not SISO. If
+        ``squeeze=False``, the output is returned as a 2D or 3D array
+        (indexed by the output [if multi-input], trace [if multi-trace]
+        and time) even if the system is SISO. The default value can be
+        set using config.defaults['control.squeeze_time_response'].
 
     transpose : bool, optional
         If True, transpose all input and output arrays (for backward
@@ -183,6 +181,9 @@ class TimeResponseData():
          t, y = step_response(sys)
          t, y, x = step_response(sys, return_x=True)
 
+       When using this (legacy) interface, the state vector is not affected by
+       the `squeeze` parameter.
+
     2. For backward compatibility with earlier version of python-control,
        this class has ``__getitem__`` and ``__len__`` methods that allow the
        return value to be indexed:
@@ -191,11 +192,16 @@ class TimeResponseData():
          response[1]: returns the output vector
          response[2]: returns the state vector
 
+       When using this (legacy) interface, the state vector is not affected by
+       the `squeeze` parameter.
+
     3. The default settings for ``return_x``, ``squeeze`` and ``transpose``
        can be changed by calling the class instance and passing new values:
 
          response(tranpose=True).input
 
+       See :meth:`TimeResponseData.__call__` for more information.
+
     """
 
     def __init__(
@@ -251,8 +257,8 @@ def __init__(
             the system is SISO. The default value can be set using
             config.defaults['control.squeeze_time_response'].
 
-        Additional parameters
-        ---------------------
+        Other parameters
+        ----------------
         transpose : bool, optional
             If True, transpose all input and output arrays (for backward
             compatibility with MATLAB and :func:`scipy.signal.lsim`).
@@ -391,8 +397,10 @@ def __init__(
             raise ValueError("unknown squeeze value")
         self.squeeze = squeeze
 
-        # Store legacy keyword values (only needed for legacy interface)
+        # Keep track of whether to transpose for MATLAB/scipy.signal
         self.transpose = transpose
+
+        # Store legacy keyword values (only needed for legacy interface)
         self.return_x = return_x
 
     def __call__(self, **kwargs):
@@ -405,13 +413,13 @@ def __call__(self, **kwargs):
         Parameters
         ----------
         squeeze : bool, optional
-            If squeeze=True, access to the output response will
-            remove single-dimensional entries from the shape of the inputs
-            and outputs even if the system is not SISO. If squeeze=False,
-            keep the input as a 2D or 3D array (indexed by the input (if
-            multi-input), trace (if single input) and time) and the output
-            as a 3D array (indexed by the output, trace, and time) even if
-            the system is SISO.
+            If squeeze=True, access to the output response will remove
+            single-dimensional entries from the shape of the inputs, outputs,
+            and states even if the system is not SISO. If squeeze=False, keep
+            the input as a 2D or 3D array (indexed by the input (if
+            multi-input), trace (if single input) and time) and the output and
+            states as a 3D array (indexed by the output/state, trace, and
+            time) even if the system is SISO.
 
         transpose : bool, optional
             If True, transpose all input and output arrays (for backward
@@ -421,13 +429,15 @@ def __call__(self, **kwargs):
         return_x : bool, optional
             If True, return the state vector when enumerating result by
             assigning to a tuple (default = False).
+
         """
         # Make a copy of the object
         response = copy(self)
 
         # Update any keywords that we were passed
         response.transpose = kwargs.pop('transpose', self.transpose)
         response.squeeze = kwargs.pop('squeeze', self.squeeze)
+        response.return_x = kwargs.pop('return_x', self.squeeze)
 
         # Make sure no unknown keywords were passed
         if len(kwargs) != 0:
@@ -452,32 +462,40 @@ def outputs(self):
 
         Output response of the system, indexed by either the output and time
         (if only a single input is given) or the output, trace, and time
-        (for multiple traces).
+        (for multiple traces).  See :attr:`TimeResponseData.squeeze` for a
+        description of how this can be modified using the `squeeze` keyword.
 
         :type: 1D, 2D, or 3D array
+
         """
         t, y = _process_time_response(
             self.t, self.y, issiso=self.issiso,
             transpose=self.transpose, squeeze=self.squeeze)
         return y
 
-    # Getter for state (implements non-standard squeeze processing)
+    # Getter for states (implements squeeze processing)
     @property
     def states(self):
         """Time response state vector.
 
         Time evolution of the state vector, indexed indexed by either the
-        state and time (if only a single trace is given) or the state,
-        trace, and time (for multiple traces).
+        state and time (if only a single trace is given) or the state, trace,
+        and time (for multiple traces).  See :attr:`TimeResponseData.squeeze`
+        for a description of how this can be modified using the `squeeze`
+        keyword.
 
         :type: 2D or 3D array
-        """
 
+        """
         if self.x is None:
             return None
 
+        elif self.squeeze is True:
+            x = self.x.squeeze()
+
         elif self.ninputs == 1 and self.noutputs == 1 and \
-             self.ntraces == 1 and self.x.ndim == 3:
+             self.ntraces == 1 and self.x.ndim == 3 and \
+             self.squeeze is not False:
             # Single-input, single-output system with single trace
             x = self.x[:, 0, :]
 
@@ -491,7 +509,7 @@ def states(self):
 
         return x
 
-    # Getter for state (implements squeeze processing)
+    # Getter for inputs (implements squeeze processing)
     @property
     def inputs(self):
         """Time response input vector.
@@ -504,7 +522,12 @@ def inputs(self):
         the two.  If a 3D vector is passed, then it represents a multi-trace,
         multi-input signal, indexed by input, trace, and time.
 
+        See :attr:`TimeResponseData.squeeze` for a description of how the
+        dimensions of the input vector can be modified using the `squeeze`
+        keyword.
+
         :type: 1D or 2D array
+
         """
         if self.u is None:
             return None
@@ -514,11 +537,45 @@ def inputs(self):
             transpose=self.transpose, squeeze=self.squeeze)
         return u
 
+    # Getter for legacy state (implements non-standard squeeze processing)
+    @property
+    def _legacy_states(self):
+        """Time response state vector (legacy version).
+
+        Time evolution of the state vector, indexed indexed by either the
+        state and time (if only a single trace is given) or the state,
+        trace, and time (for multiple traces).
+
+        The `legacy_states` property is not affected by the `squeeze` keyword
+        and hence it will always have these dimensions.
+
+        :type: 2D or 3D array
+
+        """
+
+        if self.x is None:
+            return None
+
+        elif self.ninputs == 1 and self.noutputs == 1 and \
+             self.ntraces == 1 and self.x.ndim == 3:
+            # Single-input, single-output system with single trace
+            x = self.x[:, 0, :]
+
+        else:
+            # Return the full set of data
+            x = self.x
+
+        # Transpose processing
+        if self.transpose:
+            x = np.transpose(x, np.roll(range(x.ndim), 1))
+
+        return x
+
     # Implement iter to allow assigning to a tuple
     def __iter__(self):
         if not self.return_x:
             return iter((self.time, self.outputs))
-        return iter((self.time, self.outputs, self.states))
+        return iter((self.time, self.outputs, self._legacy_states))
 
     # Implement (thin) getitem to allow access via legacy indexing
     def __getitem__(self, index):
@@ -533,7 +590,7 @@ def __getitem__(self, index):
         if index == 1:
             return self.outputs
         if index == 2:
-            return self.states
+            return self._legacy_states
         raise IndexError
 
     # Implement (thin) len to emulate legacy testing interface
diff --git a/doc/classes.rst b/doc/classes.rst
@@ -17,6 +17,7 @@ these directly.
    TransferFunction
    StateSpace
    FrequencyResponseData
+   TimeResponseData
 
 Input/output system subclasses
 ==============================
@@ -47,4 +48,3 @@ Additional classes
    flatsys.SystemTrajectory
    optimal.OptimalControlProblem
    optimal.OptimalControlResult
-   TimeResponseData
diff --git a/doc/control.rst b/doc/control.rst
@@ -70,7 +70,6 @@ Time domain simulation
     input_output_response
     step_response
     phase_plot
-    TimeResponseData
 
 Block diagram algebra
 =====================
