def load_output_data_by_draw( draw_id: int, scenario: str, data_interface: ForecastDataInterface ) -> Tuple[pd.Series, pd.Series, pd.Series]: draw_df = data_interface.load_raw_outputs(scenario, draw_id) draw_df = draw_df.set_index(['location_id', 'date']).sort_index() deaths = draw_df.reset_index().set_index( ['location_id', 'date', 'observed'])['deaths'].rename(draw_id) infections = draw_df['infections'].rename(draw_id) r_effective = draw_df['r_effective'].rename(draw_id) return deaths, infections, r_effective
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)