Create utility for putting 3D objects into Pandas DataFrames

[jdebacker]

For those looking to use the model output in other software (e.g., for plotting), it can be helpful to be able to put three dimensional objects into a DataFrame as panel dataset.

We could add a utility to utils.py to simplify this. The code could be something like

def D3_to_df(tpi_vars, var="c_path", start_year=2025):

    data = tpi_vars[var]
    # First, get the dimensions of your array
    T, S, J = data.shape

    # Create index values for each dimension
    idx_t = [f'{start_year + i}' for i in range(T)]  # You can customize these labels
    idx_s = [f'{i}' for i in range(S)]
    idx_j = [f'{i}' for i in range(J)]

    # Create MultiIndex from the first two dimensions
    multi_idx = pd.MultiIndex.from_product([idx_t, idx_s], names=['Year', 'Age'])

    # Reshape the 3D array to 2D
    reshaped_data = data.reshape(T * S, J)

    # Create the DataFrame with MultiIndex
    df = pd.DataFrame(reshaped_data, index=multi_idx, columns=idx_j)

    return df

Then use this function with:

df = D3_to_df(tpi_vars_bas, var="c_path", start_year=p_bas.start_year)
df.to_csv("c_path.csv")