How to assign prediction to timeseries in NBeats?

[DonBeo]

Hi, thanks for the great library. I have an usability question.

I have trained an NBeats model on a dataframe containing 16 time series as in the example (https://pytorch-forecasting.readthedocs.io/en/latest/tutorials/ar.html)

Now I am predicting new data from a dataframe as:

pr, x_pr, idx = best_model.predict(data_ts, mode="prediction", return_x=True, return_index=True)

as expected I get prediction for 16 time series and 192 points (this was my prediction length)

pr.shape #= torch.Size([16, 192])

Now how do I assign each row in the prediction to the correct time series? In other words I would like to know to which of my 16 time series each row refers to. In my data the time series is identified by the series column as in the attached figure

[jdb78]

You can use idx here to match them to the timeseries.

[DonBeo]

thanks for answering. So each row of the idx dataframe returned by the function corresponds to a row of pr

I really suggest to add it to the doc as it was not very clear.

Thanks for the work you are doing!!

[jdb78]

Adding in the latest PR.