tfci_tfp_example.py

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import tensorflow_probability as tfp


x = np.random.rand(30)
y = 2.2 * x + np.random.rand(30)
data = pd.DataFrame({'X': x, 'y': y}, dtype=np.float32)
obs_data = data['y'].iloc[:20]

level = tfp.sts.LocalLevel(observed_time_series=obs_data)
linear = tfp.sts.LinearRegression(design_matrix=data['X'][..., np.newaxis])
model = tfp.sts.Sum([level, linear], observed_time_series=obs_data)

samples, _ = tfp.sts.fit_with_hmc(model, obs_data)

dist = tfp.sts.forecast(model, obs_data, samples, 10)
mean, std = dist.mean(), dist.stddev()

fig = plt.figure(figsize=(12, 10))
ax = plt.subplot(1, 1 ,1)
ax.plot(np.arange(30), data['X'], label='X', lw=1)
ax.plot(np.arange(30), data['y'], label='y', lw=2, color='orangered')
ax.scatter(np.arange(30), data['y'], lw=2, color='red')
ax.plot(np.arange(20, 30), mean, color='k', ls='--', label='forecast mean')
ax.fill_between(np.arange(20, 30), np.squeeze(mean - 1.96 * std),
                np.squeeze(mean + 1.96 * std), color='orange', alpha=0.3,
                label='95% CI')
plt.legend(loc='best')