Perfect hospital (seekers) and death observations

epiforecast/epiplots.py

@@ -63,3 +63,101 @@ def plot_master_eqns(states, t, axes = None, xlims = None, reduced_system = True
     plt.tight_layout()
 
     return axes
+
+def plot_ensemble_states(states, t, axes = None,
+                                    xlims = None,
+                                    reduced_system = True,
+                                    leave = False,
+                                    figsize = (15, 4),
+                                    a_min = None,
+                                    a_max = None):
+    if axes is None:
+        fig, axes = plt.subplots(1, 2, figsize = figsize)
+
+    ensemble_size = states.shape[0]
+    if reduced_system:
+        N_eqns = 5
+        statuses = np.arange(N_eqns)
+        statuses_colors = ['C0', 'C1', 'C2', 'C4', 'C6']
+    else:
+        N_eqns = 6
+        statuses = np.arange(N_eqns)
+        statuses_colors = ['C0', 'C3', 'C1', 'C2', 'C4', 'C6']
+    population   = states.shape[1]/N_eqns
+
+    states_sum  = states.reshape(ensemble_size, N_eqns, -1, len(t)).sum(axis = 2)
+    states_perc = np.percentile(states_sum, q = [1, 10, 25, 50, 75, 90, 99], axis = 0)
+
+    for status in statuses:
+        if (reduced_system and status in [0, 3]) or (not reduced_system and status in [0, 2, 4]):
+            axes[0].fill_between(t, np.clip(states_perc[0,status], a_min, a_max), np.clip(states_perc[-1,status], a_min, a_max), alpha = .2, color = statuses_colors[status], linewidth = 0.)
+            axes[0].fill_between(t, np.clip(states_perc[1,status], a_min, a_max), np.clip(states_perc[-2,status], a_min, a_max), alpha = .2, color = statuses_colors[status], linewidth = 0.)
+            axes[0].fill_between(t, np.clip(states_perc[2,status], a_min, a_max), np.clip(states_perc[-3,status], a_min, a_max), alpha = .2, color = statuses_colors[status], linewidth = 0.)
+            axes[0].plot(t, states_perc[3,status], color = statuses_colors[status])
+
+        if (reduced_system and status in [1]) or (not reduced_system and status in [2, 3, 5]):
+            axes[1].fill_between(t, np.clip(states_perc[0,status], a_min, a_max), np.clip(states_perc[-1,status], a_min, a_max), alpha = .2, color = statuses_colors[status], linewidth = 0.)
+            axes[1].fill_between(t, np.clip(states_perc[1,status], a_min, a_max), np.clip(states_perc[-2,status], a_min, a_max), alpha = .2, color = statuses_colors[status], linewidth = 0.)
+            axes[1].fill_between(t, np.clip(states_perc[2,status], a_min, a_max), np.clip(states_perc[-3,status], a_min, a_max), alpha = .2, color = statuses_colors[status], linewidth = 0.)
+            axes[1].plot(t, states_perc[3,status], color = statuses_colors[status])
+
+        if (reduced_system and status in [2, 4]) or (not reduced_system and status in [3, 5]):
+            axes[2].fill_between(t, np.clip(states_perc[0,status], a_min, a_max), np.clip(states_perc[-1,status], a_min, a_max), alpha = .2, color = statuses_colors[status], linewidth = 0.)
+            axes[2].fill_between(t, np.clip(states_perc[1,status], a_min, a_max), np.clip(states_perc[-2,status], a_min, a_max), alpha = .2, color = statuses_colors[status], linewidth = 0.)
+            axes[2].fill_between(t, np.clip(states_perc[2,status], a_min, a_max), np.clip(states_perc[-3,status], a_min, a_max), alpha = .2, color = statuses_colors[status], linewidth = 0.)
+            axes[2].plot(t, states_perc[3,status], color = statuses_colors[status])
+
+    if reduced_system:
+        residual_state = population - states_sum.sum(axis = 1)
+        residual_state = np.percentile(residual_state, q = [1, 10, 25, 50, 75, 90, 99], axis = 0)
+        axes[0].fill_between(t, np.clip(residual_state[0], a_min, a_max), np.clip(residual_state[-1], a_min, a_max), alpha = .2, color = 'C3', linewidth = 0.)
+        axes[0].fill_between(t, np.clip(residual_state[1], a_min, a_max), np.clip(residual_state[-2], a_min, a_max), alpha = .2, color = 'C3', linewidth = 0.)
+        axes[0].fill_between(t, np.clip(residual_state[2], a_min, a_max), np.clip(residual_state[-3], a_min, a_max), alpha = .2, color = 'C3', linewidth = 0.)
+        axes[0].plot(t, np.clip(residual_state[3], a_min, a_max), color = 'C3')
+
+    axes[0].legend(['Susceptible', 'Resistant', 'Exposed'],
+        bbox_to_anchor=(0., 1.02, 1., .102), loc=3,
+        ncol=3, mode="expand", borderaxespad=0.);
+    axes[1].legend(['Infected'],
+        bbox_to_anchor=(0., 1.02, 1., .102), loc=3,
+        ncol=2, mode="expand", borderaxespad=0.);
+    axes[2].legend(['Hospitalized', 'Death'],
+        bbox_to_anchor=(0., 1.02, 1., .102), loc=3,
+        ncol=4, mode="expand", borderaxespad=0.);
+
+    for kk, ax in enumerate(axes):
+        ax.set_xlim(xlims)
+
+    plt.tight_layout()
+
+    return axes
+
+def plot_kinetic_model_data(kinetic_model, axes):
+    statuses_name   = kinetic_model.return_statuses
+    statuses_colors = ['C0', 'C3', 'C1', 'C2', 'C4', 'C6']
+    colors_dict = dict(zip(statuses_name, statuses_colors))
+
+    data = kinetic_model.statuses
+    axes[1].scatter(kinetic_model.current_time, data['I'][-1], c = colors_dict['I'], marker = 'x')
+    axes[2].scatter(kinetic_model.current_time, data['H'][-1], c = colors_dict['H'], marker = 'x')
+    axes[2].scatter(kinetic_model.current_time, data['D'][-1], c = colors_dict['D'], marker = 'x')
+
+    return axes
+
+def plot_ensemble_transmission_latent_fraction(community_transmission_rate_trace, latent_periods_trace, time_horizon):
+    transmission_perc = np.percentile(community_transmission_rate_trace, q = [1, 25, 50, 75, 99], axis = 0)
+    latent_periods_perc = np.percentile(latent_periods_trace, q = [1, 25, 50, 75, 99], axis = 0)
+
+    fig, axes = plt.subplots(1, 2, figsize = (12, 4))
+
+    axes[0].fill_between(time_horizon, transmission_perc[0], transmission_perc[-1], alpha = .2, color = 'C0')
+    axes[0].fill_between(time_horizon, transmission_perc[1], transmission_perc[-2], alpha = .2, color = 'C0')
+    axes[0].plot(time_horizon, transmission_perc[2])
+    axes[0].set_title(r'Transmission rate: $\beta$');
+
+    axes[1].fill_between(time_horizon, latent_periods_perc[0], latent_periods_perc[-1], alpha = .2, color = 'C0')
+    axes[1].fill_between(time_horizon, latent_periods_perc[1], latent_periods_perc[-2], alpha = .2, color = 'C0')
+    axes[1].plot(time_horizon, latent_periods_perc[2])
+    axes[1].set_title(r'Hospitalization fraction: $h$');
+
+    return axes

epiforecast/measurements.py

@@ -39,8 +39,8 @@ def _set_ppv(self, scale = 'log'):
              ((1 - self.sensitivity) * self.prevalence + self.specificity * (1 - self.prevalence))
 
         if scale == 'log':
-            logit_ppv  = np.log(PPV/(1 - PPV))
-            logit_for  = np.log(FOR/(1 - FOR))
+            logit_ppv  = np.log(PPV/(1 - PPV + 1e-8))
+            logit_for  = np.log(FOR/(1 - FOR + 1e-8))
 
             self.logit_ppv_mean = logit_ppv.mean()
             self.logit_ppv_var  = logit_ppv.var()
@@ -117,10 +117,10 @@ def __init__(self,
         #number of nodes in the graph
         self.N = N
         #number of different states a node can be in
-      
+
         if reduced_system == True:
             self.status_catalog = dict(zip(['S', 'I', 'H', 'R', 'D'], np.arange(5)))
-  
+
         else:
             self.status_catalog = dict(zip(['S', 'E', 'I', 'H', 'R', 'D'], np.arange(6)))
         self.n_status = len(self.status_catalog.keys())
@@ -158,7 +158,7 @@ def find_observation_states(self, state):
         else: #The value is too small
             self.obs_states=np.array([],dtype=int)
             print("no observation was above the threshold")
-    
+
 
 #combine them together
 class Observation(StateInformedObservation, TestMeasurement):
@@ -173,9 +173,9 @@ def __init__(self,
                  reduced_system=True,
                  sensitivity = 0.80,
                  specificity = 0.99):
-        
+
         self.name=obs_name
-        
+
         StateInformedObservation.__init__(self,
                                           N,
                                           obs_frac,
@@ -189,8 +189,8 @@ def __init__(self,
                                  sensitivity,
                                  specificity,
                                  reduced_system)
-        
-    #State is a numpy array of size [self.N * n_status]    
+
+    #State is a numpy array of size [self.N * n_status]
     #def find_observation_states(self, state):
         # obtain where one should make an observation based on the
         # current state, and the contact network