def fit_all_case_data(num_procs=4): pool = multi.Pool(num_procs) print(f"Pooling with {num_procs} processors") case_counts = parse_tsv() scenario_data = load_population_data() age_distributions = load_distribution() params = [] for region in case_counts: if region in scenario_data: params.append([ region, case_counts[region], scenario_data.get(region, None), age_distributions[scenario_data[region]['ages']], False ]) results = pool.map(fit_population, params) results_dict = {} for k, params in results: if params is None: results_dict[k] = None elif np.isfinite(params['logInitial']): results_dict[k] = params else: results_dict[k] = None return results_dict
with open(output_json, "w+") as fd: output = ScenarioArray(scenarios) output.marshalJSON(fd) if __name__ == '__main__': generate('test.json', recalculate=False) from scripts.test_fitting_procedure import generate_data, check_fit from scripts.model import trace_ages, get_IFR from matplotlib import pyplot as plt case_counts = parse_tsv() scenario_data = load_population_data() age_distributions = load_distribution() # region = 'JPN-Kagawa' region = 'United States of America' # region = 'Germany' region = 'Switzerland' region = 'USA-Texas' age_dis = age_distributions[scenario_data[region]['ages']] region, p, fit_params = fit_population( (region, case_counts[region], scenario_data[region], age_dis, True)) model_data = generate_data(fit_params) model_cases = model_data['cases'][7:] - model_data['cases'][:-7] model_deaths = model_data['deaths'][7:] - model_data['deaths'][:-7] model_time = fit_params.time[7:] cases = cumulative_to_rolling_average(