<jupyter_start><jupyter_text>Load the simulation inputs<jupyter_code>params = load_params("baseline")

pmsm = get_parallelizable_msm_criterion(
    prefix="gridsearch",
    fall_start_date=FALL_SIM_START,
    fall_end_date=FALL_SIM_END,
    spring_start_date=SPRING_SIM_START,
    spring_end_date=SPRING_SIM_END,
    mode="fall",
    debug=DEBUG,
)<jupyter_output><empty_output><jupyter_text>Modify Params<jupyter_code>index_bundles = get_index_bundles(params)
hh_probs = index_bundles["hh"]
school_probs = index_bundles["school"]
young_educ_probs = index_bundles["young_educ"]
work_probs = index_bundles["work"]
other_probs = index_bundles["other"]

b117_rate = [("events", "b117_cases_per_100_000", "2021-01-31")]

vacation_probs = params.query("category == 'additional_other_vacation_contact'").index

rapid_test_level = [
    ("rapid_test_demand", "private_demand", "2021-03-31"),
    ("rapid_test_demand", "private_demand", "2021-04-06"),
]

free_probs = [
    bundle[0]
    for bundle in list(index_bundles.values())
    + [vacation_probs, b117_rate, rapid_test_level]
]
params.loc[other_probs, "value"] = 0.15875
params.loc[young_educ_probs, "value"] = 0.005
params.loc[school_probs, "value"] = 0.012
params.loc[work_probs, "value"] = 0.1475
params.loc[hh_probs, "value"] = 0.1
params.loc[vacation_probs, "value"] = 0.5
params.loc[free_probs]
params<jupyter_output><empty_output><jupyter_text>Run estimation<jupyter_code>dimensions = "1d"

n_gridpoints = 16
loc1 = b117_rate
gridspec1 = (0.97, 1.0, n_gridpoints)
# only used if 2d
loc2 = work_probs
gridspec2 = (0.14, 0.16, n_gridpoints)

n_seeds = 20
n_cores = 40
mask = None
if dimensions == "2d":
    results, grid, best_index, fig = run_2d_gridsearch(
        func=pmsm,
        params=params,
        loc1=loc1,
        gridspec1=gridspec1,
        loc2=loc2,
        gridspec2=gridspec2,
        n_seeds=n_seeds,
        n_cores=n_cores,
        mask=mask,
    )
else:
    results, grid, best_index, fig = run_1d_gridsearch(
        func=pmsm,
        params=params,
        loc=loc1,
        gridspec=gridspec1,
        n_seeds=n_seeds,
        n_cores=n_cores,
    )
pd.to_pickle(results, "results.pkl")
fig
ix = best_index
plot_estimation_moment(results[ix], "aggregated_infections_not_log")
plot_estimation_moment(results[ix], "aggregated_b117_share")
plot_estimation_moment(results[ix], "infections_by_age_group")
# plot_infection_channels(results[ix], aggregate=True)
grid[ix]
best_index
np.mean([res["value"] for res in results[best_index]])<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.12672176308539945, 'fraction_numerical': 0.04132231404958678, 'mean_word_length': 5.261083743842365, 'pattern_counts': {'":': 1, '<': 13, '<?xml version=': 0, '>': 13, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '3', 'dir': 'jupyter-structured-clean-dedup', 'id': '494470', 'license': 'MIT', 'n_tokens_mistral': 1072, 'n_tokens_neox': 997, 'n_words': 145, 'path': 'src/exploration/janos_simulation_notebook.ipynb', 'repo_name': 'covid-19-impact-lab/sid-germany', 'size': '2541'}