Add list of systems functionality to time response functions

murrayrm · murrayrm · commit 8b416eb82c6b · 2024-06-29T12:41:18.000-07:00
diff --git a/control/ctrlplot.py b/control/ctrlplot.py
@@ -125,11 +125,10 @@ def _update_suptitle(fig, title, rcParams=None, frame='axes'):
     if fig is not None and isinstance(title, str):
         # Get the current title, if it exists
         old_title = None if fig._suptitle is None else fig._suptitle._text
-        new_title = title
 
         if old_title is not None:
             # Find the common part of the titles
-            common_prefix = commonprefix([old_title, new_title])
+            common_prefix = commonprefix([old_title, title])
 
             # Back up to the last space
             last_space = common_prefix.rfind(' ')
@@ -138,9 +137,9 @@ def _update_suptitle(fig, title, rcParams=None, frame='axes'):
             common_len = len(common_prefix)
 
             # Add the new part of the title (usually the system name)
-            if old_title[common_len:] != new_title[common_len:]:
+            if old_title[common_len:] != title[common_len:]:
                 separator = ',' if len(common_prefix) > 0 else ';'
-                new_title = old_title + separator + new_title[common_len:]
+                title = old_title + separator + title[common_len:]
 
         # Add the title
         suptitle(title, fig=fig, rcParams=rcParams, frame=frame)
diff --git a/control/frdata.py b/control/frdata.py
@@ -653,6 +653,14 @@ def plot(self, plot_type=None, *args, **kwargs):
 
     # Convert to pandas
     def to_pandas(self):
+        """Convert response data to pandas data frame.
+
+        Creates a pandas data frame for the value of the frequency
+        response at each `omega`.  The frequency response values are
+        labeled in the form "H_{<out>, <in>}" where "<out>" and "<in>"
+        are replaced with the output and input labels for the system.
+
+        """
         if not pandas_check():
             ImportError('pandas not installed')
         import pandas
diff --git a/control/nlsys.py b/control/nlsys.py
@@ -27,8 +27,8 @@
 from . import config
 from .iosys import InputOutputSystem, _parse_spec, _process_iosys_keywords, \
     _process_signal_list, common_timebase, isctime, isdtime
-from .timeresp import TimeResponseData, _check_convert_array, \
-    _process_time_response
+from .timeresp import _check_convert_array, _process_time_response, \
+    TimeResponseData, TimeResponseList
 
 __all__ = ['NonlinearIOSystem', 'InterconnectedSystem', 'nlsys',
            'input_output_response', 'find_eqpt', 'linearize',
@@ -1448,14 +1448,17 @@ def input_output_response(
     if kwargs:
         raise TypeError("unrecognized keyword(s): ", str(kwargs))
 
-    # Convert the first argument to a list
-    syslist = sysdata if isinstance(sysdata, (list, tuple)) else [sysdata]
-
-    # TODO: implement step responses for multiple systems
-    if len(syslist) > 1:
-        raise NotImplementedError(
-            "step responses for multiple systems not yet implemented")
-    sys = syslist[0]
+    # If passed a list, recursively call individual responses with given T
+    if isinstance(sysdata, (list, tuple)):
+        responses = []
+        for sys in sysdata:
+            responses.append(input_output_response(
+                sys, T, U=U, X0=X0, params=params, transpose=transpose,
+                return_x=return_x, squeeze=squeeze, t_eval=t_eval,
+                solve_ivp_kwargs=solve_ivp_kwargs, **kwargs))
+        return TimeResponseList(responses)
+    else:
+        sys = sysdata
 
     # Sanity checking on the input
     if not isinstance(sys, NonlinearIOSystem):
diff --git a/control/tests/kwargs_test.py b/control/tests/kwargs_test.py
@@ -308,6 +308,7 @@ def test_response_plot_kwargs(data_fcn, plot_fcn, mimo):
     'StateSpace.sample': test_unrecognized_kwargs,
     'TimeResponseData.__call__': trdata_test.test_response_copy,
     'TimeResponseData.plot': timeplot_test.test_errors,
+    'TimeResponseList.plot': timeplot_test.test_errors,
     'TransferFunction.__init__': test_unrecognized_kwargs,
     'TransferFunction.sample': test_unrecognized_kwargs,
     'optimal.OptimalControlProblem.__init__':
diff --git a/control/tests/timeplot_test.py b/control/tests/timeplot_test.py
@@ -313,14 +313,52 @@ def test_combine_time_responses():
         combresp6 = ct.combine_time_responses([resp1, resp])
 
 
-@pytest.mark.xfail(
-    reason="step responses for multiple systems not yet implemented")
-def test_list_responses():
-    sys1 = ct.rss(2, 2, 2)
-    sys2 = ct.rss(2, 2, 2)
+@pytest.mark.parametrize("resp_fcn", [
+    ct.step_response, ct.initial_response, ct.impulse_response,
+    ct.forced_response, ct.input_output_response])
+def test_list_responses(resp_fcn):
+    sys1 = ct.rss(2, 2, 2, strictly_proper=True)
+    sys2 = ct.rss(2, 2, 2, strictly_proper=True)
+
+    # Figure out the expected shape of the system
+    match resp_fcn:
+        case ct.step_response | ct.impulse_response:
+            shape = (2, 2)
+            kwargs = {}
+        case ct.initial_response:
+            shape = (2, 1)
+            kwargs = {}
+        case ct.forced_response | ct.input_output_response:
+            shape = (4, 1)      # outputs and inputs both plotted
+            T = np.linspace(0, 10)
+            U = [np.sin(T), np.cos(T)]
+            kwargs = {'T': T, 'U': U}
+
+    resp1 = resp_fcn(sys1, **kwargs)
+    resp2 = resp_fcn(sys2, **kwargs)
+
+    # Sequential plotting results in colors rotating
+    plt.figure()
+    out1 = resp1.plot()
+    out2 = resp2.plot()
+    assert out1.shape == shape
+    assert out2.shape == shape
+    for row in range(2):        # just look at the outputs
+        for col in range(shape[1]):
+            assert out1[row, col][0].get_color() == 'tab:blue'
+            assert out2[row, col][0].get_color() == 'tab:orange'
 
-    resp = ct.step_response([sys1, sys2]).plot()
-    assert resp.ntraces == 2
+    plt.figure()
+    resp_combined = resp_fcn([sys1, sys2], **kwargs)
+    assert isinstance(resp_combined, ct.timeresp.TimeResponseList)
+    assert resp_combined[0].time[-1] == max(resp1.time[-1], resp2.time[-1])
+    assert resp_combined[1].time[-1] == max(resp1.time[-1], resp2.time[-1])
+    out = resp_combined.plot()
+    assert out.shape == shape
+    for row in range(2):        # just look at the outputs
+        for col in range(shape[1]):
+            assert out[row, col][0].get_color() == 'tab:blue'
+            assert out[row, col][1].get_color() == 'tab:orange'
 
 
 @slycotonly
@@ -421,6 +459,12 @@ def test_errors():
             out = stepresp.plot('k-', **propkw)
             assert out[0, 0][0].get_color() == 'k'
 
+    # Make sure TimeResponseLists also work
+    stepresp = ct.step_response([sys, sys])
+    with pytest.raises(AttributeError,
+                       match="(has no property|unexpected keyword)"):
+        stepresp.plot(unknown=None)
+
 if __name__ == "__main__":
     #
     # Interactive mode: generate plots for manual viewing
@@ -519,3 +563,28 @@ def test_errors():
     input_props=[{'color': c} for c in ['red', 'green']],
     trace_props=[{'linestyle': s} for s in ['-', '--']])
     plt.savefig('timeplot-mimo_step-linestyle.png')
+
+    sys1 = ct.rss(4, 2, 2)
+    sys2 = ct.rss(4, 2, 2)
+    resp_list = ct.step_response([sys1, sys2])
+
+    fig = plt.figure()
+    ct.combine_time_responses(
+        [ct.step_response(sys1, resp_list[0].time),
+         ct.step_response(sys2, resp_list[1].time)]
+    ).plot(overlay_traces=True)
+    ct.suptitle("[Combine] " + fig._suptitle._text)
+
+    fig = plt.figure()
+    ct.step_response(sys1).plot()
+    ct.step_response(sys2).plot()
+    ct.suptitle("[Sequential] " + fig._suptitle._text)
+
+    fig = plt.figure()
+    ct.step_response(sys1).plot(color='b')
+    ct.step_response(sys2).plot(color='r')
+    ct.suptitle("[Seq w/color] " + fig._suptitle._text)
+
+    fig = plt.figure()
+    ct.step_response([sys1, sys2]).plot()
+    ct.suptitle("[List] " + fig._suptitle._text)
diff --git a/control/timeplot.py b/control/timeplot.py
@@ -44,6 +44,7 @@
     'timeplot.time_label': "Time [s]",
 }
 
