add {control, disturbance}_indices to create_estimator_iosystem

murrayrm · murrayrm · commit e165dd5d7d76 · 2023-03-18T10:36:19.000-07:00
diff --git a/control/namedio.py b/control/namedio.py
@@ -22,8 +22,8 @@
     'namedio.sampled_system_name_prefix': '',
     'namedio.sampled_system_name_suffix': '$sampled'
 }
-    
-    
+
+
 class NamedIOSystem(object):
     def __init__(
             self, name=None, inputs=None, outputs=None, states=None, **kwargs):
@@ -586,21 +586,46 @@ def _process_signal_list(signals, prefix='s'):
         raise TypeError("Can't parse signal list %s" % str(signals))
 
 
+#
 # Utility function to process signal indices
-def _process_indices(arg, name, labels, default=None):
-    arg = default if arg is None else arg
-    if arg is None:
-        return None;
+#
+# Signal indices can be specified in one of four ways:
+#
+# 1. As a positive integer 'm', in which case we return a list
+#    corresponding to the first 'm' elements of a range of a given length
+#
+# 2. As a negative integer '-m', in which case we return a list
+#    corresponding to the last 'm' elements of a range of a given length
+#
+# 3. As a slice, in which case we return the a list corresponding to the
+#    indices specified by the slice of a range of a given length
+#
+# 4. As a list of ints or strings specifying specific indices.  Strings are
+#    compared to a list of labels to determine the index.
+#
+def _process_indices(arg, name, labels, length):
+    # Default is to return indices up to a certain length
+    arg = length if arg is None else arg
 
     if isinstance(arg, int):
-        return range(arg)
+        # Return the start or end of the list of possible indices
+        return list(range(arg)) if arg > 0 else list(range(length))[arg:]
+
     elif isinstance(arg, slice):
-        return arg
+        # Return the indices referenced by the slice
+        return list(range(length))[arg]
+
     elif isinstance(arg, list):
+        # Make sure the length is OK
+        if len(arg) > length:
+            raise ValueError(
+                f"{name}_indices list is too long; max length = {length}")
+
+        # Return the list, replacing strings with corresponding indices
         arg=arg.copy()
         for i, idx in enumerate(arg):
             if isinstance(idx, str):
                 arg[i] = labels.index(arg[i])
         return arg
-    else:
-        raise ValueError(f"invalid argument for {name}_indices")
+
+    raise ValueError(f"invalid argument for {name}_indices")
diff --git a/control/stochsys.py b/control/stochsys.py
@@ -22,7 +22,7 @@
 
 from .iosys import InputOutputSystem, LinearIOSystem, NonlinearIOSystem
 from .lti import LTI
-from .namedio import isctime, isdtime
+from .namedio import isctime, isdtime, _process_indices
 from .mateqn import care, dare, _check_shape
 from .statesp import StateSpace, _ssmatrix
 from .exception import ControlArgument, ControlNotImplemented
