<jupyter_start><jupyter_text>Sample run<jupyter_code>sim = Simulator(graph_dict=graph_dict)
sim.run()
# Code to show graph

aggregated = np.sum(sim.state_var, axis=2)

fig, ax = plt.subplots()
ax.plot(sim.time_series, aggregated[:, sim.c_idx["S"]], "g", label="Susceptibles", linewidth=2)
ax.plot(sim.time_series, aggregated[:, sim.c_idx["E"]], "orange", label="Exposed", linewidth=2)
ax.plot(sim.time_series, aggregated[:, sim.c_idx["I"]], "r", label="Infected", linewidth=2)
ax.plot(sim.time_series, aggregated[:, sim.c_idx["R"]], "b", label="Recovered", linewidth=2)
ax.plot(sim.time_series, aggregated[:, sim.c_idx["D"]], "gray", label="Dead", linewidth=2)
ax.legend()
plt.xlabel("Days")
plt.ylabel("Size of compartments")
plt.title("Time series of the compartments")
plt.tight_layout()
os.makedirs('./plots', exist_ok=True)
plt.savefig("./plots/compartments.pdf",dpi=300)
plt.show()
# Code to show graph
colors = cm.Greens(0.5+np.linspace(0, 0.5, 3))

fig, ax = plt.subplots()
ax.plot(sim.time_series, sim.state_var[:, sim.c_idx["S"], 0],
        color=colors[0], label="Age group 0-14", linewidth=2)
ax.plot(sim.time_series, sim.state_var[:, sim.c_idx["S"], 1],
        color=colors[1], label="Age group 15-59", linewidth=2)
ax.plot(sim.time_series, sim.state_var[:, sim.c_idx["S"], 2],
        color=colors[2], label="Age group 60+", linewidth=2)
ax.legend()
plt.xlabel("Days")
plt.ylabel("Number of susceptibles")
plt.title("Susceptible compartments by age groups")
plt.tight_layout()
os.makedirs('./plots', exist_ok=True)
plt.savefig("./plots/susceptible_age_groups.pdf",dpi=300)
plt.show()<jupyter_output><empty_output><jupyter_text>Multiple runs from a given point<jupyter_code>sim_original = Simulator(graph_dict=graph_dict)
sim_original.run_after_peak(days=10)
time = np.array(sim_original.time_series).searchsorted(35)
time1 = np.array(sim_original.time_series).searchsorted(25)
ts_list = []
i_list = []
for i in range(3):
    sim_1 = Simulator(graph_dict=graph_dict)
    sim_1.reset_for_new_simulation(init=sim_original.state_var[time])
    sim_1.run_after_peak(days=10)
    
    aggregated = np.sum(sim_1.state_var, axis=2)
    
    ts_list.append(sim_1.time_series[:])
    i_list.append(aggregated[:, sim_1.c_idx["I"]])
aggregated = np.sum(sim_original.state_var, axis=2)
i_time_series = sim_original.time_series[time1:time]
i_values = aggregated[time1:time, sim_original.c_idx["I"]]
days, day_i = generate_daily_data(time_data=i_time_series, 
                                  values=i_values, 
                                  day=26)
fig, ax = plt.subplots()
aggregated = np.sum(sim_original.state_var, axis=2)
ax.plot(days, day_i, "m.", linewidth=2, label="Original data")

color = cm.Reds(0.5+np.linspace(0, 0.5, len(ts_list)))
for i, c in zip(range(len(ts_list)), color):
    ax.plot(sim_original.time_series[time]+np.array(ts_list[i]), i_list[i],
            color=c, linewidth=1, label="Simulation #"+str(i+1))
plt.xlabel("Days")
plt.ylabel("Number of infected people")
plt.title("Halting condition: 10 days after the first peak")
plt.legend()
plt.tight_layout()
os.makedirs('./plots', exist_ok=True)
plt.savefig("./plots/stop_after_peak.pdf",dpi=300)
plt.show()<jupyter_output><empty_output><jupyter_text>Contact reduction<jupyter_code>sim_original = Simulator(graph_dict=graph_dict)
sim_original.run_after_peak(days=20)
time1 = np.array(sim_original.time_series).searchsorted(30)
time2 = np.array(sim_original.time_series).searchsorted(45)
ts_list = []
i_list = []
for i in range(1,4):
    sim_1 = Simulator(graph_dict=graph_dict)
    sim_1.reset_for_new_simulation(init=sim_original.state_var[time2])
    sim_1.model.contact_matrix = sim_1.model.contact_matrix * (1 - 0.1 * i)
    sim_1.run_after_peak(days=10)
    
    aggregated = np.sum(sim_1.state_var, axis=2)
    
    ts_list.append(sim_1.time_series[:])
    i_list.append(aggregated[:, sim_1.c_idx["I"]])
fig, ax = plt.subplots()
aggregated = np.sum(sim_original.state_var, axis=2)
i_lista = aggregated[:, sim_original.c_idx["I"]]
ax.plot(sim_original.time_series[time1:], i_lista[time1:], "r",
        linewidth=2, label="No contact reduction")

