def test_forecast_io(self, tmpdir, components, beta_scales, forecast_outputs): forecast_paths = ForecastPaths( root_dir=Path(tmpdir), scenarios=['happy'], ) di = ForecastDataInterface( forecast_paths=None, regression_paths=None, covariate_paths=None, regression_marshall=None, forecast_marshall=CSVMarshall.from_paths(forecast_paths), ) # Step 1: save files di.save_components(components, scenario="happy", draw_id=4, strict=True) di.save_beta_scales(beta_scales, scenario="happy", draw_id=4) di.save_raw_outputs(forecast_outputs, scenario="happy", draw_id=4, strict=True) # Step 2: test save location # this is sort of cheating, but it ensures that scenario things are # nicely nested as they should be assert (Path(tmpdir) / "happy" / "component_draws" / "draw_4.csv").exists() assert (Path(tmpdir) / "happy" / "beta_scaling" / "draw_4.csv").exists() assert (Path(tmpdir) / "happy" / "raw_outputs" / "draw_4.csv").exists() # Step 3: load those files loaded_components = di.load_components(scenario="happy", draw_id=4) # Load components now does some formatting, which broke the tests. # Back out these changes here. loaded_components = loaded_components.reset_index() loaded_components['date'] = loaded_components['date'].astype(str) loaded_components = loaded_components[ components.columns] # Use the same sort order. loaded_beta_scales = di.load_beta_scales(scenario="happy", draw_id=4) loaded_forecast_outputs = di.load_raw_outputs(scenario="happy", draw_id=4) # Step 4: test files pandas.testing.assert_frame_equal(components, loaded_components) pandas.testing.assert_frame_equal(beta_scales, loaded_beta_scales) pandas.testing.assert_frame_equal(forecast_outputs, loaded_forecast_outputs)
def load_betas_by_draw(draw_id: int, scenario: str, data_interface: ForecastDataInterface) -> pd.Series: components = data_interface.load_components(scenario, draw_id) draw_betas = (components.sort_index()['beta'].rename(draw_id)) return draw_betas