Ejemplo n.º 1
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.º 2
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.º 3
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.º 4
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))