Simulating an epidemic

%23%20%2F%2F%2F%20script%0A%23%20requires-python%20%3D%20%22%3E%3D3.10%2C%3C3.15%22%0A%23%0A%23%20%5Btool.marimo-studio%5D%0A%23%20default%20%3D%20%22animation%22%0A%23%20runtimes%20%3D%20%5B%22server%22%5D%0A%23%20%2F%2F%2F%0A%0Aimport%20marimo%0A%0A__generated_with%20%3D%20%220.25.1%22%0Aapp%20%3D%20marimo.App(width%3D%22medium%22%2C%20app_title%3D%22Simulating%20an%20epidemic%22)%0A%0A%0A%40app.cell%0Adef%20_()%3A%0A%20%20%20%20import%20marimo%20as%20mo%0A%20%20%20%20import%20matplotlib.pyplot%20as%20plt%0A%20%20%20%20import%20numpy%20as%20np%0A%0A%20%20%20%20return%20mo%2C%20np%2C%20plt%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20introduction(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%20Simulating%20an%20epidemic%0A%0A%20%20%20%20Two%20hundred%20people%20wander%20around%20a%20city.%20Each%20is%20susceptible%2C%20infectious%2C%0A%20%20%20%20or%20recovered.%20An%20infectious%20person%20infects%20each%20susceptible%20neighbour%0A%20%20%20%20within%20the%20infection%20radius%20at%20a%20fixed%20rate%20per%20day%20and%20recovers%20after%0A%20%20%20%20five%20days.%20People%20who%20keep%20their%20distance%20push%20away%20from%20their%20ten%0A%20%20%20%20nearest%20neighbours.%0A%0A%20%20%20%20The%20classic%20SIR%20model%20averages%20this%20into%20three%20differential%20equations.%0A%20%20%20%20Simulating%20every%20person%20keeps%20what%20the%20averages%20hide%3A%20who%20infected%20whom%2C%0A%20%20%20%20and%20how%20a%20local%20habit%20bends%20the%20curve%20of%20the%20whole%20city.%0A%0A%20%20%20%20Adapted%20from%203Blue1Brown's%0A%20%20%20%20%5B*Simulating%20an%20epidemic*%5D(https%3A%2F%2Fwww.youtube.com%2Fwatch%3Fv%3DgxAaO2rsdIs)%0A%20%20%20%20and%20its%20%60SIRSimulation%60%20in%0A%20%20%20%20%5B3b1b%2Fvideos%5D(https%3A%2F%2Fgithub.com%2F3b1b%2Fvideos%2Fblob%2Fmaster%2F_2020%2Fsir.py).%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20controls(mo)%3A%0A%20%20%20%20distancing%20%3D%20mo.ui.slider(%0A%20%20%20%20%20%20%20%200%2C%0A%20%20%20%20%20%20%20%200.75%2C%0A%20%20%20%20%20%20%20%20step%3D0.25%2C%0A%20%20%20%20%20%20%20%20value%3D0%2C%0A%20%20%20%20%20%20%20%20label%3D%22Share%20keeping%20their%20distance%22%2C%0A%20%20%20%20%20%20%20%20show_value%3DTrue%2C%0A%20%20%20%20)%0A%20%20%20%20mo.hstack(%5Bdistancing%5D%2C%20justify%3D%22start%22)%0A%20%20%20%20return%20(distancing%2C)%0A%0A%0A%40app.cell%0Adef%20city()%3A%0A%20%20%20%20population%20%3D%20200%0A%20%20%20%20city_size%20%3D%207%0A%20%20%20%20infection_radius%20%3D%200.5%0A%20%20%20%20infection_rate%20%3D%200.2%20%20%23%20per%20day%2C%20per%20infectious%20neighbour%0A%20%20%20%20infectious_days%20%3D%205%0A%20%20%20%20days%20%3D%2024%0A%20%20%20%20%23%20A%20static%20export%20keeps%20the%20samples%20of%20every%20distancing%20level%2C%20so%20the%20city%0A%20%20%20%20%23%20records%20positions%20twice%20a%20day%20and%20the%20scene%20draws%20curves%20between%20them.%0A%20%20%20%20samples_per_day%20%3D%202%0A%20%20%20%20return%20(%0A%20%20%20%20%20%20%20%20city_size%2C%0A%20%20%20%20%20%20%20%20days%2C%0A%20%20%20%20%20%20%20%20infection_radius%2C%0A%20%20%20%20%20%20%20%20infection_rate%2C%0A%20%20%20%20%20%20%20%20infectious_days%2C%0A%20%20%20%20%20%20%20%20population%2C%0A%20%20%20%20%20%20%20%20samples_per_day%2C%0A%20%20%20%20)%0A%0A%0A%40app.cell%0Adef%20model(%0A%20%20%20%20city_size%2C%0A%20%20%20%20days%2C%0A%20%20%20%20infection_radius%2C%0A%20%20%20%20infection_rate%2C%0A%20%20%20%20infectious_days%2C%0A%20%20%20%20np%2C%0A%20%20%20%20population%2C%0A%20%20%20%20samples_per_day%2C%0A)%3A%0A%20%20%20%20def%20simulate(distancing%2C%20seed%3D1%2C%20steps_per_day%3D20)%3A%0A%20%20%20%20%20%20%20%20%22%22%22Run%20the%20city%20for%20%60%60days%60%60%20and%20return%20every%20person's%20path%20and%20status.%0A%0A%20%20%20%20%20%20%20%20Each%20day%2C%20every%20person%20drifts%20toward%20a%20new%20gravity%20well%20one%20unit%20away%20in%0A%20%20%20%20%20%20%20%20a%20random%20direction%2C%20bounces%20off%20the%20city%20walls%2C%20and%20caps%20their%20speed%2C%20as%0A%20%20%20%20%20%20%20%20in%203Blue1Brown's%20model.%0A%20%20%20%20%20%20%20%20%22%22%22%0A%20%20%20%20%20%20%20%20rng%20%3D%20np.random.default_rng(seed)%0A%20%20%20%20%20%20%20%20half%2C%20dt%20%3D%20city_size%20%2F%202%2C%201%20%2F%20steps_per_day%0A%20%20%20%20%20%20%20%20position%20%3D%20rng.uniform(-half%20%2B%200.3%2C%20half%20-%200.3%2C%20(population%2C%202))%0A%20%20%20%20%20%20%20%20velocity%20%3D%20np.zeros_like(position)%0A%20%20%20%20%20%20%20%20well%20%3D%20position.copy()%0A%20%20%20%20%20%20%20%20keeps_distance%20%3D%20rng.random(population)%20%3C%20distancing%0A%20%20%20%20%20%20%20%20infected%20%3D%20np.full(population%2C%20np.inf)%0A%20%20%20%20%20%20%20%20source%20%3D%20np.full(population%2C%20-1)%0A%20%20%20%20%20%20%20%20infected%5Bnp.argmin(np.linalg.norm(position%2C%20axis%3D1))%5D%20%3D%200%0A%20%20%20%20%20%20%20%20paths%2C%20history%20%3D%20%5B%5D%2C%20%5B%5D%0A%0A%20%20%20%20%20%20%20%20for%20step%20in%20range(days%20*%20steps_per_day%20%2B%201)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20time%20%3D%20step%20*%20dt%0A%20%20%20%20%20%20%20%20%20%20%20%20sick%20%3D%20(infected%20%3C%3D%20time)%20%26%20(time%20%3C%20infected%20%2B%20infectious_days)%0A%20%20%20%20%20%20%20%20%20%20%20%20if%20step%20%25%20(steps_per_day%20%2F%2F%20samples_per_day)%20%3D%3D%200%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20paths.append(position.copy())%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20recovered%20%3D%20infected%20%2B%20infectious_days%20%3C%3D%20time%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20history.append(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%5B1%20-%20np.mean(sick%20%7C%20recovered)%2C%20sick.mean()%2C%20recovered.mean()%5D%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20)%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20if%20step%20%25%20steps_per_day%20%3D%3D%200%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20angle%20%3D%20rng.uniform(0%2C%202%20*%20np.pi%2C%20population)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20well%20%3D%20position%20%2B%20np.column_stack(%5Bnp.cos(angle)%2C%20np.sin(angle)%5D)%0A%20%20%20%20%20%20%20%20%20%20%20%20to_well%20%3D%20well%20-%20position%0A%20%20%20%20%20%20%20%20%20%20%20%20gap%20%3D%20np.maximum(np.linalg.norm(to_well%2C%20axis%3D1%2C%20keepdims%3DTrue)%2C%200.05)%0A%20%20%20%20%20%20%20%20%20%20%20%20force%20%3D%200.2%20*%20to_well%20%2F%20gap**3%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20offset%20%3D%20position%5BNone%2C%20%3A%2C%20%3A%5D%20-%20position%5B%3A%2C%20None%2C%20%3A%5D%0A%20%20%20%20%20%20%20%20%20%20%20%20distance%20%3D%20np.linalg.norm(offset%2C%20axis%3D2)%20%2B%20np.diag(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20np.full(population%2C%20np.inf)%0A%20%20%20%20%20%20%20%20%20%20%20%20)%0A%20%20%20%20%20%20%20%20%20%20%20%20if%20keeps_distance.any()%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20rows%20%3D%20np.flatnonzero(keeps_distance)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20nearest%20%3D%20np.argpartition(distance%5Brows%5D%2C%2010%2C%20axis%3D1)%5B%3A%2C%20%3A10%5D%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20near%20%3D%20offset%5Brows%5B%3A%2C%20None%5D%2C%20nearest%5D%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20near_gap%20%3D%20np.maximum(distance%5Brows%5B%3A%2C%20None%5D%2C%20nearest%5D%2C%200.05)%5B...%2C%20None%5D%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20force%5Brows%5D%20-%3D%202%20*%20(near%20%2F%20near_gap**3).sum(axis%3D1)%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20force%20%2B%3D%20np.maximum(1%20%2F%20np.maximum(position%20%2B%20half%2C%201e-3)%20-%201%2C%200)%0A%20%20%20%20%20%20%20%20%20%20%20%20force%20-%3D%20np.maximum(1%20%2F%20np.maximum(half%20-%20position%2C%201e-3)%20-%201%2C%200)%0A%20%20%20%20%20%20%20%20%20%20%20%20velocity%20%2B%3D%20force%20*%20dt%0A%20%20%20%20%20%20%20%20%20%20%20%20speed%20%3D%20np.linalg.norm(velocity%2C%20axis%3D1%2C%20keepdims%3DTrue)%0A%20%20%20%20%20%20%20%20%20%20%20%20velocity%20*%3D%20np.minimum(1%2C%200.5%20%2F%20np.maximum(speed%2C%201e-9))%0A%20%20%20%20%20%20%20%20%20%20%20%20position%20%3D%20np.clip(position%20%2B%20velocity%20*%20dt%2C%20-half%2C%20half)%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20contact%20%3D%20(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20(distance%20%3C%20infection_radius)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%26%20sick%5BNone%2C%20%3A%5D%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%26%20~np.isfinite(infected)%5B%3A%2C%20None%5D%0A%20%20%20%20%20%20%20%20%20%20%20%20)%0A%20%20%20%20%20%20%20%20%20%20%20%20caught%20%3D%20contact%20%26%20(rng.random(contact.shape)%20%3C%20infection_rate%20*%20dt)%0A%20%20%20%20%20%20%20%20%20%20%20%20new%20%3D%20np.flatnonzero(caught.any(axis%3D1))%0A%20%20%20%20%20%20%20%20%20%20%20%20source%5Bnew%5D%20%3D%20caught%5Bnew%5D.argmax(axis%3D1)%0A%20%20%20%20%20%20%20%20%20%20%20%20infected%5Bnew%5D%20%3D%20time%0A%0A%20%20%20%20%20%20%20%20return%20np.array(paths)%2C%20infected%2C%20source%2C%20np.array(history)%0A%0A%20%20%20%20return%20(simulate%2C)%0A%0A%0A%40app.cell%0Adef%20simulation(days%2C%20distancing%2C%20infectious_days%2C%20np%2C%20population%2C%20simulate)%3A%0A%20%20%20%20positions%2C%20infected%2C%20source%2C%20history%20%3D%20simulate(distancing.value)%0A%20%20%20%20caught%20%3D%20np.isfinite(infected)%0A%20%20%20%20cases%20%3D%20np.bincount(source%5Bsource%20%3E%3D%200%5D%2C%20minlength%3Dpopulation)%0A%20%20%20%20%23%20Each%20person's%20path%20as%20x0%2C%20y0%2C%20x1%2C%20y1%2C%20and%20so%20on.%0A%20%20%20%20paths%20%3D%20positions.round(2).transpose(1%2C%200%2C%202).reshape(population%2C%20-1).tolist()%0A%0A%20%20%20%20def%20day(time)%3A%0A%20%20%20%20%20%20%20%20return%20round(float(time)%2C%202)%20if%20np.isfinite(time)%20else%20None%0A%0A%20%20%20%20people%20%3D%20%5B%0A%20%20%20%20%20%20%20%20%7B%0A%20%20%20%20%20%20%20%20%20%20%20%20%22infected%22%3A%20day(infected%5Bk%5D)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%22recovered%22%3A%20day(infected%5Bk%5D%20%2B%20infectious_days)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%22source%22%3A%20int(source%5Bk%5D)%20if%20source%5Bk%5D%20%3E%3D%200%20else%20None%2C%0A%20%20%20%20%20%20%20%20%7D%0A%20%20%20%20%20%20%20%20for%20k%20in%20range(population)%0A%20%20%20%20%5D%0A%20%20%20%20peak%20%3D%20int(history%5B%3A%2C%201%5D.argmax())%0A%20%20%20%20first_cases%20%3D%20np.argsort(infected)%5B%3A15%5D%0A%20%20%20%20summary%20%3D%20%7B%0A%20%20%20%20%20%20%20%20%22peak%22%3A%20float(history%5Bpeak%2C%201%5D)%2C%0A%20%20%20%20%20%20%20%20%22peak_day%22%3A%20peak%20*%20days%20%2F%20(len(history)%20-%201)%2C%0A%20%20%20%20%20%20%20%20%22infected%22%3A%20float(caught.mean())%2C%0A%20%20%20%20%20%20%20%20%22reproduction%22%3A%20float(cases%5Bfirst_cases%5Bcaught%5Bfirst_cases%5D%5D%5D.mean())%2C%0A%20%20%20%20%7D%0A%20%20%20%20history%20%3D%20history.round(4).tolist()%0A%20%20%20%20return%20(summary%2C)%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20outbreak(mo%2C%20summary)%3A%0A%20%20%20%20mo.md(f%22%22%22%0A%20%20%20%20At%20its%20peak%20on%20day%20**%7Bsummary%5B%22peak_day%22%5D%3A.1f%7D**%2C%20**%7Bsummary%5B%22peak%22%5D%3A.0%25%7D**%0A%20%20%20%20of%20the%20city%20is%20infectious%20at%20once.%20**%7Bsummary%5B%22infected%22%5D%3A.0%25%7D**%20catch%20it%20in%0A%20%20%20%20total%2C%20and%20each%20of%20the%20first%20cases%20infects%20**%7Bsummary%5B%22reproduction%22%5D%3A.1f%7D**%0A%20%20%20%20others%20on%20average.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20flattening(days%2C%20distancing%2C%20mo%2C%20np%2C%20plt%2C%20simulate)%3A%0A%20%20%20%20levels%20%3D%20%5B0%2C%200.25%2C%200.5%2C%200.75%5D%0A%20%20%20%20flattening_figure%2C%20axes%20%3D%20plt.subplots(figsize%3D(7%2C%203.4))%0A%20%20%20%20for%20level%20in%20levels%3A%0A%20%20%20%20%20%20%20%20_%2C%20_%2C%20_%2C%20curve%20%3D%20simulate(level)%0A%20%20%20%20%20%20%20%20current%20%3D%20np.isclose(level%2C%20distancing.value)%0A%20%20%20%20%20%20%20%20axes.plot(%0A%20%20%20%20%20%20%20%20%20%20%20%20np.linspace(0%2C%20days%2C%20len(curve))%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20curve%5B%3A%2C%201%5D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20color%3D%22%23d3453d%22%20if%20current%20else%20%22%239aa5b1%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20linewidth%3D2.4%20if%20current%20else%201.4%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20label%3Df%22%7Blevel%3A.0%25%7D%20keep%20their%20distance%22%2C%0A%20%20%20%20%20%20%20%20)%0A%20%20%20%20axes.set(xlabel%3D%22Day%22%2C%20ylabel%3D%22Share%20infectious%22%2C%20xlim%3D(0%2C%20days)%2C%20ylim%3D(0%2C%20None))%0A%20%20%20%20axes.yaxis.set_major_formatter(lambda%20share%2C%20_%3A%20f%22%7Bshare%3A.0%25%7D%22)%0A%20%20%20%20axes.legend(frameon%3DFalse)%0A%20%20%20%20axes.spines%5B%5B%22top%22%2C%20%22right%22%5D%5D.set_visible(False)%0A%20%20%20%20flattening_figure.tight_layout()%0A%20%20%20%20mo.vstack(%0A%20%20%20%20%20%20%20%20%5B%0A%20%20%20%20%20%20%20%20%20%20%20%20flattening_figure%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20mo.md(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22Each%20line%20runs%20the%20same%20city%20with%20a%20different%20share%20of%20people%20%22%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22keeping%20their%20distance.%20The%20highlighted%20run%20is%20the%20one%20above.%22%0A%20%20%20%20%20%20%20%20%20%20%20%20)%2C%0A%20%20%20%20%20%20%20%20%5D%0A%20%20%20%20)%0A%20%20%20%20return%0A%0A%0Aif%20__name__%20%3D%3D%20%22__main__%22%3A%0A%20%20%20%20app.run()%0A