def do_knfit(df,
outputDir='./knfit',
errorbudget=1.0,
n_live_points=100,
evidence_tolerance=0.5):
if len(df.index) < 2:
print('Not enough data for candidate %s... continuing.',
df.candname.values[0])
return
if not os.path.isdir(outputDir):
os.makedirs(outputDir)
candname = df.name.values[0].decode()
filters = list(set(df.filtname))
mint = 0.0
maxt = 7.0
dt = 0.05
tt = np.arange(mint, maxt, dt)
doWaveformExtrapolate = True
ZPRange = 5.0
T0Range = 2.0
ModelPath = "../gwemlightcurves/output/svdmodels/"
T0 = np.inf
mag_min = np.inf
data_out = {}
for index, row in df.iterrows():
filt = row.filtname
mag = row.magpsf
dmag = row.sigmapsf
mjd = Time(row.jd, format='jd').mjd
magzp = row.magzpsci
if 99.0 == mag:
mag = magzp
dmag = np.inf
else:
T0 = np.min([T0, mjd])
mag_min = np.min([mag_min, mag])
if not filt in data_out:
data_out[filt] = np.empty((0, 3), float)
data_out[filt] = np.append(data_out[filt],
np.array([[mjd, mag, dmag]]),
axis=0)
distmod = mag_min - -16
distance = 10**((distmod / 5.0) + 1.0) / 1e6
for ii, key in enumerate(data_out.keys()):
if key == "t":
continue
else:
data_out[key][:, 0] = data_out[key][:, 0] - T0
data_out[key][:, 1] = data_out[key][:, 1] - 5 * (
np.log10(distance * 1e6) - 1)
for ii, key in enumerate(data_out.keys()):
if key == "t":
continue
else:
idxs = np.intersect1d(
np.where(data_out[key][:, 0] >= mint)[0],
np.where(data_out[key][:, 0] <= maxt)[0])
data_out[key] = data_out[key][idxs, :]
for ii, key in enumerate(data_out.keys()):
idxs = np.where(~np.isnan(data_out[key][:, 2]))[0]
if key == "t":
continue
else:
data_out[key] = data_out[key][idxs, :]
for ii, key in enumerate(data_out.keys()):
if not key in filters:
del data_out[key]
for ii, key in enumerate(data_out.keys()):
if ii == 0:
samples = data_out[key].copy()
else:
samples = np.vstack((samples, data_out[key].copy()))
idx = np.argmin(samples[:, 0])
samples = samples[idx, :]
Global.data_out = data_out
Global.errorbudget = errorbudget
Global.ZPRange = ZPRange
Global.T0Range = T0Range
Global.doLightcurves = 1
Global.filters = filters
Global.doWaveformExtrapolate = doWaveformExtrapolate
modelfile = os.path.join(ModelPath, 'Bu2019inc_mag.pkl')
with open(modelfile, 'rb') as handle:
svd_mag_model = pickle.load(handle)
Global.svd_mag_model = svd_mag_model
modelfile = os.path.join(ModelPath, 'Bu2019inc_lbol.pkl')
with open(modelfile, 'rb') as handle:
svd_lbol_model = pickle.load(handle)
Global.svd_lbol_model = svd_lbol_model
plotDir = os.path.join(outputDir, candname)
if not os.path.isdir(plotDir):
os.makedirs(plotDir)
max_iter = -1
best = []
parameters = ["t0", "mej", "phi", "theta", "zp"]
labels = [
r"$T_0$", r"${\rm log}_{10} (M_{\rm ej})$", r"$\Phi$", r"$\Theta$",
"ZP"
]
n_params = len(parameters)
pymultinest.run(myloglike_Bu2019inc_ejecta,
myprior_Bu2019inc_ejecta,
n_params,
importance_nested_sampling=False,
resume=True,
verbose=True,
sampling_efficiency='parameter',
n_live_points=n_live_points,
outputfiles_basename='%s/2-' % plotDir,
evidence_tolerance=evidence_tolerance,
multimodal=False,
max_iter=max_iter)
multifile = lightcurve_utils.get_post_file(plotDir)
data = np.loadtxt(multifile)
t0, mej, phi, theta, zp, loglikelihood = data[:,
0], 10**data[:,
1], data[:,
2], data[:,
3], data[:,
4], data[:,
5]
idx = np.argmax(loglikelihood)
t0_best, mej_best, phi_best, theta_best, zp_best = data[idx, 0], 10**data[
idx, 1], data[idx, 2], data[idx, 3], data[idx, 4]
zp_mu, zp_std = 0.0, Global.ZPRange
zp_best = scipy.stats.norm(zp_mu, zp_std).ppf(zp_best)
tmag, lbol, mag = Bu2019inc_model_ejecta(mej_best, phi_best, theta_best)
pcklFile = os.path.join(plotDir, "data.pkl")
f = open(pcklFile, 'wb')
pickle.dump((data_out, data, tmag, lbol, mag, t0_best, zp_best, n_params,
labels, best), f)
f.close()
title_fontsize = 30
label_fontsize = 30
plotName = "%s/corner.pdf" % (plotDir)
figure = corner.corner(data[:, :-1],
labels=labels,
quantiles=[0.16, 0.5, 0.84],
show_titles=True,
title_kwargs={"fontsize": title_fontsize},
label_kwargs={"fontsize": label_fontsize},
title_fmt=".2f",
smooth=3,
color="coral")
figure.set_size_inches(14.0, 14.0)
plt.savefig(plotName)
plt.close()
tmag = tmag + t0_best
#colors=cm.rainbow(np.linspace(0,1,len(filters)))
colors = cm.Spectral(np.linspace(0, 1, len(filters)))[::-1]
color2 = 'coral'
color1 = 'cornflowerblue'
plotName = "%s/models_panels.pdf" % (plotDir)
#plt.figure(figsize=(20,18))
plt.figure(figsize=(20, 28))
cnt = 0
for filt, color in zip(filters, colors):
cnt = cnt + 1
if cnt == 1:
ax1 = plt.subplot(len(filters), 1, cnt)
else:
ax2 = plt.subplot(len(filters), 1, cnt, sharex=ax1, sharey=ax1)
if not filt in data_out: continue
samples = data_out[filt]
t, y, sigma_y = samples[:, 0], samples[:, 1], samples[:, 2]
idx = np.where(~np.isnan(y))[0]
t, y, sigma_y = t[idx], y[idx], sigma_y[idx]
if len(t) == 0: continue
idx = np.where(np.isfinite(sigma_y))[0]
plt.errorbar(t[idx],
y[idx],
sigma_y[idx],
fmt='o',
c=color,
markersize=16,
label='%s-band' % filt)
idx = np.where(~np.isfinite(sigma_y))[0]
plt.errorbar(t[idx],
y[idx],
sigma_y[idx],
fmt='v',
c=color,
markersize=16)
magave = lightcurve_utils.get_mag(mag, filt)
ii = np.where(~np.isnan(magave))[0]
f = interp.interp1d(tmag[ii], magave[ii], fill_value='extrapolate')
maginterp = f(tt)
#plt.plot(tt,maginterp+zp_best,'--',c=color,linewidth=3)
#plt.fill_between(tt,maginterp+zp_best-errorbudget,maginterp+zp_best+errorbudget,facecolor=color,alpha=0.2)
plt.plot(tt, maginterp + zp_best, '--', c=color2, linewidth=3)
plt.fill_between(tt,
maginterp + zp_best - errorbudget,
maginterp + zp_best + errorbudget,
facecolor=color2,
alpha=0.2)
plt.ylabel('%s' % filt, fontsize=48, rotation=0, labelpad=40)
plt.xlim([0.0, 10.0])
plt.ylim([-22.0, -8.0])
plt.gca().invert_yaxis()
plt.grid()
if cnt == 1:
ax1.set_yticks([-22, -18, -14, -10])
plt.setp(ax1.get_xticklabels(), visible=False)
#l = plt.legend(loc="upper right",prop={'size':36},numpoints=1,shadow=True, fancybox=True)
elif not cnt == len(filters):
plt.setp(ax2.get_xticklabels(), visible=False)
plt.xticks(fontsize=36)
plt.yticks(fontsize=36)
ax1.set_zorder(1)
plt.xlabel('Time [days]', fontsize=48)
plt.tight_layout()
plt.savefig(plotName)
plt.close()
return
ax2 = plt.subplot(len(filts),1,cnt,sharex=ax1,sharey=ax1)
if not filt in data_out: continue
samples = data_out[filt]
t, y, sigma_y = samples[:,0], samples[:,1], samples[:,2]
idx = np.where(~np.isnan(y))[0]
t, y, sigma_y = t[idx], y[idx], sigma_y[idx]
if len(t) == 0: continue
idx = np.where(np.isfinite(sigma_y))[0]
plt.errorbar(t[idx],y[idx],sigma_y[idx],fmt='o',c=color, markersize=16)
idx = np.where(~np.isfinite(sigma_y))[0]
plt.errorbar(t[idx],y[idx],sigma_y[idx],fmt='v',c=color, markersize=16)
magave1 = lightcurve_utils.get_mag(mag1,filt)
magave2 = lightcurve_utils.get_mag(mag2,filt)
magave3 = lightcurve_utils.get_mag(mag3,filt)
ii = np.where(~np.isnan(magave1))[0]
f = interp.interp1d(tmag1[ii], magave1[ii], fill_value='extrapolate')
maginterp1 = f(tt)
plt.plot(tt,maginterp1+zp_best1,'--',c=color1,linewidth=2,label='Kilonova')
plt.plot(tt,maginterp1+zp_best1-errorbudget,'-',c=color1,linewidth=2)
plt.plot(tt,maginterp1+zp_best1+errorbudget,'-',c=color1,linewidth=2)
plt.fill_between(tt,maginterp1+zp_best1-errorbudget,maginterp1+zp_best1+errorbudget,facecolor=color1,alpha=0.2)
ii = np.where(~np.isnan(magave2))[0]
f = interp.interp1d(tmag2[ii], magave2[ii], fill_value='extrapolate')
maginterp2 = f(tt)
plt.plot(tt,maginterp2+zp_best2,'--',c=color2,linewidth=2,label='Afterglow')
def calc_prob(tmag, lbol, mag, t0, zp):
if Global.doLuminosity:
if np.sum(lbol) == 0.0:
prob = -np.inf
return prob
tmag = tmag + t0
count = 0
chisquare = np.nan
t = Global.data_out["tt"]
y = Global.data_out["Lbol"]
sigma_y = Global.data_out["Lbol_err"]
idx = np.where(~np.isnan(y))[0]
t = t[idx]
y = y[idx]
sigma_y = sigma_y[idx]
ii = np.where(~np.isnan(lbol))[0]
if len(ii) == 0:
lbolinterp = np.nan*np.ones(t.shape)
else:
f = interp.interp1d(tmag[ii], np.log10(lbol[ii]), fill_value='extrapolate')
lbolinterp = 10**f(t)
zp_factor = 10**(zp/-2.5)
lbolinterp = lbolinterp*zp_factor
sigma_y = np.abs(sigma_y/(y*np.log(10)))
sigma = np.sqrt((np.log10(1+Global.errorbudget))**2 + sigma_y**2)
y = np.log10(y)
lbolinterp = np.log10(lbolinterp)
chisquarevals = ((y-lbolinterp)/sigma)**2
chisquaresum = np.sum(chisquarevals)
if np.isnan(chisquaresum):
chisquare = np.nan
return chisquare
if not float(len(chisquarevals)-1) == 0:
chisquaresum = (1/float(len(chisquarevals)-1))*chisquaresum
chisquare = chisquaresum
if np.isnan(chisquare):
prob = -np.inf
else:
#prob = scipy.stats.chi2.logpdf(chisquare, count, loc=0, scale=1)
#prob = -chisquare/2.0
#prob = chisquare
prob = scipy.stats.chi2.logpdf(chisquare, 1, loc=0, scale=1)
if np.isnan(prob):
prob = -np.inf
#if np.isfinite(prob):
# print t0, zp, prob
return prob
elif Global.doLightcurves:
if len(np.isfinite(lbol)) == 0:
prob = -np.inf
return prob
if np.sum(lbol) == 0.0:
prob = -np.inf
return prob
tmag = tmag + t0
count = 0
chisquare = np.nan
gaussprob = np.nan
for key in Global.data_out:
samples = Global.data_out[key]
t = samples[:,0]
y = samples[:,1]
sigma_y = samples[:,2]
idx = np.where(~np.isnan(y))[0]
t = t[idx]
y = y[idx]
sigma_y = sigma_y[idx]
if len(idx) == 0: continue
if not key in Global.filters: continue
keyslist = ["u","g","r","i","z","y","J","H","K"]
if key in keyslist:
idx = keyslist.index(key)
ii = np.where(np.isfinite(mag[idx]))[0]
if len(ii) == 0:
maginterp = np.nan*np.ones(t.shape)
else:
if Global.doWaveformExtrapolate:
f = interp.interp1d(tmag[ii], mag[idx][ii], fill_value='extrapolate')
else:
f = interp.interp1d(tmag[ii], mag[idx][ii], fill_value=np.nan, bounds_error = False)
maginterp = f(t)
elif key in ["w","c","o","V","B","R","I","F606W","F160W","F814W"]:
magave = lightcurve_utils.get_mag(mag,key)
ii = np.where(np.isfinite(magave))[0]
if len(ii) == 0:
maginterp = np.nan*np.ones(t.shape)
else:
if Global.doWaveformExtrapolate:
f = interp.interp1d(tmag[ii], magave[ii], fill_value='extrapolate')
else:
f = interp.interp1d(tmag[ii], magave[ii], fill_value=np.nan, bounds_error = False)
maginterp = f(t)
else:
continue
maginterp = maginterp + zp
sigma = np.sqrt(Global.errorbudget**2 + sigma_y**2)
chisquarevals = np.zeros(y.shape)
chisquarevals = ((y-maginterp)/sigma)**2
idx = np.where(~np.isfinite(sigma))[0]
if len(idx) > 0:
gaussprobvals = 1-scipy.stats.norm.cdf(y[idx], maginterp[idx], Global.errorbudget)
gaussprobsum = np.sum(np.log(gaussprobvals))
else:
gaussprobsum = 0.0
chisquaresum = np.sum(chisquarevals)
if np.isnan(chisquaresum):
chisquare = np.nan
break
if not float(len(chisquarevals)-1) == 0:
chisquaresum = (1/float(len(chisquarevals)-1))*chisquaresum
if count == 0:
chisquare = chisquaresum
gaussprob = gaussprobsum
else:
chisquare = chisquare + chisquaresum
gaussprob = gaussprob + gaussprobsum
#count = count + len(chisquarevals)
count = count + 1
if np.isnan(chisquare):
prob = -np.inf
else:
#prob = scipy.stats.chi2.logpdf(chisquare, count, loc=0, scale=1)
#prob = -chisquare/2.0
#prob = chisquare
chiprob = scipy.stats.chi2.logpdf(chisquare, 1, loc=0, scale=1)
prob = chiprob + gaussprob
if np.isnan(prob):
prob = -np.inf
#if np.isfinite(prob):
# print t0, zp, prob
return prob
else:
print "Enable doLuminosity or doLightcurves..."
exit(0)
tini, tmax, dt = opts.tmin, opts.tmax, 0.1
tt = np.arange(tini, tmax, dt)
plotName = "%s/lightcurve.pdf" % (plotDir)
plt.figure(figsize=(10, 8))
for filt, color in zip(filters, colors):
if not filt in filters: continue
if not filt in data_out: continue
samples = data_out[filt]
t, y, sigma_y = samples[:, 0], samples[:, 1], samples[:, 2]
idx = np.where(~np.isnan(y))[0]
t, y, sigma_y = t[idx], y[idx], sigma_y[idx]
if len(t) == 0: continue
plt.errorbar(t, y, sigma_y, fmt='o', c=color, label='%s-band' % filt)
magave = lightcurve_utils.get_mag(mag, filt)
ii = np.where(~np.isnan(magave))[0]
f = interp.interp1d(tmag[ii], magave[ii], fill_value='extrapolate')
maginterp = f(tt)
plt.plot(tt, maginterp + zp_best, 'k--', linewidth=2)
if opts.filters == "c,o":
plt.xlim([opts.tmin - 2, opts.tmax + 2])
plt.ylim([-20.0, -5.0])
elif opts.model == "SN":
plt.xlim([0.0, 10.0])
plt.ylim([-20.0, -5.0])
elif opts.model == "Me2017":
plt.xlim([0.0, 18.0])
plt.ylim([-20.0, -5.0])
