def test_printing(): ''' Test printing functions ''' example = sc.prettyobj() example.data = sc.vectocolor(10) print('sc.pr():') sc.pr(example) print('sc.pp():') sc.pp(example.data) string = sc.pp(example.data, doprint=False) return string
def test_printing(): sc.heading('test_printing()') example = sc.prettyobj() example.data = sc.vectocolor(10) print('sc.pr():') sc.pr(example) print('sc.pp():') sc.pp(example.data) string = sc.pp(example.data, doprint=False) print('sc.printdata():') sc.printdata(example.data) return string
def computation(seed=0, n=1000): # Make graph pl.seed(int(seed)) fig = pl.figure() ax = fig.add_subplot(111) xdata = pl.randn(n) ydata = pl.randn(n) colors = sc.vectocolor(pl.sqrt(xdata**2+ydata**2)) ax.scatter(xdata, ydata, c=colors) # Convert to FE graphjson = sw.mpld3ify(fig, jsonify=False) # Convert to dict return graphjson # Return the JSON representation of the Matplotlib figure
def pairplotpars(df, inds=None, keys=None, color_column=None, bounds=None, cmap='parula', bins=None, edgecolor='w', facecolor='#F8A493', figsize=(20, 16)): ''' Plot scatterplots, histograms, and kernel densities ''' if inds is not None: df = df.iloc[inds, :].copy() # Choose the colors if color_column: colors = sc.vectocolor(df[color_column].values, cmap=cmap) else: colors = [facecolor for i in range(len(df))] df['color_column'] = [sc.rgb2hex(rgba[:-1]) for rgba in colors] if keys is not None: df = df.loc[:, keys + ['color_column']].copy() # Make the plot grid = sns.PairGrid(df) grid = grid.map_lower(pl.scatter, **{'facecolors': df['color_column']}) grid = grid.map_diag(pl.hist, bins=bins, edgecolor=edgecolor, facecolor=facecolor) grid = grid.map_upper(pl.hexbin, cmap="Blues", edgecolor="none", gridsize=25) grid.fig.set_size_inches(figsize) grid.fig.tight_layout() # Set bounds if bounds: for ax in grid.axes.flatten(): xlabel = ax.get_xlabel() ylabel = ax.get_ylabel() if xlabel in bounds: ax.set_xlim(bounds[xlabel]) if ylabel in bounds: ax.set_ylim(bounds[ylabel]) return grid
def pairplotpars(data, inds=None, color_column=None, bounds=None, cmap='parula', bins=None, edgecolor='w', facecolor='#F8A493', figsize=(20, 16)): ''' Plot scatterplots, histograms, and kernel densities ''' import seaborn as sns # Optional import data = sc.odict(sc.dcp(data)) # Create the dataframe df = pd.DataFrame.from_dict(data) if inds is not None: df = df.iloc[inds, :].copy() # Choose the colors if color_column: colors = sc.vectocolor(df[color_column].values, cmap=cmap) else: colors = [facecolor for i in range(len(df))] df['color_column'] = [sc.rgb2hex(rgba[:-1]) for rgba in colors] # Make the plot grid = sns.PairGrid(df) grid = grid.map_lower(pl.scatter, **{'facecolors': df['color_column']}) grid = grid.map_diag(pl.hist, bins=bins, edgecolor=edgecolor, facecolor=facecolor) grid = grid.map_upper(sns.kdeplot) grid.fig.set_size_inches(figsize) grid.fig.tight_layout() # Set bounds if bounds: for ax in grid.axes.flatten(): xlabel = ax.get_xlabel() ylabel = ax.get_ylabel() if xlabel in bounds: ax.set_xlim(bounds[xlabel]) if ylabel in bounds: ax.set_ylim(bounds[ylabel]) return grid
import pylab as pl import sciris as sc import hiptool as hp sc.heading('Initializing...') sc.tic() dosave = True missing_data = ['remove', 'assumption'][1] # Choose how to handle missing data spendings = [0.1, 0.3, 1, 3, 10, 30, 100, 300, 1000]#, 3000, 10000] nspendings = len(spendings) colors = sc.vectocolor(len(spendings)) # Load input files D = sc.odict() # Data R = sc.odict() # Results bod_data = sc.loadobj('gbd-data.dat') country_data = sc.loadspreadsheet('country-data.xlsx') baseline_factor = country_data.findrow('Zambia', asdict=True)['icer_multiplier'] # Zambia was used for this # Create default P = hp.Project() P.loadburden(filename='rapid_BoD.xlsx') P.loadinterventions(filename='rapid_interventions.xlsx') ninterventions = P.intervsets[0].data.nrows # Load data missing_data_adjustment_factor = 2 sc.heading('Loading data...') for c,country in enumerate(country_data['name'].tolist()): print(f' Working on {country}...')
def plot(): fig = pl.figure(num='Fig. 2: Transmission dynamics', figsize=(20,14)) piey, tsy, r3y = 0.68, 0.50, 0.07 piedx, tsdx, r3dx = 0.2, 0.9, 0.25 piedy, tsdy, r3dy = 0.2, 0.47, 0.35 pie1x, pie2x = 0.12, 0.65 tsx = 0.07 dispx, cumx, sympx = tsx, 0.33+tsx, 0.66+tsx ts_ax = pl.axes([tsx, tsy, tsdx, tsdy]) pie_ax1 = pl.axes([pie1x, piey, piedx, piedy]) pie_ax2 = pl.axes([pie2x, piey, piedx, piedy]) symp_ax = pl.axes([sympx, r3y, r3dx, r3dy]) disp_ax = pl.axes([dispx, r3y, r3dx, r3dy]) cum_ax = pl.axes([cumx, r3y, r3dx, r3dy]) off = 0.06 txtdispx, txtcumx, txtsympx = dispx-off, cumx-off, sympx-off+0.02 tsytxt = tsy+tsdy r3ytxt = r3y+r3dy labelsize = 40-wf pl.figtext(txtdispx, tsytxt, 'a', fontsize=labelsize) pl.figtext(txtdispx, r3ytxt, 'b', fontsize=labelsize) pl.figtext(txtcumx, r3ytxt, 'c', fontsize=labelsize) pl.figtext(txtsympx, r3ytxt, 'd', fontsize=labelsize) #%% Fig. 2A -- Time series plot layer_keys = list(sim.layer_keys()) layer_mapping = {k:i for i,k in enumerate(layer_keys)} n_layers = len(layer_keys) colors = sc.gridcolors(n_layers) layer_counts = np.zeros((sim.npts, n_layers)) for source_ind, target_ind in tt.count_transmissions(): dd = tt.detailed[target_ind] date = dd['date'] layer_num = layer_mapping[dd['layer']] layer_counts[date, layer_num] += sim.rescale_vec[date] mar12 = cv.date('2020-03-12') mar23 = cv.date('2020-03-23') mar12d = sim.day(mar12) mar23d = sim.day(mar23) labels = ['Household', 'School', 'Workplace', 'Community', 'LTCF'] for l in range(n_layers): ts_ax.plot(sim.datevec, layer_counts[:,l], c=colors[l], lw=3, label=labels[l]) sc.setylim(ax=ts_ax) sc.boxoff(ax=ts_ax) ts_ax.set_ylabel('Transmissions per day') ts_ax.set_xlim([sc.readdate('2020-01-18'), sc.readdate('2020-06-09')]) ts_ax.xaxis.set_major_formatter(mdates.DateFormatter('%b-%d')) ts_ax.set_xticks([sim.date(d, as_date=True) for d in np.arange(0, sim.day('2020-06-09'), 14)]) ts_ax.legend(frameon=False, bbox_to_anchor=(0.85,0.1)) color = [0.2, 0.2, 0.2] ts_ax.axvline(mar12, c=color, linestyle='--', alpha=0.4, lw=3) ts_ax.axvline(mar23, c=color, linestyle='--', alpha=0.4, lw=3) yl = ts_ax.get_ylim() labely = yl[1]*1.015 ts_ax.text(mar12, labely, 'Schools close ', color=color, alpha=0.9, style='italic', horizontalalignment='center') ts_ax.text(mar23, labely, ' Stay-at-home', color=color, alpha=0.9, style='italic', horizontalalignment='center') #%% Fig. 2A inset -- Pie charts pre_counts = layer_counts[0:mar12d, :].sum(axis=0) post_counts = layer_counts[mar23d:, :].sum(axis=0) pre_counts = pre_counts/pre_counts.sum()*100 post_counts = post_counts/post_counts.sum()*100 lpre = [ f'Household\n{pre_counts[0]:0.1f}%', f'School\n{pre_counts[1]:0.1f}% ', f'Workplace\n{pre_counts[2]:0.1f}% ', f'Community\n{pre_counts[3]:0.1f}%', f'LTCF\n{pre_counts[4]:0.1f}%', ] lpost = [ f'Household\n{post_counts[0]:0.1f}%', f'School\n{post_counts[1]:0.1f}%', f'Workplace\n{post_counts[2]:0.1f}%', f'Community\n{post_counts[3]:0.1f}%', f'LTCF\n{post_counts[4]:0.1f}%', ] pie_ax1.pie(pre_counts, colors=colors, labels=lpre, **pieargs) pie_ax2.pie(post_counts, colors=colors, labels=lpost, **pieargs) pie_ax1.text(0, 1.75, 'Transmissions by layer\nbefore schools closed', style='italic', horizontalalignment='center') pie_ax2.text(0, 1.75, 'Transmissions by layer\nafter stay-at-home', style='italic', horizontalalignment='center') #%% Fig. 2B -- histogram by overdispersion # Process targets n_targets = tt.count_targets(end_day=mar12) # Handle bins max_infections = n_targets.max() edges = np.arange(0, max_infections+2) # Analysis counts = np.histogram(n_targets, edges)[0] bins = edges[:-1] # Remove last bin since it's an edge norm_counts = counts/counts.sum() raw_counts = counts*bins total_counts = raw_counts/raw_counts.sum()*100 n_bins = len(bins) index = np.linspace(0, 100, len(n_targets)) sorted_arr = np.sort(n_targets) sorted_sum = np.cumsum(sorted_arr) sorted_sum = sorted_sum/sorted_sum.max()*100 change_inds = sc.findinds(np.diff(sorted_arr) != 0) pl.set_cmap('Spectral_r') sscolors = sc.vectocolor(n_bins) width = 1.0 for i in range(n_bins): disp_ax.bar(bins[i], total_counts[i], width=width, facecolor=sscolors[i]) disp_ax.set_xlabel('Number of transmissions per case') disp_ax.set_ylabel('Proportion of transmissions (%)') sc.boxoff() disp_ax.set_xlim([0.5, 32.5]) disp_ax.set_xticks(np.arange(0, 32.5, 4)) sc.boxoff(ax=disp_ax) dpie_ax = pl.axes([dispx+0.05, 0.20, 0.2, 0.2]) trans1 = total_counts[1:3].sum() trans2 = total_counts[3:5].sum() trans3 = total_counts[5:8].sum() trans4 = total_counts[8:].sum() labels = [ f'1-2:\n{trans1:0.0f}%', f' 3-4:\n {trans2:0.0f}%', f'5-7: \n{trans3:0.0f}%\n', f'>7: \n{trans4:0.0f}%\n', ] dpie_args = sc.mergedicts(pieargs, dict(labeldistance=1.2)) # Slightly smaller label distance dpie_ax.pie([trans1, trans2, trans3, trans4], labels=labels, colors=sscolors[[0,4,7,12]], **dpie_args) #%% Fig. 2C -- cumulative distribution function rev_ind = 100 - index n_change_inds = len(change_inds) change_bins = bins[counts>0][1:] for i in range(n_change_inds): ib = int(change_bins[i]) ci = change_inds[i] ici = index[ci] sci = sorted_sum[ci] color = sscolors[ib] if i>0: cim1 = change_inds[i-1] icim1 = index[cim1] scim1 = sorted_sum[cim1] cum_ax.plot([icim1, ici], [scim1, sci], lw=4, c=color) cum_ax.scatter([ici], [sci], s=150, zorder=50-i, c=[color], edgecolor='w', linewidth=0.2) if ib<=6 or ib in [8, 10, 25]: xoff = 5 - 2*(ib==1) + 3*(ib>=10) + 1*(ib>=20) yoff = 2*(ib==1) cum_ax.text(ici-xoff, sci+yoff, ib, fontsize=18-wf, color=color) cum_ax.set_xlabel('Proportion of primary infections (%)') cum_ax.set_ylabel('Proportion of transmissions (%)') xmin = -2 ymin = -2 cum_ax.set_xlim([xmin, 102]) cum_ax.set_ylim([ymin, 102]) sc.boxoff(ax=cum_ax) # Draw horizontal lines and annotations ancol1 = [0.2, 0.2, 0.2] ancol2 = sscolors[0] ancol3 = sscolors[6] i01 = sc.findlast(sorted_sum==0) i20 = sc.findlast(sorted_sum<=20) i50 = sc.findlast(sorted_sum<=50) cum_ax.plot([xmin, index[i01]], [0, 0], '--', lw=2, c=ancol1) cum_ax.plot([xmin, index[i20], index[i20]], [20, 20, ymin], '--', lw=2, c=ancol2) cum_ax.plot([xmin, index[i50], index[i50]], [50, 50, ymin], '--', lw=2, c=ancol3) # Compute mean number of transmissions for 80% and 50% thresholds q80 = sc.findfirst(np.cumsum(total_counts)>20) # Count corresponding to 80% of cumulative infections (100-80) q50 = sc.findfirst(np.cumsum(total_counts)>50) # Count corresponding to 50% of cumulative infections n80, n50 = [sum(bins[q:]*norm_counts[q:]/norm_counts[q:].sum()) for q in [q80, q50]] # Plot annotations kw = dict(bbox=dict(facecolor='w', alpha=0.9, lw=0), fontsize=20-wf) cum_ax.text(2, 3, f'{index[i01]:0.0f}% of infections\ndo not transmit', c=ancol1, **kw) cum_ax.text(8, 23, f'{rev_ind[i20]:0.0f}% of infections cause\n80% of transmissions\n(mean: {n80:0.1f} per infection)', c=ancol2, **kw) cum_ax.text(14, 53, f'{rev_ind[i50]:0.0f}% of infections cause\n50% of transmissions\n(mean: {n50:0.1f} per infection)', c=ancol3, **kw) #%% Fig. 2D -- histogram by date of symptom onset # Calculate asymp_count = 0 symp_counts = {} minind = -5 maxind = 15 for _, target_ind in tt.transmissions: dd = tt.detailed[target_ind] date = dd['date'] delta = sim.rescale_vec[date] # Increment counts by this much if dd['s']: if tt.detailed[dd['source']]['date'] <= date: # Skip dynamical scaling reinfections sdate = dd['s']['date_symptomatic'] if np.isnan(sdate): asymp_count += delta else: ind = int(date - sdate) if ind not in symp_counts: symp_counts[ind] = 0 symp_counts[ind] += delta # Convert to an array xax = np.arange(minind-1, maxind+1) sympcounts = np.zeros(len(xax)) for i,val in symp_counts.items(): if i<minind: ind = 0 elif i>maxind: ind = -1 else: ind = sc.findinds(xax==i)[0] sympcounts[ind] += val # Plot total_count = asymp_count + sympcounts.sum() sympcounts = sympcounts/total_count*100 presymp = sc.findinds(xax<=0)[-1] colors = ['#eed15b', '#ee943a', '#c3211a'] asymp_frac = asymp_count/total_count*100 pre_frac = sympcounts[:presymp].sum() symp_frac = sympcounts[presymp:].sum() symp_ax.bar(xax[0]-2, asymp_frac, label='Asymptomatic', color=colors[0]) symp_ax.bar(xax[:presymp], sympcounts[:presymp], label='Presymptomatic', color=colors[1]) symp_ax.bar(xax[presymp:], sympcounts[presymp:], label='Symptomatic', color=colors[2]) symp_ax.set_xlabel('Days since symptom onset') symp_ax.set_ylabel('Proportion of transmissions (%)') symp_ax.set_xticks([minind-3, 0, 5, 10, maxind]) symp_ax.set_xticklabels(['Asymp.', '0', '5', '10', f'>{maxind}']) sc.boxoff(ax=symp_ax) spie_ax = pl.axes([sympx+0.05, 0.20, 0.2, 0.2]) labels = [f'Asymp-\ntomatic\n{asymp_frac:0.0f}%', f' Presymp-\n tomatic\n {pre_frac:0.0f}%', f'Symp-\ntomatic\n{symp_frac:0.0f}%'] spie_ax.pie([asymp_frac, pre_frac, symp_frac], labels=labels, colors=colors, **pieargs) return fig
tts = sc.objdict() for key, sim in sims.items(): sim.run() sim.people.make_detailed_transtree() tts[key] = sim.people.transtree.detailed if plot_sim: to_plot = cv.get_sim_plots() to_plot['Total counts'] = [ 'cum_infections', 'cum_diagnoses', 'cum_quarantined', 'n_quarantined' ] sim.plot(to_plot=to_plot) #%% Plotting colors = sc.vectocolor(sim.n, cmap='parula') msize = 10 suscol = [0.5, 0.5, 0.5] plargs = dict(lw=2, alpha=0.5) idelay = 0.05 daydelay = 0.3 pl.rcParams['font.size'] = 18 F = sc.objdict() T = sc.objdict() D = sc.objdict() Q = sc.objdict() for key in sims.keys():
for ind, person in datadict.items(): mapping[person.Case] = ind contacts = {} for ind in datadict.keys(): contacts[ind] = [] for ind, person in datadict.items(): if person.exposure_source != -1: contacts[mapping[person.exposure_source]].append(ind) # Handle colors ages_for_color = data.age.to_numpy() for i, person in enumerate(datadict.values()): assert ages_for_color[i] == person.age colors = sc.vectocolor(ages_for_color, cmap=sc.parulacolormap()) # Process into days plotdata = {} for i in range(N): row = data.loc[i] if row.day not in plotdata: plotdata[row.day] = [] plotdata[row.day].append(row) #%% Plotting # Plot settings delta = 1.05 yoff = 0.01 markersize = 300
def animate_transtree(tt, **kwargs): ''' Plot an animation of the transmission tree; see TransTree.animate() for documentation ''' # Settings animate = kwargs.get('animate', True) verbose = kwargs.get('verbose', False) msize = kwargs.get('markersize', 10) sus_color = kwargs.get('sus_color', [0.5, 0.5, 0.5]) fig_args = kwargs.get('fig_args', dict(figsize=(24, 16))) axis_args = kwargs.get( 'axis_args', dict(left=0.10, bottom=0.05, right=0.85, top=0.97, wspace=0.25, hspace=0.25)) plot_args = kwargs.get('plot_args', dict(lw=2, alpha=0.5)) delay = kwargs.get('delay', 0.2) font_size = kwargs.get('font_size', 18) colors = kwargs.get('colors', None) cmap = kwargs.get('cmap', 'parula') pl.rcParams['font.size'] = font_size if colors is None: colors = sc.vectocolor(tt.pop_size, cmap=cmap) # Initialization n = tt.n_days + 1 frames = [list() for i in range(n)] tests = [list() for i in range(n)] diags = [list() for i in range(n)] quars = [list() for i in range(n)] # Construct each frame of the animation for i, entry in enumerate(tt.detailed): # Loop over every person frame = sc.objdict() tdq = sc.objdict() # Short for "tested, diagnosed, or quarantined" # This person became infected if entry: source = entry['source'] target = entry['target'] target_date = entry['date'] if source: # Seed infections and importations won't have a source source_date = tt.detailed[source]['date'] else: source = 0 source_date = 0 # Construct this frame frame.x = [source_date, target_date] frame.y = [source, target] frame.c = colors[source] frame.i = True # If this person is infected frames[target_date].append(frame) # Handle testing, diagnosis, and quarantine tdq.t = target tdq.d = target_date tdq.c = colors[target] date_t = entry.t.date_tested date_d = entry.t.date_diagnosed date_q = entry.t.date_known_contact if ~np.isnan(date_t) and date_t < n: tests[int(date_t)].append(tdq) if ~np.isnan(date_d) and date_d < n: diags[int(date_d)].append(tdq) if ~np.isnan(date_q) and date_q < n: quars[int(date_q)].append(tdq) # This person did not become infected else: frame.x = [0] frame.y = [i] frame.c = sus_color frame.i = False frames[0].append(frame) # Configure plotting fig = pl.figure(**fig_args) pl.subplots_adjust(**axis_args) ax = fig.add_subplot(1, 1, 1) # Create the legend ax2 = pl.axes([0.85, 0.05, 0.14, 0.9]) ax2.axis('off') lcol = colors[0] na = np.nan # Shorten pl.plot(na, na, '-', c=lcol, **plot_args, label='Transmission') pl.plot(na, na, 'o', c=lcol, markersize=msize, **plot_args, label='Source') pl.plot(na, na, '*', c=lcol, markersize=msize, **plot_args, label='Target') pl.plot(na, na, 'o', c=lcol, markersize=msize * 2, fillstyle='none', **plot_args, label='Tested') pl.plot(na, na, 's', c=lcol, markersize=msize * 1.2, **plot_args, label='Diagnosed') pl.plot(na, na, 'x', c=lcol, markersize=msize * 2.0, label='Known contact') pl.legend() # Plot the animation pl.sca(ax) for day in range(n): pl.title(f'Day: {day}') pl.xlim([0, n]) pl.ylim([0, tt.pop_size]) pl.xlabel('Day') pl.ylabel('Person') flist = frames[day] tlist = tests[day] dlist = diags[day] qlist = quars[day] if verbose: print(i, flist) for f in flist: if verbose: print(f) pl.plot(f.x[0], f.y[0], 'o', c=f.c, markersize=msize, **plot_args) # Plot sources pl.plot(f.x, f.y, '-', c=f.c, **plot_args) # Plot transmission lines if f.i: # If this person is infected pl.plot(f.x[1], f.y[1], '*', c=f.c, markersize=msize, **plot_args) # Plot targets for tdq in tlist: pl.plot(tdq.d, tdq.t, 'o', c=tdq.c, markersize=msize * 2, fillstyle='none') # Tested; No alpha for this for tdq in dlist: pl.plot(tdq.d, tdq.t, 's', c=tdq.c, markersize=msize * 1.2, **plot_args) # Diagnosed for tdq in qlist: pl.plot(tdq.d, tdq.t, 'x', c=tdq.c, markersize=msize * 2.0) # Quarantine; no alpha for this pl.plot([0, day], [0.5, 0.5], c='k', lw=5) # Plot the endless march of time if animate: # Whether to animate pl.pause(delay) return fig
def test_plot_pop(): plotconnections = True doclear = False pause = 0.2 n = 20000 alpha = 0.5 # indices = pl.arange(1000) pl.seed(1) indices = pl.randint(0, n, 20) max_contacts = {'S': 20, 'W': 10} population = sp.make_population(n=n, max_contacts=max_contacts, as_objdict=True) nside = np.ceil(np.sqrt(n)) x, y = np.meshgrid(np.arange(nside), np.arange(nside)) x = x.flatten()[:n] y = y.flatten()[:n] people = population.values() for p, person in enumerate(people): person.loc = sc.objdict(dict(x=x[p], y=y[p])) ages = np.array([person.age for person in people]) f_inds = [ind for ind, person in enumerate(people) if not person.sex] m_inds = [ind for ind, person in enumerate(people) if person.sex] # import matplotlib.pyplot as plt # import matplotlib.colors as colors # colors_undersea = plt.cm.terrain(np.linspace(0, 0.17, 256)) # colors_land = plt.cm.terrain(np.linspace(0.25, 1, 256)) # all_colors = np.vstack((colors_undersea, colors_land)) # terrain_map = colors.LinearSegmentedColormap.from_list('terrain_map', # all_colors) # pl.set_cmap(terrain_map) fig = pl.figure(figsize=(24, 18)) ax = pl.subplot(111) # sc.turbocolormap(apply=True) minval = 0 # ages.min() maxval = 100 # ages.min() colors = sc.vectocolor(ages, minval=minval, maxval=maxval) for i, inds in enumerate([f_inds, m_inds]): pl.scatter(x[inds], y[inds], marker='os'[i], c=colors[inds]) pl.clim([minval, maxval]) pl.colorbar() if plotconnections: lcols = dict(H=[0, 0, 0], S=[0, 0.5, 1], W=[0, 0.7, 0], C=[1, 1, 0]) for index in indices: person = people[index] contacts = person.contacts lines = [] for lkey in lcols.keys(): for contactkey in contacts[lkey]: contact = population[contactkey] tmp = pl.plot([person.loc.x, contact.loc.x], [person.loc.y, contact.loc.y], c=lcols[lkey], alpha=alpha) lines.append(tmp) # pl.title(f'Index: {index}') # pl.pause(pause) # if doclear: # ax.lines = [] return fig
sc.colorize(showhelp=True) sc.colorize('green', 'hi') # Simple example sc.colorize(['yellow', 'bgblack']) print('Hello world') print('Goodbye world') sc.colorize('reset') # Colorize all output in between bluearray = sc.colorize(color='blue', string=str(range(5)), output=True) print("c'est bleu: " + bluearray) sc.colorize('magenta') # Now type in magenta for a while print('this is magenta') sc.colorize('reset') # Stop typing in magenta # Test printing functions if 'printing' in torun: example = sc.prettyobj() example.data = sc.vectocolor(10) print('sc.pr():') sc.pr(example) print('sc.pp():') sc.pp(example.data) string = sc.pp(example.data, doprint=False) # Test profiling functions if 'profile' in torun: def slow_fn(): n = 10000 int_list = [] int_dict = {} for i in range(n): int_list.append(i)
n_days=30, pop_infected=100, rand_seed=25857 + p * 241, #, 29837*(p+298), pop_type='random', verbose=0, ) sims[key].run() results.append(sims[key].results['cum_infections'].values) #%% Plotting pl.figure(figsize=(18, 6), dpi=200) # pl.rcParams['font.size'] = 18 pl.subplot(1, 3, 1) colors = sc.vectocolor(pl.log10(popsizes), cmap='parula') for k, key in enumerate(keys): label = f'{int(float(key[1:]))/1000}k: {results[k][-1]:0.0f}' pl.plot(results[k], label=label, lw=3, color=colors[k]) print(label) # pl.legend() pl.title('Total number of infections') pl.xlabel('Day') pl.ylabel('Number of infections') sc.commaticks(axis='y') pl.subplot(1, 3, 2) for k, key in enumerate(keys): label = f'{int(float(key[1:]))/1000}k: {results[k][-1]/popsizes[k]*100:0.1f}' pl.plot(results[k] / popsizes[k] * 100, label=label, lw=3, color=colors[k]) print(label)
def plot_histogram(self, bins=None, fig_args=None, width=0.8, font_size=18): ''' Plots a histogram of the number of transmissions ''' if bins is None: max_infections = self.n_targets.max() bins = np.arange(0, max_infections + 2) # Analysis counts = np.histogram(self.n_targets, bins)[0] bins = bins[:-1] # Remove last bin since it's an edge total_counts = counts * bins # counts = counts*100/counts.sum() # total_counts = total_counts*100/total_counts.sum() n_bins = len(bins) n_trans = sum(total_counts) index = np.linspace(0, 100, len(self.n_targets)) sorted_arr = np.sort(self.n_targets) sorted_sum = np.cumsum(sorted_arr) sorted_sum = sorted_sum / sorted_sum.max() * 100 change_inds = sc.findinds(np.diff(sorted_arr) != 0) # Plotting fig_args = sc.mergedicts(dict(figsize=(24, 15))) pl.rcParams['font.size'] = font_size fig = pl.figure(**fig_args) pl.set_cmap('Spectral') pl.subplots_adjust(left=0.08, right=0.92, bottom=0.08, top=0.92) colors = sc.vectocolor(n_bins) pl.subplot(1, 2, 1) w05 = width * 0.5 w025 = w05 * 0.5 pl.bar(bins - w025, counts, width=w05, facecolor='k', label='Number of events') for i in range(n_bins): label = 'Number of transmissions (events × transmissions per event)' if i == 0 else None pl.bar(bins[i] + w025, total_counts[i], width=w05, facecolor=colors[i], label=label) pl.xlabel('Number of transmissions per person') pl.ylabel('Count') pl.xticks(ticks=bins) pl.legend() pl.title('Numbers of events and transmissions') pl.subplot(2, 2, 2) total = 0 for i in range(n_bins): new = total_counts[i] / n_trans * 100 pl.bar(bins[i:], new, width=width, bottom=total, facecolor=colors[i]) total += new pl.xticks(ticks=bins) pl.xlabel('Number of transmissions per person') pl.ylabel('Proportion of infections caused (%)') pl.title('Proportion of transmissions, by number of transmissions') pl.subplot(2, 2, 4) pl.plot(index, sorted_sum, lw=3, c='k', alpha=0.5) for i in range(len(change_inds)): pl.scatter([index[change_inds[i]]], [sorted_sum[change_inds[i]]], s=150, zorder=10, c=[colors[i]], label=f'Transmitted to {i+1} people') pl.xlabel( 'Proportion of population, ordered by the number of people they infected (%)' ) pl.ylabel('Proportion of infections caused (%)') pl.legend() pl.ylim([0, 100]) pl.title('Proportion of transmissions, by proportion of population') pl.axes([0.25, 0.65, 0.2, 0.2]) berry = [0.8, 0.1, 0.2] pl.plot(self.sim_results.t, self.sim_results.cum_infections, lw=2, c=berry) pl.xlabel('Day') pl.ylabel('Cumulative infections') return fig
def animate(self, *args, **kwargs): ''' Animate the transmission tree. Args: animate (bool): whether to animate the plot (otherwise, show when finished) verbose (bool): print out progress of each frame markersize (int): size of the markers sus_color (list): color for susceptibles fig_args (dict): arguments passed to pl.figure() axis_args (dict): arguments passed to pl.subplots_adjust() plot_args (dict): arguments passed to pl.plot() delay (float): delay between frames in seconds font_size (int): size of the font colors (list): color of each person cmap (str): colormap for each person (if colors is not supplied) Returns: fig: the figure object ''' # Settings animate = kwargs.get('animate', True) verbose = kwargs.get('verbose', False) msize = kwargs.get('markersize', 10) sus_color = kwargs.get('sus_color', [0.5, 0.5, 0.5]) fig_args = kwargs.get('fig_args', dict(figsize=(24, 16))) axis_args = kwargs.get( 'axis_args', dict(left=0.10, bottom=0.05, right=0.85, top=0.97, wspace=0.25, hspace=0.25)) plot_args = kwargs.get('plot_args', dict(lw=2, alpha=0.5)) delay = kwargs.get('delay', 0.2) font_size = kwargs.get('font_size', 18) colors = kwargs.get('colors', None) cmap = kwargs.get('cmap', 'parula') pl.rcParams['font.size'] = font_size if colors is None: colors = sc.vectocolor(self.pop_size, cmap=cmap) # Initialization n = self.n_days + 1 frames = [list() for i in range(n)] tests = [list() for i in range(n)] diags = [list() for i in range(n)] quars = [list() for i in range(n)] # Construct each frame of the animation for ddict in self.detailed: # Loop over every person if ddict is None: continue # Skip the 'None' node corresponding to seeded infections frame = sc.objdict() tdq = sc.objdict() # Short for "tested, diagnosed, or quarantined" target = ddict.t target_ind = ddict['target'] if not np.isnan(ddict['date']): # If this person was infected source_ind = ddict[ 'source'] # Index of the person who infected the target target_date = ddict['date'] if source_ind is not None: # Seed infections and importations won't have a source source_date = self.detailed[source_ind]['date'] else: source_ind = 0 source_date = 0 # Construct this frame frame.x = [source_date, target_date] frame.y = [source_ind, target_ind] frame.c = colors[source_ind] frame.i = True # If this person is infected frames[int(target_date)].append(frame) # Handle testing, diagnosis, and quarantine tdq.t = target_ind tdq.d = target_date tdq.c = colors[int(target_ind)] date_t = target['date_tested'] date_d = target['date_diagnosed'] date_q = target['date_known_contact'] if ~np.isnan(date_t) and date_t < n: tests[int(date_t)].append(tdq) if ~np.isnan(date_d) and date_d < n: diags[int(date_d)].append(tdq) if ~np.isnan(date_q) and date_q < n: quars[int(date_q)].append(tdq) else: frame.x = [0] frame.y = [target_ind] frame.c = sus_color frame.i = False frames[0].append(frame) # Configure plotting fig = pl.figure(**fig_args) pl.subplots_adjust(**axis_args) ax = fig.add_subplot(1, 1, 1) # Create the legend ax2 = pl.axes([0.85, 0.05, 0.14, 0.9]) ax2.axis('off') lcol = colors[0] na = np.nan # Shorten pl.plot(na, na, '-', c=lcol, **plot_args, label='Transmission') pl.plot(na, na, 'o', c=lcol, markersize=msize, **plot_args, label='Source') pl.plot(na, na, '*', c=lcol, markersize=msize, **plot_args, label='Target') pl.plot(na, na, 'o', c=lcol, markersize=msize * 2, fillstyle='none', **plot_args, label='Tested') pl.plot(na, na, 's', c=lcol, markersize=msize * 1.2, **plot_args, label='Diagnosed') pl.plot(na, na, 'x', c=lcol, markersize=msize * 2.0, label='Known contact') pl.legend() # Plot the animation pl.sca(ax) for day in range(n): pl.title(f'Day: {day}') pl.xlim([0, n]) pl.ylim([0, len(self)]) pl.xlabel('Day') pl.ylabel('Person') flist = frames[day] tlist = tests[day] dlist = diags[day] qlist = quars[day] for f in flist: if verbose: print(f) pl.plot(f.x[0], f.y[0], 'o', c=f.c, markersize=msize, **plot_args) # Plot sources pl.plot(f.x, f.y, '-', c=f.c, **plot_args) # Plot transmission lines if f.i: # If this person is infected pl.plot(f.x[1], f.y[1], '*', c=f.c, markersize=msize, **plot_args) # Plot targets for tdq in tlist: pl.plot(tdq.d, tdq.t, 'o', c=tdq.c, markersize=msize * 2, fillstyle='none') # Tested; No alpha for this for tdq in dlist: pl.plot(tdq.d, tdq.t, 's', c=tdq.c, markersize=msize * 1.2, **plot_args) # Diagnosed for tdq in qlist: pl.plot(tdq.d, tdq.t, 'x', c=tdq.c, markersize=msize * 2.0) # Quarantine; no alpha for this pl.plot([0, day], [0.5, 0.5], c='k', lw=5) # Plot the endless march of time if animate: # Whether to animate pl.pause(delay) return fig
def plot_pop(do_show=False, pause=0.2): ''' Plot an example population ''' plotconnections = True n = 5000 alpha = 0.5 # indices = pl.arange(1000) pl.seed(1) indices = pl.randint(0, n, 20) max_contacts = {'S': 20, 'W': 10} population = sp.make_population(n=n, max_contacts=max_contacts) nside = np.ceil(np.sqrt(n)) x, y = np.meshgrid(np.arange(nside), np.arange(nside)) x = x.flatten()[:n] y = y.flatten()[:n] people = list(population.values()) for p, person in enumerate(people): person['loc'] = dict(x=x[p], y=y[p]) ages = np.array([person['age'] for person in people]) f_inds = [ind for ind, person in enumerate(people) if not person['sex']] m_inds = [ind for ind, person in enumerate(people) if person['sex']] if do_show: use_terrain = False if use_terrain: import matplotlib.pyplot as plt import matplotlib.colors as colors colors_undersea = plt.cm.terrain(np.linspace(0, 0.17, 256)) colors_land = plt.cm.terrain(np.linspace(0.25, 1, 256)) all_colors = np.vstack((colors_undersea, colors_land)) terrain_map = colors.LinearSegmentedColormap.from_list( 'terrain_map', all_colors) pl.set_cmap(terrain_map) fig = pl.figure(figsize=(24, 18)) pl.subplot(111) minval = 0 # ages.min() maxval = 100 # ages.min() colors = sc.vectocolor(ages, minval=minval, maxval=maxval) for i, inds in enumerate([f_inds, m_inds]): pl.scatter(x[inds], y[inds], marker='os'[i], c=colors[inds]) pl.clim([minval, maxval]) pl.colorbar() if plotconnections: lcols = dict(H=[0, 0, 0], S=[0, 0.5, 1], W=[0, 0.7, 0], C=[1, 1, 0]) for index in indices: person = people[index] contacts = person['contacts'] lines = [] for lkey in lcols.keys(): for contactkey in contacts[lkey]: contact = population[contactkey] tmp = pl.plot( [person['loc']['x'], contact['loc']['x']], [person['loc']['y'], contact['loc']['y']], c=lcols[lkey], alpha=alpha) lines.append(tmp) if pause: pl.pause(pause) return fig
# transition_steps = 0 # heat+freeze transition_steps = 4 # increase if closest ends of the color maps are far apart, values to try: 4, 8, 16 transition = mplt.colors.LinearSegmentedColormap.from_list( "transition", [cmap1(1.), cmap2(0)])(np.linspace(0, 1, transition_steps)) colors = np.vstack((colors1, transition, colors2)) colors = np.flipud(colors) new_cmap = mplt.colors.LinearSegmentedColormap.from_list(new_cmap_name, colors) cmap = new_cmap # Assign colors to age groups age_cutoffs = np.arange( 0, 101, 10) # np.array([0, 4, 6, 18, 22, 30, 45, 65, 80, 90, 100]) if discrete: raw_colors = sc.vectocolor(len(age_cutoffs), cmap=cmap) colors = [] for age in sim.people.age: ind = sc.findinds(age_cutoffs <= age)[-1] colors.append(raw_colors[ind]) colors = np.array(colors) else: age_map = sim.people.age * 0.1 + np.sqrt(sim.people.age) colors = sc.vectocolor(age_map, cmap=cmap) # Create the legend if plot_stacked: ax = fig.add_axes([0.85, 0.05, 0.14, 0.93]) elif len(keys_to_plot) % 2 != 0: ax = fig.add_axes([0.85, 0.05, 0.14, 0.93]) else:
def plot_surface(ax, dfr, col=0, colval=0): ''' Plot one of the surfaces ''' all_tests = dfr['cum_tests'] quar_tests = cs.d_calcs.quar_test * dfr['cum_quarantined'] non_quar_tests = all_tests - quar_tests scaled_nq_tests = non_quar_tests * 1000 / kcpop / n_days x = scaled_nq_tests y = np.array(dfr['trprob']) z = np.array(dfr['r_eff']) min_x = 0 min_y = 0 max_x = 6 max_y = 1.0 min_z = 0.25 max_z = 1.75 eps = 0.08 npts = 100 xi = np.linspace(min_x, max_x, npts) yi = np.linspace(min_y, max_y, npts) xx, yy = np.meshgrid(xi, yi) zz = gauss2d(x, y, z, xi, yi, eps=eps, xscale=max_x - min_x, yscale=max_y - min_y) im = ax.contourf(xx, yy, zz, cmap=colormap, levels=np.linspace(min_z, max_z, 300)) if col == 0: ax.set_ylabel('Contact tracing probability at home and work', labelpad=20) ax.set_yticks([0.0, 0.2, 0.4, 0.6, 0.8, 1.0]) if col == 1: ax.set_xlabel('Number of routine tests per 1,000 people per day') axhandle = ax.contour(xx, yy, zz, levels=[1], colors='k', linestyles=':', linewidths=2) if col == 0: ax.clabel(axhandle, fmt=r'$R_{e}=%2.1d$', colors='k', fontsize=18, manual=[(1, 0.3)]) elif col == 1: ax.clabel(axhandle, fmt=r'$R_{e}=%2.1d$', colors='k', fontsize=18, manual=[(2, 0.4)]) else: ax.clabel(axhandle, fmt=r'$R_{e}=%2.1d$', colors='k', fontsize=18, manual=[(3, 0.5)]) ax.set_title('%i%% mobility' % (colval * 100)) scolors = sc.vectocolor(z, cmap=colormap, minval=min_z, maxval=max_z) ax.scatter(x, y, marker='o', c=scolors, edgecolor=[0.3] * 3, s=15, linewidth=0.1, alpha=0.5) # Add the measles plots ax.set_xlim([min_x, max_x]) ax.set_ylim([min_y, max_y]) if verbose: print( f'Plot: {col}, min zz: {zz.min():0.2f}; max zz: {zz.max():0.2f}; min z: {z.min():0.2f}; max z: {z.max():0.2f}' ) return im
'pop_size': pop_size, # start with a small pool 'pop_type': pop_type, # synthpops, hybrid 'contacts': contacts[pop_type], 'n_days': 1, # 'rand_seed': None, } # Create sim sim = cv.Sim(pars=pars) sim.initialize() mnet = pymnet.MultilayerNetwork(aspects=1) # fig = pl.figure(figsize=(16,16), dpi=120) mapping = dict(a='All', h='Households', s='Schools', w='Work', c='Community') colors = sc.vectocolor(sim.people.age, cmap='turbo') keys = list(contacts[pop_type].keys()) keys.remove('c') # nrowcol = np.ceil(np.sqrt(len(keys))) G = nx.MultiGraph() node_set = set() home_set = set() school_set = set() work_set = set() sample_size = 1 for p in range(sample_size):
def plot_histograms(self, start_day=None, end_day=None, bins=None, width=0.8, fig_args=None, font_size=18): ''' Plots a histogram of the number of transmissions. Args: start_day (int/str): the day on which to start counting people who got infected end_day (int/str): the day on which to stop counting people who got infected bins (list): bin edges to use for the histogram width (float): width of bars fig_args (dict): passed to pl.figure() font_size (float): size of font ''' # Process targets n_targets = self.count_targets(start_day, end_day) # Handle bins if bins is None: max_infections = n_targets.max() bins = np.arange(0, max_infections+2) # Analysis counts = np.histogram(n_targets, bins)[0] bins = bins[:-1] # Remove last bin since it's an edge total_counts = counts*bins n_bins = len(bins) index = np.linspace(0, 100, len(n_targets)) sorted_arr = np.sort(n_targets) sorted_sum = np.cumsum(sorted_arr) sorted_sum = sorted_sum/sorted_sum.max()*100 change_inds = sc.findinds(np.diff(sorted_arr) != 0) max_labels = 15 # Maximum number of ticks and legend entries to plot # Plotting fig_args = sc.mergedicts(dict(figsize=(24,15)), fig_args) pl.rcParams['font.size'] = font_size fig = pl.figure(**fig_args) pl.set_cmap('Spectral') pl.subplots_adjust(left=0.08, right=0.92, bottom=0.08, top=0.92) colors = sc.vectocolor(n_bins) pl.subplot(1,2,1) w05 = width*0.5 w025 = w05*0.5 pl.bar(bins-w025, counts, width=w05, facecolor='k', label='Number of events') for i in range(n_bins): label = 'Number of transmissions (events × transmissions per event)' if i==0 else None pl.bar(bins[i]+w025, total_counts[i], width=w05, facecolor=colors[i], label=label) pl.xlabel('Number of transmissions per person') pl.ylabel('Count') if n_bins<max_labels: pl.xticks(ticks=bins) pl.legend() pl.title('Numbers of events and transmissions') pl.subplot(2,2,2) total = 0 for i in range(n_bins): pl.bar(bins[i:], total_counts[i], width=width, bottom=total, facecolor=colors[i]) total += total_counts[i] if n_bins<max_labels: pl.xticks(ticks=bins) pl.xlabel('Number of transmissions per person') pl.ylabel('Number of infections caused') pl.title('Number of transmissions, by transmissions per person') pl.subplot(2,2,4) pl.plot(index, sorted_sum, lw=3, c='k', alpha=0.5) n_change_inds = len(change_inds) label_inds = np.linspace(0, n_change_inds, max_labels).round() # Don't allow more than this many labels for i in range(n_change_inds): if i in label_inds: # Don't plot more than this many labels label = f'Transmitted to {bins[i+1]:n} people' else: label = None pl.scatter([index[change_inds[i]]], [sorted_sum[change_inds[i]]], s=150, zorder=10, c=[colors[i]], label=label) pl.xlabel('Proportion of population, ordered by the number of people they infected (%)') pl.ylabel('Proportion of infections caused (%)') pl.legend() pl.ylim([0, 100]) pl.grid(True) pl.title('Proportion of transmissions, by proportion of population') pl.axes([0.30, 0.65, 0.15, 0.2]) berry = [0.8, 0.1, 0.2] dirty_snow = [0.9, 0.9, 0.9] start_day = self.day(start_day, which='start') end_day = self.day(end_day, which='end') pl.axvspan(start_day, end_day, facecolor=dirty_snow) pl.plot(self.sim_results['t'], self.sim_results['cum_infections'], lw=2, c=berry) pl.xlabel('Day') pl.ylabel('Cumulative infections') return fig