def test_forecast_features_num_threads(self): output_time_points = 5 forecasts = feature_utils.forecast_features( self.covariates, self.feature_specs, num_forecast_steps=output_time_points, num_threads=1, ) forecasts_parallel = feature_utils.forecast_features( self.covariates, self.feature_specs, num_forecast_steps=output_time_points, num_threads=2, ) _, n_locations, n_features = self.covariates.shape self.assertTupleEqual(forecasts.shape, (output_time_points, n_locations, n_features)) np.testing.assert_allclose( forecasts, self.expected_output[:output_time_points, :, :]) np.testing.assert_allclose(forecasts, forecasts_parallel)
def test_forecast_features_forecast_shorter(self): output_time_points = 3 _, n_locations, n_features = self.covariates.shape expected_output = np.zeros( (output_time_points, n_locations, n_features)) expected_output[:, :, 1] = 27 expected_output[:, :, 2] = np.arange(21, 24).reshape(-1, 1) expected_output[:, :, 3] = 1.0 forecasts = feature_utils.forecast_features( self.covariates, self.feature_specs, num_forecast_steps=output_time_points, ) _, n_locations, n_features = self.covariates.shape self.assertTupleEqual(forecasts.shape, (output_time_points, n_locations, n_features)) np.testing.assert_allclose( forecasts, self.expected_output[:output_time_points, :, :])