color = cm.Reds(0.5+np.linspace(0, 0.5, len(ts_list)))
for i, c in zip(range(len(ts_list)), color):
    ax.plot(sim_original.time_series[time2]+np.array(ts_list[i])[time1:], i_list[i][time1:], color=c,
            linewidth=1, label=str((i+1)*10) + "% contact reduction")
plt.xlabel("Days")
plt.ylabel("Number of infected people")
plt.title("Effect of contact reduction")
plt.legend()
plt.tight_layout()
os.makedirs('./plots', exist_ok=True)
plt.savefig("./plots/contact_reduction.pdf",dpi=300)
plt.show()<jupyter_output><empty_output><jupyter_text>Simulation with confidence calculation<jupyter_code>do_conf_simulation = False
graph_dict = json.load(open("params/graph_dict.json"))
sim = Simulator(graph_dict=graph_dict)
if do_conf_simulation:
    sim.run_conf()
# Number of simulations
n_simulations = 104
# List for storing infection curves for all simulations
to_plot = []
# List for storing peak sizes/times from all simulations
peak_sizes = []
peak_times = []
# List for storing daily sampled data
daily_data = []
y_data = []

for sim_cnt in range(n_simulations):
    # Load data
    data=np.load("./data/simulation_" + str(sim_cnt) + ".npz")
    time_series = data['t']
    state_var = data['y']
    aggregated = np.sum(state_var, axis=2)
    
    # Save data for multiple time series plot
    to_plot.append((time_series, aggregated[:, sim.c_idx["I"]]))

    # Generate daily data
    days, day_i = generate_daily_data(time_data=time_series, values=aggregated[:, sim.c_idx["I"]])
    daily_data.append((days, day_i))
    y_data.append(day_i)

    # Save data for confidence calculation
    i_agg_max = np.amax(aggregated[:, sim.c_idx["I"]])
    peak_sizes.append(i_agg_max)
    peak_times.append(time_series[np.argmax(aggregated[:, sim.c_idx["I"]])])
y_data = np.array(y_data)

# Standard deviation calculated for confidence band
m = np.mean(np.array(peak_sizes))
s = np.std(np.array(peak_sizes))
# Plot multiple loaded time series

fig, ax = plt.subplots()
for ts, sv in to_plot:
    ax.plot(ts, sv)
plt.show()
# Plot multiple daily time series

fig, ax = plt.subplots()
for ts, sv in daily_data:
    ax.plot(ts, sv)
plt.show()
os.makedirs('./plots', exist_ok=True)
fig.savefig('./plots/conf_plot.pdf')
# Plot confidence band 1.

y_mean = np.mean(y_data, axis=0)

fig, ax = plt.subplots()
ax.fill_between(daily_data[0][0], y_mean - 2 * s, y_mean + 2 * s, alpha=0.3)
ax.plot(daily_data[0][0], y_mean, "r")
plt.show()
os.makedirs('./plots', exist_ok=True)
fig.savefig('./plots/conf_plot_2.pdf')
# Plot confidence band 2.

y_mean = np.mean(y_data, axis=0)
y_min = np.amin(y_data, axis=0)
y_max = np.amax(y_data, axis=0)

fig, ax = plt.subplots()
ax.fill_between(daily_data[0][0], y_min, y_max, alpha=0.3)
ax.plot(daily_data[0][0], y_mean, "r")
plt.show()
os.makedirs('./plots', exist_ok=True)
fig.savefig('./plots/conf_plot_3.pdf')
# Plot confidence band 3.

t_vec = np.array(daily_data[0][0])
peak = m * np.ones_like(t_vec)
filt = (min(peak_times) < t_vec) & (t_vec < max(peak_times))

fig, ax = plt.subplots()
ax.fill_between(t_vec[filt], peak[filt] - 2 * s, peak[filt] + 2 * s, alpha=0.9)
for ts, sv in daily_data:
    ax.plot(t_vec[filt], np.array(sv)[filt])
ax.set_ylim(28000, 30000)
plt.show()
os.makedirs('./plots', exist_ok=True)
fig.savefig('./plots/conf_plot_4.pdf')
# Plot confidence band 4.

y_mean = np.mean(y_data, axis=0)
peak = np.amax(y_mean) * np.ones_like(y_mean)
t_vec = np.array(daily_data[0][0])
filt = (min(peak_times) < t_vec) & (t_vec < max(peak_times))

fig, ax = plt.subplots()
ax.fill_between(t_vec[filt], peak[filt] - 2 * s, peak[filt] + 2 * s, alpha=0.9)
ax.plot(t_vec[filt], y_mean[filt], "r")
ax.set_ylim(22000, 30000)
plt.show()
os.makedirs('./plots', exist_ok=True)
fig.savefig('./plots/conf_plot_5.pdf')
print(np.array([len(v) for _, v in daily_data]))
print(np.array([len(t) for t, _ in daily_data]))<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1707682698313554, 'fraction_numerical': 0.02236102435977514, 'mean_word_length': 5.394568690095847, 'pattern_counts': {'":': 0, '<': 21, '<?xml version=': 0, '>': 17, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '338988', 'license': 'MIT', 'n_tokens_mistral': 3325, 'n_tokens_neox': 3148, 'n_words': 550, 'path': 'demo.ipynb', 'repo_name': 'zsvizi/propensity-matrix-epidemics', 'size': '8005'}