def SI(cause: str) -> DiseaseModel:
healthy = SusceptibleState(cause)
infected = DiseaseState(cause)
healthy.allow_self_transitions()
healthy.add_transition(infected, source_data_type='rate')
infected.allow_self_transitions()
return DiseaseModel(cause, states=[healthy, infected])
def NeonatalSWC_without_incidence(cause):
with_condition_data_functions = {'birth_prevalence':
lambda cause, builder: builder.data.load(f"cause.{cause}.birth_prevalence")}
healthy = SusceptibleState(cause)
with_condition = DiseaseState(cause, get_data_functions=with_condition_data_functions)
healthy.allow_self_transitions()
with_condition.allow_self_transitions()
return DiseaseModel(cause, states=[healthy, with_condition])
def test_dwell_time(assign_cause_mock, base_config, base_plugins, disease,
base_data):
time_step = 10
assign_cause_mock.side_effect = lambda population, *args: pd.DataFrame(
{'condition_state': 'healthy'}, index=population.index)
base_config.update(
{
'time': {
'step_size': time_step
},
'population': {
'population_size': 10
}
}, **metadata(__file__))
healthy_state = BaseDiseaseState('healthy')
data_function = base_data(0)
data_function['dwell_time'] = lambda _, __: pd.Timedelta(days=28)
data_function['disability_weight'] = lambda _, __: 0.0
event_state = DiseaseState('event', get_data_functions=data_function)
done_state = BaseDiseaseState('sick')
healthy_state.add_transition(event_state)
event_state.add_transition(done_state)
model = DiseaseModel(disease,
initial_state=healthy_state,
states=[healthy_state, event_state, done_state])
simulation = InteractiveContext(components=[TestPopulation(), model],
configuration=base_config,
plugin_configuration=base_plugins)
# Move everyone into the event state
simulation.step()
event_time = simulation._clock.time
assert np.all(simulation.get_population()[disease] == 'event')
simulation.step()
simulation.step()
# Not enough time has passed for people to move out of the event state, so they should all still be there
assert np.all(simulation.get_population()[disease] == 'event')
simulation.step()
# Now enough time has passed so people should transition away
assert np.all(simulation.get_population()[disease] == 'sick')
assert np.all(simulation.get_population().event_event_time ==
pd.to_datetime(event_time))
assert np.all(simulation.get_population().event_event_count == 1)
def test_prevalence_multiple_sequelae(base_config, base_plugins, disease,
base_data, test_prevalence_level):
year_start = base_config.time.start.year
year_end = base_config.time.end.year
healthy = BaseDiseaseState('healthy')
sequela = dict()
for i, p in enumerate(test_prevalence_level):
data_funcs = base_data(p)
data_funcs.update({'disability_weight': lambda _, __: 0.0})
sequela[i] = DiseaseState('sequela' + str(i),
get_data_functions=data_funcs)
model = DiseaseModel(disease,
initial_state=healthy,
states=[healthy, sequela[0], sequela[1], sequela[2]])
base_config.update({'population': {
'population_size': 100000
}}, **metadata(__file__))
simulation = InteractiveContext(components=[TestPopulation(), model],
configuration=base_config,
plugin_configuration=base_plugins)
error_message = "initial sequela status of simulants should be matched to the prevalence data."
assert np.allclose([
get_test_prevalence(simulation, 'sequela0'),
get_test_prevalence(simulation, 'sequela1'),
get_test_prevalence(simulation, 'sequela2')
], test_prevalence_level, .02), error_message
def test__assign_event_time_for_prevalent_cases():
pop_data = pd.DataFrame(index=range(100))
random_func = lambda index: pd.Series(0.4, index=index)
current_time = pd.Timestamp(2017, 1, 10, 12)
dwell_time_func = lambda index: pd.Series(10, index=index)
# 10* 0.4 = 4 ; 4 days before the current time
expected = pd.Series(pd.Timestamp(2017, 1, 6, 12), index=pop_data.index)
assert expected.equals(
DiseaseState._assign_event_time_for_prevalent_cases(
pop_data, current_time, random_func, dwell_time_func))
def test_prevalence_birth_prevalence_initial_assignment(
base_config, base_plugins, disease):
healthy = BaseDiseaseState('healthy')
data_funcs = {
'prevalence': lambda _, __: 1,
'birth_prevalence': lambda _, __: 0.5,
'disability_weight': lambda _, __: 0
}
with_condition = DiseaseState('with_condition',
get_data_functions=data_funcs)
model = DiseaseModel(disease,
initial_state=healthy,
states=[healthy, with_condition])
base_config.update(
{
'population': {
'population_size': 1000,
'age_start': 0,
'age_end': 5
}
}, **metadata(__file__))
simulation = InteractiveContext(components=[TestPopulation(), model],
configuration=base_config,
plugin_configuration=base_plugins)
# prevalence should be used for assigning initial status at sim start
assert np.isclose(get_test_prevalence(simulation, "with_condition"), 1)
# birth prevalence should be used for assigning initial status to newly-borns on time steps
simulation._clock.step_forward()
simulation.simulant_creator(1000,
population_configuration={
'age_start': 0,
'age_end': 0,
'sim_state': 'time_step'
})
assert np.isclose(get_test_prevalence(simulation, "with_condition"), 0.75,
0.01)
# and prevalence should be used for ages not start = end = 0
simulation._clock.step_forward()
simulation.simulant_creator(1000,
population_configuration={
'age_start': 0,
'age_end': 5,
'sim_state': 'time_step'
})
assert np.isclose(get_test_prevalence(simulation, "with_condition"), 0.83,
0.01)
def test_prevalence_single_state_with_migration(base_config, base_plugins,
disease, base_data,
test_prevalence_level):
"""
Test the prevalence for the single state over newly migrated population.
Start with the initial population, check the prevalence for initial assignment.
Then add new simulants and check whether the initial status is
properly assigned to new simulants based on the prevalence data and pre-existing simulants status
"""
year_start = base_config.time.start.year
year_end = base_config.time.end.year
healthy = BaseDiseaseState('healthy')
data_funcs = base_data(test_prevalence_level)
data_funcs.update({'disability_weight': lambda _, __: 0.0})
sick = DiseaseState('sick', get_data_functions=data_funcs)
model = DiseaseModel(disease,
initial_state=healthy,
states=[healthy, sick])
base_config.update({'population': {
'population_size': 50000
}}, **metadata(__file__))
simulation = InteractiveContext(components=[TestPopulation(), model],
configuration=base_config,
plugin_configuration=base_plugins)
error_message = "initial status of simulants should be matched to the prevalence data."
assert np.isclose(get_test_prevalence(simulation, 'sick'),
test_prevalence_level, 0.01), error_message
simulation._clock.step_forward()
assert np.isclose(get_test_prevalence(simulation, 'sick'),
test_prevalence_level, .01), error_message
simulation.simulant_creator(50000,
population_configuration={
'age_start': 0,
'age_end': 5,
'sim_state': 'time_step'
})
assert np.isclose(get_test_prevalence(simulation, 'sick'),
test_prevalence_level, .01), error_message
simulation._clock.step_forward()
simulation.simulant_creator(50000,
population_configuration={
'age_start': 0,
'age_end': 5,
'sim_state': 'time_step'
})
assert np.isclose(get_test_prevalence(simulation, 'sick'),
test_prevalence_level, .01), error_message
def SIS_fixed_duration(cause: str, duration: str) -> DiseaseModel:
duration = pd.Timedelta(days=float(duration) // 1, hours=(float(duration) % 1) * 24.0)
healthy = SusceptibleState(cause)
infected = DiseaseState(cause, get_data_functions={'dwell_time': lambda _, __: duration})
healthy.allow_self_transitions()
healthy.add_transition(infected, source_data_type='rate')
infected.add_transition(healthy)
infected.allow_self_transitions()
return DiseaseModel(cause, states=[healthy, infected])
def NeonatalSIS(cause):
with_condition_data_functions = {'birth_prevalence':
lambda cause, builder: builder.data.load(f"cause.{cause}.birth_prevalence")}
healthy = SusceptibleState(cause)
with_condition = DiseaseState(cause, get_data_functions=with_condition_data_functions)
healthy.allow_self_transitions()
healthy.add_transition(with_condition, source_data_type='rate')
with_condition.allow_self_transitions()
with_condition.add_transition(healthy, source_data_type='rate')
return DiseaseModel(cause, states=[healthy, with_condition])
def test_mortality_rate(base_config, base_plugins, disease):
year_start = base_config.time.start.year
year_end = base_config.time.end.year
time_step = pd.Timedelta(days=base_config.time.step_size)
healthy = BaseDiseaseState('healthy')
mort_get_data_funcs = {
'dwell_time':
lambda _, __: pd.Timedelta(days=0),
'disability_weight':
lambda _, __: 0.0,
'prevalence':
lambda _, __: build_table(0.000001, year_start - 1, year_end,
['age', 'year', 'sex', 'value']),
'excess_mortality_rate':
lambda _, __: build_table(0.7, year_start - 1, year_end),
}
mortality_state = DiseaseState('sick',
get_data_functions=mort_get_data_funcs)
healthy.add_transition(mortality_state)
model = DiseaseModel(disease,
initial_state=healthy,
states=[healthy, mortality_state])
simulation = InteractiveContext(
components=[TestPopulation(), model,
Mortality()],
configuration=base_config,
plugin_configuration=base_plugins)
mortality_rate = simulation._values.get_value('mortality_rate')
simulation.step()
# Folks instantly transition to sick so now our mortality rate should be much higher
assert np.allclose(
from_yearly(0.7, time_step),
mortality_rate(simulation.get_population().index)['sick'])
def Stroke(stroke_name):
stroke_types = [
'ischemic_stroke', 'subarachnoid_hemorrhage',
'intracerebral_hemorrhage'
]
if stroke_name not in stroke_types:
raise ValueError(
f'Stroke name must be one of {stroke_types}. You supplied {stroke_name}'
)
susceptible = SusceptibleState(stroke_name)
data_funcs = {'dwell_time': lambda *args: pd.Timedelta(days=28)}
acute = DiseaseState(f'acute_{stroke_name}',
cause_type='sequela',
get_data_functions=data_funcs)
#data_funcs = {'dwell_time': lambda *args: pd.Timedelta(days=365.25)}
post = RelapseState(f'post_{stroke_name}', cause_type='sequela')
susceptible.allow_self_transitions()
data_funcs = {
'incidence_rate':
lambda _, builder: builder.data.load(
f'cause.{stroke_name}.incidence_rate')
}
susceptible.add_transition(acute,
source_data_type='rate',
get_data_functions=data_funcs)
acute.allow_self_transitions()
acute.add_transition(post)
post.allow_self_transitions()
data_funcs = {
'relapse_rate':
lambda _, builder: builder.data.load(
f'cause.{stroke_name}.incidence_rate')
}
post.add_transition(acute,
source_data_type='rate',
get_data_functions=data_funcs)
#post.add_transition(susceptible, source_data_type='time')
return DiseaseModel(stroke_name, states=[susceptible, acute, post])
def IschemicHeartDisease():
susceptible = SusceptibleState('ischemic_heart_disease')
data_funcs = {'dwell_time': lambda *args: pd.Timedelta(days=28)}
acute_mi = DiseaseState('acute_myocardial_infarction',
cause_type='sequela',
get_data_functions=data_funcs)
#data_funcs = {'dwell_time': lambda *args: pd.Timedelta(days=365.25)}
post_mi = RelapseState(
'post_myocardial_infarction',
cause_type='sequela',
)
susceptible.allow_self_transitions()
data_funcs = {
'incidence_rate':
lambda _, builder: builder.data.load(
'cause.ischemic_heart_disease.incidence_rate')
}
susceptible.add_transition(acute_mi,
source_data_type='rate',
get_data_functions=data_funcs)
acute_mi.allow_self_transitions()
acute_mi.add_transition(post_mi)
post_mi.allow_self_transitions()
data_funcs = {
'relapse_rate':
lambda _, builder: builder.data.load(
'cause.ischemic_heart_disease.incidence_rate')
}
post_mi.add_transition(acute_mi,
source_data_type='rate',
get_data_functions=data_funcs)
#post_mi.add_transition(susceptible, source_data_type='time')
return DiseaseModel('ischemic_heart_disease',
states=[susceptible, acute_mi, post_mi])
def test_dwell_time_with_mortality(base_config, base_plugins, disease):
year_start = base_config.time.start.year
year_end = base_config.time.end.year
time_step = 10
pop_size = 100
base_config.update(
{
'time': {
'step_size': time_step
},
'population': {
'population_size': pop_size
}
}, **metadata(__file__))
healthy_state = BaseDiseaseState('healthy')
mort_get_data_funcs = {
'dwell_time':
lambda _, __: pd.Timedelta(days=14),
'excess_mortality_rate':
lambda _, __: build_table(0.7, year_start - 1, year_end),
'disability_weight':
lambda _, __: 0.0
}
mortality_state = DiseaseState('event',
get_data_functions=mort_get_data_funcs)
done_state = BaseDiseaseState('sick')
healthy_state.add_transition(mortality_state)
mortality_state.add_transition(done_state)
model = DiseaseModel(disease,
initial_state=healthy_state,
states=[healthy_state, mortality_state, done_state])
mortality = Mortality()
simulation = InteractiveContext(
components=[TestPopulation(), model, mortality],
configuration=base_config,
plugin_configuration=base_plugins)
# Move everyone into the event state
simulation.step()
assert np.all(simulation.get_population()[disease] == 'event')
simulation.step()
# Not enough time has passed for people to move out of the event state, so they should all still be there
assert np.all(simulation.get_population()[disease] == 'event')
simulation.step()
# Make sure some people have died and remained in event state
assert (simulation.get_population()['alive'] == 'alive').sum() < pop_size
assert ((simulation.get_population()['alive'] == 'dead').sum() == (
simulation.get_population()[disease] == 'event').sum())
# enough time has passed so living people should transition away to sick
assert ((simulation.get_population()['alive'] == 'alive').sum() == (
simulation.get_population()[disease] == 'sick').sum())