def _init_meeting_patterns(self):
self.mean_stochastic_interactions = self.config.get(
'mean_stochastic_interactions')
superspreader_mask = cp.random.random(
self._size) < self.config.get('suprespreader_ratio')
self.stochastic_interactions = cp.random.poisson(
self.mean_stochastic_interactions, self._size)
self.stochastic_interactions[
superspreader_mask] = cp.random.negative_binomial(
15, 0.1, int(superspreader_mask.astype(int).sum()))
check_plot_dir = self.config.get('check_plots')
ensure_dir(check_plot_dir)
if cuda:
interactions = cp.asnumpy(self.stochastic_interactions)
else:
interactions = self.stochastic_interactions
interaction_multiplicities(
interactions,
filepath=f'{check_plot_dir}/interaction-multiplicities.png')
interaction_multiplicities(
interactions,
filepath=f'{check_plot_dir}/interaction-multiplicities-log.png',
log_scale=True)
def to_json(self, filepath: str):
"""
Saves results as a .json file to the given path
"""
ensure_dir('/'.join(filepath.split('/')[:-1]))
with open(filepath, 'w') as f:
json.dump(self._data, f)
def to_json(self, filepath: str):
"""
Saves results as a .json file to the given path
"""
ensure_dir('/'.join(filepath.split('/')[:-1]))
with open(filepath, 'w') as f:
json.dump({
'days': self.days,
'infected': self.infected,
'unaffected': self.unaffected,
'new_cases': self.new_cases,
'immune': self.immune,
'dead': self.dead,
'hospitalized': self.hospitalized
}, f)
def to_json(self, filepath: str):
"""
Saves results as a .json file to the given path
"""
ensure_dir('/'.join(filepath.split('/')[:-1]))
with open(filepath, 'w') as f:
json.dump(
{
'data': self._data,
'city_ids': self._city_ids,
'city_names': self._city_names,
'city_sizes': self._city_sizes,
'longitudes': self._longitudes,
'latitudes': self._latitudes,
'parameters': self._parameters,
'config': self._config,
'R_eff': self._r_eff
}, f)
def _plot_time_ranges(self):
"""
Plots histograms of time ranges as used in the simulation.
Figure will be saved to the directory specified in the config file (check_plots)
"""
check_plot_dir = self.config.get('check_plots')
ensure_dir(check_plot_dir)
need_hospital = self._health_condition <= self._hospitalization_percentage
need_critical_care = self._health_condition <= self._critical_care_percentage
healing = self._infectious_start + self._healing_days
healing[need_hospital] = self._hospitalization_finish[need_hospital]
healing[need_critical_care] = self._critical_care_finish[
need_critical_care]
if cuda:
infectious = cp.asnumpy(self._infectious_start)
healing = cp.asnumpy(healing)
hospitalization = cp.asnumpy(
self._hospitalization_start[need_hospital])
critical_care = cp.asnumpy(
self._critical_care_start[need_critical_care])
else:
infectious = self._infectious_start
healing = healing
hospitalization = self._hospitalization_start[need_hospital]
critical_care = self._critical_care_start[need_critical_care]
time_ranges(infectious=infectious,
healing=healing,
hospitalization=hospitalization,
critical_care=critical_care,
filepath=f'{check_plot_dir}/time-ranges.png')
def plot_grid(self,
longitudes: np.ndarray,
latitudes: np.ndarray,
values: np.ndarray,
filepath=None,
cmap='gnuplot',
alpha=0.7,
values_range=None,
title=None,
hide_ticks=True,
show_coords=True):
"""
Makes a 2D contour plot
"""
fig = pl.figure(figsize=(15, 9))
ax1 = fig.add_subplot(1, 1, 1)
ax1.imshow(self.map_image)
X, Y, Z = self._get_grid(longitudes, latitudes, values)
X, Y = self._transform_xy(X, Y)
if self.borders is not None:
xc, yc = self.borders[0], self.borders[1]
pth = Path(np.vstack((xc, yc)).T, closed=False)
mask = pth.contains_points(np.vstack((X.ravel(), Y.ravel())).T)
mask = mask.reshape(X.shape)
X = np.ma.masked_array(X, ~mask)
Y = np.ma.masked_array(Y, ~mask)
Z = np.ma.masked_array(Z, ~mask)
if values_range is None:
vmin = min(values)
vmax = max(values)
else:
vmin = values_range[0]
vmax = values_range[1]
cf = ax1.contourf(X,
Y,
Z,
alpha=alpha,
cmap=cmap,
vmin=vmin,
vmax=vmax)
cbar = fig.colorbar(cf, ax=ax1)
cbar.set_ticks(np.linspace(vmin, vmax, 10))
cbar.set_ticklabels(np.linspace(vmin, vmax, 10))
# cbar.set_clim(vmin, vmax)
if hide_ticks:
pl.xticks([])
pl.yticks([])
if show_coords:
x_lin = np.linspace(0, self.image_width, 10)
y_lin = np.linspace(0, self.image_height, 10)
lon_lin, lat_lin = self._xy_to_coords(x_lin, y_lin)
pl.xticks(x_lin, np.around(lon_lin, 2))
pl.yticks(y_lin, np.around(lat_lin, 2))
pl.tight_layout()
if title is not None:
pl.title(title, fontsize=18)
if filepath is None:
pl.show()
else:
ensure_dir('/'.join(filepath.split('/')[:-1]))
pl.savefig(filepath)
pl.close()
logging.info(f'Result plot saved to: {filepath}')
def plot_pandemic(data: TimeSeriesResult, filepath=None, logscale=False):
config = Config()
virus = Virus.from_string(config.get('virus', 'name'), )
fig, (ax1, ax2, ax3) = pl.subplots(3, sharex=True, figsize=(8, 10))
ax1.set_title('COVID19 simulation', fontsize=20)
p1 = ax1.plot(data.days, data.susceptible, 'tab:blue', linewidth=2)
p2 = ax1.plot(data.days, data.infected, 'tab:red', linewidth=2)
p3 = ax1.plot(data.days, data.immune, 'tab:green', linewidth=2)
ax1.legend((p1[0], p2[0], p3[0]), ('Susceptible', 'Infected', 'Immune'))
if logscale:
ax1.set_yscale('log')
ax1.grid(True)
p4 = ax2.plot(data.days, data.hospitalized, 'tab:red', linewidth=2)
p5 = ax2.plot(data.days, data.critical_care, 'tab:blue', linewidth=2)
ax2.legend((p4[0], p5[0]), ('Hospitalized', 'Critical Care'))
if logscale:
ax2.set_yscale('log')
ax2.grid(True)
infected = np.array(data.infected)
p6 = ax3.plot(data.days,
infected * virus.asymptomatic_ratio,
'tab:blue',
linewidth=2)
p7 = ax3.plot(data.days,
infected * virus.mild_symptoms_ratio,
'tab:green',
linewidth=2)
p8 = ax3.plot(data.days, data.new_cases, 'tab:red', linewidth=2)
p9 = ax3.plot(data.days, data.dead, 'k', linewidth=2)
ax3.legend(
(p6[0], p7[0], p8[0], p9[0]),
('Asymptomatic', 'Mild symptoms', 'New cases', 'Dead cumulative'))
ax3.set_xlabel('Days', fontsize=16)
if logscale:
ax3.set_yscale('log')
ax3.grid(True)
pl.tight_layout()
if filepath is None:
pl.show()
else:
ensure_dir('/'.join(filepath.split('/')[:-1]))
pl.savefig(filepath)
pl.close()
logging.info(f'Result plot saved to: {filepath}')
def _init_households(self, household_data: dict):
"""
Initializes parameters for households.
Seniors living in pairs are simulated separately.
"""
dice = cp.random.random(self._size)
elderly_indexes = self._indices[(self.age >= 60) *
(dice <= household_data['elderly'][2])]
current_city_ids = self.city_id[elderly_indexes]
meetings = [[], []]
for city_id in self.city_ids:
city_indexes = elderly_indexes[current_city_ids == city_id]
split_indexes = cp.split(
city_indexes[:int(len(city_indexes) / 2) * 2], 2)
for i, s in enumerate(split_indexes):
meetings[i].append(s)
self._household_meetings_elderly = {
'first': cp.hstack(meetings[0]),
'second': cp.hstack(meetings[1])
}
# === split the rest of the population into households:
self._household_sizes = cp.ones(self._size)
splitting_dice = cp.random.random(self._size)
current_threshold = 0
self._max_household_size = max(
[s for s in household_data['young'].keys()])
for s, ratio in household_data['young'].items():
mask = (current_threshold <= splitting_dice) * (
splitting_dice < current_threshold + ratio)
self._household_sizes[mask] = s
current_threshold += ratio
self._household_sizes[cp.hstack([
self._household_meetings_elderly['first'],
self._household_meetings_elderly['second'],
])] = -1
for household_size in range(2, self._max_household_size + 1):
current_indexes = self._indices[self._household_sizes ==
household_size]
if len(current_indexes) == 0:
continue
current_city_ids = self.city_id[current_indexes]
meetings = [[] for i in range(household_size)]
for city_id in self.city_ids:
city_indexes = current_indexes[current_city_ids == city_id]
split_indexes = cp.split(
city_indexes[:int(len(city_indexes) / household_size) *
household_size], household_size)
for i, s in enumerate(split_indexes):
meetings[i].append(s)
self._household_meetings[household_size] = [
cp.hstack(m) for m in meetings
]
self._household_sizes[self._household_sizes == -1] = 2
# === plot the distribution of households:
check_plot_dir = self.config.get('check_plots')
ensure_dir(check_plot_dir)
if cuda:
ages = cp.asnumpy(self.age)
household_sizes = cp.asnumpy(self._household_sizes)
else:
ages = self.age
household_sizes = self._household_sizes
household_age_distribution(
ages=ages,
household_sizes=household_sizes,
filepath=f'{check_plot_dir}/household-distributions.png')
age_distribution(ages,
filepath=f'{check_plot_dir}/age-distributions.png')