def test_records(self, jhu_data, population_data):
scenario = Scenario(jhu_data, population_data, country="Italy")
record_df = scenario.records(show_figure=False)
warnings.filterwarnings("ignore", category=UserWarning)
record_df = scenario.records(show_figure=True)
assert isinstance(record_df, pd.DataFrame)
assert set(record_df.columns) == set(Term.NLOC_COLUMNS)
def test_estimate_delay(self, jhu_data, population_data, oxcgrt_data, indicator, target):
warnings.simplefilter("ignore", category=UserWarning)
for country in jhu_data.countries():
snl = Scenario(jhu_data, population_data, country)
delay, df = snl.estimate_delay(
oxcgrt_data, indicator=indicator, target=target)
assert isinstance(delay, int)
assert isinstance(df, pd.DataFrame)
def test_retrospective_before_estimate(self, jhu_data):
scenario = Scenario(country="Japan")
scenario.register(jhu_data)
scenario.retrospective(
beginning_date="01Jan2021", model=SIRF,
control="Control", target="Retro", timeout=1, timeout_iteration=1)
scenario.simulate(name="Control")
scenario.simulate(name="Retro")
def snl(jhu_data, population_data, japan_data, oxcgrt_data, vaccine_data,
pcr_data):
snl = Scenario(country="Italy",
province=None,
tau=None,
auto_complement=True)
snl.register(jhu_data,
population_data,
extras=[japan_data, oxcgrt_data, pcr_data, vaccine_data])
return snl
def test_start(self, jhu_data, population_data, country, province, tau):
if country == "Italy" and province == "Tokyo":
with pytest.raises(KeyError):
Scenario(
jhu_data, population_data, country, province=province)
return
if tau == 1000:
with pytest.raises(ValueError):
Scenario(
jhu_data, population_data, country, province=province, tau=tau)
return
Scenario(
jhu_data, population_data, country, province=province, tau=tau)
def test_retrospective(self, jhu_data, population_data, country):
# Setting
snl = Scenario(jhu_data, population_data, country)
snl.first_date = "01Apr2020"
snl.last_date = "19Oct2020"
snl.trend(show_figure=False)
# Retrospective analysis
snl.retrospective(
"01Sep2020", model=SIR, control="Main", target="Retrospective")
def test_retrospective(self, jhu_data, population_data, country):
# Setting
snl = Scenario(jhu_data, population_data, country)
snl.first_date = "01Apr2020"
snl.last_date = "01Jun2020"
snl.trend(show_figure=False)
# Retrospective analysis
snl.retrospective(
"01May2020", model=SIRF, control="Main", target="Retrospective",
timeout=1, timeout_iteration=1)
def test_interactive(self, jhu_data, population_data, country):
with pytest.raises(ValueError):
Scenario()
# Setting
scenario = Scenario(country=country)
scenario.register(jhu_data, population_data)
# Force interactive
scenario._interactive = True
warnings.filterwarnings("ignore", category=UserWarning)
scenario.records(show_figure=True)
def test_add_phase_dep(self, jhu_data, population_data, country):
# Setting
snl = Scenario(jhu_data, population_data, country)
snl.first_date = "01Apr2020"
snl.last_date = "01Aug2020"
# Test
warnings.simplefilter("error")
with pytest.raises(DeprecationWarning):
snl.add_phase(end_date="01May2020")
warnings.simplefilter("ignore")
snl.add_phase(end_date="01May2020")
def test_complement(self, jhu_data, population_data, country):
max_ignored = 100
# Auto complement
snl = Scenario(jhu_data,
population_data,
country,
auto_complement=True)
recovered = snl.records(show_figure=False)[Term.R]
assert recovered[recovered > max_ignored].value_counts().max() == 1
# Restore the raw data
snl.complement_reverse()
recovered = snl.records(show_figure=False)[Term.R]
assert recovered[recovered > max_ignored].value_counts().max() != 1
# Complement
snl.complement()
recovered = snl.records(show_figure=False)[Term.R]
assert recovered[recovered > max_ignored].value_counts().max() == 1
def test_class_as_dict(self, jhu_data, population_data, country):
# Setting
snl = Scenario(jhu_data, population_data, country)
# Create a phase series
population = population_data.value(country)
series = PhaseSeries("01Apr2020", "01Aug2020", population)
with pytest.raises(ScenarioNotFoundError):
assert snl["New"]
# Add scenario
snl["New"] = series
# Get scenario
assert snl["New"] == series
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_class_as_dict(self, jhu_data, population_data, country):
# Setting
snl = Scenario(jhu_data, population_data, country)
# Create a phase series
population = population_data.value(country)
sr_df = jhu_data.to_sr(country=country, population=population)
series = PhaseSeries("01Apr2020", "01Aug2020", population)
series.trend(sr_df, show_figure=False)
# Add scenario
snl["New"] = series
# Get scenario
assert snl["New"] == series
assert len(snl["New"]) == len(series)
def test_start_record_range(self, jhu_data, population_data, country):
# Setting
snl = Scenario(jhu_data, population_data, country)
# Test
snl.first_date = "01Apr2020"
assert snl.first_date == "01Apr2020"
snl.last_date = "01May2020"
assert snl.last_date == "01May2020"
with pytest.raises(ValueError):
snl.first_date = "01Jan2019"
with pytest.raises(ValueError):
tomorrow = Term.tomorrow(datetime.now().strftime(Term.DATE_FORMAT))
snl.last_date = tomorrow
def test_records(self, jhu_data, population_data, country):
warnings.simplefilter("ignore", category=UserWarning)
# Setting
snl = Scenario(jhu_data, population_data, country)
snl.first_date = "01Apr2020"
snl.last_date = "01Aug2020"
# Test
df = snl.records(show_figure=False)
assert isinstance(df, pd.DataFrame)
assert set(df.columns) == set(Term.NLOC_COLUMNS)
dates = df[Term.DATE]
assert dates.min() == Term.date_obj(snl.first_date)
assert dates.max() == Term.date_obj(snl.last_date)
df2 = snl.records(show_figure=True)
assert isinstance(df2, pd.DataFrame)
assert set(df2.columns) == set(Term.NLOC_COLUMNS)
def test_score(self, jhu_data, population_data, country):
snl = Scenario(jhu_data, population_data, country, tau=360)
snl.trend(show_figure=False)
snl.estimate(SIRF, timeout=10)
metrics_list = ["MAE", "MSE", "MSLE", "RMSE", "RMSLE"]
for metrics in metrics_list:
score = snl.score(metrics=metrics)
assert isinstance(score, float)
# Selected phases
df = snl.summary()
all_phases = df.index.tolist()
sel_score = snl.score(phases=all_phases[-2:])
# Selected past days (when the begging date is a start date)
beginning_date = df.loc[df.index[-2], Term.START]
past_days = Term.steps(beginning_date, snl.last_date, tau=1440)
assert snl.score(past_days=past_days) == sel_score
# Selected past days
snl.score(past_days=60)
def test_simulate(self, jhu_data, population_data, country):
warnings.simplefilter("ignore", category=UserWarning)
warnings.simplefilter("ignore", category=DeprecationWarning)
# Setting
snl = Scenario(jhu_data, population_data, country)
snl.first_date = "01Apr2020"
snl.last_date = "01Aug2020"
snl.trend(show_figure=False)
# Parameter estimation
with pytest.raises(ValueError):
# Deprecated
snl.param_history(["rho"])
snl.estimate(SIR)
# Simulation
snl.simulate()
# Parameter history (Deprecated)
snl.param_history([Term.RT], divide_by_first=False)
snl.param_history(["rho"])
snl.param_history(["rho"], show_figure=False)
snl.param_history(["rho"], show_box_plot=False)
with pytest.raises(KeyError):
snl.param_history(["feeling"])
# Comparison of scenarios
snl.describe()
snl.history(target="Rt")
snl.history(target="sigma")
snl.history(target="Infected")
with pytest.raises(KeyError):
snl.history(target="temperature")
# Change rate of parameters
snl.history_rate(name="Main")
with pytest.raises(TypeError):
snl.history_rate(params="", name="Main")
# Add new scenario
snl.add(end_date="01Sep2020", name="New")
snl.describe()
def test_estimate_tau(self, jhu_data, population_data, country):
# Setting
snl = Scenario(jhu_data, population_data, country)
snl.trend(show_figure=False)
with pytest.raises(ValueError):
snl.estimate(SIR, tau=1440)
def test_analysis(self, jhu_data, population_data):
scenario = Scenario(jhu_data, population_data, country="Italy")
with pytest.raises(KeyError):
scenario.simulate(name="Main", show_figure=False)
with pytest.raises(ValueError):
scenario.estimate(model=SIRF)
# S-R trend analysis
scenario.trend(show_figure=False)
warnings.filterwarnings("ignore", category=UserWarning)
scenario.trend(show_figure=True)
# Parameter estimation of SIR-F model
with pytest.raises(ValueError):
scenario.param_history(targets=["Rt"], show_figure=False)
with pytest.raises(ValueError):
scenario.estimate(model=SIRF, tau=1440)
scenario.estimate(model=SIRF)
# History of estimation
scenario.estimate_history(phase="1st")
with pytest.raises(KeyError):
scenario.estimate_history(phase="0th")
# Accuracy of estimation
scenario.estimate_accuracy(phase="1st")
with pytest.raises(KeyError):
scenario.estimate_accuracy(phase="0th")
# Prediction
scenario.add(name="Main", days=100)
scenario.simulate(name="Main", show_figure=False)
scenario.simulate(name="Main", show_figure=True)
scenario.param_history(targets=["Rt"], show_figure=False)
scenario.param_history(targets=["Rt"], divide_by_first=False)
scenario.param_history(targets=["Rt"], show_box_plot=False)
with pytest.raises(KeyError):
scenario.param_history(targets=["Rt", "Value"])
with pytest.raises(KeyError):
scenario.param_history(targets=["Rt"], box_plot=False)
# New scenario
sigma_new = scenario.get("sigma", phase="last") * 2
with pytest.raises(KeyError):
scenario.get("value")
warnings.filterwarnings("ignore", category=DeprecationWarning)
scenario.add_phase(name="New medicines", days=100, sigma=sigma_new)
# Summarize scenarios
summary_df = scenario.summary()
assert isinstance(summary_df, pd.DataFrame)
desc_df = scenario.describe()
assert isinstance(desc_df, pd.DataFrame)
# Estimation errors
with pytest.raises(TypeError):
scenario.estimate(SIRF, phases="1st")
with pytest.raises(KeyError):
scenario.estimate(SIRF, phases=["100th"])
def test_summary(self, jhu_data, population_data, country):
# Setting
snl = Scenario(jhu_data, population_data, country)
snl.first_date = "01Apr2020"
snl.last_date = "01Aug2020"
snl.trend(show_figure=False)
# One scenario
assert set(snl.summary().columns) == set(
[Term.TENSE, Term.START, Term.END, Term.N])
# Show two scenarios
snl.clear(name="New")
cols = snl.summary().reset_index().columns
assert set([Term.SERIES, Term.PHASE]).issubset(set(cols))
# Show selected scenario
cols_sel = snl.summary(name="New").reset_index().columns
assert not set([Term.SERIES, Term.PHASE]).issubset(set(cols_sel))
# Columns to show
show_cols = [Term.N, Term.START]
assert set(snl.summary(columns=show_cols).columns) == set(show_cols)
with pytest.raises(TypeError):
snl.summary(columns=Term.N)
with pytest.raises(KeyError):
snl.summary(columns=[Term.N, "Temperature"])
# To markdown
snl.summary().to_markdown()
def test_score_error(self, jhu_data, population_data, country):
snl = Scenario(jhu_data, population_data, country)
with pytest.raises(ValueError):
snl.score()
snl.trend(show_figure=False)
with pytest.raises(NameError):
snl.score()
all_phases = snl.summary().index.tolist()
snl.disable(phases=all_phases[:-2])
snl.estimate(SIRF, timeout=10)
with pytest.raises(TypeError):
snl.score(phases="0th")
with pytest.raises(KeyError):
snl.score(phases=["100th"])
with pytest.raises(TypeError):
snl.score(variables="Infected")
with pytest.raises(KeyError):
snl.score(variables=["Susceptible"])
with pytest.raises(ValueError):
snl.score(metrics="Subjective evaluation")
with pytest.raises(ValueError):
snl.score(phases=["0th"], past_days=100)
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_fit_predict(self, jhu_data, population_data, oxcgrt_data, country):
snl = Scenario(jhu_data, population_data, country)
snl.trend(show_figure=False)
with pytest.raises(UnExecutedError):
snl.fit(oxcgrt_data)
snl.estimate(SIRF, timeout=1, timeout_iteration=1)
# Fitting
with pytest.raises(UnExecutedError):
snl.predict()
info_dict = snl.fit(oxcgrt_data)
assert isinstance(info_dict, dict)
# Prediction
snl.predict()
assert Term.FUTURE in snl.summary()[Term.TENSE].unique()
# Fitting & predict
snl.fit_predict(oxcgrt_data)
assert Term.FUTURE in snl.summary()[Term.TENSE].unique()
def test_add_past_phases(self, jhu_data, population_data):
scenario = Scenario(jhu_data, population_data, country="India")
scenario.delete()
# Phase series
scenario.clear(name="Medicine")
scenario.add(days=100)
scenario.delete(name="Medicine")
with pytest.raises(TypeError):
scenario.delete(phase="0th")
with pytest.raises(TypeError):
scenario.summary(columns="Population")
with pytest.raises(KeyError):
scenario.summary(columns=["Value"])
# Range of past phases
scenario.first_date = "01Mar2020"
scenario.first_date
scenario.last_date = "16Jul2020"
scenario.last_date
with pytest.raises(ValueError):
scenario.first_date = "01Aug2020"
with pytest.raises(ValueError):
scenario.last_date = "01Feb2020"
# With trend analysis
scenario.trend(set_phases=True)
with pytest.raises(ValueError):
scenario.trend(set_phases=False, n_points=3)
scenario.combine(phases=["3rd", "4th"])
scenario.separate(date="30May2020", phase="1st")
scenario.delete(phases=["1st"])
scenario.trend(set_phases=False)
trend_df = scenario.summary()
assert len(trend_df) == 4
# add scenarios one by one
scenario.clear(include_past=True)
scenario.add(end_date="29May2020")
scenario.add(end_date="05Jun2020").delete(phases=["0th"])
scenario.add(end_date="15Jun2020")
scenario.add(end_date="04Jul2020")
scenario.add()
one_df = scenario.summary()
assert len(one_df) == 4
# With 0th phase
scenario.use_0th = True
scenario.trend(set_phases=False, include_init_phase=True)
scenario.use_0th = False
scenario.trend(set_phases=True, include_init_phase=True)
scenario.delete(phases=["0th"])
assert len(scenario.summary()) == 5
with pytest.raises(TypeError):
scenario.delete(phases="1st")
def test_edit_series(self, jhu_data, population_data, country):
# Setting
snl = Scenario(jhu_data, population_data, country)
snl.first_date = "01Apr2020"
snl.last_date = "01Aug2020"
# Add and clear
assert snl.summary().empty
snl.add(end_date="05May2020")
snl.add(days=20)
snl.add()
snl.add(end_date="01Sep2020")
assert len(snl["Main"]) == 4
snl.clear(include_past=True)
snl.add(end_date="01Sep2020", name="New")
assert len(snl["New"]) == 2
# Delete
snl.delete(name="Main")
assert len(snl["Main"]) == 0
with pytest.raises(TypeError):
snl.delete(phases="1st", name="New")
snl.delete(phases=["1st"], name="New")
assert len(snl["New"]) == 1
snl.delete(name="New")
with pytest.raises(KeyError):
assert len(snl["New"]) == 1
def test_estimate(self, jhu_data, population_data, country):
warnings.simplefilter("ignore", category=UserWarning)
# Setting
snl = Scenario(jhu_data, population_data, country)
snl.first_date = "01Apr2020"
snl.last_date = "01Aug2020"
with pytest.raises(ValueError):
snl.estimate(SIR)
snl.trend(include_init_phase=True, show_figure=False)
snl.disable(phases=["0th"])
with pytest.raises(AttributeError):
snl.estimate_history(phase="1th")
# Parameter estimation
with pytest.raises(KeyError):
snl.estimate(SIR, phases=["30th"])
with pytest.raises(ValueError):
snl.estimate(model=SIR, tau=1440)
snl.enable(phases=["0th"])
with pytest.raises(TypeError):
snl.estimate(model=SIR, phases="1st")
with pytest.raises(ValueError):
snl.estimate(model=SIR, phases=["0th"])
snl.clear(include_past=True)
snl.trend(show_figure=False)
snl.estimate(SIR)
# Estimation history
snl.estimate_history(phase="1st")
# Estimation accuracy
snl.estimate_accuracy(phase="1st")
# Get a value
snl.get(Term.RT)
with pytest.raises(KeyError):
snl.get("feeling")
def test_edit(self, jhu_data, population_data, country):
# Setting
snl = Scenario(jhu_data, population_data, country)
snl.first_date = "01Apr2020"
snl.last_date = "01Aug2020"
snl.trend(show_figure=False)
# Combine
length = len(snl["Main"])
snl.combine(["1st", "2nd"])
n_changed = int(population_data.value(country) * 0.98)
snl.combine(["2nd", "3rd"], population=n_changed)
assert len(snl["Main"]) == length - 2
# Separate
with pytest.raises(IndexError):
snl.separate(date="01Dec2020")
snl.separate(date="01May2020")
assert len(snl["Main"]) == length - 1
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_adjust_end(self, jhu_data, population_data, country):
# Setting
scenario = Scenario(country=country)
scenario.register(jhu_data, population_data)
scenario.timepoints(first_date="01Dec2020", today="01Feb2021")
# Main scenario
scenario.add(end_date="01Apr2021", name="Main")
# New scenario
scenario.clear(name="New", include_past=True)
scenario.add(end_date="01Jan2021", name="New")
# Adjust end date
scenario.adjust_end()
# Check output
assert scenario.get(Term.END, phase="last", name="Main") == "01Apr2021"
assert scenario.get(Term.END, phase="last", name="New") == "01Apr2021"
def test_trend(self, jhu_data, population_data, country):
# Setting
snl = Scenario(jhu_data, population_data, country)
snl.first_date = "01Apr2020"
snl.last_date = "01Aug2020"
snl.trend(show_figure=False)
assert snl["Main"]
with pytest.raises(ValueError):
snl.trend(show_figure=False, n_points=3)
# Disable/enable
length = len(snl["Main"])
snl.disable(phases=["0th"], name="Main")
assert len(snl["Main"]) == length - 1
snl.enable(phases=["0th"], name="Main")
assert len(snl["Main"]) == length
with pytest.raises(TypeError):
snl.enable(phases="0th", name="Main")
with pytest.raises(TypeError):
snl.disable(phases="1st", name="Main")