statsmodels_armina_prediction.py

from collections import OrderedDict
import pandas as pd
from dateutil import parser


def forecast_ts_model(res, start, end):
    """
    Take a statsmodels.tsa.arima_model.ARIMAResults object
    and predict (forecast) the timeseries values
    between the "start" and "end" times, which
    must be in the future relative to the training data
    for the timeseries.
    """
    
    if type(start) == str:
        start = parser.parse(start)
        
    if type(end) == str:
        end = parser.parse(end)
    
    training_end = res.data.dates.max()
    time_delta = res.data.dates.freq.delta
    
    if end <= start:
        raise Exception("End: {} must be after start: {}", end, start)
    
    if start < training_end:
        raise Exception("Start: {} must be after the end of the training data: {}".format(start, training_end))
    
    idx = 0
    begin_idx = None
    begin_dt = None
    end_idx = None
    end_dt = None
    
    current_dt = training_end
    
    while True:
        
        idx += 1
        current_dt = current_dt + time_delta
        
        if begin_idx is None and current_dt > start:
            begin_idx = idx
            begin_dt = current_dt
            
        if end_idx is None and current_dt + time_delta > end:
            end_idx = idx
            end_dt = current_dt
            break
    
    index = pd.date_range(begin_dt, end_dt, freq=time_delta)
        
    vals, std_err, conf_int =  res.forecast(end_idx)
    
    conf_int_low = [x[0] for x in conf_int]
    conf_int_high = [x[1] for x in conf_int]

    res = OrderedDict()
    res['forecast'] = vals[begin_idx-1:]
    res['stderr'] = std_err[begin_idx-1:]
    res['conf_int_low'] = conf_int_low[begin_idx-1:]
    res['conf_int_high'] = conf_int_high[begin_idx-1:]
    
    return pd.DataFrame(res, index=index)