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)
