def test_interpolated_tables_without_uninterpolated_columns(base_config):
year_start = base_config.time.start.year
year_end = base_config.time.end.year
years = build_table(lambda age, sex, year: year, year_start, year_end)
del years['sex']
years = years.drop_duplicates()
base_config.population.update({'population_size': 10000})
simulation = setup_simulation([TestPopulation()], input_config=base_config)
manager = simulation.tables
years = manager.build_table(years,
key_columns=(),
parameter_columns=(
'year',
'age',
))
result_years = years(simulation.population.population.index)
fractional_year = simulation.clock.time.year
fractional_year += simulation.clock.time.timetuple().tm_yday / 365.25
assert np.allclose(result_years, fractional_year)
simulation.clock._time += pd.Timedelta(30.5 * 125, unit='D')
result_years = years(simulation.population.population.index)
fractional_year = simulation.clock.time.year
fractional_year += simulation.clock.time.timetuple().tm_yday / 365.25
assert np.allclose(result_years, fractional_year)
def test_side_effects():
class DoneState(State):
def setup(self, builder):
super().setup(builder)
self.population_view = builder.population.get_view(['count'])
def _transition_side_effect(self, index, event_time):
pop = self.population_view.get(index)
self.population_view.update(pop['count'] + 1)
done_state = DoneState('done')
start_state = State('start')
start_state.add_transition(done_state)
done_state.add_transition(start_state)
machine = Machine('state', states=[start_state, done_state])
simulation = setup_simulation([machine, _population_fixture('state', 'start'), _population_fixture('count', 0)])
event_time = simulation.clock.time + simulation.clock.step_size
machine.transition(simulation.population.population.index, event_time)
assert np.all(simulation.population.population['count'] == 1)
machine.transition(simulation.population.population.index, event_time)
assert np.all(simulation.population.population['count'] == 1)
machine.transition(simulation.population.population.index, event_time)
assert np.all(simulation.population.population['count'] == 2)
def test_transition():
done_state = State('done')
start_state = State('start')
start_state.add_transition(done_state)
machine = Machine('state', states=[start_state, done_state])
simulation = setup_simulation([machine, _population_fixture('state', 'start')])
event_time = simulation.clock.time + simulation.clock.step_size
machine.transition(simulation.population.population.index, event_time)
assert np.all(simulation.population.population.state == 'done')
def test_interpolated_tables(base_config):
year_start = base_config.time.start.year
year_end = base_config.time.end.year
years = build_table(lambda age, sex, year: year, year_start, year_end)
ages = build_table(lambda age, sex, year: age, year_start, year_end)
one_d_age = ages.copy()
del one_d_age['year']
one_d_age = one_d_age.drop_duplicates()
base_config.population.update({'population_size': 10000})
simulation = setup_simulation([TestPopulation()], input_config=base_config)
manager = simulation.tables
years = manager.build_table(years,
key_columns=('sex', ),
parameter_columns=(
'age',
'year',
))
ages = manager.build_table(ages,
key_columns=('sex', ),
parameter_columns=(
'age',
'year',
))
one_d_age = manager.build_table(one_d_age,
key_columns=('sex', ),
parameter_columns=('age', ))
result_years = years(simulation.population.population.index)
result_ages = ages(simulation.population.population.index)
result_ages_1d = one_d_age(simulation.population.population.index)
fractional_year = simulation.clock.time.year
fractional_year += simulation.clock.time.timetuple().tm_yday / 365.25
assert np.allclose(result_years, fractional_year)
assert np.allclose(result_ages, simulation.population.population.age)
assert np.allclose(result_ages_1d, simulation.population.population.age)
simulation.clock._time += pd.Timedelta(30.5 * 125, unit='D')
simulation.population._population.age += 125 / 12
result_years = years(simulation.population.population.index)
result_ages = ages(simulation.population.population.index)
result_ages_1d = one_d_age(simulation.population.population.index)
fractional_year = simulation.clock.time.year
fractional_year += simulation.clock.time.timetuple().tm_yday / 365.25
assert np.allclose(result_years, fractional_year)
assert np.allclose(result_ages, simulation.population.population.age)
assert np.allclose(result_ages_1d, simulation.population.population.age)
def test_null_transition(base_config):
base_config.population.update({'population_size': 10000})
a_state = State('a')
start_state = State('start')
start_state.add_transition(a_state, probability_func=lambda agents: np.full(len(agents), 0.5))
start_state.allow_self_transitions()
machine = Machine('state', states=[start_state, a_state])
simulation = setup_simulation([machine, _population_fixture('state', 'start')], base_config)
event_time = simulation.clock.time + simulation.clock.step_size
machine.transition(simulation.population.population.index, event_time)
a_count = (simulation.population.population.state == 'a').sum()
assert round(a_count/len(simulation.population.population), 1) == 0.5
def test_no_null_transition(base_config):
base_config.population.update({'population_size': 10000})
a_state = State('a')
b_state = State('b')
start_state = State('start')
a_transition = Transition(start_state, a_state)
b_transition = Transition(start_state, b_state)
start_state.transition_set.allow_null_transition = False
start_state.transition_set.extend((a_transition, b_transition))
machine = Machine('state')
machine.states.extend([start_state, a_state, b_state])
simulation = setup_simulation([machine, _population_fixture('state', 'start')], base_config)
event_time = simulation.clock.time + simulation.clock.step_size
machine.transition(simulation.population.population.index, event_time)
a_count = (simulation.population.population.state == 'a').sum()
assert round(a_count/len(simulation.population.population), 1) == 0.5
def test_interpolated_tables__exact_values_at_input_points(base_config):
year_start = base_config.time.start.year
year_end = base_config.time.end.year
years = build_table(lambda age, sex, year: year, year_start, year_end)
input_years = years.year.unique()
base_config.population.update({'population_size': 10000})
simulation = setup_simulation([TestPopulation()], input_config=base_config)
manager = simulation.tables
years = manager.build_table(years,
key_columns=('sex', ),
parameter_columns=(
'age',
'year',
))
for year in input_years:
simulation.clock._time = pd.Timestamp(year, 1, 1)
assert np.allclose(years(simulation.population.population.index),
simulation.clock.time.year + 1 / 365)