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)