Predicting out-of-sample for autoregressive models

I believe I have a basic working version of this. Assumes your trace has a rho and scale. This also allows for multiple observations for the same date, which is often a case I’m working with. But if you want to just have one observation per date you could just remove the date_idx parts.

 def predict_outofsample(trace, date_idx):
        """
        trace: a pymc3 MultiTrace object
        date_idx: np.ndarray with shape (N_obs), indicating for each observation what date it corresponds to
            (so you can have multiple observations on the same day that will have the same prediction)
        """
        samples = []
        horizon = np.max(date_idx)
        for point in enumerate(trace.points()):
            rho, scale = point['rho'], point['scale']
            thetas = [np.random.normal(loc=0, scale=scale)]
            for i in range(horizon):
                thetas.append(rho*thetas[-1] + np.random.normal(loc=0, scale=scale))
            samples.append(thetas)
        return np.array(samples)[:, date_idx]