import control as ct
import numpy as np

model = ct.rss(states=['x0', 'x1'], outputs=['y0', 'y1'], inputs=['u0', 'u1'], name='My Model')

T = np.linspace(0, 10, 100, endpoint=False)
X0 = np.zeros(model.nstates)

res = ct.step_response(model, T=T, X0=X0, input=0)

df = res.to_pandas()

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[140], line 6
      3 T = np.linspace(0, 10, 100, endpoint=False)
      4 X0 = np.zeros(model.nstates)
----> 6 res = ct.step_response(model, T=T, X0=X0, input=0).to_pandas()

File ~.env/lib/python3.10/site-packages/control/timeresp.py:723, in TimeResponseData.to_pandas(self)
    719 if self.nstates > 0:
    720     data.update(
    721         {name: self.x[i] for i, name in enumerate(self.state_labels)})
--> 723 return pandas.DataFrame(data)

File ~.env/lib/python3.10/site-packages/pandas/core/frame.py:778, in DataFrame.__init__(self, data, index, columns, dtype, copy)
    772     mgr = self._init_mgr(
    773         data, axes={"index": index, "columns": columns}, dtype=dtype, copy=copy
    774     )
    776 elif isinstance(data, dict):
    777     # GH#38939 de facto copy defaults to False only in non-dict cases
--> 778     mgr = dict_to_mgr(data, index, columns, dtype=dtype, copy=copy, typ=manager)
    779 elif isinstance(data, ma.MaskedArray):
    780     from numpy.ma import mrecords

File ~.env/lib/python3.10/site-packages/pandas/core/internals/construction.py:503, in dict_to_mgr(data, index, columns, dtype, typ, copy)
    499     else:
    500         # dtype check to exclude e.g. range objects, scalars
    501         arrays = [x.copy() if hasattr(x, "dtype") else x for x in arrays]
--> 503 return arrays_to_mgr(arrays, columns, index, dtype=dtype, typ=typ, consolidate=copy)

File ~.env/lib/python3.10/site-packages/pandas/core/internals/construction.py:114, in arrays_to_mgr(arrays, columns, index, dtype, verify_integrity, typ, consolidate)
    111 if verify_integrity:
    112     # figure out the index, if necessary
    113     if index is None:
--> 114         index = _extract_index(arrays)
    115     else:
    116         index = ensure_index(index)

File ~.env/lib/python3.10/site-packages/pandas/core/internals/construction.py:664, in _extract_index(data)
    662         raw_lengths.append(len(val))
    663     elif isinstance(val, np.ndarray) and val.ndim > 1:
--> 664         raise ValueError("Per-column arrays must each be 1-dimensional")
    666 if not indexes and not raw_lengths:
    667     raise ValueError("If using all scalar values, you must pass an index")

ValueError: Per-column arrays must each be 1-dimensional

import matplotlib.pyplot as plt
import control as ct
import numpy as np

def step_response_to_pandas(step_response):
    return pd.DataFrame(
        {'trace_label': np.array([[label] * (len(res.time)) for label in res.trace_labels]).ravel()} |
        {'time': res.time.repeat(len(res.trace_labels))} |
        {label: res.inputs[i].ravel() for i,label in enumerate(res.input_labels)} |
        {label: res.outputs[i].ravel() for i,label in enumerate(res.output_labels)} |
        {label: res.states[i].ravel() for i,label in enumerate(res.state_labels)}
    )

def plot_step_response_dataframe(df):
    grouped = df.groupby(level='trace_label')
    row_size = 1

    for trace_label, group in grouped:
        fig, axes = plt.subplots(len(group.columns), 1, figsize=(6.4, len(group.columns) * row_size), sharex=True)
        fig.suptitle(f'Trace: {trace_label}', fontsize=16)
        
        if len(group.columns) == 1:
            axes = [axes]
        
        for ax, (signal_name, signal_data) in zip(axes, group.items()):
            ax.plot(group.index.get_level_values('time'), signal_data, label=signal_name)
            ax.grid(True)
            ax.set_ylabel(signal_name)
        
        axes[-1].set_xlabel('Time')
        
        plt.tight_layout()
        plt.show()


model = ct.rss(states=['x0', 'x1'], outputs=['y0', 'y1'], inputs=['u0', 'u1'], name='My Model')

T = np.linspace(0, 10, 100, endpoint=False)
X0 = np.zeros(model.nstates)

res = ct.step_response(model, T=T, X0=X0)

df = step_response_to_pandas(res)
df = df.set_index(['trace_label', 'time'])

plot_step_response_dataframe(df)

display(df)