compute time averaged r_eff via keyword argument
@@ -441,14 +441,19 @@ class Sim(cvbase.BaseSim):
return
- def compute_r_eff(self):
+ def compute_r_eff(self, window=1):
'''
Effective reproductive number based on number of people each person infected.
+
+ Args:
+ window (int): the size of the window used (larger values are more accurate but less precise)
+
'''
# Initialize arrays to hold sources and targets infected each day
sources = np.zeros(self.npts)
targets = np.zeros(self.npts)
+ window = int(window)
for t in self.tvec:
@@ -464,8 +469,21 @@ class Sim(cvbase.BaseSim):
# Populate the array -- to avoid divide-by-zero, skip indices that are 0
inds = sc.findinds(sources>0)
- r_eff = targets[inds]/sources[inds]
- self.results['r_eff'].values[inds] = r_eff
+ r_eff = np.zeros(self.npts)*np.nan
+ r_eff[inds] = targets[inds]/sources[inds]
+
+ # use stored weights calculate the moving average over the window of timesteps, n
+ num = np.nancumsum(r_eff * sources)
+ num[window:] = num[window:] - num[:-window]
+ den = np.cumsum(sources)
+ den[window:] = den[window:] - den[:-window]
+
+ # avoid dividing by zero
+ values = np.zeros(num.shape)*np.nan
+ ind = den > 0
+ values[ind] = num[ind]/den[ind]
+
+ self.results['r_eff'].values = values
return
