Ejemplo n.º 1
0
    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)
Ejemplo n.º 2
0
 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
Ejemplo n.º 3
0
 def test_no_overwriting(self, instance, parameters):
     self.assert_no_accidental_overwrites(instance,
                                          parameters,
                                          key=Keys.parameter(4))
Ejemplo n.º 4
0
 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))
Ejemplo n.º 5
0
 def test_date_marshall(self, instance, dates):
     self.assert_load_dump_workflow_correct(instance,
                                            dates,
                                            key=Keys.date(4))
Ejemplo n.º 6
0
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))
Ejemplo n.º 7
0
 def save_regression_betas(self, df: pd.DataFrame, draw_id: int) -> None:
     self.regression_marshall.dump(df, key=MKeys.regression_beta(draw_id))
Ejemplo n.º 8
0
 def save_date_file(self, df: pd.DataFrame, draw_id: int) -> None:
     self.regression_marshall.dump(df, key=MKeys.date(draw_id))
Ejemplo n.º 9
0
 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)
Ejemplo n.º 10
0
 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()
Ejemplo n.º 11
0
 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
Ejemplo n.º 12
0
 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)
Ejemplo n.º 13
0
 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
Ejemplo n.º 14
0
 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
Ejemplo n.º 15
0
 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))
Ejemplo n.º 16
0
 def save_beta_param_file(self, df: pd.DataFrame, draw_id: int) -> None:
     self.regression_marshall.dump(df, key=MKeys.parameter(draw_id))
Ejemplo n.º 17
0
 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'])
Ejemplo n.º 18
0
 def save_regression_coefficients(self, coefficients: pd.DataFrame,
                                  draw_id: int) -> None:
     self.regression_marshall.dump(coefficients,
                                   key=MKeys.coefficient(draw_id))
Ejemplo n.º 19
0
 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))
Ejemplo n.º 20
0
 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))
Ejemplo n.º 21
0
 def load_beta_scales(self, scenario: str, draw_id: int):
     return self.forecast_marshall.load(
         MKeys.beta_scales(scenario=scenario, draw_id=draw_id))
Ejemplo n.º 22
0
 def test_parameters_marshall(self, instance, parameters):
     self.assert_load_dump_workflow_correct(instance,
                                            parameters,
                                            key=Keys.parameter(4))
Ejemplo n.º 23
0
 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)
Ejemplo n.º 24
0
 def test_coefficients_marshall(self, instance, coefficients):
     self.assert_load_dump_workflow_correct(instance,
                                            coefficients,
                                            key=Keys.coefficient(4))
Ejemplo n.º 25
0
 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))
Ejemplo n.º 26
0
 def test_components_marshall(self, instance, components):
     self.assert_load_dump_workflow_correct(
         instance,
         components,
         key=Keys.components(scenario='happy', draw_id=4))
Ejemplo n.º 27
0
 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))
Ejemplo n.º 28
0
 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))
Ejemplo n.º 29
0
 def load_regression_coefficients(self, draw_id: int) -> pd.DataFrame:
     return self.regression_marshall.load(MKeys.coefficient(draw_id))