def _lngrid_from_trace(self, trace, make_plot): extents = {'x0': self.par_limits[1], 'x1': self.par_limits[2]} samps = MCSamples(samples=trace, names=['x0', 'x1'], ranges=extents) density = samps.get2DDensity('x0', 'x1') # set up the grid on which to evaluate the likelihood x_bins, y_bins = self.get_x0_x1 xx, yy = np.meshgrid(x_bins, y_bins) pos = np.vstack([xx.ravel(), yy.ravel()]).T # Evalaute density on a grid prob = np.array([density.Prob(*ele) for ele in pos]) prob[prob < 0] = 1e-50 ln_prob = np.log(prob) ln_prob -= ln_prob.max() ln_prob = ln_prob.reshape(xx.shape) if make_plot: plt.pcolormesh(x_bins, y_bins, ln_prob) plt.colorbar() plt.clim(0, -5) return ln_prob
def mcmcSkewer(bundleObj, logdef=3, binned=False, niter=2500, do_mcmc=True, return_sampler=False, evalgrid=True, in_axes=None, viz=False, VERBOSITY=False, seed=None, truths=[0.002, 3.8]): """ Script to fit simple flux model on each restframe wavelength skewer Parameters: ----------- bundleObj : A list of [z, f, ivar] with the skewer_index logdef : Which model to use niter : The number of iterations to run the mcmc (40% for burn-in) do_mcmc : Flag whether to perform mcmc plt_pts : Plot the data along with best fit from scipy and mcmc return_sampler : Whether to return the raw sampler without flatchaining triangle : Display triangle plot of the parameters evalgrid : Whether to compute loglikelihood on a specified grid in_axes : axes over which to draw the plots xx_viz : draw marginalized contour in modifed space VERBOSITY : print extra information seed : how to seed the random state truths : used with logdef=4, best-fit values of tau0 and gamma Returns: mcmc_chains if return_sampler, else None """ z, f, ivar = bundleObj[0].T ind = (ivar > 0) & (np.isfinite(f)) z, f, sigma = z[ind], f[ind], 1.0 / np.sqrt(ivar[ind]) # ------------------------------------------------------------------------- # continuum flux estimate given a value of (tau0, gamma) if logdef == 4: if VERBOSITY: print('Continuum estimates using optical depth parameters:', truths) chisq4 = lambda *args: -outer(*truths)(*args) opt_res = minimize(chisq4, 1.5, args=(z, f, sigma), method='Nelder-Mead') return opt_res['x'] if VERBOSITY: print('Carrying analysis for skewer', bundleObj[1]) if logdef == 1: nll, names, labels, guess = chisq1, names1, labels1, guess1 ndim, kranges, lnlike = 4, kranges1, lnlike1 elif logdef == 2: nll, names, labels, guess = chisq2, names2, labels2, guess2 ndim, kranges, lnlike = 5, kranges2, lnlike2 elif logdef == 3: nll, names, labels, guess = chisq3, names3, labels3, guess3 ndim, kranges, lnlike = 3, kranges3, lnlike3 # Try to fit with scipy optimize routine opt_res = minimize(nll, guess, args=(z, f, sigma), method='Nelder-Mead') print('Scipy optimize results:') print('Success =', opt_res['success'], 'params =', opt_res['x'], '\n') if viz: if in_axes is None: fig, in_axes = plt.subplots(1) in_axes.errorbar(z, f, sigma, fmt='o', color='gray', alpha=0.2) in_axes.plot( zline, opt_res['x'][0] * np.exp(-np.exp(opt_res['x'][1]) * (1 + zline)**opt_res['x'][2])) if binned: mu = binned_statistic(z, f, bins=binx).statistic sig = binned_statistic(z, f, bins=binx, statistic=sig_func).statistic ixs = sig > 0 z, f, sigma = centers[ixs], mu[ixs], sig[ixs] if viz: in_axes.errorbar(z, f, sigma, fmt='o', color='r') nll, names, labels, guess = lsq, names3, labels3, guess3 ndim, kranges, lnlike = 3, kranges3, simpleln # -------------------------------------------------------------------------- if do_mcmc: np.random.seed() nwalkers = 100 p0 = [guess + 1e-4 * np.random.randn(ndim) for i in range(nwalkers)] # configure the sampler sampler = emcee.EnsembleSampler(nwalkers, ndim, lnlike, args=(z, f, sigma)) # burn-in time - Is this enough? p0, __, __ = sampler.run_mcmc(p0, 500) sampler.reset() # Production step sampler.run_mcmc(p0, niter) print("Burn-in and production completed \n") if return_sampler: return sampler.chain else: # pruning 40 percent of the samples as extra burn-in lInd = int(niter * 0.4) samps = sampler.chain[:, lInd:, :].reshape((-1, ndim)) # using percentiles as confidence intervals CenVal = np.median(samps, axis=0) # print BIC at the best estimate point, BIC = - 2 * ln(L_0) + k ln(n) print('CHISQ_R', -2 * lnlike(CenVal, z, f, sigma) / (len(z) - 3)) print('BIC:', -2 * lnlike(CenVal, z, f, sigma) + ndim * np.log(len(z))) # Rotate the points to the other basis and 1D estimates # and write them to the file # Format : center, top error, bottom error tg_est = list( map(lambda v: (v[1], v[2] - v[1], v[1] - v[0]), zip(*np.percentile(samps, [16, 50, 84], axis=0)))) xx = xfm(samps[:, 1:], shift, tilt, direction='up') xx_est = list( map(lambda v: (v[1], v[2] - v[1], v[1] - v[0]), zip(*np.percentile(xx, [16, 50, 84], axis=0)))) f_name2 = 'tg_est_' + str(bundleObj[1]) + '.dat' np.savetxt(f_name2, tg_est) f_name3 = 'xx_est_' + str(bundleObj[1]) + '.dat' np.savetxt(f_name3, xx_est) if viz: in_axes.plot( zline, CenVal[0] * np.exp(-np.exp(CenVal[1]) * (1 + zline)**CenVal[2]), '-g') # instantiate a getdist object MC = MCSamples(samples=samps, names=names, labels=labels, ranges=kranges) # MODIFY THIS TO BE PRETTIER if viz: g = plots.getSubplotPlotter() g.triangle_plot(MC) # Evaluate the pdf on a rotated grid for better estimation if evalgrid: print('Evaluating on the grid specified \n') pdist = MC.get2DDensity('t0', 'gamma') # Evalaute density on a grid pgrid = np.array([pdist.Prob(*ele) for ele in modPos]) # Prune to remove negative densities pgrid[pgrid < 0] = 1e-50 # Convert to logLikelihood logP = np.log(pgrid) logP -= logP.max() logP = logP.reshape(x0.shape) # Visualize the contour in modified space per skewer if viz: fig, ax2 = plt.subplots(1) ax2.contour(x0, x1, cts(logP), levels=[ 0.683, 0.955, ], colors='k') ax2.axvline(xx_est[0][0] + xx_est[0][1]) ax2.axvline(xx_est[0][0] - xx_est[0][2]) ax2.axhline(xx_est[1][0] + xx_est[1][1]) ax2.axhline(xx_est[1][0] - xx_est[1][2]) ax2.set_xlabel(r'$x_0$') ax2.set_ylabel(r'$x_1$') plt.show() # fileName1: the log-probability evaluated in the tilted grid f_name1 = 'gridlnlike_' + str(bundleObj[1]) + '.dat' np.savetxt(f_name1, logP)