add option to change Rt

covid19/models.py

@@ -181,6 +181,16 @@ def init_param_dist(cls, param_init):
         return dist
 
 
+    @classmethod
+    def init_r0_dist(cls, param_init):
+        if isinstance(param_init, tuple):
+            raise NotImplementedError('r0_dist can not be initiated with a tuple')
+        elif isinstance(param_init, rv_frozen):
+            dist = param_init
+        else:
+            dist = EmpiricalDistribution(param_init)
+        return dist
+
     def sample(self, size=1, return_param_samples=False):
         '''Sample from model.
         Args:
@@ -228,12 +238,13 @@ def sample(self, size=1, return_param_samples=False):
                       for _ in range(4)]
         S[0, ], E[0, ], I[0, ], R[0, ] = self._params['init_conditions']
 
-        r0 = self._params['r0_dist'].rvs(size)
+        _r0 = self._params['r0_dist'].rvs(self._params['t_max'])
+        r0 = np.full((size, len(_r0)),_r0)
         gamma = 1/self._params['gamma_inv_dist'].rvs(size)
         alpha = 1/self._params['alpha_inv_dist'].rvs(size)
-        beta = r0*gamma
 
         for t in t_space[1:]:
+            beta = r0[:,t]*gamma
             SE = npr.binomial(S[t-1, ].astype(int),
                               expon(scale=1/(beta*I[t-1, ]/N)).cdf(1))
             EI = npr.binomial(E[t-1, ].astype(int),

simulator/pages/seir.py

@@ -241,6 +241,35 @@ def make_r0_widgets(defaults=DEFAULT_PARAMS):
     return (r0_inf, r0_sup, .95, 'lognorm')
 
 
+def make_r0_array(t_max, r0_mean, rt_mean, rt_t):
+    if (rt_t==0) | (rt_t>=t_max):
+        r0_len = len(range(t_max))
+        return np.full(r0_len, r0_mean)
+    else: 
+        r0_len = len(range(0,rt_t))
+        rt_len = len(range(rt_t, t_max))
+        return np.concatenate([np.full(r0_len, r0_mean), np.full(rt_len, rt_mean)])
+
+
+def make_rt_widgets(w_params, defaults=DEFAULT_PARAMS):
+    t_max = w_params['t_max']
+    r0_mean = st.number_input(
+             'Número básico de reprodução (R0)',
+             min_value=0.01, max_value=10.0, step=0.25,
+             value=defaults['r0_dist'][0])
+
+    rt_mean = st.number_input(
+            'Número R(t) de reproducao ',
+            min_value=0.01, max_value=10.0, step=0.25,
+            value=defaults['r0_dist'][1])
+
+    rt_t = st.number_input('Atraso de R(t)',
+                           value=0,
+                           min_value=0,
+                           max_value=179)
+    return make_r0_array(t_max, r0_mean, rt_mean, rt_t)
+
+
 def write():
     st.markdown("## Modelo Epidemiológico (SEIR-Bayes)")
     st.sidebar.markdown(texts.PARAMETER_SELECTION)
@@ -267,7 +296,10 @@ def write():
                                   options=options_date,
                                   index=len(options_date)-1)
     NEIR0 = make_NEIR0(cases_df, population_df, w_place, w_date)
-
+
+    # Param Widgets
+    w_params = make_param_widgets(NEIR0)
+
     # Estimativa R0
     st.markdown(texts.r0_ESTIMATION_TITLE)
     should_estimate_r0 = st.checkbox(
@@ -292,11 +324,15 @@ def write():
                     f'${np.quantile(r0_dist, 0.01):.03}$ e ${np.quantile(r0_dist, 0.99):.03}$*')
         st.markdown(texts.r0_CITATION)
     else:
-        r0_dist = make_r0_widgets()
-        st.markdown(texts.r0_ESTIMATION_DONT)
+        should_use_rt = st.checkbox('Utilizar R(t) com atraso de t dias',
+                                    value=True)
+        if should_use_rt:
+            r0_dist = make_rt_widgets(w_params, defaults=DEFAULT_PARAMS)
+        else:
+            r0_dist = make_r0_widgets()
+            st.markdown(texts.r0_ESTIMATION_DONT)
 
     # Previsão de infectados
-    w_params = make_param_widgets(NEIR0)
     model = SEIRBayes(**w_params, r0_dist=r0_dist)
     model_output = model.sample(SAMPLE_SIZE)
     ei_df = make_EI_df(model, model_output, SAMPLE_SIZE)