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
