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 test_generate_population_initial_age(age_bounds_mock, initial_age_mock): creation_time = pd.Timestamp(1990, 7, 2) step_size = pd.Timedelta(days=1) age_params = {'age_start': 0, 'age_end': 0} pop_data = dt.assign_demographic_proportions( make_uniform_pop_data(age_bin_midpoint=True)) r = { k: get_randomness() for k in ['general_purpose', 'bin_selection', 'age_smoothing'] } sims = make_base_simulants() simulant_ids = sims.index bp.generate_population(simulant_ids, creation_time, step_size, age_params, pop_data, r, lambda *args, **kwargs: None) initial_age_mock.assert_called_once() mock_args = initial_age_mock.call_args[0] assert mock_args[0].equals(sims) assert mock_args[1].equals(pop_data) assert mock_args[2] == float(age_params['age_start']) assert mock_args[3] == step_size assert mock_args[4] == r age_bounds_mock.assert_not_called()
def test_BasePopulation(config, base_plugins, generate_population_mock): num_days = 600 time_step = 100 # Days sims = make_full_simulants() start_population_size = len(sims) generate_population_mock.return_value = sims.drop(columns=['tracked']) base_pop = bp.BasePopulation() components = [base_pop] config.update( { 'population': { 'population_size': start_population_size }, 'time': { 'step_size': time_step } }, layer='override') simulation = InteractiveContext(components=components, configuration=config, plugin_configuration=base_plugins) time_start = simulation._clock.time pop_structure = simulation._data.load('population.structure') pop_structure['location'] = simulation.configuration.input_data.location uniform_pop = dt.assign_demographic_proportions(pop_structure) assert base_pop.population_data.equals(uniform_pop) age_params = { 'age_start': config.population.age_start, 'age_end': config.population.age_end } sub_pop = bp.BasePopulation.select_sub_population_data( uniform_pop, time_start.year) generate_population_mock.assert_called_once() # Get a dictionary of the arguments used in the call mock_args = generate_population_mock.call_args[1] assert mock_args[ 'creation_time'] == time_start - simulation._clock.step_size assert mock_args['age_params'] == age_params assert mock_args['population_data'].equals(sub_pop) assert mock_args['randomness_streams'] == base_pop.randomness pop = simulation.get_population() for column in pop: assert pop[column].equals(sims[column]) final_ages = pop.age + num_days / utilities.DAYS_PER_YEAR simulation.run_for(duration=pd.Timedelta(days=num_days)) pop = simulation.get_population() assert np.allclose(pop.age, final_ages, atol=0.5 / utilities.DAYS_PER_YEAR) # Within a half of a day.
def test__get_bins_and_proportions_with_youngest_bin(): pop_data = dt.assign_demographic_proportions( make_uniform_pop_data(age_bin_midpoint=True)) pop_data = pop_data[(pop_data.year_start == 1990) & (pop_data.location == 1) & (pop_data.sex == 'Male')] age = dt.AgeValues(current=2.5, young=0, old=7.5) endpoints, proportions = dt._get_bins_and_proportions(pop_data, age) assert endpoints.left == 0 assert endpoints.right == 5 bin_width = endpoints.right - endpoints.left assert proportions.current == 1 / len(pop_data) / bin_width assert proportions.young == 1 / len(pop_data) / bin_width assert proportions.old == 1 / len(pop_data) / bin_width
def test_rescale_binned_proportions_full_range(): pop_data = dt.assign_demographic_proportions( make_uniform_pop_data(age_bin_midpoint=True)) pop_data = pop_data[pop_data.year_start == 1990] pop_data_scaled = dt.rescale_binned_proportions(pop_data, age_start=0, age_end=100) pop_data_scaled = pop_data_scaled[pop_data_scaled.age.isin( pop_data.age.unique())] assert np.allclose(pop_data['P(sex, location, age| year)'], pop_data_scaled['P(sex, location, age| year)'])
def test_assign_demographic_proportions(): pop_data = dt.assign_demographic_proportions( make_uniform_pop_data(age_bin_midpoint=True)) assert np.allclose(pop_data['P(sex, location, age| year)'], len(pop_data.year_start.unique()) / len(pop_data)) assert np.allclose(pop_data['P(sex, location | age, year)'], (len(pop_data.year_start.unique()) * len(pop_data.age.unique()) / len(pop_data))) assert np.allclose( pop_data['P(age | year, sex, location)'], (len(pop_data.year_start.unique()) * len(pop_data.sex.unique()) * len(pop_data.location.unique()) / len(pop_data)))
def test__assign_demography_with_age_bounds_error(): pop_data = dt.assign_demographic_proportions( make_uniform_pop_data(age_bin_midpoint=True)) simulants = make_base_simulants() age_start, age_end = 110, 120 r = { k: get_randomness() for k in ['general_purpose', 'bin_selection', 'age_smoothing'] } with pytest.raises(ValueError): bp._assign_demography_with_age_bounds(simulants, pop_data, age_start, age_end, r, lambda *args, **kwargs: None)
def test_rescale_binned_proportions_age_bin_edges(): pop_data = dt.assign_demographic_proportions( make_uniform_pop_data(age_bin_midpoint=True)) pop_data = pop_data[pop_data.year_start == 1990] # Test edge case where age_start/age_end fall on age bin boundaries. pop_data_scaled = dt.rescale_binned_proportions(pop_data, age_start=5, age_end=10) assert len(pop_data_scaled.age.unique()) == len(pop_data.age.unique()) + 2 assert 7.5 in pop_data_scaled.age.unique() correct_data = ([1 / len(pop_data)] * (len(pop_data_scaled) // 2 - 2) + [0, 0]) * 2 assert np.allclose(pop_data_scaled['P(sex, location, age| year)'], correct_data)
def test__assign_demography_with_initial_age_error(): 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 = 200 r = { k: get_randomness() for k in ['general_purpose', 'bin_selection', 'age_smoothing'] } step_size = pd.Timedelta(days=1) with pytest.raises(ValueError): bp._assign_demography_with_initial_age(simulants, pop_data, initial_age, step_size, r, lambda *args, **kwargs: None)
def test_rescale_binned_proportions_clipped_ends(): pop_data = dt.assign_demographic_proportions( make_uniform_pop_data(age_bin_midpoint=True)) pop_data = pop_data[pop_data.year_start == 1990] scale = len(pop_data.location.unique()) * len(pop_data.sex.unique()) pop_data_scaled = dt.rescale_binned_proportions(pop_data, age_start=2, age_end=7) base_p = 1 / len(pop_data) p_scaled = [ base_p * 7 / 5, base_p * 3 / 5, base_p * 2 / 5, base_p * 8 / 5 ] + [base_p] * (len(pop_data_scaled) // scale - 5) + [0] for group, sub_population in pop_data_scaled.groupby(['sex', 'location']): assert np.allclose(sub_population['P(sex, location, age| year)'], p_scaled)
def test_smooth_ages(): pop_data = dt.assign_demographic_proportions( make_uniform_pop_data(age_bin_midpoint=True)) pop_data = pop_data[pop_data.year_start == 1990] simulants = pd.DataFrame({ 'age': [22.5] * 10000 + [52.5] * 10000, 'sex': ['Male', 'Female'] * 10000, 'location': [1, 2] * 10000 }) randomness = get_randomness() smoothed_simulants = dt.smooth_ages(simulants, pop_data, randomness) assert math.isclose(len(smoothed_simulants.age.unique()), len(smoothed_simulants.index), abs_tol=1) # Tolerance is 3*std_dev of the sample mean assert math.isclose(smoothed_simulants.age.mean(), 37.5, abs_tol=3 * math.sqrt(13.149778198**2 / 2000))
def test__assign_demography_with_age_bounds(): 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() age_start, age_end = 0, 180 r = { k: get_randomness(k) for k in [ 'general_purpose', 'bin_selection', 'age_smoothing', 'age_smoothing_age_bounds' ] } simulants = bp._assign_demography_with_age_bounds( simulants, pop_data, age_start, age_end, r, lambda *args, **kwargs: None) 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) ages = np.sort(simulants.age.values) age_deltas = ages[1:] - ages[:-1] age_bin_width = 5 # See `make_uniform_pop_data` num_bins = len(pop_data.age.unique()) n = len(simulants) assert math.isclose(age_deltas.mean(), age_bin_width * num_bins / n, rel_tol=1e-3) assert age_deltas.max( ) < 100 * age_bin_width * num_bins / n # Make sure there are no big age gaps.