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)
