Create For Loop To Predict Next Value Over Group (Python)

user_2813606 0
Viewed 953

I am working on a project where I need to take groups of data and predict the next value for that group using a time series model. In my data, I have a grouping variable and a numeric variable.

Here is an example of my data:

import pandas as pd
data = [
    ["A", 10],
    ["B", 10],
    ["C", 15],
    ["D", 12],
    ["A", 18],
    ["B", 19],
    ["C", 14],
    ["D", 22],
    ["A", 20],
    ["B", 25],
    ["C", 12],
    ["D", 30],
    ["A", 36],
    ["B", 27],
    ["C", 10],
    ["D", 45]
]


data = pd.DataFrame(
    data,
    columns=[
        "group",
        "value"
    ],
)

What I want to do is to create a for loop that iterates over the groups and predicts the next value for A, B, C, and D. Essentially, my end result would be a new data frame with 4 rows, one for each new predicted value. It would look something like this:

group pred_value
A     40
B     36
C     8
D     42

Here is my attempt at that so far:

from statsmodels.tsa.ar_model import AutoReg
final=pd.DataFrame()
for i in data['group']:
    group = data[data['group']==i]
    model = AutoReg(group['value'], lags=1)
    model_fit = model.fit()
    yhat = model_fit.predict(len(group), len(group))
    final = final.append(yhat,ignore_index=True)

Unfortunately, this produces a data frame with 15 rows and I'm not sure how to get the end result that I described above.

Can anyone help point me in the right direction? Any help would be appreciated! Thank you!

pandas time-series for-loop python
1 Answers

You can groupby first and then iterate. We can store the results in a dict and after the loop convert it to a DataFrame:

# will hold the predictions
forecasts = {}

# in each turn e.g., group == "A", values are [10, 18, 20, 36]
for group, values in data.groupby("group").value:
    # form the model and fit
    model = AutoReg(values, lags=1)
    result = model.fit()

    # predict
    prediction = result.forecast(steps=1)
    
    # store
    forecasts[group] = prediction

# after `for` ends, convert to DataFrame
all_predictions = pd.DataFrame(forecasts)

to get

>>> all_predictions

           A          B         C          D
4  51.809524  28.561404  7.285714  62.110656

We can also do this all with apply:

>>> data.groupby("group").value.apply(lambda x: AutoReg(x, lags=1).fit().forecast(1))

group
A      4    51.809524
B      4    28.561404
C      4     7.285714
D      4    62.110656
Name: value, dtype: float64

However, we potentially lose the ability to hold references to the fitted models, whereas in explicit for, we could keep them aside. But if that is not wanted anyway, this can be used.

user_93321870
Related