+
 # Plot the input/output response of a system
 def time_response_plot(
         data, *fmt, ax=None, plot_inputs=None, plot_outputs=True,
@@ -364,6 +365,16 @@ def _make_line_label(signal_index, signal_labels, trace_index):
 
         return label
 
+    #
+    # Store the color offsets with the figure to allow color/style cycling
+    #
+    # To allow repeated calls to time_response_plot() to cycle through
+    # colors, we store an offset in the figure object that we can
+    # retrieve at a later date, if needed.
+    #
+    output_offset = fig._output_offset = getattr(fig, '_output_offset', 0)
+    input_offset = fig._input_offset = getattr(fig, '_input_offset', 0)
+
     # Go through each trace and each input/output
     for trace in range(ntraces):
         # Plot the output
@@ -373,7 +384,8 @@ def _make_line_label(signal_index, signal_labels, trace_index):
             # Set up line properties for this output, trace
             if len(fmt) == 0:
                 line_props = output_props[
-                    i % oprop_len if overlay_signals else 0].copy()
+                    (i + output_offset) % oprop_len if overlay_signals
+                    else output_offset].copy()
                 line_props.update(
                     trace_props[trace % tprop_len if overlay_traces else 0])
                 line_props.update(kwargs)
@@ -397,7 +409,8 @@ def _make_line_label(signal_index, signal_labels, trace_index):
             # Set up line properties for this output, trace
             if len(fmt) == 0:
                 line_props = input_props[
-                    i % iprop_len if overlay_signals else 0].copy()
+                    (i + input_offset) % iprop_len if overlay_signals
+                    else input_offset].copy()
                 line_props.update(
                     trace_props[trace % tprop_len if overlay_traces else 0])
                 line_props.update(kwargs)
@@ -407,6 +420,12 @@ def _make_line_label(signal_index, signal_labels, trace_index):
             out[input_map[i, trace]] += ax_array[input_map[i, trace]].plot(
                 x, y, *fmt, label=label, **line_props)
 
+    # Update the offsets so that we start at a new color/style the next time
+    fig._output_offset = (
+        output_offset + (noutputs if overlay_signals else 1)) % oprop_len
+    fig._input_offset = (
+        input_offset + (ninputs if overlay_signals else 1)) % iprop_len
+
     # Stop here if the user wants to control everything
     if not relabel:
         return out
diff --git a/control/timeresp.py b/control/timeresp.py
diff --git a/doc/conventions.rst b/doc/conventions.rst
diff --git a/doc/plotting.rst b/doc/plotting.rst
