Add missing labels when returning TimeResponseData

joaoantoniocardoso · joaoantoniocardoso · commit 76748132c009 · 2023-05-20T16:36:59.000-03:00
diff --git a/control/timeresp.py b/control/timeresp.py
@@ -1111,6 +1111,8 @@ def forced_response(sys, T=None, U=0., X0=0., transpose=False,
 
     return TimeResponseData(
         tout, yout, xout, U, issiso=sys.issiso(),
+        output_labels=sys.output_index, input_labels=sys.input_index,
+        state_labels=sys.state_index,
         transpose=transpose, return_x=return_x, squeeze=squeeze)
 
 
@@ -1374,8 +1376,14 @@ 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 = sys.input_index if input is None else {k: 0 for k, v in sys.input_index.items() if v == input}
+    output_index = sys.output_index if output is None else {k: 0 for k, v in sys.output_index.items() if v == output}
+
     return TimeResponseData(
         response.time, yout, xout, uout, issiso=issiso,
+        output_labels=output_index, input_labels=input_index,
+        state_labels=sys.state_index,
         transpose=transpose, return_x=return_x, squeeze=squeeze)
 
 
@@ -1704,9 +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 output, if any
+    output_index = sys.output_index if output is None else {k: 0 for k, v in sys.output_index.items() if v == output}
+
     # Store the response without an input
     return TimeResponseData(
         response.t, response.y, response.x, None, issiso=issiso,
+        output_labels=output_index, input_labels=None,
+        state_labels=sys.state_index,
         transpose=transpose, return_x=return_x, squeeze=squeeze)
 
 
@@ -1869,8 +1882,14 @@ 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 = sys.input_index if input is None else {k: 0 for k, v in sys.input_index.items() if v == input}
+    output_index = sys.output_index if output is None else {k: 0 for k, v in sys.output_index.items() if v == output}
+
     return TimeResponseData(
         response.time, yout, xout, uout, issiso=issiso,
+        output_labels=output_index, input_labels=input_index,
+        state_labels=sys.state_index,
         transpose=transpose, return_x=return_x, squeeze=squeeze)
 
 
