def test_datetime(self, instance, regression_beta): """ Re-use regression_Beta fixture but cast date as a datetime. """ regression_beta['date'] = pandas.to_datetime(regression_beta['date']) key = Keys.regression_beta(4) instance.dump(regression_beta, key=key) loaded = instance.load(key) pandas.testing.assert_frame_equal(regression_beta, loaded)
def load_transition_date(self, draw_id: int) -> pd.Series: dates_df = self.regression_marshall.load(key=MKeys.date(draw_id)) dates_df['end_date'] = pd.to_datetime(dates_df['end_date']) transition_date = dates_df.set_index( 'location_id').sort_index()['end_date'].rename('date') return transition_date
def test_no_overwriting(self, instance, parameters): self.assert_no_accidental_overwrites(instance, parameters, key=Keys.parameter(4))
def test_beta_scales_marshall(self, instance, beta_scales): self.assert_load_dump_workflow_correct( instance, beta_scales, key=Keys.beta_scales(scenario='happy', draw_id=4))
def test_date_marshall(self, instance, dates): self.assert_load_dump_workflow_correct(instance, dates, key=Keys.date(4))
def test_CSVMarshall_uses_same_paths_as_paths_module(tmpdir): """ Runs the gamut of paths produced by paths.py and compares to CSVMarshall path computation. """ from covid_model_seiir_pipeline.paths import ( RegressionPaths, ForecastPaths, ) root = pathlib.Path(tmpdir) m = CSVMarshall(root) r = RegressionPaths(root) f = ForecastPaths(root, scenarios=["happy"]) dargs = {'draw_id': 4} sdargs = {'scenario': "happy", 'draw_id': 4} # forecasting/data.py # load_regression_coefficients assert r.get_coefficient_file(**dargs) == m.resolve_key( Keys.coefficient(**dargs)) # load_dates_df assert r.get_date_file(**dargs) == m.resolve_key(Keys.date(**dargs)) # load_beta_regression assert r.get_beta_regression_file(**dargs) == m.resolve_key( Keys.regression_beta(**dargs)) # load_infection_data assert r.get_data_file(**dargs) == m.resolve_key( Keys.location_data(**dargs)) # load_beta_params assert r.get_beta_param_file(**dargs) == m.resolve_key( Keys.parameter(**dargs)) # save_components assert f.get_components_path(**sdargs) == m.resolve_key( Keys.components(**sdargs)) # save_beta_scales assert f.get_beta_scaling_path(**sdargs) == m.resolve_key( Keys.beta_scales(**sdargs)) # save_outputs assert f.get_outputs_path(**sdargs) == m.resolve_key( Keys.forecast_raw_outputs(**sdargs)) # regression/data.py # save_beta_param_file assert r.get_beta_param_file(**dargs) == m.resolve_key( Keys.parameter(**dargs)) # save_date_file assert r.get_date_file(**dargs) == m.resolve_key(Keys.date(**dargs)) # save_regression_coefficients assert r.get_coefficient_file(**dargs) == m.resolve_key( Keys.coefficient(**dargs)) # save_regression_betas assert r.get_beta_regression_file(**dargs) == m.resolve_key( Keys.regression_beta(**dargs)) # save_location_data assert r.get_data_file(**dargs) == m.resolve_key( Keys.location_data(**dargs))
def save_regression_betas(self, df: pd.DataFrame, draw_id: int) -> None: self.regression_marshall.dump(df, key=MKeys.regression_beta(draw_id))
def save_date_file(self, df: pd.DataFrame, draw_id: int) -> None: self.regression_marshall.dump(df, key=MKeys.date(draw_id))
def save_components(self, forecasts: pd.DataFrame, scenario: str, draw_id: int, strict: bool): self.forecast_marshall.dump(forecasts, key=MKeys.components(scenario=scenario, draw_id=draw_id), strict=strict)
def load_beta_params(self, draw_id: int) -> Dict[str, float]: df = self.regression_marshall.load(key=MKeys.parameter( draw_id=draw_id)) return df.set_index('params')['values'].to_dict()
def load_raw_covariates(self, scenario: str, draw_id: int): covariates = self.forecast_marshall.load( key=MKeys.forecast_raw_covariates(scenario=scenario, draw_id=draw_id)) covariates['date'] = pd.to_datetime(covariates['date']) return covariates
def save_raw_covariates(self, covariates: pd.DataFrame, scenario: str, draw_id: int, strict: bool): self.forecast_marshall.dump(covariates, key=MKeys.forecast_raw_covariates( scenario=scenario, draw_id=draw_id), strict=strict)
def load_infection_data(self, draw_id: int) -> pd.DataFrame: infection_data = self.regression_marshall.load( key=MKeys.location_data(draw_id)) infection_data['date'] = pd.to_datetime(infection_data['date']) return infection_data
def load_beta_regression(self, draw_id: int) -> pd.DataFrame: beta_regression = self.regression_marshall.load( key=MKeys.regression_beta(draw_id)) beta_regression['date'] = pd.to_datetime(beta_regression['date']) return beta_regression
def load_reimposition_dates(self, scenario: str, reimposition_number: int): return self.forecast_marshall.load(key=MKeys.reimposition_dates( scenario=scenario, reimposition_number=reimposition_number))
def save_beta_param_file(self, df: pd.DataFrame, draw_id: int) -> None: self.regression_marshall.dump(df, key=MKeys.parameter(draw_id))
def load_components(self, scenario: str, draw_id: int) -> pd.DataFrame: components = self.forecast_marshall.load( key=MKeys.components(scenario=scenario, draw_id=draw_id)) components['date'] = pd.to_datetime(components['date']) return components.set_index(['location_id', 'date'])
def save_regression_coefficients(self, coefficients: pd.DataFrame, draw_id: int) -> None: self.regression_marshall.dump(coefficients, key=MKeys.coefficient(draw_id))
def save_beta_scales(self, scales: pd.DataFrame, scenario: str, draw_id: int): self.forecast_marshall.dump(scales, key=MKeys.beta_scales(scenario=scenario, draw_id=draw_id))
def save_location_data(self, df: pd.DataFrame, draw_id: int) -> None: # quasi-inverse of load_all_location_data, except types are different self.regression_marshall.dump(df, key=MKeys.location_data(draw_id))
def load_beta_scales(self, scenario: str, draw_id: int): return self.forecast_marshall.load( MKeys.beta_scales(scenario=scenario, draw_id=draw_id))
def test_parameters_marshall(self, instance, parameters): self.assert_load_dump_workflow_correct(instance, parameters, key=Keys.parameter(4))
def save_raw_outputs(self, raw_outputs: pd.DataFrame, scenario: str, draw_id: int, strict: bool): self.forecast_marshall.dump(raw_outputs, key=MKeys.forecast_raw_outputs( scenario=scenario, draw_id=draw_id), strict=strict)
def test_coefficients_marshall(self, instance, coefficients): self.assert_load_dump_workflow_correct(instance, coefficients, key=Keys.coefficient(4))
def load_raw_outputs(self, scenario: str, draw_id: int) -> pd.DataFrame: return self.forecast_marshall.load( key=MKeys.forecast_raw_outputs(scenario=scenario, draw_id=draw_id))
def test_components_marshall(self, instance, components): self.assert_load_dump_workflow_correct( instance, components, key=Keys.components(scenario='happy', draw_id=4))
def save_output_miscellaneous(self, output_miscellaneous: pd.DataFrame, scenario: str, measure: str): self.forecast_marshall.dump(output_miscellaneous, key=MKeys.forecast_output_miscellaneous( scenario=scenario, measure=measure))
def save_reimposition_dates(self, reimposition_dates: pd.DataFrame, scenario: str, reimposition_number: int): self.forecast_marshall.dump( reimposition_dates, key=MKeys.reimposition_dates( scenario=scenario, reimposition_number=reimposition_number))
def load_regression_coefficients(self, draw_id: int) -> pd.DataFrame: return self.regression_marshall.load(MKeys.coefficient(draw_id))