fixup: Add missing labels when returning TimeResponseData

joaoantoniocardoso · web-flow · commit 7eeaa6d980aa · 2023-05-20T15:10:39.000-03:00
diff --git a/control/timeresp.py b/control/timeresp.py
@@ -1376,9 +1376,13 @@ def step_response(sys, T=None, X0=0., input=None, output=None, T_num=None,
     # Figure out if the system is SISO or not
     issiso = sys.issiso() or (input is not None and output is not None)
 
+    # Select only the given input and output, if any
+    input_index = {k: v for k, v in sys.input_index.items() if v is input} if input else sys.input_index
+    output_index = {k: v for k, v in sys.output_index.items() if v is output} if output else sys.output_index
+
     return TimeResponseData(
         response.time, yout, xout, uout, issiso=issiso,
-        output_labels=sys.output_index, input_labels=sys.input_index,
+        output_labels=output_index, input_labels=input_index,
         state_labels=sys.state_index,
         transpose=transpose, return_x=return_x, squeeze=squeeze)
 
@@ -1708,10 +1712,14 @@ def initial_response(sys, T=None, X0=0., input=0, output=None, T_num=None,
     # Figure out if the system is SISO or not
     issiso = sys.issiso() or (input is not None and output is not None)
 
+    # Select only the given input and output, if any
+    input_index = {k: v for k, v in sys.input_index.items() if v is input} if input else sys.input_index
+    output_index = {k: v for k, v in sys.output_index.items() if v is output} if output else sys.output_index
+
     # Store the response without an input
     return TimeResponseData(
         response.t, response.y, response.x, None, issiso=issiso,
-        output_labels=sys.output_index, input_labels=sys.input_index,
+        output_labels=output_index, input_labels=input_index,
         state_labels=sys.state_index,
         transpose=transpose, return_x=return_x, squeeze=squeeze)
 
@@ -1875,9 +1883,13 @@ def impulse_response(sys, T=None, X0=0., input=None, output=None, T_num=None,
     # Figure out if the system is SISO or not
     issiso = sys.issiso() or (input is not None and output is not None)
 
+    # Select only the given input and output, if any
+    input_index = {k: v for k, v in sys.input_index.items() if v is input} if input else sys.input_index
+    output_index = {k: v for k, v in sys.output_index.items() if v is output} if output else sys.output_index
+
     return TimeResponseData(
         response.time, yout, xout, uout, issiso=issiso,
-        output_labels=sys.output_index, input_labels=sys.input_index,
+        output_labels=output_index, input_labels=input_index,
         state_labels=sys.state_index,
         transpose=transpose, return_x=return_x, squeeze=squeeze)
 