@@ -314,6 +314,7 @@ def dlqe(*args, **kwargs):
 #
 def create_estimator_iosystem(
         sys, QN, RN, P0=None, G=None, C=None,
+        control_indices=None, disturbance_indices=None,
         state_labels='xhat[{i}]', output_labels='xhat[{i}]',
         covariance_labels='P[{i},{j}]', sensor_labels=None):
     r"""Create an I/O system implementing a linear quadratic estimator
@@ -347,9 +348,10 @@ def create_estimator_iosystem(
 
     Parameters
     ----------
-    sys : InputOutputSystem
-        The I/O system that represents the process dynamics.  If no estimator
-        is given, the output of this system should represent the full state.
+    sys : LinearIOSystem
+        The linear I/O system that represents the process dynamics.  If no
+        estimator is given, the output of this system should represent the
+        full state.
     QN, RN : ndarray
         Process and sensor noise covariance matrices.
     P0 : ndarray, optional
@@ -362,14 +364,6 @@ def create_estimator_iosystem(
         If the system has full state output, define the measured values to
         be used by the estimator.  Otherwise, use the system output as the
         measured values.
-    {state, covariance, sensor, output}_labels : str or list of str, optional
-        Set the name of the signals to use for the internal state, covariance,
-        sensors, and outputs (state estimate).  If a single string is
-        specified, it should be a format string using the variable `i` as an
-        index (or `i` and `j` for covariance).  Otherwise, a list of
-        strings matching the size of the respective signal should be used.
-        Default is ``'xhat[{i}]'`` for state and output labels, ``'y[{i}]'``
-        for output labels and ``'P[{i},{j}]'`` for covariance labels.
 
     Returns
     -------
@@ -378,6 +372,47 @@ def create_estimator_iosystem(
         the system output y and input u and generates the estimated state
         xhat.
 
+    Other Parameters
+    ----------------
+    control_indices : int, slice, or list of int or string, optional
+        Specify the indices in the system input vector that correspond to
+        the control inputs.  These inputs will be used as known control
+        inputs for the estimator. If value is an integer `m`, the first `m`
+        system inputs are used.  Otherwise, the value should be a slice or
+        a list of indices.  The list of indices can be specified as either
+        integer offsets or as system input signal names.  If not specified,
+        defaults to the system inputs.
+    disturbance_indices : int, list of int, or slice, optional
+        Specify the indices in the system input vector that correspond to
+        the unknown disturbances.  These inputs are assumed to be white
+        noise with noise intensity QN.  If value is an integer `m`, the
+        last `m` system inputs are used.  Otherwise, the value should be a
+        slice or a list of indices.  The list of indices can be specified
+        as either integer offsets or as system input signal names.  If not
+        specified, the disturbances are assumed to be added to the system
+        inputs.
+    state_labels : str or list of str, optional
+        Set the names of the internal state estimate variables.  If a
+        single string is specified, it should be a format string using the
+        variable `i` as an index.  Otherwise, a list of strings matching
+        the number of system states should be used.  Default is "xhat[{i}]".
+    covariance_labels : str or list of str, optional
+        Set the name of the the covariance state variables.  If a single
+        string is specified, it should be a format string using the
+        variables `i` and `j` as indices.  Otherwise, a list of strings
+        matching the size of the covariance matrix should be used.  Default
+        is "P[{i},{j}]".
+    sensor_labels : str or list of str, optional
+        Set the name of the sensor signals (estimator inputs).  If
+        specified, it should be a format string using the variable `i` as
+        an index.  Otherwise, a list of strings matching the size of the
+        measured system outputs should be used.  Default is "y[{i}]".
+    output_labels : str or list of str, optional
+        Set the name of the estimator outputs (state estimate).  If a
+        single string is specified, it should be a format string using the
+        variable `i` as an index.  Otherwise, a list of strings matching
+        the size of the system state should be used.  Default is "xhat[{i}]".
+
     Notes
     -----
     This function can be used with the ``create_statefbk_iosystem()`` function
@@ -403,11 +438,45 @@ def create_estimator_iosystem(
     if not isinstance(sys, LinearIOSystem):
         raise ControlArgument("Input system must be a linear I/O system")
 
-    # Extract the matrices that we need for easy reference
-    A, B = sys.A, sys.B
+    # Set the state matrix for later use
+    A = sys.A
+
+    # Set the disturbance matrices (indices take priority over G)
+    ctrl_idx = _process_indices(
+        control_indices, 'control', sys.input_labels, sys.ninputs)
+
+    if disturbance_indices is None and control_indices is not None:
+        # Disturbance indices are the complement of control indices
+        dist_idx = [i for i in range(sys.ninputs) if i not in ctrl_idx]
+        if G is not None:
+            warn("'control_indices' and 'G' both specified; ignoring 'G'")
+        G = sys.B[:, dist_idx]
+
+    elif disturbance_indices is not None:
+        if G is not None:
+            warn("'disturbance_indices' and 'G' both specified; ignoring 'G'")
+
+        # If passed an integer, count from the end of the input vector
+        arg = -disturbance_indices if isinstance(disturbance_indices, int) \
+            else disturbance_indices
 
-    # Set the disturbance and output matrices
-    G = sys.B if G is None else G
+        dist_idx = _process_indices(
+            arg, 'disturbance', sys.input_labels, sys.ninputs)
+        G = sys.B[:, dist_idx]
+
+        # Set control indices to complement disturbance indices, if needed
+        if control_indices is None:
+            ctrl_idx = [i for i in range(sys.ninputs) if i not in dist_idx]
+
+    elif G is None:
+        G = sys.B
+
+    # Set the input and direct matrices
+    B = sys.B[:, ctrl_idx]
+    if not np.allclose(sys.D, 0):
+        raise NotImplemented("nonzero 'D' matrix not yet implemented")
+
+    # Set the output matrices
     if C is not None:
         # Make sure that we have the full system output
         if not np.array_equal(sys.C, np.eye(sys.nstates)):
@@ -425,7 +494,7 @@ def create_estimator_iosystem(
     # Initialize the covariance matrix
     if P0 is None:
         # Initalize P0 to the steady state value
-        L0, P0, _ = lqe(A, G, C, QN, RN)
+        _, P0, _ = lqe(A, G, C, QN, RN)
 
     # Figure out the labels to use
     if isinstance(state_labels, str):
@@ -447,6 +516,10 @@ def create_estimator_iosystem(
         # Generate the list of labels using the argument as a format string
         sensor_labels = [sensor_labels.format(i=i) for i in range(C.shape[0])]
 
+    # Set the input labels based on the system input
+    # TODO: allow these to be overriden
+    input_labels = [sys.input_labels[i] for i in ctrl_idx]
+
     if isctime(sys):
         # Create an I/O system for the state feedback gains
         # Note: reshape vectors into column vectors for legacy np.matrix
@@ -470,7 +543,7 @@ def _estim_update(t, x, u, params):
             L = P @ C.T @ R_inv
 
             # Update the state estimate
-            dxhat = A @ xhat + B @ u            # prediction
+            dxhat = A @ xhat + B @ u                    # prediction
             if correct:
                 dxhat -= L @ (C @ xhat - y)     # correction
 
@@ -500,7 +573,7 @@ def _estim_update(t, x, u, params):
             L = A @ P @ C.T @ Reps_inv
 
             # Update the state estimate
-            dxhat = A @ xhat + B @ u            # prediction
+            dxhat = A @ xhat + B @ u                    # prediction
             if correct:
                 dxhat -= L @ (C @ xhat - y)     # correction
 
@@ -518,7 +591,7 @@ def _estim_output(t, x, u, params):
     # Define the estimator system
     return NonlinearIOSystem(
         _estim_update, _estim_output, states=state_labels + covariance_labels,
-        inputs=sensor_labels + sys.input_labels, outputs=output_labels,
+        inputs=sensor_labels + input_labels, outputs=output_labels,
         dt=sys.dt)
 
 
diff --git a/control/tests/stochsys_test.py b/control/tests/stochsys_test.py
@@ -319,6 +319,7 @@ def test_correlation():
         T = np.logspace(0, 2, T.size)
         tau, Rtau = ct.correlation(T, V)
 
+@pytest.mark.slow
 @pytest.mark.parametrize('dt', [0, 1])
 def test_oep(dt):
     # Define the system to test, with additional input
@@ -455,6 +456,7 @@ def test_oep(dt):
         est3.states[:, -1], res3.states[:, -1], atol=1e-1, rtol=1e-2)
 
 
+@pytest.mark.slow
 def test_mhe():
     # Define the system to test, with additional input
     csys = ct.ss(
@@ -495,3 +497,62 @@ def test_mhe():
 
     # Make sure the estimated state is close to the actual state
     np.testing.assert_allclose(estp.outputs, resp.states, atol=1e-2, rtol=1e-4)
+
+@pytest.mark.parametrize("ctrl_indices, dist_indices", [
+    (slice(0, 3), None),
+    (3, None),
+    (None, 2),
+    ([0, 1, 4], None),
+    (['u[0]', 'u[1]', 'u[4]'], None),
+    (['u[0]', 'u[1]', 'u[4]'], ['u[1]', 'u[3]']),
+    (slice(0, 3), slice(3, 5))
+])
+def test_indices(ctrl_indices, dist_indices):
+    # Define a system with inputs (0:3), disturbances (3:5), and noise (5, 7)
+    ninputs = 3
+    nstates = ninputs + 1
+    ndisturbances = 2
+    noutputs = 2
+    nnoises = 0
+    # TODO: remove strictly proper
+    sys = ct.rss(nstates, noutputs, ninputs + ndisturbances + nnoises, strictly_proper=True)
+
+    # Create a system whose state we want to estimate
+    if ctrl_indices is not None:
+        ctrl_idx = ct.namedio._process_indices(
+            ctrl_indices, 'control', sys.input_labels, sys.ninputs)
+    else:
+        arg = -dist_indices if isinstance(dist_indices, int) else dist_indices
+        dist_idx = ct.namedio._process_indices(
+            arg, 'disturbance', sys.input_labels, sys.ninputs)
+        ctrl_idx = [i for i in range(sys.ninputs) if i not in dist_idx]
+    sysm = ct.ss(sys.A, sys.B[:, ctrl_idx], sys.C, sys.D[:, ctrl_idx])
+
+    # Set the simulation time based on the slowest system pole
+    from math import log
+    T = 10 / min(-sys.poles().real)
+
+    # Generate a system response with no disturbances
+    timepts = np.linspace(0, T, 50)
+    U = np.vstack([np.sin(timepts + i) for i in range(ninputs)])
+    resp = ct.input_output_response(
+        sysm, timepts, U, np.zeros(nstates),
+        solve_ivp_kwargs={'method': 'RK45', 'max_step': 0.01,
+                          'atol': 1, 'rtol': 1})
+    Y = resp.outputs
+
+    # Create an estimator
+    QN = np.eye(ndisturbances)
+    RN = np.eye(noutputs)
+    P0 = np.eye(nstates)
+    estim = ct.create_estimator_iosystem(
+        sys, QN, RN, control_indices=ctrl_indices,
+        disturbance_indices=dist_indices)
+
+    # Run estimator (no prediction + same solve_ivp params => should be exact)
+    resp_estim = ct.input_output_response(
+        estim, timepts, [Y, U], [np.zeros(nstates), P0],
+        solve_ivp_kwargs={'method': 'RK45', 'max_step': 0.01,
+                          'atol': 1, 'rtol': 1},
+            params={'correct': False})
+    np.testing.assert_allclose(resp.states, resp_estim.outputs, rtol=1e-2)
