def show_ratio_box(self, ratio, area, start, end, copy_to_clipboard=True): """Displays a text box that shows the ratio and of areas between two spectra within start and end. """ if self.anchored_box: self.anchored_box.set_visible(False) self.anchored_box = None child_boxes = [] if ratio is None: text = "Invalid selection, \nno ratio could \nbe calculated" child_boxes.append( offsetbox.TextArea(text, textprops=dict(color="k", size=12))) else: ratio_round = 9 # Round decimal number text = f"Difference: {round(area, 2)}\n" \ f"Ratio: {round(ratio, ratio_round)}\n" \ f"Interval: [{round(start, 2)}, {round(end, 2)}]" child_boxes.append( offsetbox.TextArea(text, textprops=dict(color="k", size=12))) if copy_to_clipboard: self.clipboard.setText(str(round(ratio, ratio_round))) text_2 = "\nRatio copied to clipboard." child_boxes.append( offsetbox.TextArea(text_2, textprops=dict(color="k", size=10))) box = offsetbox.VPacker(children=child_boxes, align="center", pad=0, sep=0) self.anchored_box = offsetbox.AnchoredOffsetbox( loc=2, child=box, pad=0.5, frameon=False, bbox_to_anchor=(1.0, 1.0), bbox_transform=self.axes.transAxes, borderpad=0., ) self.axes.add_artist(self.anchored_box) self.axes.add_artist(self.leg) self.canvas.draw_idle()
def legend_outside(ncol, extra_text): handles, labels = plt.gca().get_legend_handles_labels() lgd = plt.legend(handles, labels, loc='upper center', ncol=ncol, bbox_to_anchor=(-0.15, -0.3)) import matplotlib.offsetbox as offsetbox txt=offsetbox.TextArea(extra_text, textprops=dict(size=lgd.get_texts()[0].get_fontsize())) box = lgd._legend_box box.get_children().append(txt) box.set_figure(box.figure) return lgd
def __init__(self, length=1, extent = 0.03, label="", loc=2, ax=None, pad=0.4, borderpad=0.5, ppad = 0, sep=2, prop=None, frameon=True, linekw={}, **kwargs): if not ax: ax = plt.gca() trans = ax.get_xaxis_transform() size_bar = offbox.AuxTransformBox(trans) line = Line2D([0,length],[0,0], **linekw) size_bar.add_artist(line) txt = offbox.TextArea(label, minimumdescent=False, textprops=dict(color="black",size=14, fontweight='bold')) self.vpac = offbox.VPacker(children=[size_bar,txt], align="center", pad=ppad, sep=sep) offbox.AnchoredOffsetbox.__init__(self, loc, pad=pad, borderpad=borderpad, child=self.vpac, prop=prop, frameon=frameon, **kwargs)
def _ano_fid_note(self, ax, hitf): hM = hitf['hit'].max() for fid in hitf.index: hv = hitf.ix[fid].hit # and hitf.ix[fid].nonhit hva = hM * (fid % 5) * 0.05 hva = hva if fid % 2 is 0 else -hva oft = ofb.TextArea("{}[{}]".format(fid, hitf.ix[fid].sid)) bbprop = dict(boxstyle='round,pad=0.1', ec='g') ab = ofb.AnnotationBbox(oft, (fid, hv), xycoords='data', xybox=(fid, hv + hva), boxcoords=("data"), bboxprops=bbprop, arrowprops=dict(arrowstyle="->", color='darkblue')) ax.add_artist(ab)
def legend(ax, legend=None, legend_text=None, **kwargs): """ Place legend on an axis Additional kwargs are ignored to facilitate argument passing. Parameters ---------- legend : str either a mpl name for the location or tuple coordinates legend_text : str additional text to place above legend semicolons are replaced with newlines """ if legend is None: return bbox = {} # Specify x,y coordinates for upper left corner if '(' in legend: bbox['bbox_to_anchor'] = tuple( float(x) for x in legend[1:-1].split(',')) bbox['bbox_transform'] = ax.transAxes legend = 2 handles, labels = ax.get_legend_handles_labels() legend = ax.legend(handles, labels, frameon=False, loc=legend, prop={'size': 18}, labelspacing=0.25, ncol=1 if len(labels) < 6 else 2, **bbox) # Add extra text above legend. if legend_text: for sub in legend_text.split(";")[::-1]: txt = offsetbox.TextArea(sub, {'size': 18}) box = legend._legend_box box.get_children().insert(0, txt) box.set_figure(box.figure) return legend
ax1.plot(xf, expfit(xf, *popt1), 'k', label=fitlabel) ax1.plot(xf, expfit(xf, *popt1) + 41.75, 'r--', label='confidence') ax1.plot(xf, expfit(xf, *popt1) - 41.75, 'r--') ax1.plot(lit_vec, interval, 'ko', color='b', label='lithology to time interval') ax1.errorbar(lit_vec, interval, fmt='o', yerr=16, color='b') resiudal = interval - expfit(lit_vec, *popt1) ss_res = np.sum(resiudal**2) ss_tot = np.sum((interval - np.mean(interval))**2) r_squared = 1 - (ss_res / ss_tot) leg = ax1.legend(loc=[0.05, .8]) leg_r_squerd = "r^2 = {0:.3f}".format(r_squared) txt = offsetbox.TextArea(leg_r_squerd) box = leg._legend_box box.get_children().append(txt) box.set_figure(box.figure) ax1.set_xlim(0, 1.1) ax1.set_xlabel('$\\frac{gravel\:percentage}{gravel\:percentage\:maximum}$', fontsize=25) ax1.set_ylabel('days', fontsize=22) ax2.semilogy(xf, expfit(xf, *popt1), 'k', label=fitlabel) ax2.semilogy(xf, expfit(xf, *popt1) + 41.75, 'r--', label='confidence') ax2.semilogy(xf, expfit(xf, *popt1) - 41.75, 'r--') ax2.semilogy(lit_vec, interval, 'ko', color='b',
for idx_nn, nn in enumerate(input_trainings): data = datas[idx_nn] plt.scatter(data['epoch'], data['%s' % met], linestyle='-', label='Training, lr = %.1e' % (param_dicts[idx_nn]['lr']), color=colors[idx_nn]) axes = plt.gca() axes.set_xlim(-1, 50) ymin, ymax = axes.get_ylim() xmin, xmax = axes.get_xlim() current_dict = param_dicts[0] # plt.text(xmin+abs(xmin)*0.1,ymax*0.95,'%s'%(str(current_dict['architecture'])), fontsize=20) # plt.text(xmin+abs(xmin)*0.1,ymax*0.93,'%s'%(current_dict['loss']), fontsize=20) # plt.text(xmin+abs(xmin)*0.1,ymax*0.91,'dropout %s'%(current_dict['dropout']), fontsize=20) text0 = str(current_dict['architecture']) text1 = str(current_dict['loss']) text2 = 'dropout %s' % (current_dict['dropout']) textbox = text0 + '\n' + text1 + '\n' + text2 plt.xlabel('# epochs') plt.ylabel('%s' % met) leg = plt.legend() txt = offsetbox.TextArea("%s" % textbox, textprops=dict(size=20)) box = leg._legend_box box.get_children().append(txt) box.set_figure(box.figure) pltsavename = options.out_dir + '/learning_rate_loss_vs_epoch_%s' % (met) plt.savefig(pltsavename + '.png') plt.savefig(pltsavename + '.pdf') plt.clf()
def plot_lines(data, labels, xlabel, ylabel, title, suptitle, conf=None, additionals=None, path=None, ls='-', begin=1): fig, ax = plt.subplots() fig.set_size_inches(15, 10) for i, d in enumerate(data): ax.plot(range(begin, len(d) + begin), d, color=colors[i], ls=ls, label=labels[i]) legend = ax.legend(framealpha=.3) if conf is not None: import matplotlib.offsetbox as offsetbox keys = [] vals = [] for k, v in conf.items(): if k != 'features': if v is None: v = "none" elif type(v) == bool: if v == True: v = "true" else: v = "false" if k == 'n_input': keys.append(offsetbox.TextArea(r"$features$")) vals.append( offsetbox.TextArea(r"${}$".format(v), textprops={'size': 10})) elif k == 'unit_range': if v is not None: keys.append(offsetbox.TextArea(r"$unit range$")) vals.append( offsetbox.TextArea(r"${}-{}$".format(v[0], v[1]), textprops={'size': 10})) elif k == 'units': keys.append(offsetbox.TextArea(r"$hidden\_layers$")) vals.append( offsetbox.TextArea(r"${}$".format(len(v)), textprops={'size': 10})) for l, layer in enumerate(v): keys.append( offsetbox.TextArea(r"$units_{HL{%d}}$" % (l))) vals.append( offsetbox.TextArea(r"${}$".format(layer), textprops={'size': 10})) elif k == 'learning_rate': keys.append(offsetbox.TextArea(r"$\eta$")) vals.append( offsetbox.TextArea(r"${}$".format(v), textprops={'size': 10})) else: keys.append( offsetbox.TextArea(r"${}$".format(k.replace("_", "\_")))) vals.append( offsetbox.TextArea(r"${}$".format(v), textprops={'size': 10})) if additionals is not None: keys.append( offsetbox.TextArea(r"$E_{min, train}$", textprops={'color': colors[0]})) keys.append( offsetbox.TextArea(r"$E_{min, valid}$", textprops={'color': colors[1]})) vals.append( offsetbox.TextArea(r"${:.2e} {} (epoch {})$".format( additionals[0][1], " ", additionals[0][0]), textprops={ 'color': colors[0], 'size': 10, 'weight': 'bold' })) vals.append( offsetbox.TextArea(r"${:.2e} {} (epoch {})$".format( additionals[1][1], " ", additionals[1][0]), textprops={ 'color': colors[1], 'size': 10, 'weight': 'bold' })) vp1 = offsetbox.VPacker(children=keys, align="left", pad=0, sep=3) vp2 = offsetbox.VPacker(children=vals, align="right", pad=0, sep=5) hp = offsetbox.HPacker(children=(vp1, vp2), align="right", pad=5.76, sep=28.8) box = legend._legend_box box.get_children()[1].get_children()[0].get_children().append(hp) box.set_figure(box.figure) if additionals is not None: add = np.array(additionals).T plt.vlines(x=add[0], ymin=0, ymax=add[1], colors=colors, linestyles='dotted') plt.xlabel(xlabel) plt.ylabel(ylabel if ylabel not in ylabels else ylabels[ylabel]) plt.xlim(xmax=len(data[0])) plt.title(title) if path is not None: plt.savefig(path, dpi=300) plt.show()
def plot_double(data1, data2, label1, label2, era, output_dir, is_cumul, is_sequential, split_by_era, is_recorded = True): if split_by_era and not is_cumul: # Applies only when is_cumul is set to true return fig, ax = plt.subplots(figsize = (10, 8), dpi = 120) figure_type = 'recorded' if is_recorded else 'delivered' subject = '{}{}'.format('cumulative_' if is_cumul else '', figure_type) if split_by_era and is_cumul: subject += '_era' ylim = ( 10 ** math.floor(math.log10(min(min(map(lambda x: x[subject], data1)), min(map(lambda x: x[subject], data2))))), 10 ** math.ceil (math.log10(max(max(map(lambda x: x[subject], data1)), max(map(lambda x: x[subject], data2))))) ) ax.set_ylim(ylim) run_count = 0 runmap = {} expected_cumul = 0. has_labeled = False for era_idx, era_key in enumerate(runs[era]): if not has_labeled: has_labeled = True label1_on_plot = label1 label2_on_plot = label2 else: label1_on_plot = '' label2_on_plot = '' run_set_1 = set(map(lambda x: x['run'], filter(lambda y: y['era'] == era_key, data1))) run_set_2 = set(map(lambda x: x['run'], filter(lambda y: y['era'] == era_key, data2))) run_set = run_set_1 | run_set_2 xrange_run_real = list(sorted(list(run_set))) run_len = len(run_set) xrange_run_count = list(range(run_count, run_count + run_len)) xrange_plot = xrange_run_count if is_sequential else xrange_run_real yvals_1_filtered = list(sorted(filter(lambda y: y['era'] == era_key, data1), key = lambda record: record['run'])) yvals_2_filtered = list(sorted(filter(lambda y: y['era'] == era_key, data2), key = lambda record: record['run'])) yvals_1, yvals_2 = [], [] for xval in xrange_run_real: if xval not in run_set_1: if yvals_1: yvals_1.append(yvals_1[-1]) else: yvals_1.append(0.) else: yval_candidate = [ yval for yval in yvals_1_filtered if yval['run'] == xval ] assert(len(yval_candidate) == 1) yvals_1.append(yval_candidate[0][subject]) if xval not in run_set_2: if yvals_2: yvals_2.append(yvals_2[-1]) else: yvals_2.append(0.) else: yval_candidate = [ yval for yval in yvals_2_filtered if yval['run'] == xval ] assert(len(yval_candidate) == 1) yvals_2.append(yval_candidate[0][subject]) ax.plot(xrange_plot, yvals_2, c = '#1f77b4', ls = '-', label = label2_on_plot) ax.plot(xrange_plot, yvals_1, c = '#ff7f0e', ls = '--', label = label1_on_plot) if is_cumul and xrange_plot: if split_by_era: expected_cumul = runs[era][era_key][figure_type] * 1e3 else: expected_cumul += runs[era][era_key][figure_type] * 1e3 ax.hlines( y = expected_cumul, xmin = xrange_plot[0], xmax=xrange_plot[-1], color = 'black', linestyle=':', ) runmap[era_key] = (run_count, run_count + run_len) run_count += run_len if not is_sequential: runmap = { era_key : runs[era][era_key]['ranges'] for era_key in runs[era] } era_colors = ['gray', 'olive'] for era_idx, era_key in enumerate(runs[era]): ax.fill_between(runmap[era_key], *ylim, facecolor = era_colors[era_idx % 2], alpha = 0.2) runb_avg = (runmap[era_key][0] + runmap[era_key][1]) / 2 ab = obox.AnnotationBbox(obox.TextArea(era_key), xy = (runb_avg, 10 * ylim[0]), xycoords = 'data') ax.add_artist(ab) plt.xlabel('Run {}'.format('index' if is_sequential else 'number')) plt.ylabel('{} [/pb]'.format(subject.replace('_', ' ').capitalize())) plt.yscale('log') plt.grid(True) plt.legend(loc = 'best') plt.title('Run {}'.format(era)) out_path = os.path.join(output_dir, '{}_vs_{}_{}'.format(label1, label2, subject)) if is_sequential: out_path += '_seq' out_path += '.png' plt.savefig(out_path, bbox_inches = 'tight') logging.info('Created plot {}'.format(out_path)) plt.close()
def plot_ratio(data1, data2, label1, label2, era, output_dir, is_cumul, is_recorded = True): fig, ax = plt.subplots(figsize = (10, 8), dpi = 120) subject = '{}{}'.format('cumulative_' if is_cumul else '', 'recorded' if is_recorded else 'delivered') run_count = 0 runmap = {} min_ratio, max_ratio = 1e+6, -1 has_labeled = False for era_idx, era_key in enumerate(runs[era]): run_len_1 = len(list(map(lambda x: x['run'], filter(lambda y: y['era'] == era_key, data1)))) run_len_2 = len(list(map(lambda x: x['run'], filter(lambda y: y['era'] == era_key, data2)))) assert(run_len_1 == run_len_2) run_len = run_len_1 data2_vals = list(map(lambda x: x[subject], filter(lambda y: y['era'] == era_key, data2))) data1_vals = list(map(lambda x: x[subject], filter(lambda y: y['era'] == era_key, data1))) data2_data1_ratio = [data2_vals[run_idx] / data1_vals[run_idx] for run_idx in range(run_len)] if not data2_data1_ratio: continue data2_data1_ratio_min = min(data2_data1_ratio) data2_data1_ratio_max = max(data2_data1_ratio) if data2_data1_ratio_min < min_ratio: min_ratio = data2_data1_ratio_min if data2_data1_ratio_max > max_ratio: max_ratio = data2_data1_ratio_max ratio_mean = np.mean(data2_data1_ratio) if abs(1. - ratio_mean) > 1e-3 and not is_cumul: ax.hlines( y = ratio_mean, xmin = run_count, xmax = run_count + run_len, color = 'black', linestyle='--', label = 'Average prescale: %.1f' % ratio_mean ) if not has_labeled: has_labeled = True label_on_plot = '{} / {}'.format(label2, label1) else: label_on_plot = '' ax.plot( list(range(run_count, run_count + run_len)), data2_data1_ratio, c = '#1f77b4', ls = '-', label = label_on_plot ) runmap[era_key] = (run_count, run_count + run_len) run_count += run_len ylim = (min_ratio * (0.7 if is_cumul else 0), max_ratio / 0.7) ax.set_ylim(ylim) era_colors = ['gray', 'olive'] for era_idx, era_key in enumerate(runs[era]): if era_key not in runmap: continue ax.fill_between(runmap[era_key], *ylim, facecolor = era_colors[era_idx % 2], alpha = 0.2) runb_avg = (runmap[era_key][0] + runmap[era_key][1]) / 2 ab = obox.AnnotationBbox(obox.TextArea(era_key), xy = (runb_avg, ylim[1] * 0.8), xycoords = 'data') ax.add_artist(ab) plt.xlabel('run (indexed)') plt.ylabel('Ratio of {}'.format(subject.replace('_', ' '))) plt.grid(True) plt.legend(loc = 'best') plt.title('Run {}'.format(era)) out_path = os.path.join(output_dir, '{}_vs_{}_{}_ratio.png'.format(label1, label2, subject)) plt.savefig(out_path, bbox_inches = 'tight') logging.info('Created plot {}'.format(out_path)) plt.close()
def plot_single(data, label, era, output_dir, is_cumul, is_sequential, split_by_era, is_recorded = True): if split_by_era and not is_cumul: # Applies only when is_cumul is set to true return fig, ax = plt.subplots(figsize = (10, 8), dpi = 120) figure_type = 'recorded' if is_recorded else 'delivered' subject = '{}{}'.format('cumulative_' if is_cumul else '', figure_type) if split_by_era and is_cumul: subject += '_era' ylim = ( 10 ** math.floor(math.log10(min(map(lambda x: x[subject], data)))), 10 ** math.ceil (math.log10(max(map(lambda x: x[subject], data)))) ) run_count = 0 runmap = {} expected_cumul = 0. has_labeled = False for era_idx, era_key in enumerate(runs[era]): if not has_labeled: has_labeled = True label_on_plot = label else: label_on_plot = '' run_len = len(list(map(lambda x: x['run'], filter(lambda y: y['era'] == era_key, data)))) xrange_run_count = list(range(run_count, run_count + run_len)) xrange_run_real = list(map(lambda x: x['run'], filter(lambda y: y['era'] == era_key, data))) xrange_plot = xrange_run_count if is_sequential else xrange_run_real ax.plot( xrange_plot, list(map(lambda x: x[subject], filter(lambda y: y['era'] == era_key, data))), c = '#1f77b4', ls = '-', label = label_on_plot ) if is_cumul and xrange_plot: if split_by_era: expected_cumul = runs[era][era_key][figure_type] * 1e3 else: expected_cumul += runs[era][era_key][figure_type] * 1e3 ax.hlines( y = expected_cumul, xmin = xrange_plot[0], xmax=xrange_plot[-1], color = 'black', linestyle=':', ) if expected_cumul > ylim[1]: ylim = (ylim[0], expected_cumul * 5) runmap[era_key] = (run_count, run_count + run_len) run_count += run_len ax.set_ylim(ylim) if not is_sequential: runmap = { era_key : runs[era][era_key]['ranges'] for era_key in runs[era] } era_colors = ['gray', 'olive'] for era_idx, era_key in enumerate(runs[era]): ax.fill_between(runmap[era_key], *ylim, facecolor = era_colors[era_idx % 2], alpha = 0.2) runb_avg = (runmap[era_key][0] + runmap[era_key][1]) / 2 ab = obox.AnnotationBbox(obox.TextArea(era_key), xy = (runb_avg, 10 * ylim[0]), xycoords = 'data') ax.add_artist(ab) plt.xlabel('Run {}'.format('index' if is_sequential else 'number')) plt.ylabel('{} [/pb]'.format(subject.replace('_', ' ').capitalize())) plt.yscale('log') plt.grid(True) plt.legend(loc = 'best') plt.title('Run {}'.format(era)) out_path = os.path.join(output_dir, '{}_{}'.format(label, subject)) if is_sequential: out_path += '_seq' out_path += '.png' plt.savefig(out_path, bbox_inches = 'tight') logging.info('Created plot {}'.format(out_path)) plt.close()
def PlotSingleNeuCompositionCompare(E,body,param,sparam = PC.PhysicsConstants()): """ Plots the composition of a single mass neutrino state. E : neutrino energy [eV] body : body with the asociated density profile. param : set of physical parameters used to make the plot. param can be a list. sparam : standard parameters """ fig = plt.figure() ax = plt.subplot(111) mpl.rcParams['axes.labelsize'] = "x-large" mpl.rcParams['xtick.labelsize'] = "x-large" mpl.rcParams['legend.fontsize'] = "small" mpl.rcParams['font.size'] = 18 colors = ['b', 'g', 'r', 'c', 'm', 'y', 'k'] linestyles = ['--','-.',':','-','-'] #B Initializing variables param.Refresh() fM2 = no.flavorM2(param) sparam.Refresh() fM2STD = no.flavorM2(sparam) R = np.arange(1.0,0.01,-0.001) Rho = map(lambda r : body.rdensity(r), R) #E Initializing variables #B Estimating Energy Scale if(E/param.MeV <= 500.0) : scale = param.MeV scalen = "MeV" elif(E/param.GeV <= 500.0) : scale = param.GeV scalen = "GeV" else : scale = param.TeV scalen = "TeV" #E Estimating Energy Scale #B Adding title #tit = "Energy : "+str(E/scale)+" "+scalen+ " Parameters :"#+" $\\th_{12}$ = " + str(param.th12) + " $\\th_{23}$ = " + str(param.th23) + " $\\th_{13}$ = "+str(param.th13) #atit = [] #[[ atit.append(" $\\theta_{"+str(j)+str(i)+"}$ = "+format(param.th[j][i],'.4f')) for i in range(1,param.numneu+1) if i>j] for j in range(1,param.numneu+1) ] #[[ atit.append(" $\\Delta m^2_{"+str(i)+str(j)+"}$ = "+format(param.dm2[j][i],'.4f')) for i in range(1,param.numneu+1) if i>j and j == 1] for j in range(1,param.numneu+1) ] ##[[ atit.append(" $\\Delta m^2_{"+str(j)+str(i)+"}$ = "+format(param.dm2[j][i],'.4f')) for i in range(1,param.numneu+1) if i>j and j == 1] for j in range(1,param.numneu+1) ] #for i in range(len(atit)): # tit = tit + atit[i] #plt.suptitle(tit,horizontalalignment='center') #E Adding title ##B PLOTTING MASS BASIS AS FUNCTION OF FLAVOR BASIS for i in [1]: #fig.add_subplot(2,param.numneu,i+1) flavor = False NeuComp = map(lambda x : no.NeuComposition(i,E, x, body, fM2, param,flavor),R) plt.xlabel(r"$\rho \mathrm{[g/cm^{3}]}$") pp = [] for k in range(param.numneu): kNeuComp = map(lambda x: x[k],NeuComp) # log interpolator rholog = gt.LogSpaceEnergies(float(Rho[0]),float(Rho[-1]),100) #print rholog , float(Rho[0]),float(Rho[-1]) rholog[-1] = Rho[-1] inter_neu = interpolate.interp1d(Rho,kNeuComp) logkNeuComp = map(inter_neu,rholog) if k == 3: #pp.append(plt.plot(Rho,kNeuComp,'o-',color = 'r',markevery = 10,markeredgewidth = 0.0, ms = 2.0)) pp.append(plt.plot(rholog,logkNeuComp,'x-',color = 'r',markevery = 10,markeredgewidth = 0.0, ms = 2.0,aa = True,solid_joinstyle = 'bevel')) elif k == 4: pp.append(plt.plot(Rho,kNeuComp,'o-',color = 'r',markevery = 10,markeredgewidth = 0.0, ms = 2.0)) else : pp.append(plt.plot(Rho,kNeuComp,linestyle = linestyles[k] ,color = 'r')) if i<=2 : NeuCompSTD = map(lambda x : no.NeuComposition(i,E, x, body, fM2STD, sparam,flavor),R) for k in range(sparam.numneu): kNeuCompSTD = map(lambda x: x[k],NeuCompSTD) plt.plot(Rho,kNeuCompSTD, color = 'k', linestyle = linestyles[k]) #Solar density #ps = plt.vlines(150, 0.0, 1.0, linestyle = "dashed", label = r"$\rho_S$") #B plt format plt.title(r"Composition of $\nu_"+str(i+1)+"$") plt.semilogx() plt.ylim(0.0,1.0) plt.xlim(1.0,150.0) plt.yticks(np.arange(0.0,1.1,0.1)) xtt = [1.0,5.0,10.0,30.0,100.0]#,150.0] #plt.xticks(xtt) ax.set_xticks(xtt) ax.set_xticklabels(['$1$','$5$','$10$','$30$','$100$'])#,'$\\rho_\\odot = 150$']) #plt.xscale() #B LEGEND plots = [] for e in pp : plots.append(e[0]) #plots.append(ps) leg = ["$\\nu_e$","$\\nu_\mu$","$\\nu_\\tau$"] ss = ["$\\nu_{s"+str(i)+"}$" for i in np.arange(1,param.numneu-3+1,1)] if ss != []: leg.extend(ss) leg = plt.legend(plots,leg,loc = 6,fancybox=True,bbox_to_anchor = (0.05, 0.75)) leg.get_frame().set_alpha(0.25) #E LEGEND #E plt format #B EXTRA LEGEND box1t = osb.TextArea("STD", textprops=dict(color="k")) box1d = osb.DrawingArea(60, 20, 0, 0) el1 = ptc.Ellipse((10, 10), width=5, height=5, angle=0, fc="k", edgecolor = 'none') box1d.add_artist(el1) box2t = osb.TextArea("2+3", textprops=dict(color="k")) box2d = osb.DrawingArea(60, 20, 0, 0) el2 = ptc.Ellipse((10, 10), width=5, height=5, angle=0, fc="r", edgecolor = 'none') box2d.add_artist(el2) box1 = osb.HPacker(children=[box1t, box1d], align="center", pad=0, sep=1) box2 = osb.HPacker(children=[box2t, box2d], align="center", pad=0, sep=5) anchored_box1 = osb.AnchoredOffsetbox(loc=9, child=box1, pad=5.0, frameon=False, #bbox_to_anchor=(0., 1.02), #bbox_transform=ax.transAxes, borderpad=0., ) anchored_box2 = osb.AnchoredOffsetbox(loc=9, child=box2, pad=6.0, frameon=False, #bbox_to_anchor=(0., 1.02), #bbox_transform=ax.transAxes, borderpad=0., ) ax.add_artist(anchored_box1) ax.add_artist(anchored_box2) #E EXTRA LEGEND ##E PLOTTING MASS BASIS AS FUNCTION OF FLAVOR BASIS #plt.suptitle("*Dashed colored lines are 3-flavor standard oscillations.", x = 0.15, y = 0.03) path = "../plots/" filename = "PlotNeuComposition_E_"+str(E/scale)+"_"+scalen+"_FIG2.eps" plt.savefig(path + filename, dpi = 1200)
def PlotOscProb(iineu,Enumin,Enumax,param, datapath = "../data/SunOscProbabilities/",plotpath = "../plots",plot_survival_probability = False,filename_append = '', fmt = 'eps'): """Plots P(neu_ineu -> neu_fneu) as a function of the Energy from an initial flavor state (ineu) to all final flavor states (fneu) on the sun # iineu : 0 (electron), 1 (muon), 2 (tau) # Enumin : minimum neutrino energy [eV] # Enumax : maximum neutrino energy [eV] # param : physics parameter set list [param_1,param_2,...,param_n] """ plt.cla() plt.clf() plt.close() #fig = plt.figure(figsize = (4*3+2,6)) fig = plt.figure(figsize = (10,7.5)) ax = plt.subplot(111) mpl.rcParams['axes.labelsize'] = "xx-large" mpl.rcParams['xtick.labelsize'] = "xx-large" mpl.rcParams['legend.fontsize'] = "small" mpl.rcParams['font.size'] = 30 ordmag = np.log10(Enumax)-np.log10(Enumin) npoints = 1000.0*ordmag # Figuring out best energy scale try: if(Enumax/param[0].MeV <= 500.0) : scale = param[0].MeV scalen = "MeV" elif(Enumax/param[0].GeV <= 1000.0) : scale = param[0].GeV scalen = "GeV" else : scale = param[0].GeV#param[0].TeV scalen = "GeV"#"TeV" except (TypeError,AttributeError): if(Enumax/param.MeV <= 500.0) : scale = param.MeV scalen = "MeV" elif(Enumax/param.GeV <= 1000.0) : scale = param.GeV scalen = "GeV" else : scale = param.TeV scalen = "TeV" try : Emin = Enumin/param[0].GeV Emax = Enumax/param[0].GeV except : Emin = Enumin/param.GeV Emax = Enumax/param.GeV neulabel = {0 : "e",1 : "\\mu", 2 : "\\tau",3 : "{s_1}",4 : "{s_2}",5 : "{s_3}"} sneulabel = {0 : "e",1 : "mu", 2 : "tau",3 : "s1",4 : "s2",5 : "s3"} # RK points #ERKstep = (np.log10(Enumax)-np.log10(Enumin))/(20.0) #ERK = np.arange(np.log10(Enumin),np.log10(Enumax),ERKstep) #ERK = map(lambda E : (10**E)/scale,ERK) #ERK.append(1000.0) ERK = gt.MidPoint(gt.LogSpaceEnergies(Enumin/scale,Enumax/scale, binnum = 200)) ERK = [ERK[i] for i in range(len(ERK)-1)] ERK = [ERK[i] for i in range(len(ERK)-1)] Estep = (Enumax-Enumin)/npoints Enu = np.arange(Enumin,Enumax,Estep)/scale # Generating plot #totalPRK = [0.0]*len(ERK) #colors = ['orange', 'r', 'k', 'c', 'm', 'y', 'k'] colors = ['b', 'r', 'g','c', 'm', 'y', 'k'] linestyles = ['--','-.',':','-..','-'] for ineu,fneu in [[iineu,0],[iineu,1],[iineu,2],[iineu,5]]:#[[0,0],[1,1],[2,2],[2,1],[2,0],[0,0]]: for i,p in enumerate(param): p.Refresh() plt.xlabel(r"$\mathrm{E}_\nu\mathrm{["+scalen+"]}$") plt.ylabel("$ \mathrm{Probability}$") if (p.name == "STD" or p.name == "STD_XXX" or p.numneu == 3) and fneu > 2: pass else : if fneu == 5: fneu = 3 PRK_3 = map(lambda E : float(no.InterOscProb(ineu,fneu,E*scale,p,datapath,Emin,Emax,filename_append = filename_append)),ERK) fneu = 4 PRK_4 = map(lambda E : float(no.InterOscProb(ineu,fneu,E*scale,p,datapath,Emin,Emax,filename_append = filename_append)),ERK) if p.numneu > 5 : fneu = 5 PRK_5 = map(lambda E : float(no.InterOscProb(ineu,fneu,E*scale,p,datapath,Emin,Emax,filename_append = filename_append)),ERK) else : PRK_5 = map(lambda E : 0.0*E ,ERK) #PRK = [PRK_3[i] + PRK_4[i] for i in range(len(PRK_3))] PRK = [PRK_3[i] + PRK_4[i] + PRK_5[i] for i in range(len(PRK_3))] if p.neutype == "neutrino": plt.plot(ERK,PRK,linestyle = linestyles[-1],label='$ \mathrm{P}(\\nu_'+neulabel[ineu]+'\\rightarrow \\nu_s)$',color = p.colorstyle, lw = 6,solid_joinstyle = 'bevel') elif p.neutype == "antineutrino": plt.plot(ERK,PRK,linestyle = linestyles[-1],label='$ \mathrm{P}(\\bar{\\nu}_'+neulabel[ineu]+'\\rightarrow \\bar{\\nu}_s)$',color = p.colorstyle, lw = 6,solid_joinstyle = 'bevel') else: print "Wrong neutrino type." quit() else : PRK = map(lambda E : no.InterOscProb(ineu,fneu,E*scale,p,datapath,Emin,Emax,filename_append = filename_append),ERK) if p.name == "STD" or p.name == "STD_XXX" or p.numneu == 3: plt.plot(ERK,PRK,linestyle = linestyles[fneu],color = p.colorstyle, lw = 4) else : if p.neutype == "neutrino": plt.plot(ERK,PRK,linestyle = linestyles[fneu],label='$ \mathrm{P}(\\nu_'+neulabel[ineu]+'\\rightarrow \\nu_'+neulabel[fneu]+')$', color = p.colorstyle, lw = 4) elif p.neutype == "antineutrino": plt.plot(ERK,PRK,linestyle = linestyles[fneu],label='$ \mathrm{P}(\\bar{\\nu}_'+neulabel[ineu]+'\\rightarrow \\bar{\\nu}_'+neulabel[fneu]+')$', color = p.colorstyle, lw = 4) else: print "Wrong neutrino type." quit() if plot_survival_probability : for p in param: P_surival = map(lambda E : no.InterOscProb(iineu,p.numneu,E*scale,p,datapath,Emin,Emax,filename_append = filename_append),ERK) plt.plot(ERK,P_surival,linestyle = 'solid', lw = 4, color = p.colorstyle,solid_joinstyle = 'bevel') plt.semilogx() #plt.loglog() plt.axis([Enumin/scale,Enumax/scale,0.0,1.0]) #fig.subplots_adjust(left=0.05, right=0.8,wspace = 0.35, top = 0.85, bottom = 0.15) #mpl.rcParams['legend.fontsize'] = "small" #plt.legend(bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.,fancybox = True) plt.legend(loc='upper right',fancybox = True) ######################### begin extra legend ########################## for i,p in enumerate(param): # removing everything after an _ (underscore) try: paramname = re.search('.*(?=_)',p.name).group(0) except AttributeError: paramname = p.name # create text box boxt = osb.TextArea(paramname, textprops=dict(color="k")) boxd = osb.DrawingArea(60, 20, 0, 0) el = ptc.Ellipse((10, 10), width=5, height=5, angle=0, fc=p.colorstyle, edgecolor = 'none') boxd.add_artist(el) box = osb.HPacker(children=[boxt, boxd], align="center", pad=0, sep=1) # anchor boxes anchored_box = osb.AnchoredOffsetbox(loc=2, child=box, pad=0.25, frameon=False, bbox_to_anchor=(0.0, 1.0-0.06*i), bbox_transform=ax.transAxes, borderpad=0., ) ax.add_artist(anchored_box) ########################## end extra legend ############################## fig.subplots_adjust(bottom = 0.12, top = 0.95, left = 0.12, right = 0.95) path = plotpath mpl.rcParams['font.size'] = 30 try: filename = path+"PlotOscProbability_ineu_"+str(iineu) for p in param: filename = filename + +"_" + p.name+"_"+p.neutype filename = filename + "."+fmt except TypeError: filename = path+"PlotOscProbability_ineu_"+str(iineu)+"_"+param.name+"_"+param.neutype+"."+fmt plt.savefig(filename, dpi = 1200) plt.clf() plt.close()
def set_marginal_histogram_title(ax, fmt, color, label=None, rotated=False): """ Sets the title of the marginal histograms. Parameters ---------- ax : Axes The `Axes` instance for the plot. fmt : str The string to add to the title. color : str The color of the text to add to the title. label : str If title does not exist, then include label at beginning of the string. rotated : bool If `True` then rotate the text 270 degrees for sideways title. """ # get rotation angle of the title rotation = 270 if rotated else 0 # get how much to displace title on axes xscale = 1.05 if rotated else 0.0 if rotated: yscale = 1.0 elif len(ax.get_figure().axes) > 1: yscale = 1.15 else: yscale = 1.05 # get class that packs text boxes vertical or horizonitally packer_class = offsetbox.VPacker if rotated else offsetbox.HPacker # if no title exists if not hasattr(ax, "title_boxes"): # create a text box title = "{} = {}".format(label, fmt) tbox1 = offsetbox.TextArea(title, textprops=dict(color=color, size=15, rotation=rotation, ha='left', va='bottom')) # save a list of text boxes as attribute for later ax.title_boxes = [tbox1] # pack text boxes ybox = packer_class(children=ax.title_boxes, align="bottom", pad=0, sep=5) # else append existing title else: # delete old title ax.title_anchor.remove() # add new text box to list tbox1 = offsetbox.TextArea(" {}".format(fmt), textprops=dict(color=color, size=15, rotation=rotation, ha='left', va='bottom')) ax.title_boxes = ax.title_boxes + [tbox1] # pack text boxes ybox = packer_class(children=ax.title_boxes, align="bottom", pad=0, sep=5) # add new title and keep reference to instance as an attribute anchored_ybox = offsetbox.AnchoredOffsetbox(loc=2, child=ybox, pad=0., frameon=False, bbox_to_anchor=(xscale, yscale), bbox_transform=ax.transAxes, borderpad=0.) ax.title_anchor = ax.add_artist(anchored_ybox)
def plot_annote(self, ax, raw, gnd, nodes, ffrom=0, fto=2000): """ `ax`: axes for plot area `raw`: raw data for observation `gnd`: ground truth data `nodes`: nodes dict contained (index array, value array, text array) for marking important data. `section`: data range to display `window_size`: data size to display """ def add_artist(ax, fid, pval, bounce=(.5, .5), img=None): if img is None: img = vid.get_frame(by='id', value=fid) oft = ofb.OffsetImage(img, zoom=0.1) boxp = dict(boxstyle='round,pad=0.1', ec='g') arrp = dict(arrowstyle="->", color='g') ab = ofb.\ AnnotationBbox(oft, (fid, pval), xycoords='data', xybox=bounce, boxcoords=('axes fraction', "axes fraction"), bboxprops=boxp, arrowprops=arrp) ax.add_artist(ab) vid = Video(self.root, self.name) slid = Slider(self.name, self.root) xlim = (ffrom, fto) # plot diff value ln_raw = ax.plot(raw, color='black') ln_peak = ax.plot([ii['index'] for ii in nodes], [vv['value'] for vv in nodes], 'ro') # plot ground truth ax2 = ax.twinx() ln_gnd = ax2.plot(gnd['fid'], gnd['sid'], 'b-') rng_sid = gnd[(gnd['fid'] > ffrom) & (gnd['fid'] < fto)].sid ax2.set_ylim(rng_sid.min() - 1, rng_sid.max() + 1) ax.set_xlim(xlim) # make legend lns = ln_raw + ln_peak + ln_gnd # labs = [l.get_label() for l in lns] ax.legend(lns, ['raw', 'peak', 'ground'], loc=0) for nod in nodes: inx = nod['index'] nov = nod['value'] gsid = int(nod['gsid']) if inx < xlim[0]: continue if inx > xlim[1]: break frac = (inx - ffrom) * 1.0 / (fto - ffrom) add_artist(ax, inx - 1, nov, (frac - 0.05, .65)) # last add_artist(ax, inx, nov, (frac, .95)) add_artist(ax, inx, nov, (frac, .05), img=slid.get_slide(gsid)) oft = ofb.TextArea("{}[{}]".format(inx, gsid)) boxp = dict(boxstyle='round,pad=0.1', ec='g') ab = ofb.\ AnnotationBbox(oft, (inx, nov), xycoords='data', xybox=(inx, 0.90), bboxprops=boxp, boxcoords=('data', "axes fraction"), arrowprops=dict(arrowstyle="->", color='g')) ax.add_artist(ab) # last add_artist(ax, inx + 1, nov, (frac + 0.05, .8))
plt.text(X_tsne[i, 0], X_tsne[i, 1], X_txt, color=plt.cm.Dark2(y_pred[i] / 9.), fontdict={ 'weight': 'bold', 'size': 6 }) plt.xticks([]), plt.yticks([]) if hasattr(offsetbox, 'AnnotationBbox'): # only print thumbnails with matplotlib > 1.0 shown_images = np.array([[1., 1.]]) # just something big for i in range(X_tsne.shape[0]): dist = np.sum((X_tsne[i] - shown_images)**2, 1) if np.min(dist) < args.dist: # don't show points that are too close continue shown_images = np.r_[shown_images, [X_tsne[i]]] X_txt = shorten(X_data[i], width=args.width, placeholder='...') imagebox = offsetbox.AnnotationBbox( offsetbox.TextArea(X_txt, textprops={'size': 8}), X_tsne[i]) ax.add_artist(imagebox) fig.savefig(args.output) elif args.logdir: X = torch.from_numpy(X).float() print("Writing tensorboard embedding...") writer = SummaryWriter(log_dir=args.logdir) writer.add_embedding(X, global_step=0, metadata=X_data) writer.close()
def plot_model(model, save_file=None, ax=None, show=False, fig_title="Demographic Model", pop_labels=None, nref=None, draw_ancestors=True, draw_migrations=True, draw_scale=True, arrow_size=0.01, transition_size=0.05, gen_time=0, gen_time_units="Years", reverse_timeline=False, fig_bg_color='#ffffff', plot_bg_color='#ffffff', text_color='#002b36', gridline_color='#586e75', pop_color='#268bd2', arrow_color='#073642', label_size=16, tick_size=12, grid=True): """ Plots a demographic model based on information contained within a _ModelInfo object. See the matplotlib docs for valid entries for the color parameters. model : A _ModelInfo object created using generate_model(). save_file : If not None, the figure will be saved to this location. Otherwise the figure will be displayed to the screen. fig_title : Title of the figure. pop_labels : If not None, should be a list of strings of the same length as the total number of final populations in the model. The string at index i should be the name of the population along axis i in the model's SFS. nref : If specified, this will update the time and population size labels to use units based on an ancestral population size of nref. See the documentation for details. draw_ancestors : Specify whether the ancestral populations should be drawn in beginning of plot. Will fade off with a gradient. draw_migrations : Specify whether migration arrows are drawn. draw_scale : Specify whether scale bar should be shown in top-left corner. arrow_size : Float to control the size of the migration arrows. transition_size : Float specifying size of the "transitional periods" between populations. gen_time : If greater than 0, and nref given, timeline will be adjusted to show absolute time values, using this value as the time elapsed per generation. gen_time_units : Units used for gen_time (e.g. Years, Thousand Years, etc.). reverse_timeline : If True, the labels on the timeline will be reversed, so that "0 time" is the present time period, rather than the time of the original population. fig_bg_color : Background color of figure (i.e. border surrounding the drawn model). plot_bg_color : Background color of the actual plot area. text_color : Color of text in the figure. gridline_color : Color of the plot gridlines. pop_color : Color of the populations. arrow_color : Color of the arrows showing migrations between populations. """ # Set up the plot with a title and axis labels fig_kwargs = { 'figsize': (9.6, 5.4), 'dpi': 200, 'facecolor': fig_bg_color, 'edgecolor': fig_bg_color } if ax == None: fig = plt.figure(**fig_kwargs) ax = fig.add_subplot(111) ax.set_facecolor(plot_bg_color) ax.set_title(fig_title, color=text_color, fontsize=24) xlabel = "Time Ago" if reverse_timeline else "Time" if nref: if gen_time > 0: xlabel += " ({})".format(gen_time_units) else: xlabel += " (Generations)" ylabel = "Population Sizes" else: xlabel += " (Genetic Units)" ylabel = "Relative Population Sizes" ax.set_xlabel(xlabel, color=text_color, fontsize=label_size) ax.set_ylabel(ylabel, color=text_color, fontsize=label_size) # Determine various maximum values for proper scaling within the plot xmax = model.tp_list[-1].time[-1] ymax = sum(model.tp_list[0].framesizes) ax.set_xlim([-1 * xmax * 0.1, xmax]) ax.set_ylim([0, ymax]) mig_max = 0 for tp in model.tp_list: if tp.migrations is None: continue mig = np.amax(tp.migrations) mig_max = mig_max if mig_max > mig else mig # Configure axis border colors ax.spines['top'].set_color(text_color) ax.spines['right'].set_color(text_color) ax.spines['bottom'].set_color(text_color) ax.spines['left'].set_color(text_color) # Major ticks along x-axis (time) placed at each population split xticks = [tp.time[0] for tp in model.tp_list] xticks.append(xmax) ax.xaxis.set_major_locator(mticker.FixedLocator(xticks)) ax.xaxis.set_minor_locator(mticker.NullLocator()) ax.tick_params(which='both', axis='x', labelcolor=text_color, labelsize=tick_size, top=False) # Choose correct time labels based on nref, gen_time, and reverse_timeline if reverse_timeline: xticks = [xmax - x for x in xticks] if nref: if gen_time > 0: xticks = [2 * nref * gen_time * x for x in xticks] else: xticks = [2 * nref * x for x in xticks] ax.set_xticklabels(['{:.0f}'.format(x) for x in xticks]) else: ax.set_xticklabels(['{:.2f}'.format(x) for x in xticks]) # Gridlines along y-axis (population size) spaced by nref size if grid: ax.yaxis.set_major_locator(mticker.FixedLocator(np.arange(ymax))) ax.yaxis.set_minor_locator(mticker.NullLocator()) ax.grid(b=True, which='major', axis='y', color=gridline_color) ax.tick_params(which='both', axis='y', colors='none', labelsize=tick_size) else: ax.set_yticks([]) # Add scale in top-left corner displaying ancestral population size (Nref) if draw_scale: # Bidirectional arrow of height Nref arrow = mbox.AuxTransformBox(ax.transData) awidth = xmax * arrow_size * 0.2 alength = ymax * arrow_size arrow_kwargs = { 'width': awidth, 'head_width': awidth * 3, 'head_length': alength, 'color': text_color, 'length_includes_head': True } arrow.add_artist( plt.arrow(0, 0.25, 0, 0.75, zorder=100, **arrow_kwargs)) arrow.add_artist( plt.arrow(0, 0.75, 0, -0.75, zorder=100, **arrow_kwargs)) # Population bar of height Nref bar = mbox.AuxTransformBox(ax.transData) bar.add_artist( mpatches.Rectangle((0, 0), xmax / ymax, 1, color=pop_color)) # Appropriate label depending on scale label = mbox.TextArea(str(nref) if nref else "Nref") label.get_children()[0].set_color(text_color) bars = mbox.HPacker(children=[label, arrow, bar], pad=0, sep=2, align="center") scalebar = mbox.AnchoredOffsetbox(2, pad=0.25, borderpad=0.25, child=bars, frameon=False) ax.add_artist(scalebar) # Add ancestral populations using a gradient fill. if draw_ancestors: time = -1 * xmax * 0.1 for i, ori in enumerate(model.tp_list[0].origins): # Draw ancestor for each initial pop xlist = np.linspace(time, 0.0, model.precision) dx = xlist[1] - xlist[0] low, mid, top = (ori[1], ori[1] + 1.0, ori[1] + model.tp_list[0].popsizes[i][0]) tsize = int(transition_size * model.precision) y1list = np.array([low] * model.precision) y2list = np.array([mid] * (model.precision - tsize)) y2list = np.append(y2list, np.linspace(mid, top, tsize)) # Custom color map runs from bg color to pop color cmap = mcolors.LinearSegmentedColormap.from_list( "custom_map", [plot_bg_color, pop_color]) colors = np.array( cmap(np.linspace(0.0, 1.0, model.precision - tsize))) # Gradient created by drawing multiple small rectangles for x, y1, y2, color in zip(xlist[:-1 * tsize], y1list[:-1 * tsize], y2list[:-1 * tsize], colors): rect = mpatches.Rectangle((x, y1), dx, y2 - y1, color=color) ax.add_patch(rect) ax.fill_between(xlist[-1 * tsize:], y1list[-1 * tsize:], y2list[-1 * tsize:], color=pop_color, edgecolor=pop_color) # Iterate through time periods and populations to draw everything for tp_index, tp in enumerate(model.tp_list): # Keep track of migrations to evenly space arrows across time period total_migrations = np.count_nonzero(tp.migrations) num_migrations = 0 for pop_index in range(len(tp.popsizes)): # Draw current population origin = tp.origins[pop_index] popsize = tp.popsizes[pop_index] direc = tp.direcs[pop_index] y1 = origin[1] y2 = origin[1] + (direc * popsize) ax.fill_between(tp.time, y1, y2, color=pop_color, edgecolor=pop_color) # Draw connections to next populations if necessary if tp.descendants is not None and tp.descendants[pop_index] != -1: desc = tp.descendants[pop_index] tp_next = model.tp_list[tp_index + 1] # Split population case if isinstance(desc, tuple): # Get origins connect_below = tp_next.origins[desc[0]][1] connect_above = tp_next.origins[desc[1]][1] # Get popsizes subpop_below = tp_next.popsizes[desc[0]][0] subpop_above = tp_next.popsizes[desc[1]][0] # Determine correct connection location connect_below -= direc * subpop_below connect_above += direc * subpop_above # Single population case else: connect_below = tp_next.origins[desc][1] subpop = tp_next.popsizes[desc][0] connect_above = connect_below + direc * subpop # Draw the connections tsize = int(transition_size * model.precision) cx = tp.time[-1 * tsize:] cy_below_1 = [origin[1]] * tsize cy_above_1 = origin[1] + direc * popsize[-1 * tsize:] cy_below_2 = np.linspace(cy_below_1[0], connect_below, tsize) cy_above_2 = np.linspace(cy_above_1[0], connect_above, tsize) ax.fill_between(cx, cy_below_1, cy_below_2, color=pop_color, edgecolor=pop_color) ax.fill_between(cx, cy_above_1, cy_above_2, color=pop_color, edgecolor=pop_color) # Draw migrations if necessary if draw_migrations and tp.migrations is not None: # Iterate through migrations for current population for mig_index, mig_val in enumerate(tp.migrations[pop_index]): # If no migration, continue if mig_val == 0: continue # Calculate proper offset for arrow within this period num_migrations += 1 offset = int(tp.precision * num_migrations / (total_migrations + 1.0)) x = tp.time[offset] dx = 0 # Determine which sides of populations are closest y1 = origin[1] y2 = y1 + direc * popsize[offset] mig_y1 = tp.origins[mig_index][1] mig_y2 = mig_y1 + (tp.direcs[mig_index] * tp.popsizes[mig_index][offset]) y = y1 if abs(mig_y1 - y1) < abs(mig_y1 - y2) else y2 dy = mig_y1-y if abs(mig_y1 - y) < abs(mig_y2 - y) \ else mig_y2-y # Scale arrow to proper size mig_scale = max(0.1, mig_val / mig_max) awidth = xmax * arrow_size * mig_scale alength = ymax * arrow_size ax.arrow(x, y, dx, dy, width=awidth, head_width=awidth * 3, head_length=alength, color=arrow_color, length_includes_head=True) # Label populations if proper labels are given tp_last = model.tp_list[-1] if pop_labels and len(pop_labels) == len(tp_last.popsizes): ax2 = ax.twinx() ax2.set_xlim(ax.get_xlim()) ax2.set_ylim(ax.get_ylim()) # Determine placement of ticks yticks = [ tp_last.origins[i][1] + 0.5 * tp_last.direcs[i] * tp_last.popsizes[i][-1] for i in range(len(tp_last.popsizes)) ] ax2.yaxis.set_major_locator(mticker.FixedLocator(yticks)) ax2.set_yticklabels(pop_labels) ax2.tick_params(which='both', color='none', labelcolor=text_color, labelsize=label_size, left=False, top=False, right=False) ax2.spines['top'].set_color(text_color) ax2.spines['left'].set_color(text_color) ax2.spines['right'].set_color(text_color) ax2.spines['bottom'].set_color(text_color) # Display figure if save_file: plt.savefig(save_file, **fig_kwargs) else: if show == True: plt.show() if ax == None: plt.close(fig)