Temporal Fusion Transformer: how to join the prediction with the original data

Question

I'm using Temporal Fusion Transformer to do sales forecast. For validation, how to join the prediction with the original data? I guess I need both the 'time_idx' and the group_ids to join the prediction with the original data. I have get the 'time_idx' from x, how to get the group_ids=['STORE_NUMBER'] information of the predictions.

Thanks

    return TimeSeriesDataSet(
        data,
        time_idx="time_idx",
        target="SALES",
        group_ids=['STORE_NUMBER'],
        min_encoder_length=1,  # keep encoder length long (as it is in the validation set)
        max_encoder_length=max_encoder_length,
        min_prediction_length=1,
        max_prediction_length=max_prediction_length,
        static_categoricals=['CITY'],  
        static_reals=['STORE_AGE'],        
time_varying_known_categoricals=['day','week','holiday'], 
        time_varying_known_reals=["time_idx"],
        time_varying_unknown_categoricals=[],
        time_varying_unknown_reals=["SALES"],
        add_relative_time_idx=True,
        add_target_scales=True,
        add_encoder_length=True,
        allow_missing_timesteps=True,
        },
    )

print("Keys in x:", x.keys())
print("Keys in raw_predictions:", raw_predictions.keys())

actuals = x['decoder_target']
times = x['decoder_time_idx']```

Answer 1

Use return_index in your predict statement:

predictions = best_tft.predict(data, mode="raw", return_index=True, return_x=True)

predictions.index will contain the group ids