def setUpClass(self):
np.random.seed(8675309)
np.random.seed(0xDEADBEEF)
fake_data = np.random.randn(100, 10)
fake_coefficients = np.array([5, -4, 3] + [0] * 7)
random_walk = np.cumsum(np.random.randn(100))
noise = np.random.randn(100) * 1.5
y = random_walk + fake_data @ fake_coefficients + noise
data = pd.DataFrame(
fake_data,
columns=["V" + str(x + 1) for x in range(10)],
)
data["y"] = y
model = Bsts()
model.add_state(LocalLevelStateModel(y))
model.train(formula=f"y ~ {dot(data, omit='y')}", niter=100, data=data)
self._regression_model = model
data = np.log(AirPassengers)
model = Bsts()
model.add_state(LocalLinearTrendStateModel(data))
model.add_state(SeasonalStateModel(data, nseasons=12))
model.train(data, niter=100)
self._model = model
model = Bsts()
model.add_state(LocalLinearTrendStateModel(data))
model.train(data, niter=100)
self._model2 = model
class TestSeasonalStateModel(unittest.TestCase):
def setUp(self):
np.random.seed(8675309)
self.data = np.random.randn(100)
self.model = SeasonalStateModel(self.data, nseasons=4)
self.model.set_state_index(1)
def tearDown(self):
if os.path.exists("seas.pkl") and os.path.isfile("seas.pkl"):
os.remove("seas.pkl")
def test_state_dimension(self):
self.assertEqual(3, self.model.state_dimension)
def test_storage(self):
time_dimension = 20
niter = 10
self.model.allocate_space(niter, time_dimension)
self.assertEqual(self.model.sigma_draws.shape, (niter, ))
self.assertEqual(self.model._state_contribution.shape,
(niter, time_dimension))
self.model.set_state_index(2)
state_matrix = np.random.randn(3, time_dimension)
iteration = 4
self.model.sigma_draws[4] = 1.9
self.model.restore_state(iteration)
self.assertAlmostEqual(1.9, self.model._state_model.sigma)
iteration = 8
self.model.record_state(iteration, state_matrix)
self.assertAlmostEqual(self.model.sigma_draws[iteration], 1.9)
np.testing.assert_array_equal(
self.model._state_contribution[iteration, :],
state_matrix[self.model._state_index, :])
with open("seas.pkl", "wb") as pkl:
pickle.dump(self.model, pkl)
with open("seas.pkl", "rb") as pkl:
_ = pickle.load(pkl)
def test_auto_ar(self):
model = Bsts()
model.add_state(LocalLinearTrendStateModel(self._y))
model.add_state(SeasonalStateModel(self._y, nseasons=12))
model.add_state(AutoArStateModel(self._y, lags=1))
model.train(data=self._y, niter=1000)
fname = "ar.pkl"
with open(fname, "wb") as pkl:
pickle.dump(model, pkl)
with open(fname, "rb") as pkl:
m2 = pickle.load(pkl)
self.assertEqual(model.time_dimension, m2.time_dimension)
self.assertIsInstance(m2, Bsts)
def test_basic_structural_model(self):
model = Bsts(family="student")
y = np.log(AirPassengers)
model.add_state(LocalLinearTrendStateModel(y))
model.add_state(SeasonalStateModel(y, nseasons=12))
model.train(data=y, niter=1000)
fig = model.plot("comp")
if _show_figs:
fig.show()
horizon = 24
predictions = model.predict(horizon)
target_quantiles = [0.025] + list(np.linspace(.05, 0.95, 19)) + [0.975]
prediction_quantiles = np.quantile(predictions.distribution,
target_quantiles,
axis=0)
self.assertIsInstance(prediction_quantiles, np.ndarray)
with open("bsm.pkl", "wb") as pkl:
pickle.dump(model, pkl)
with open("bsm.pkl", "rb") as pkl:
m2 = pickle.load(pkl)
pred2 = m2.predict(horizon)
fig, ax = plt.subplots(1, 2, figsize=(10, 5))
predictions.plot(ax=ax[0], original_series=60)
pred_plot_fig, pred_plot_ax = pred2.plot(ax=ax[1], original_series=60)
if _show_figs:
fig.show()
self.assertIsInstance(pred_plot_ax, plt.Axes)
np.testing.assert_array_almost_equal(predictions.distribution,
pred2.distribution)
def setUp(self):
np.random.seed(8675309)
self.data = np.random.randn(100)
self.model = SeasonalStateModel(self.data, nseasons=4)
self.model.set_state_index(1)