|
| 1 | +import marimo |
| 2 | + |
| 3 | +__generated_with = "0.20.4" |
| 4 | +app = marimo.App(width="medium") |
| 5 | + |
| 6 | + |
| 7 | +@app.cell |
| 8 | +def _(): |
| 9 | + import marimo as mo |
| 10 | + import matplotlib.pyplot as plt |
| 11 | + import numpy as np |
| 12 | + |
| 13 | + return mo, np, plt |
| 14 | + |
| 15 | + |
| 16 | +@app.cell(hide_code=True) |
| 17 | +def _(mo): |
| 18 | + mo.md(r""" |
| 19 | + # Simulating Daily Incidences with Flow-Based Models |
| 20 | +
|
| 21 | + ## Introduction |
| 22 | +
|
| 23 | + In Tutorial 1, we learned how to set up and simulate an ODE-based SECIR-type model. The output of the standard simulation gave us the exact number of individuals in each compartment (e.g., Susceptible, InfectedSymptoms, Recovered) at any given time. |
| 24 | +
|
| 25 | + However, public health data usually reports *incidence* -- such as **daily new cases** or **daily new hospitalizations**. If we simply look at the difference in the `InfectedSymptoms` compartment between today and yesterday, we do not get the true number of new cases. This is because, during that same day, other individuals may have recovered or progressed to a more severe state, leaving the compartment. |
| 26 | +
|
| 27 | + To obtain the exact number of new transitions (flows) between compartments, MEmilio provides a **flow-based formulation**. |
| 28 | +
|
| 29 | + The computational overhead of this flow-based formulation is minimal. The transition rates between compartments have to be evaluated at every step in the standard compartmental formulation anyway. Therefore, calculating the cumulative flows simultaneously only adds a small amount of overhead due to the mapping of these variables and a slightly higher memory usage to store the additional values. |
| 30 | +
|
| 31 | + In this tutorial, we will show how to use `simulate_flows` to obtain these transition counts. |
| 32 | + """) |
| 33 | + return |
| 34 | + |
| 35 | + |
| 36 | +@app.cell(hide_code=True) |
| 37 | +def _(mo): |
| 38 | + mo.md(r""" |
| 39 | + ## Model Setup |
| 40 | +
|
| 41 | + We set up the exact same model as in Tutorial 1. We define the simulation timeframe, initialize the model parameters, and set the initial populations. |
| 42 | + """) |
| 43 | + return |
| 44 | + |
| 45 | + |
| 46 | +@app.cell |
| 47 | +def _(np): |
| 48 | + import memilio.simulation.osecir as osecir |
| 49 | + from memilio.simulation import AgeGroup, LogLevel, set_log_level |
| 50 | + set_log_level(LogLevel.Off) |
| 51 | + |
| 52 | + # Initialize total population, simulation start time, timeframe, and step size |
| 53 | + total_population = 100000 |
| 54 | + t0 = 0 |
| 55 | + tmax = 100 |
| 56 | + dt = 0.1 |
| 57 | + |
| 58 | + # Create model with one age group |
| 59 | + model = osecir.Model(1) |
| 60 | + group = AgeGroup(0) |
| 61 | + |
| 62 | + # Set infection state stay times (in days) |
| 63 | + model.parameters.TimeExposed[group] = 3.2 |
| 64 | + model.parameters.TimeInfectedNoSymptoms[group] = 2. |
| 65 | + model.parameters.TimeInfectedSymptoms[group] = 6. |
| 66 | + model.parameters.TimeInfectedSevere[group] = 12. |
| 67 | + model.parameters.TimeInfectedCritical[group] = 8. |
| 68 | + |
| 69 | + # Set infection state transition probabilities |
| 70 | + model.parameters.RelativeTransmissionNoSymptoms[group] = 0.67 |
| 71 | + model.parameters.TransmissionProbabilityOnContact[group] = 0.1 |
| 72 | + model.parameters.RecoveredPerInfectedNoSymptoms[group] = 0.2 |
| 73 | + model.parameters.RiskOfInfectionFromSymptomatic[group] = 0.25 |
| 74 | + model.parameters.SeverePerInfectedSymptoms[group] = 0.2 |
| 75 | + model.parameters.CriticalPerSevere[group] = 0.25 |
| 76 | + model.parameters.DeathsPerCritical[group] = 0.3 |
| 77 | + |
| 78 | + # Set contact frequency |
| 79 | + model.parameters.ContactPatterns.cont_freq_mat[0].baseline = np.ones((1, 1)) * 10 |
| 80 | + |
| 81 | + # Initialize compartments |
| 82 | + model.populations[group, osecir.InfectionState.Exposed] = 0.005 * total_population |
| 83 | + model.populations[group, osecir.InfectionState.InfectedNoSymptoms] = 0.005 * total_population |
| 84 | + model.populations.set_difference_from_total( |
| 85 | + (group, osecir.InfectionState.Susceptible), total_population) |
| 86 | + return dt, model, osecir, t0, tmax |
| 87 | + |
| 88 | + |
| 89 | +@app.cell(hide_code=True) |
| 90 | +def _(mo): |
| 91 | + mo.md(r""" |
| 92 | + Before running the simulation, we check if all initial values and parameters are within a valid range. |
| 93 | + *Note: MEmilio's `check_constraints()` returns `True` if a constraint is violated, and `False` if everything is fine.* |
| 94 | + """) |
| 95 | + return |
| 96 | + |
| 97 | + |
| 98 | +@app.cell |
| 99 | +def _(model): |
| 100 | + constraints_violated = model.check_constraints() |
| 101 | + return |
| 102 | + |
| 103 | + |
| 104 | +@app.cell(hide_code=True) |
| 105 | +def _(mo): |
| 106 | + mo.md(r""" |
| 107 | + ## Simulating Flows |
| 108 | +
|
| 109 | + Instead of using `osecir.simulate`, we now use `osecir.simulate_flows`. |
| 110 | +
|
| 111 | + This function integrates the transition rates between compartments. The result is a vector which containts the compartment states as first element, and the cumulative flows between compartments as the second element. The cumulative flows are the total number of transitions that have occurred between compartments up to that point in time. The compartment states are the same as the output of `osecir.simulate`. Note that the entries of the vectors are `TimeSeries` objects. |
| 112 | + """) |
| 113 | + return |
| 114 | + |
| 115 | + |
| 116 | +@app.cell |
| 117 | +def _(dt, model, osecir, t0, tmax): |
| 118 | + # Simulate flows instead of just compartment states |
| 119 | + results = osecir.simulate_flows(t0, tmax, dt, model) |
| 120 | + return (results,) |
| 121 | + |
| 122 | + |
| 123 | +@app.cell |
| 124 | +def _(results): |
| 125 | + compartments = results[0] # The first element of results contains the compartment data |
| 126 | + flows = results[1] # The second element of results contains the flow data |
| 127 | + return compartments, flows |
| 128 | + |
| 129 | + |
| 130 | +@app.cell(hide_code=True) |
| 131 | +def _(mo): |
| 132 | + mo.md(r""" |
| 133 | + ## Extracting Daily Incidences |
| 134 | +
|
| 135 | + The `flow_result` contains continuous cumulative data. To extract daily metrics (like daily new cases), we need to: |
| 136 | + 1. Interpolate the flow results to full integer days. |
| 137 | + 2. Calculate the difference between consecutive days to get the discrete daily incidence. |
| 138 | +
|
| 139 | + Since the flow values are stored in a `TimeSeries` object, we can directly use the interpolation method to achieve this. |
| 140 | +
|
| 141 | + Note: the adaptive integrator may make time steps which are larger than 1 day. Following, `interpolate_simulation_result` then fills missing days by linear interpolation of the cumulative flows, producing constant values in the daily differences. |
| 142 | + """) |
| 143 | + return |
| 144 | + |
| 145 | + |
| 146 | +@app.cell |
| 147 | +def _(flows, np, osecir): |
| 148 | + # 1. Interpolate flows to full days |
| 149 | + interpolated_flows = osecir.interpolate_simulation_result(flows) |
| 150 | + |
| 151 | + # Transform the interpolated flows into a NumPy array |
| 152 | + flow_array = interpolated_flows.as_ndarray() |
| 153 | + |
| 154 | + # Row 0 contains the time points |
| 155 | + time_days = flow_array[0, :] |
| 156 | + |
| 157 | + # 2. Calculate daily differences |
| 158 | + # We slice [1:, :] to remove the time row. Now row 0 is Flow 0! |
| 159 | + daily_flows = np.diff(flow_array[1:, :], axis=1) |
| 160 | + |
| 161 | + # Slice time_days to match the dimension of daily_flows (N-1) |
| 162 | + plot_time = time_days[1:] |
| 163 | + |
| 164 | + print(flow_array.shape) |
| 165 | + print("ji") |
| 166 | + return daily_flows, plot_time |
| 167 | + |
| 168 | + |
| 169 | +@app.cell(hide_code=True) |
| 170 | +def _(mo): |
| 171 | + mo.md(r""" |
| 172 | + ## Extracting Specific Transitions |
| 173 | +
|
| 174 | + To plot specific transitions, we need to extract the correct rows from our `daily_flows` array. MEmilio flattens all possible flows into a 1D list. The sorting follows two clear rules: |
| 175 | +
|
| 176 | + 1. **By Age Group:** The array contains all flows for the first age group, followed by all flows for the second age group, and so on. |
| 177 | + 2. **By Infection Progression:** Within each age group, the sequence strictly follows the chronological progression of the disease (including parallel paths for confirmed cases). |
| 178 | +
|
| 179 | + For a model with $N_G$ age groups and $N_{T_G}$ flows per group, the flat index for flow $f$ in age group $g$ is simply calculated as $g \times N_{T_G} + f$. |
| 180 | +
|
| 181 | + Since our tutorial model only has one age group ($g=0$), we just need the base indices of the flows. The SECIR model defines 15 specific flows. The first few are ordered like this: |
| 182 | + * Index 0: `Susceptible` $\rightarrow$ `Exposed` (Total new infections) |
| 183 | + * Index 1: `Exposed` $\rightarrow$ `InfectedNoSymptoms` |
| 184 | + * **Index 2: `InfectedNoSymptoms` $\rightarrow$ `InfectedSymptoms` (New symptomatic cases)** |
| 185 | + * Index 3: `InfectedNoSymptoms` $\rightarrow$ `Recovered` |
| 186 | + * ... |
| 187 | + * **Index 6: `InfectedSymptoms` $\rightarrow$ `InfectedSevere` (New hospitalizations)** |
| 188 | +
|
| 189 | + The full list of flows can be found in the [MEmilio documentation](https://memilio.readthedocs.io/en/latest/cpp/models/osecir.html#flows). |
| 190 | + We will now hardcode indices 2 and 6 to extract the daily incidences for new symptomatic cases and new hospitalizations. |
| 191 | + """) |
| 192 | + return |
| 193 | + |
| 194 | + |
| 195 | +@app.cell |
| 196 | +def _(compartments, daily_flows, osecir, plot_time, plt): |
| 197 | + # 1. Prepare compartment data for comparison |
| 198 | + # Interpolate to match the time grid of our flows |
| 199 | + comp_array = osecir.interpolate_simulation_result(compartments).as_ndarray() |
| 200 | + |
| 201 | + # Extract the "InfectedSevere" compartment (Currently Hospitalized) |
| 202 | + # Adding 1 because row 0 is the time axis |
| 203 | + state_severe_idx = 1 + int(osecir.InfectionState.InfectedSevere) |
| 204 | + current_severe = comp_array[state_severe_idx, 1:] |
| 205 | + |
| 206 | + # 2. Flow indices |
| 207 | + new_symptomatic_idx = 2 |
| 208 | + new_hospitalized_idx = 6 |
| 209 | + # EXERCISE: Identify the flow index for new ICU admissions |
| 210 | + # (the flow from InfectedSevere to InfectedCritical). |
| 211 | + # use the MEmilio documentation: |
| 212 | + # https://memilio.readthedocs.io/en/latest/cpp/models/osecir.html#flows |
| 213 | + # new_icu_idx = ??? |
| 214 | + |
| 215 | + # 3. Create the plots |
| 216 | + fig, ax = plt.subplots(1, 2, figsize=(14, 6)) |
| 217 | + |
| 218 | + # --- Subplot 1: Disease Progression (The Time Lag) --- |
| 219 | + ax[0].bar(plot_time, daily_flows[new_symptomatic_idx, :], color='coral', alpha=0.5, label='New Symptomatic Cases') |
| 220 | + ax[0].bar(plot_time, daily_flows[new_hospitalized_idx, :], color='darkred', alpha=0.8, label='New Hospitalizations') |
| 221 | + # EXERCISE: Extend subplot 1 to also show new ICU admissions. |
| 222 | + # Hint: add a third bar for daily_flows[new_icu_idx, :] with a suitable color and label, |
| 223 | + # e.g. color='maroon', alpha=0.9, label='New ICU Admissions'. |
| 224 | + ax[0].set_title('Disease Progression: Infection to Hospitalization') |
| 225 | + ax[0].set_xlabel('Time [days]') |
| 226 | + ax[0].set_ylabel('Daily New Cases [#]') |
| 227 | + ax[0].legend() |
| 228 | + |
| 229 | + # --- Subplot 2: Incidence vs. Bed Occupancy --- |
| 230 | + # Axis 1: Daily New Cases (Incidence) |
| 231 | + ax[1].bar(plot_time, daily_flows[new_hospitalized_idx, :], color='steelblue', alpha=0.5, label='Daily Admissions (Incidence)') |
| 232 | + ax[1].set_xlabel('Time [days]') |
| 233 | + ax[1].set_ylabel('Daily New Hospitalizations [#]', color='steelblue') |
| 234 | + ax[1].tick_params(axis='y', labelcolor='steelblue') |
| 235 | + |
| 236 | + # Axis 2: Current Occupancy |
| 237 | + ax2 = ax[1].twinx() |
| 238 | + ax2.plot(plot_time, current_severe, color='black', linewidth=3, label='Currently Hospitalized') |
| 239 | + ax2.set_ylabel('Total Hospital Bed Occupancy [#]', color='black') |
| 240 | + ax2.tick_params(axis='y', labelcolor='black') |
| 241 | + |
| 242 | + ax[1].set_title('Incidence vs. Bed Occupancy') |
| 243 | + |
| 244 | + # Combine legends from both axes |
| 245 | + lines_1, labels_1 = ax[1].get_legend_handles_labels() |
| 246 | + lines_2, labels_2 = ax2.get_legend_handles_labels() |
| 247 | + ax2.legend(lines_1 + lines_2, labels_1 + labels_2, loc='center right') |
| 248 | + |
| 249 | + plt.tight_layout() |
| 250 | + plt.show() |
| 251 | + return |
| 252 | + |
| 253 | + |
| 254 | +@app.cell(hide_code=True) |
| 255 | +def _(mo): |
| 256 | + mo.md(r""" |
| 257 | + ## Summary |
| 258 | +
|
| 259 | + By utilizing MEmilio's flow-based simulation, we can easily extract daily transition values, such as daily new cases or hospitalizations, without losing information to simultaneous compartmental outflows. |
| 260 | +
|
| 261 | + In **Tutorial 3**, we will explore how to influence these trajectories by applying non-pharmaceutical interventions (NPIs) using `Damping` parameters. |
| 262 | + """) |
| 263 | + return |
| 264 | + |
| 265 | + |
| 266 | +if __name__ == "__main__": |
| 267 | + app.run() |
0 commit comments