def create_data_viz()

in src/deep_demand_forecast/monitor.py [0:0]

• 20 lines of code
• 1 McCabe index (conditional complexity)

def create_data_viz(train_df, test_df, context_length, prediction_length, num_sample_ts):
    num_sample_ts = min(num_sample_ts, train_df.shape[1])
    ts_col_names = list(train_df.columns)
    selected_cols = ts_col_names[1:num_sample_ts]

    selected_train_df = train_df.loc[:, ["time"] + selected_cols]
    train_df_melt = pd.melt(
        selected_train_df.tail(context_length),
        id_vars=["time"],
        value_vars=selected_cols,
    )
    train_df_melt.rename(columns={"variable": "covariate"}, inplace=True)
    selected_test_df = test_df.loc[:, ["time"] + selected_cols]
    num_train = selected_train_df.shape[0] - 1
    test_df_melt = pd.melt(
        selected_test_df.iloc[num_train : num_train + prediction_length],
        id_vars=["time"],
        value_vars=selected_cols,
    )
    test_df_melt.rename(columns={"variable": "covariate"}, inplace=True)
    return train_df_melt, test_df_melt, selected_cols