def test_set_ode(self):
unit = PhaseUnit("01Jan2020", "01Feb2020", 1000)
unit.set_ode(model=SIR, tau=720, rho=0.2)
assert issubclass(unit.model, SIR)
summary_dict = unit.to_dict()
assert summary_dict[Term.TAU] == 720
assert summary_dict[Term.ODE] == SIR.NAME
assert summary_dict["rho"] == 0.2
assert summary_dict["sigma"] is None
def test_tau(self):
unit = PhaseUnit("01Jan2020", "01Feb2020", 1000)
assert unit.tau is None
with pytest.raises(ValueError):
unit.tau = 1000
unit.tau = 1440
assert unit.tau == 1440
with pytest.raises(AttributeError):
unit.tau = 720
def test_definition_property(self):
unit = PhaseUnit("01Jan2020", "01Feb2020", 1000)
assert unit.start_date == "01Jan2020"
with pytest.raises(AttributeError):
unit.start_date = "10Jan2020"
assert unit.end_date == "01Feb2020"
with pytest.raises(AttributeError):
unit.end_date = "10Jan2020"
assert unit.population == 1000
with pytest.raises(AttributeError):
unit.population = 100000
def test_simulate_y0_direct(self, jhu_data, population_data, country):
# Dataset
population = population_data.value(country)
# Parameter setting
unit = PhaseUnit("27May2020", "27Jun2020", population)
unit.set_ode(model=SIR, tau=240, rho=0.006, sigma=0.011)
# Simulation
y0_dict = {
"Susceptible": 126512455, "Infected": 1806, "Fatal or Recovered": 14839
}
sim_df = unit.simulate(y0_dict=y0_dict)
assert set(sim_df.columns) == set(Term.NLOC_COLUMNS)
def test_delete_phase(self, jhu_data, population_data, country):
# Setting
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)
first_len = len(series)
# Deletion of 0th phase is the same as disabling 0th phase
series.delete("0th")
series.enable("0th")
assert len(series) == first_len
assert "5th" in series.to_dict()
# Delete phase (not the last registered phase)
new_second = PhaseUnit(
series.unit("2nd").start_date,
series.unit("3rd").end_date,
series.unit("2nd").population)
series.delete("3rd")
assert len(series) == first_len - 1
assert series.unit("2nd") == new_second
# Delete the last phase
old_last = series.unit("last")
series.delete("last")
series.add()
assert series.unit("last").start_date == old_last.start_date
def test_estimate_with_fixed(self, jhu_data, population_data, country):
# Dataset
population = population_data.value(country)
record_df = jhu_data.subset(country, population=population)
# Setting
unit = PhaseUnit("27May2020", "27Jun2020", population)
unit.set_ode(model=SIR, tau=360, rho=0.01)
# Parameter estimation
unit.estimate(record_df=record_df)
assert unit.tau == unit.to_dict()[Term.TAU] == 360
assert unit.to_dict()["rho"] == 0.01
def test_sort(self):
unit1 = PhaseUnit("01Jan2020", "01Feb2020", 1000)
unit2 = PhaseUnit("02Feb2020", "01Mar2020", 1000)
unit3 = PhaseUnit("02Mar2020", "01Apr2020", 1000)
unit4 = PhaseUnit("02Mar2020", "01Apr2020", 1000)
unit5 = PhaseUnit("01Jan2020", "01Apr2020", 1000)
assert unit3 == unit4
assert unit1 != unit2
assert unit1 < unit2
assert unit3 > unit1
assert not unit3 > unit4
assert unit3 <= unit4
assert unit1 < "02Feb2020"
assert unit1 <= "01Feb2020"
assert unit1 > "31Dec2019"
assert unit1 >= "01Jan2020"
assert sorted([unit3, unit1, unit2]) == [unit1, unit2, unit3]
assert str(unit1 + unit2) == "Phase (01Jan2020 - 01Mar2020)"
assert str(unit5 - unit1) == "Phase (02Feb2020 - 01Apr2020)"
assert str(unit5 - unit4) == "Phase (01Jan2020 - 01Mar2020)"
assert set([unit1, unit3, unit4]) == set([unit1, unit3])
def test_enable(self):
unit = PhaseUnit("01Jan2020", "01Feb2020", 1000)
assert bool(unit)
unit.disable()
assert not bool(unit)
unit.enable()
assert bool(unit)
def test_add_phase(self, jhu_data, population_data, country):
# Setting
population = population_data.value(country)
series = PhaseSeries("01Apr2020", "01Aug2020", population)
# Last phase when empty
empty_phase = PhaseUnit("31Mar2020", "31Mar2020", population)
assert series.unit(phase="last") == empty_phase
# Add a phase with specified end date: 0th
series.add(end_date="22Apr2020")
# Add a phase with specified population value: 1st
with pytest.raises(ValueError):
series.add(end_date="22Apr2020")
series.add(end_date="05May2020", population=int(population * 0.98))
# Add a phase with specified the number of days: 2nd
series.add(days=21)
# Filling past phases and add a future phase: 3rd, 4th
series.add(end_date="01Sep2020")
# Add a future phase: 5th
series.add(days=30)
# Summary
df = series.summary()
base_cols = [Term.TENSE, Term.START, Term.END, Term.N]
assert set(df.columns) == set(base_cols)
assert series.to_dict()["0th"]["Type"] == Term.PAST
first_len = len(df)
assert set(df.loc["3rd", :].tolist()) == set(
[Term.PAST, "27May2020", "01Aug2020", 123998518])
assert set(df.loc["4th", :].tolist()) == set(
[Term.FUTURE, "02Aug2020", "01Sep2020", 123998518])
# Disable/enable a phase
series.disable("0th")
assert "0th" not in series.to_dict()
assert len(series) == first_len - 1
assert len([unit for unit in series]) == first_len
assert len([unit for unit in series if unit]) == first_len - 1
series.enable("0th")
assert "0th" in series.to_dict()
assert len(series) == first_len
# Clear future phases: 4th and 5th will be deleted
series.clear(include_past=False)
assert "4th" not in series.to_dict()
assert len(series) == first_len - 2
assert series
# Clear all phases
series.clear(include_past=True)
assert len(series) == 0
assert not series
# Filling past phases: 0th
series.add()
assert len(series) == 1
# Last phase
assert series.unit(phase="last")
def test_start(self):
unit = PhaseUnit("01Jan2020", "01Feb2020", 1000)
assert str(unit) == "Phase (01Jan2020 - 01Feb2020)"
summary_dict = unit.to_dict()
assert summary_dict[Term.START] == "01Jan2020"
assert summary_dict[Term.END] == "01Feb2020"
assert summary_dict[Term.N] == 1000
summary_df = unit.summary()
assert set(summary_df.columns) == set([Term.START, Term.END, Term.N])
with pytest.raises(NotImplementedError):
assert unit == "phase"
unit2 = PhaseUnit("01Jan2020", "01Feb2020", 100000)
assert unit == unit2
unit3 = PhaseUnit("01Jan2020", "01Mar2020", 1000)
assert unit != unit3
assert "31Dec2019" not in unit
assert "01Jan2020" in unit
assert "15Jan2020" in unit
assert "01Feb2020" in unit
assert "01Mar2020" not in unit
def test_replace(self, jhu_data, population_data, country):
# Setting
population = population_data.value(country)
series = PhaseSeries("01Apr2020", "01Aug2020", population)
series.add(end_date="01May2020")
series.add(end_date="01Jun2020")
series.add(end_date="01Jul2020")
series.add()
# Replace one old phase with one new phase
unit_old = series.unit("0th")
unit_new = PhaseUnit(
unit_old.start_date, unit_old.end_date, population
)
unit_new.set_ode(tau=360)
series.replace("0th", unit_new)
assert series.unit("0th") == unit_new
# Replace one old phase with two new phases
unit_old = series.unit("1st")
change_date = Term.date_change(unit_old.end_date, days=-7)
unit_pre = PhaseUnit(
unit_old.start_date, Term.yesterday(change_date), population)
unit_pre.set_ode(tau=360)
unit_fol = PhaseUnit(change_date, unit_old.end_date, population)
unit_fol.set_ode(tau=360)
series.replaces(phase="1st", new_list=[unit_pre, unit_fol])
print(series.unit("1st"), unit_pre)
assert series.unit("1st") == unit_pre
assert series.unit("2nd") == unit_fol
# TypeError of new_list
with pytest.raises(TypeError):
series.replaces(phase="3rd", new_list=[unit_pre, Term])
# ValueError with tense
with pytest.raises(ValueError):
future_unit = PhaseUnit("01Sep2020", "01Dec2020", population)
series.replaces(phase="3rd", new_list=[future_unit])
# Add phase without deletion of any phases
new1 = PhaseUnit("02Aug2020", "01Sep2020", population)
new2 = PhaseUnit("02Sep2020", "01Oct2020", population)
series.replaces(phase=None, new_list=[new1, new2])
assert series.unit("last") == new2
def test_set_tau(self):
unit = PhaseUnit("01Jan2020", "01Feb2020", 1000)
unit.set_ode(model=None, tau=240, rho=0.006)
assert unit.tau == 240
assert "rho" not in unit.to_dict()
def test_simulate(self, jhu_data, population_data, country):
# Dataset
population = population_data.value(country)
record_df = jhu_data.subset(country, population=population)
# Parameter setting
unit = PhaseUnit("27May2020", "27Jun2020", population)
with pytest.raises(UnExecutedError):
unit3 = PhaseUnit("27May2020", "27Jun2020", population)
unit3.set_ode(model=SIR, tau=240, rho=0.006)
unit3.simulate()
unit.set_ode(model=SIR, tau=240, rho=0.006, sigma=0.011)
summary_dict = unit.to_dict()
assert summary_dict[Term.RT] == 0.55
assert summary_dict["1/beta [day]"] == 27
assert summary_dict["1/gamma [day]"] == 15
# Set initial values
unit.set_y0(record_df)
# Simulation
sim_df = unit.simulate()
assert set(sim_df.columns) == set(Term.NLOC_COLUMNS)
def test_estimate(self, jhu_data, population_data, country):
# Dataset
population = population_data.value(country)
record_df = jhu_data.subset(country, population=population)
# Parameter estimation
unit = PhaseUnit("27May2020", "27Jun2020", population)
with pytest.raises(NameError):
unit.estimate()
unit.set_ode(model=SIR)
with pytest.raises(ValueError):
unit.estimate()
unit.record_df = record_df
assert set(unit.record_df.columns) == set(Term.SUB_COLUMNS)
unit.estimate(timeout=20)
# Check results
assert isinstance(unit.estimator, Estimator)
assert set(SIR.PARAMETERS).issubset(unit.to_dict())
assert set(SIR.DAY_PARAMETERS).issubset(unit.to_dict())
assert None not in unit.to_dict().values()