def test_scenario_with_model_change(self):
# Instance to save population values
population_data = PopulationData(filename=None)
# Set tau value and start date of records
example_data = ExampleData(tau=1440, start_date="01Jan2020")
# Preset of SIR-F parameters
preset_dict = SIRF.EXAMPLE["param_dict"]
# Create dataset from 01Jan2020 to 31Jan2020
area = {"country": "Theoretical"}
example_data.add(SIRF, step_n=30, **area)
# Register population value
population_data.update(SIRF.EXAMPLE["population"], **area)
# Create Scenario class
snl = Scenario(example_data, population_data, tau=1440, **area)
# Set 0th phase from 02Jan2020 to 31Jan2020 with preset parameter values
snl.clear(include_past=True)
snl.add(end_date="31Jan2020", model=SIRF, **preset_dict)
# Add main scenario
snl.add(end_date="31Dec2020", name="Main")
# Add lockdown scenario
rho_lock = snl.get("rho", phase="0th") / 2
snl.add(end_date="31Dec2020", name="Lockdown", rho=rho_lock)
# Add medicine scenario
kappa_med = snl.get("kappa", phase="0th") / 2
sigma_med = snl.get("sigma", phase="0th") * 2
snl.add(end_date="31Dec2020",
name="Medicine",
kappa=kappa_med,
sigma=sigma_med)
# Add vaccine scenario
snl.add(end_date="31Dec2020", name="Vaccine", model=SIRFV, omega=0.001)
# Summarize
snl.summary()
# Compare scenarios
snl.describe(y0_dict={"Vaccinated": 0})
def test_scenario(self):
area = {"country": "Theoretical"}
# Set-up example dataset (from 01Jan2020 to 31Jan2020)
example_data = ExampleData(tau=1440, start_date="01Jan2020")
example_data.add(SIRF, step_n=30, **area)
# Population value
population_data = PopulationData(filename=None)
population_data.update(SIRF.EXAMPLE["population"], **area)
# Set-up Scenario instance
snl = Scenario(tau=1440, **area)
snl.register(example_data, population_data)
# Check records
record_df = snl.records(variables="CFR")
assert set(record_df.columns) == set(
[Term.DATE, Term.C, Term.F, Term.R])
# Add a past phase to 31Jan2020 with parameter values
snl.add(model=SIRF, **SIRF.EXAMPLE["param_dict"])
# Check summary
df = snl.summary()
assert not df.empty
assert len(df) == 1
assert Term.RT in df
# Main scenario
snl.add(end_date="31Dec2020", name="Main")
assert snl.get(Term.RT, phase="last", name="Main") == 2.50
# Lockdown scenario
snl.clear(name="Lockdown")
rho_lock = snl.get("rho", phase="0th") * 0.5
snl.add(end_date="31Dec2020", name="Lockdown", rho=rho_lock)
assert snl.get(Term.RT, phase="last", name="Lockdown") == 1.25
# Medicine scenario
snl.clear(name="Medicine")
kappa_med = snl.get("kappa", phase="0th") * 0.5
sigma_med = snl.get("sigma", phase="0th") * 2
snl.add(end_date="31Dec2020",
name="Medicine",
kappa=kappa_med,
sigma=sigma_med)
assert snl.get(Term.RT, phase="last", name="Medicine") == 1.31
# Add vaccine scenario
snl.clear(name="Vaccine")
rho_vac = snl.get("rho", phase="0th") * 0.8
kappa_vac = snl.get("kappa", phase="0th") * 0.6
sigma_vac = snl.get("sigma", phase="0th") * 1.2
snl.add(end_date="31Dec2020",
name="Vaccine",
rho=rho_vac,
kappa=kappa_vac,
sigma=sigma_vac)
assert snl.get(Term.RT, phase="last", name="Vaccine") == 1.72
# Description
snl.describe()
# History
snl.history("Rt")
snl.history("rho")
snl.history("Infected")
snl.history_rate(name="Medicine")
snl.simulate(name="Vaccine")
def test_ode_two_phases(self, population_data):
# Setting
eg_tau = 1440
area = {"country": "Example", "province": "Example"}
# Simulation
example_data = ExampleData(tau=eg_tau)
example_data.add(SIRF, step_n=30, **area)
example_data.add(SIRD, step_n=30, **area)
dim_df = example_data.subset(**area)
assert isinstance(dim_df, pd.DataFrame)
assert set(dim_df.columns) == set(Term.NLOC_COLUMNS)
def test_r0(self, model):
eg_dict = model.EXAMPLE.copy()
# Calculate r0
model_ins = model(eg_dict["population"], **eg_dict["param_dict"])
eg_r0 = model_ins.calc_r0()
# Set-up example dataset (from 01Jan2020 to 31Jan2020)
area = {"country": "Theoretical"}
example_data = ExampleData(tau=1440, start_date="01Jan2020")
example_data.add(model, step_n=180, **area)
# Check the nature of r0
df = example_data.specialized(model, **area)
x_max = df.loc[df[Term.CI].idxmax(), Term.S] / eg_dict["population"]
assert round(x_max, 2) == round(1 / eg_r0, 2)
def test_ode(self, model):
# Setting
eg_tau = 1440
area = {"country": "Full", "province": model.NAME}
# Population
population_data = PopulationData(filename=None)
population_data.update(model.EXAMPLE["population"], **area)
# Simulation
example_data = ExampleData(tau=eg_tau, start_date="01Jan2020")
example_data.add(model, **area)
# Model-specialized records
spe_df = example_data.specialized(**area)
assert set(spe_df.columns) == set(
[*Term.STR_COLUMNS, *model.VARIABLES])
# Non-dimensional records
nondim_df = example_data.non_dim(**area)
assert set(nondim_df.columns) == set(
[Term.TS, *list(model.VAR_DICT.keys())])
# JHU-type records
jhu_df = example_data.subset(**area)
assert set(jhu_df.columns) == set(Term.NLOC_COLUMNS)
# Calculate Rt/day parameters when parameters are None
param_dict = {p: 0 for p in model.PARAMETERS}
model_instance = model(population_data.value(**area), **param_dict)
model_instance.calc_r0()
model_instance.calc_days_dict(eg_tau)
def test_ode(self, model):
# Setting
eg_tau = 1440
# Simulation
example_data = ExampleData(tau=eg_tau)
example_data.add(model)
nondim_df = example_data.non_dim(model)
assert isinstance(nondim_df, pd.DataFrame)
nondim_cols = [Term.TS, *list(model.VAR_DICT.keys())]
assert set(nondim_df.columns) == set(nondim_cols)
clean_df = example_data.cleaned()
assert isinstance(clean_df, pd.DataFrame)
assert set(clean_df.columns) == set(Term.COLUMNS)
dim_df = example_data.subset(model)
assert isinstance(dim_df, pd.DataFrame)
assert set(dim_df.columns) == set(Term.NLOC_COLUMNS)
# Estimation
population = model.EXAMPLE["population"]
estimator = Estimator(example_data,
model=model,
population=population,
country=model.NAME,
province=Term.UNKNOWN,
tau=eg_tau)
estimator.run()
estimated_df = estimator.summary(name=model.NAME)
assert isinstance(estimated_df, pd.DataFrame)
estimator.history(show_figure=False)
estimator.accuracy(show_figure=False)
def test_error(self, model):
# Setting
eg_tau = 1440
# Simulation
example_data = ExampleData(tau=eg_tau, start_date="01Jan2020")
with pytest.raises(KeyError):
assert not example_data.specialized(model=model).empty
with pytest.raises(KeyError):
assert not example_data.non_dim(model=model).empty
example_data.add(model)
# Model-specialized records
with pytest.raises(ValueError):
assert not example_data.specialized().empty
def test_ode_two_phases(self, population_data):
# Setting
eg_tau = 1440
# Simulation
example_data = ExampleData(tau=eg_tau)
example_data.add(SIRF, step_n=30)
example_data.add(SIRD, step_n=30)
nondim_df = example_data.non_dim(SIRF)
assert isinstance(nondim_df, pd.DataFrame)
nondim_cols = [Term.TS, *list(SIRF.VAR_DICT.keys())]
assert set(nondim_df.columns) == set(nondim_cols)
clean_df = example_data.cleaned()
assert isinstance(clean_df, pd.DataFrame)
assert set(clean_df.columns) == set(Term.COLUMNS)
dim_df = example_data.subset(SIRF)
assert isinstance(dim_df, pd.DataFrame)
assert set(dim_df.columns) == set(Term.NLOC_COLUMNS)
# Scenario analysis
population = SIRF.EXAMPLE["population"]
population_data.update(population, country=SIRF.NAME)
scenario = Scenario(example_data, population_data, country=SIRF.NAME)
scenario.trend()
def test_subset(self):
example_data = ExampleData()
example_data.add(SIRF, country="Japan")
example_data.subset(country="Japan")
example_data.subset_complement(country="Japan")
example_data.records(country="Japan")
def test_iso3(self):
example_data = ExampleData()
example_data.add(SIRF, country="Japan")
assert example_data.country_to_iso3("Japan") == "JPN"
example_data.add(SIRF, country="Moon")
assert example_data.country_to_iso3("Moon") == "---"
def test_two_phases(self, model):
example_data = ExampleData()
example_data.add(model)
example_data.add(model)
example_data.subset(model=model)
def test_one_phase(self, model):
example_data = ExampleData()
example_data.add(model)
with pytest.raises(ValueError):
example_data.subset()
# Cleaned dataset
clean_df = example_data.cleaned()
assert set(
clean_df.columns) == set(JHUData.CLEANED_COLS) - set([Term.ISO3])
# Subset
subset_df = example_data.subset(model=model)
assert subset_df.columns.tolist() == Term.SUB_COLUMNS
example_data.subset(country=model.NAME)
example_data.subset_complement(model=model)
example_data.records(model=model)
# Specialized
specialized_df = example_data.specialized(model=model)
assert specialized_df.columns.tolist() == [Term.DATE, *model.VARIABLES]
# Non-dimensional
nondim_df = example_data.non_dim(model=model)
assert nondim_df.columns.tolist() == model.VARIABLES
assert round(nondim_df.sum().sum()) == len(nondim_df)