def main(dm: DataManager, config):
"""
Main run function
"""
# get dataframes for each location, age_group and sex_group combination
logging.info('Loading data')
data = get_data(dm)
# get mortality pattern models
logging.info('Setting up models')
mortality_pattern_models = get_mortality_pattern_models(dm, data)
# fit mortality pattern models and predict results
for name, model in mortality_pattern_models.items():
logging.info('Running: {}'.format(name))
model.run_models()
include_ci = (config['n_draws'] > 1)
data[name][1] = model.predict(data[name][1],
col_pred="mortality_pattern",
include_ci=include_ci,
ci_bounds=config['summary_probs'],
num_samples=config['n_draws'],
return_draws=include_ci)
results = {name: dfs[1] for name, dfs in data.items()}
# save the mortality pattern results
logging.info('Saving results')
dm.write_data(results)
logging.info('Done!')
def get_group_specific_data(dm: DataManager, location: str, age_group: str,
sex_group: str) -> List[DataFrame]:
df = dm.read_data_location(location,
group_specs={
"age_name": [age_group],
"sex": [sex_group]
})
data = []
for i in range(2):
df_sub = dm.truncate_time_location(location, df, time_end_id=i)
if i == 0:
df_sub = df_sub[~df_sub.deaths.isna()].reset_index(drop=True)
df_sub["offset_0"] = np.log(df_sub.population)
data.append(df_sub)
return data
def get_group_specific_data(dm: DataManager, location: str, age_group: str,
sex_group: str) -> List[DataFrame]:
"""
Get location-age-sex-specific data
"""
df = dm.read_data_location(location,
group_specs={
"age_name": [age_group],
"sex": [sex_group]
})
data = []
for i in range(2):
df_sub = dm.truncate_time_location(location, df, time_end_id=i)
df_sub["offset_0"] = np.log(df_sub.population)
data.append(df_sub)
return data
def main(dm: DataManager):
# get dataframes for each location, age_group and sex_group combination
data = get_data(dm)
# get mortality pattern models
mortality_pattern_models = get_mortality_pattern_models(dm, data)
# fit mortality pattern models and predict results
for name, model in mortality_pattern_models.items():
model.run_models()
data[name][1] = model.predict(data[name][1],
col_pred="mortality_pattern")
results = {name: dfs[1] for name, dfs in data.items()}
# save the mortality pattern results
dm.write_data(results)
# plot results and save figures
plot_models(dm, results)
def main(dm: DataManager):
# get dataframes for each location, age_group and sex_group combination
data = get_data(dm)
# get mortality pattern models
mortality_pattern_models = get_mortality_pattern_models(dm, data)
# fit mortality pattern models and predict results
for name, model in mortality_pattern_models.items():
model.run_models()
data[name][1] = model.predict(data[name][1], col_pred="cases")
draws = model.get_draws(data[name][1],
col_pred="cases",
num_samples=1000)
for i in range(draws["cases"].shape[0]):
data[name][1][f"cases_draw_{i}"] = draws["cases"][i]
results = {name: dfs[1] for name, dfs in data.items()}
# save the mortality pattern results
dm.write_data(results)
# plot results and save figures
plot_models(dm, results)
def main(dm: DataManager):
# get dataframes for each location, age_group and sex_group combination
data = get_data(dm)
# get mortality pattern models
mortality_pattern_models = get_mortality_pattern_models(dm, data)
# fit mortality pattern models and predict results
for name, model in mortality_pattern_models.items():
model.run_models()
data[name][1] = model.predict(data[name][1],
col_pred="mortality_pattern")
results = {name: dfs[1] for name, dfs in data.items()}
# save the mortality pattern results
dm.write_data(results)
# plot results and save figures
plot_models(dm, results)
if __name__ == "__main__":
# inputs
i_folder = "examples/data"
o_folder = "examples/results"
meta_filename = "meta.yaml"
locations = ["AUT"]
main(DataManager(i_folder, o_folder, meta_filename, locations))