def test__assign_demography_with_initial_age_zero(config):
pop_data = dt.assign_demographic_proportions(
make_uniform_pop_data(age_bin_midpoint=True))
pop_data = pop_data[pop_data.year_start == 1990]
simulants = make_base_simulants()
initial_age = 0
r = {
k: get_randomness()
for k in ['general_purpose', 'bin_selection', 'age_smoothing']
}
step_size = utilities.to_time_delta(config.time.step_size)
simulants = bp._assign_demography_with_initial_age(
simulants, pop_data, initial_age, step_size, r,
lambda *args, **kwargs: None)
assert len(simulants) == len(simulants.age.unique())
assert simulants.age.min() > initial_age
assert simulants.age.max() < initial_age + utilities.to_years(step_size)
assert math.isclose(len(simulants[simulants.sex == 'Male']) /
len(simulants),
0.5,
abs_tol=0.01)
for location in simulants.location.unique():
assert math.isclose(len(simulants[simulants.location == location]) /
len(simulants),
1 / len(simulants.location.unique()),
abs_tol=0.01)
def on_time_step(self, event):
"""Ages simulants each time step.
Parameters
----------
event : vivarium.framework.population.PopulationEvent
"""
population = self.population_view.get(event.index,
query="alive == 'alive'")
population['age'] += utilities.to_years(event.step_size)
self.population_view.update(population)
def _assign_demography_with_initial_age(simulants, pop_data, initial_age,
step_size, randomness_streams,
register_simulants):
"""Assigns age, sex, and location information to the provided simulants given a fixed age.
Parameters
----------
simulants : pandas.DataFrame
Table that represents the new cohort of agents being added to the simulation.
pop_data : pandas.DataFrame
Table with columns 'age', 'age_start', 'age_end', 'sex', 'year',
'location', 'population', 'P(sex, location, age| year)', 'P(sex, location | age, year)'
initial_age : float
The age to assign the new simulants.
randomness_streams : Dict[str, vivarium.framework.randomness.RandomnessStream]
Source of random number generation within the vivarium common random number framework.
step_size : pandas.Timedelta
The size of the initial time step.
register_simulants : Callable
A function to register the new simulants with the CRN framework.
Returns
-------
pandas.DataFrame
Table with same columns as `simulants` and with the additional columns 'age', 'sex', and 'location'.
"""
pop_data = pop_data[(pop_data.age_start <= initial_age)
& (pop_data.age_end >= initial_age)]
if pop_data.empty:
raise ValueError(
'The age {} is not represented by the population data structure'.
format(initial_age))
age_fuzz = randomness_streams['age_smoothing'].get_draw(
simulants.index) * utilities.to_years(step_size)
simulants['age'] = initial_age + age_fuzz
register_simulants(simulants[['entrance_time', 'age']])
# Assign a demographically accurate location and sex distribution.
choices = pop_data.set_index(
['sex', 'location'])['P(sex, location | age, year)'].reset_index()
decisions = randomness_streams['general_purpose'].choice(
simulants.index,
choices=choices.index,
p=choices['P(sex, location | age, year)'])
simulants['sex'] = choices.loc[decisions, 'sex'].values
simulants['location'] = choices.loc[decisions, 'location'].values
return simulants
def on_time_step(self, event):
"""Adds a set number of simulants to the population each time step.
Parameters
----------
event
The event that triggered the function call.
"""
# Assume immigrants are uniformly distributed throughout the year.
step_size = utilities.to_years(event.step_size)
simulants_to_add = self.simulants_per_year * step_size + self.fractional_new_immigrations
self.fractional_new_immigrations = simulants_to_add % 1
simulants_to_add = int(simulants_to_add)
if simulants_to_add > 0:
self.simulant_creator(simulants_to_add,
population_configuration={
'age_start': 0,
'age_end': 100,
'sim_state': 'time_step_imm',
'immigrated': "Yes"
})
# XXX make sure this does not conflict with fertility XXX
new_residents = self.population_view.get(
event.index + simulants_to_add,
query='sex == "nan" and immigrated != "no_immigration"').copy()
if len(new_residents) > 0:
# sample residents using the immigration rates
sample_resident = self.asfr_data_immigration.sample(
len(new_residents), weights="mean_value", replace=True)
new_residents["sex"] = sample_resident["sex"].values.astype(float)
new_residents["ethnicity"] = sample_resident["ethnicity"].values
new_residents["location"] = sample_resident["location"].values
new_residents["age"] = sample_resident["age_start"].values.astype(
float)
new_residents["immigrated"] = "Yes"
new_residents['MSOA'] = self.assign_MSOA(new_residents)
self.population_view.update(new_residents[[
'immigrated', 'location', 'ethnicity', 'sex', 'age', 'MSOA'
]])
def on_time_step(self, event):
"""Adds a set number of simulants to the population each time step.
Parameters
----------
event
The event that triggered the function call.
"""
# Assume births are uniformly distributed throughout the year.
step_size = utilities.to_years(event.step_size)
simulants_to_add = self.simulants_per_year * step_size + self.fractional_new_births
self.fractional_new_births = simulants_to_add % 1
simulants_to_add = int(simulants_to_add)
if simulants_to_add > 0:
self.simulant_creator(simulants_to_add,
population_configuration={
'age_start': 0,
'age_end': 0,
'sim_state': 'time_step',
})
def on_time_step(self, event):
"""Adds new simulants every time step based on the Crude Birth Rate
and an assumption that birth is a Poisson process
Parameters
----------
event
The event that triggered the function call.
"""
birth_rate = self.birth_rate.at[self.clock().year]
step_size = utilities.to_years(event.step_size)
mean_births = birth_rate * step_size
# Assume births occur as a Poisson process
r = np.random.RandomState(
seed=self.randomness.get_seed('crude_birth_rate'))
simulants_to_add = r.poisson(mean_births)
if simulants_to_add > 0:
self.simulant_creator(simulants_to_add,
population_configuration={
'age_start': 0,
'age_end': 0,
'sim_state': 'time_step',
})