def load_scaling_parameters_by_draw( draw_id: int, scenario: str, data_interface: ForecastDataInterface) -> pd.Series: scaling_parameters = data_interface.load_beta_scales(scenario, draw_id) scaling_parameters = scaling_parameters.set_index( 'location_id').stack().reset_index() scaling_parameters.columns = ['location_id', 'scaling_parameter', draw_id] scaling_parameters = scaling_parameters.set_index( ['location_id', 'scaling_parameter'])[draw_id] return scaling_parameters
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)