def mcmc_emcee(self, n_walkers, n_run, n_burn, mean_start, sigma_start, mpi=False): """ returns the mcmc analysis of the parameter space """ if mpi: pool = MPIPool() if not pool.is_master(): pool.wait() sys.exit(0) sampler = emcee.EnsembleSampler(n_walkers, self.chain.param.num_param(), self.chain.X2_chain, pool=pool) else: sampler = emcee.EnsembleSampler(n_walkers, self.chain.param.num_param(), self.chain.X2_chain) p0 = emcee.utils.sample_ball(mean_start, sigma_start, n_walkers) new_pos, _, _, _ = sampler.run_mcmc(p0, n_burn) sampler.reset() store = InMemoryStorageUtil() for pos, prob, _, _ in sampler.sample(new_pos, iterations=n_run): store.persistSamplingValues(pos, prob, None) return store.samples
def get_pool(mpi=False, threads=None): """ Get a pool object to pass to emcee for parallel processing. If mpi is False and threads is None, pool is None. Parameters ---------- mpi : bool Use MPI or not. If specified, ignores the threads kwarg. threads : int (optional) If mpi is False and threads is specified, use a Python multiprocessing pool with the specified number of threads. """ if mpi: from emcee.utils import MPIPool # Initialize the MPI pool pool = MPIPool() # Make sure the thread we're running on is the master if not pool.is_master(): pool.wait() sys.exit(0) logger.debug("Running with MPI...") elif threads > 1: logger.debug( "Running with multiprocessing on {} cores...".format(threads)) pool = multiprocessing.Pool(threads) else: logger.debug("Running serial...") pool = SerialPool() return pool
def fit(Deltam, nsne, xi, redshiftterm, mpi=False, p0=None, nchain=2000, **kwargs): ndim, nwalkers = 4, 8 sig = numpy.array([0.1, 0.01, 0.01, 50 / 3e5]) if p0 is None: p0 = [ numpy.array([1, 0, 0.08, 200 / 3e5]) + numpy.random.uniform(low=-sig, high=sig) for i in range(nwalkers) ] if mpi: pool = MPIPool() if not pool.is_master(): pool.wait() sys.exit(0) else: import time starttime = time.time() print("Start {}".format(starttime)) sampler = emcee.EnsembleSampler( nwalkers, ndim, Fit.lnprob, args=[Deltam, nsne, xi, redshiftterm], pool=pool) else: import time starttime = time.time() print("Start {}".format(starttime)) sampler = emcee.EnsembleSampler( nwalkers, ndim, Fit.lnprob, args=[Deltam, nsne, xi, redshiftterm]) sampler.run_mcmc(p0, nchain) if mpi: if pool.is_master(): endtime = time.time() print("End {}".format(endtime)) print("Difference {}".format(endtime - starttime)) pool.close() else: endtime = time.time() print("End {}".format(endtime)) print("Difference {}".format(endtime - starttime)) return sampler
def mcmc(mass_bins, dot_val, initial_c, nwalkers, ndim, burn_in, steps_wlk): import numpy as np import scipy.optimize as op import emcee from emcee.utils import MPIPool import sys def lnlike(c, dot_val): loglike = np.zeros((1)) loglike[0] = sum( np.log( (1 - c) * np.sqrt(1 + (c / 2)) * (1 - c * (1 - 3 * (dot_val * dot_val / 2)))**(-1.5))) #log-likelihood return loglike def lnprior(c): if (-1.5 < c < 0.99): #Assumes a flat prior, uninformative prior return 0.0 return -np.inf def lnprob(c, dot_val): lp = lnprior(c) if not np.isfinite(lp): return -np.inf return lp + lnlike(c, dot_val) #Parallel MCMC - initiallizes pool object; if process isn't running as master, wait for instr. and exit pool = MPIPool() if not pool.is_master(): pool.wait() sys.exit(0) pos = [initial_c + 1e-2 * np.random.randn(ndim) for i in range(nwalkers) ] #initial positions for walkers "Gaussian ball" #MCMC Running sampler = emcee.EnsembleSampler(nwalkers, ndim, lnprob, args=[dot_val], pool=pool) pos, _, _ = sampler.run_mcmc(pos, burn_in) #running of emcee burn-in period sampler.reset() sampler.run_mcmc( pos, steps_wlk ) #running of emcee for steps specified, using pos as initial walker positions pool.close() chain = sampler.flatchain[:, 0] return chain
def run_pool(pC, pW, walk, step): #pCenter and pWidths steps = step nwalkers = walk ndim = len(pC) ## r in, del r, i, PA p0 = [pC[0], pC[1], pC[2], pC[3]] widths = [pW[0], pW[1], pW[2], pW[3]] p = emcee.utils.sample_ball(p0, widths, size=nwalkers) pool = MPIPool() if not pool.is_master(): pool.wait() sys.exit(0) sampler = emcee.EnsembleSampler(nwalkers, ndim, lnlike_vis_and_sed, live_dangerously=True, pool=pool) print 'Beginning the MCMC run.' start = time.clock() sampler.run_mcmc(p, steps) stop = time.clock() pool.close() print 'MCMC finished successfully.\n' print 'This was a simultaneous visibility and SED run with {} walkers and {} steps'.format( nwalkers, steps) print "Mean acor time: " + str(np.mean(sampler.acor)) print "Mean acceptance fraction: " + str( np.mean(sampler.acceptance_fraction)) print '\nRun took %r minutes' % ((stop - start) / 60.) chain = sampler.chain chi = (sampler.lnprobability) / (-0.5) whatbywhat = str(nwalkers) + 'x' + str(steps) os.system('mkdir MCMCRUNS/vis_and_sed/' + whatbywhat) chainFile = 'MCMCRUNS/vis_and_sed/' + whatbywhat + '/' + whatbywhat + '.chain.fits' chiFile = 'MCMCRUNS/vis_and_sed/' + whatbywhat + '/' + whatbywhat + '.chi.fits' infoFile = 'MCMCRUNS/vis_and_sed/' + whatbywhat + '/' + whatbywhat + '.runInfo.txt' fits.writeto(chainFile, chain) fits.writeto(chiFile, chi) #f = open('runInfo.txt','w') f = open(infoFile, 'w') f.write('run took %r minutes\n' % ((stop - start) / 60.)) f.write('walkers: %r\n' % nwalkers) f.write('steps: %r\n' % steps) f.write('initial model: %r\n' % p0) f.write('widths: %r\n' % widths) f.write("mean acor time: " + str(np.mean(sampler.acor))) f.write("\nmean acceptance fraction: " + str(np.mean(sampler.acceptance_fraction))) f.close() print 'Data written to: \n' + chainFile + '\n' + chiFile + '\n' + infoFile
def run_mcmc(self, n_walkers=100, n_iterations=100): """ Run emcee MCMC. Args: n_walkers (int): Number of walkers to pass to the MCMC. n_iteratins (int): Number of iterations to pass to the MCMC. """ # Initialize walker matrix with initial parameters walkers_matrix = [] # must be a list, not an np.array for walker in range(n_walkers): walker_params = [] for component in self.components: walker_params = walker_params + component.initial_values( self.data_spectrum) walkers_matrix.append(walker_params) global iteration_count iteration_count = 0 # Create MCMC sampler. To enable multiproccessing, set threads > 1. # If using multiprocessing, the "lnpostfn" and "args" parameters # must be pickleable. if self.mpi: # Initialize the multiprocessing pool object. from emcee.utils import MPIPool pool = MPIPool(loadbalance=True) if not pool.is_master(): pool.wait() sys.exit(0) self.sampler = emcee.EnsembleSampler( nwalkers=n_walkers, dim=len(walkers_matrix[0]), lnpostfn=ln_posterior, args=[self], pool=pool, runtime_sortingfn=sort_on_runtime) self.sampler.run_mcmc(walkers_matrix, n_iterations) pool.close() else: self.sampler = emcee.EnsembleSampler(nwalkers=n_walkers, dim=len(walkers_matrix[0]), lnpostfn=ln_posterior, args=[self], threads=1) #self.sampler_output = self.sampler.run_mcmc(walkers_matrix, n_iterations) self.sampler.run_mcmc(walkers_matrix, n_iterations)
def main(): import argparse parser = argparse.ArgumentParser() parser.add_argument('--mpi', action='store_true', default=False) parser.add_argument('--walkers', type=int, default=100) parser.add_argument('--steps', type=int, default=1000) opt = parser.parse_args() pool = None if opt.mpi: import socket import os from emcee.utils import MPIPool pool = MPIPool() print('Running in MPI. Host', socket.gethostname(), 'pid', os.getpid(), 'is master?', pool.is_master()) if not pool.is_master(): pool.wait() return ndim, nwalkers = 2, opt.walkers ivar = 1. / np.random.rand(ndim) p0 = [np.random.rand(ndim) for i in range(nwalkers)] sampler = emcee.EnsembleSampler(nwalkers, ndim, lnprob, args=[ivar], pool=pool) import time print('Running for', opt.steps, 'steps with', opt.walkers, 'walkers') t0 = time.time() sampler.run_mcmc(p0, opt.steps) print('Finished in', time.time() - t0, 'seconds') if pool: pool.close() import matplotlib matplotlib.use('Agg') import matplotlib.pyplot as plt plt.figure() for i in range(ndim): plt.subplot(1, ndim, 1 + i) plt.hist(sampler.flatchain[:, i], 100, color="k", histtype="step") plt.title("Dimension {0:d}".format(i)) plt.savefig('plot.png') print('Saved plot')
def __enter__(self): """ Setup the MPIPool, such that only the ``pool`` master returns, while the other processes wait for tasks """ # split ranks if we need to if self.comm.size > 1: ranges = [] for i, ranks in split_ranks(self.comm.size, self.nruns): ranges.append(ranks[0]) if self.comm.rank in ranks: color = i # split the global comm into pools of workers self.pool_comm = self.comm.Split(color, 0) # make the comm to communicate b/w parallel runs if self.nruns > 1: self.par_runs_group = self.comm.group.Incl(ranges) self.par_runs_comm = self.comm.Create(self.par_runs_group) # initialize the MPI pool, if the comm has more than 1 process if self.pool_comm is not None and self.pool_comm.size > 1: from emcee.utils import MPIPool kws = { 'loadbalance': True, 'comm': self.pool_comm, 'debug': self.debug } self.pool = MPIPool(**kws) # explicitly force non-master ranks in pool to wait if self.pool is not None and not self.pool.is_master(): self.pool.wait() self.logger.debug("exiting after pool closed") sys.exit(0) # log if self.pool is not None: self.logger.debug("using an MPIPool instance with %d worker(s)" % self.pool.size) self.rank = 0 if self.par_runs_comm is not None: self.rank = self.par_runs_comm.rank return self
def get_pool(self,mpi=False,nthreads=1): import emcee from emcee.utils import MPIPool from pathos.multiprocessing import ProcessingPool as PPool #from multiprocessing import Pool as PPool if mpi: pool = MPIPool(loadbalance=True) if not pool.is_master(): pool.wait() sys.exit(0) self.logger.info('Creating MPI pool with {:d} workers.'.format(pool.size+1)) elif nthreads > 1: pool = PPool(nthreads) self.logger.info('Creating multiprocessing pool with {:d} threads.'.format(nthreads)) else: pool = None return pool
def main(): ''' A parallel run. ''' pool = MPIPool(loadbalance=True) if not pool.is_master(): pool.wait() sys.exit(0) clf = hbsgc.HBSGC(pool=pool) # save start time clf.last_clock = time.clock() clf.filter_calcs() clf.data_calcs() clf.star_model_calcs() # if clf.calc_model_mags: # clf.star_model_mags() clf.gal_model_calcs() # if clf.calc_model_mags: # clf.gal_model_mags() clf.fit_calcs() clf.count_tot = 0 clf.sample() clf.save_proba() if clf.min_chi2_write: clf.save_min_chi2() pool.close()
def lnPost(theta): '''log-posterior ''' # prior calculations if prior_min[0] < theta[0] < prior_max[0] and \ prior_min[1] < theta[1] < prior_max[1] and \ prior_min[2] < theta[2] < prior_max[2] and \ prior_min[3] < theta[3] < prior_max[3] and \ prior_min[4] < theta[4] < prior_max[4]: lnPrior = 0.0 else: lnPrior = -np.inf if not np.isfinite(lnPrior): return -np.inf return lnPrior + lnLike(theta) """Initializing Walkers""" pos = [ np.array([11., np.log(.4), 11.5, 1.0, 13.5]) + 1e-3 * np.random.randn(Ndim) for i in range(Nwalkers) ] """Initializing MPIPool""" pool = MPIPool(loadbalance=True) if not pool.is_master(): pool.wait() sys.exit(0) """Initializing the emcee sampler""" sampler = emcee.EnsembleSampler(Nwalkers, Ndim, lnprob, pool=pool) # Burn in + Production sampler.run_mcmc(pos, Nchains_burn + Nchains_pro) # Production. samples = sampler.chain[:, Nchains_burn:, :].reshape((-1, Ndim)) #closing the pool pool.close() np.savetxt("mcmc_sample.dat", samples)
print('mpi_script.py using python {}.{}'.format( *sys.version.split('.')[:2] )) n_walkers = 16 p0 = np.random.uniform(0, 10, (n_walkers, 1)) my_object_int = class_with_method('dummy_data') my_object_ext = class_with_method('dummy_data', lnprobfunc=global_ln_probability) # In the case of non-concurrency, everything works for lnprob in [global_ln_probability, my_object_int.ln_probability, my_object_ext.ln_probability]: test_mpi(pool=None, the_func=lnprob) print('Everything works if MPI is not used') with MPIPool() as pool: if not pool.is_master(): pool.wait() sys.exit(0) # Try with global function test_mpi(pool=pool, the_func=global_ln_probability) # Try with externally defined method test_mpi(pool=pool, the_func=my_object_ext.ln_probability) print('Pickling a class with externally defined lnprob is fine') # Try with internally defined method try: test_mpi(pool=pool, the_func=my_object_int.ln_probability) print('Pickling class with internally defined lnprob is fine on python 3') except PicklingError:
def run_mcmc(cos_twol,sin_twol,sin_l,cos_l,sin_b,cos_b,data_pml,data_err_pml,\ data_pmb,data_err_pmb,data_plx,data_err_plx,data_vlos,data_err_vlos,data_age,\ age_number,distance_number,Number): distance = distance_list[distance_number] age = age_list[age_number] #define the objective function negativelnLikelihood = lambda *args: -lnlike(*args)[0] #initial guess for p p_0 = p0_list[age_number] #generate random values. np.random.randn provides Gaussian with mean 0 and standard deviation 1 #thus here is adding random values obeying the Gaussian like above to each values. pos = [p_0 + 1. * np.random.randn(ndim) for i in range(N_WALKERS)] #for multiprocessing pool = MPIPool(loadbalance=True) if not pool.is_master(): pool.wait() sys.exit(0) obsevlosables= cos_twol,sin_twol,sin_l,cos_l,sin_b,cos_b,data_pml,data_err_pml,data_pmb,\ data_err_pmb,data_plx,data_err_plx,data_vlos,data_err_vlos,data_age,age_number sampler = emcee.EnsembleSampler(N_WALKERS, ndim, lnprob, args=obsevlosables) #sampler = emcee.EnsembleSampler(N_WALKERS, ndim, lnprob, pool=pool, \ # args=(cos_twol,sin_twol,sin_l,cos_l,sin_b,cos_b,data_pml,data_err_pml,data_pmb,\ # data_err_pmb,data_plx,data_err_plx,data_vlos,data_err_vlos)) sampler.run_mcmc(pos, Nrun) pool.close() print('Done.') #--- # store the results burnin = Nburn samples = sampler.chain[:, burnin:, :].reshape((-1, ndim)) plt.clf() hight_fig_inch = np.int((ndim + 1) * 3.0) fig, axes = plt.subplots(ndim + 1, 1, sharex=True, figsize=(8, hight_fig_inch)) for i in range(ndim): axes[i].plot(sampler.chain[:, :, i].T, color='k', alpha=0.5) axes[i].set_ylabel(_list_labels[i]) # last panel shows the evolution of ln-likelihood for the ensemble of walkers axes[-1].plot(sampler.lnprobability.T, color='k', alpha=0.5) axes[-1].set_ylabel('ln(L)') maxlnlike = np.max(sampler.lnprobability) axes[-1].set_ylim(maxlnlike - 3 * ndim, maxlnlike) fig.tight_layout(h_pad=0.) filename_pre = 'newModel_1/'+distance+'/line-time_walker%dNrun%dNburn%d_withscatter_'\ +age+'Gyr_'+distance+'_%dstars_newModel_1' filename = filename_pre % (N_WALKERS, Nrun, Nburn, Number) fig.savefig(filename + '.png') # Make a triangle plot burnin = Nburn samples = sampler.chain[:, burnin:, :].reshape((-1, ndim)) #convert scatters to exp(scatters) #samples[:,-3] = np.exp(samples[:,-3]) #samples[:,-2] = np.exp(samples[:,-2]) #samples[:,-1] = np.exp(samples[:,-1]) fig = corner.corner( samples[:, :-3], labels=_list_labels, label_kwargs={'fontsize': 20}, # truths=_list_answer, quantiles=[0.16, 0.5, 0.84], plot_datapoints=True, show_titles=True, title_args={'fontsize': 20}, title_fmt='.3f', ) filename_pre = 'newModel_1/'+distance+'/trinagle_walker%dNrun%dNburn%d_withscatter_'\ +age+'Gyr_'+distance+'_%dstars_newModel_1' filename = filename_pre % (N_WALKERS, Nrun, Nburn, Number) fig.savefig(filename + '.png') #fig.savefig(filename+'.pdf') p = np.mean(samples, axis=0) e = np.var(samples, axis=0)**0.5 filename = 'newModel_1/result_' + age + 'Gyr_' + distance + '_' + str( Number) + 'stars_newModel_1' + '.txt' np.savetxt(filename, (p, e), fmt="%.3f", delimiter=',') va,vR2,sigmaphi,meanvR,sigmaR =\ lnlike(p,cos_twol,sin_twol,sin_l,cos_l,sin_b,cos_b,data_pml,data_err_pml,data_pmb,\ data_err_pmb,data_plx,data_err_plx,data_vlos,data_err_vlos,data_age,age_number)[1],\ lnlike(p,cos_twol,sin_twol,sin_l,cos_l,sin_b,cos_b,data_pml,data_err_pml,data_pmb,\ data_err_pmb,data_plx,data_err_plx,data_vlos,data_err_vlos,data_age,age_number)[2],\ lnlike(p,cos_twol,sin_twol,sin_l,cos_l,sin_b,cos_b,data_pml,data_err_pml,data_pmb,\ data_err_pmb,data_plx,data_err_plx,data_vlos,data_err_vlos,data_age,age_number)[3],\ lnlike(p,cos_twol,sin_twol,sin_l,cos_l,sin_b,cos_b,data_pml,data_err_pml,data_pmb,\ data_err_pmb,data_plx,data_err_plx,data_vlos,data_err_vlos,data_age,age_number)[4],\ lnlike(p,cos_twol,sin_twol,sin_l,cos_l,sin_b,cos_b,data_pml,data_err_pml,data_pmb,\ data_err_pmb,data_plx,data_err_plx,data_vlos,data_err_vlos,data_age,age_number)[5] f = open('va_newModel_1.txt', 'a') printline = '%s, %f, %f, %f, %f, %f, %f\n' % (filename, np.mean(va), vR2, sigmaphi, np.mean(sigmaR) / 80., meanvR, sigmaR) f.write(printline) f.close() va_list = [] sigmaR = [] print(filename) return None
def main(): # firstly, read one file and get k bins we want for the fitting range; It's dark matter power spectrum in redshift space parameter_file = sys.argv[1] with open(parameter_file, 'r') as fr: input_params = yaml.load(fr) data_m = np.genfromtxt( input_params['Pwig_ifile'], dtype='f8', comments='#', delimiter='', skip_header=11 ) # skip the first data row, the first 10 rows are comments. #print(data_m) num_kbin = np.size(data_m, axis=0) kk = data_m[:, 0] # get indices based on the ascending k indices_sort = [i[0] for i in sorted(enumerate(kk), key=lambda x: x[1])] # sort out the indices whose k<=0.3 h/Mpc for i in range(num_kbin): if kk[indices_sort[i]] > 0.3: break print(indices_sort[i - 1]) indices_p = indices_sort[0:i] k_p = kk[indices_p] N_fitbin = len(k_p) mu_p, Pwig = data_m[indices_p, 1], data_m[indices_p, 2] print(k_p, N_fitbin) # input Pnow, note the (k, mu) indices have the same order as those of Pwig data file data_m = np.genfromtxt(input_params['Pnow_ifile'], dtype='f8', comments='#', delimiter='', skip_header=11) Pnow = data_m[indices_p, 2] Pwnw_diff_obs = Pwig - Pnow # input diagonal terms of covariance matrix of (Pwig-Pnow) diag_Cov_Pwnw = np.loadtxt(input_params['diag_Cov_Pwnw_ifile'], dtype='f8', comments='#', usecols=(2, )) ivar_Pk_wnow = 1.0 / diag_Cov_Pwnw #print(ivar_Pk_wnow) # input (theoretical) linear power spectrum k_wiggle, Pk_wiggle = np.loadtxt(input_params['Pwig_linear'], dtype='f8', comments='#', unpack=True) tck_Pk_linw = interpolate.splrep(k_wiggle, Pk_wiggle) k_smooth, Pk_smooth = np.loadtxt(input_params['Pnow_linear'], dtype='f8', comments='#', unpack=True) tck_Pk_sm = interpolate.splrep(k_smooth, Pk_smooth) q_max = 110.0 # Mpc/h, BAO radius Sigma2_sm_array = np.linspace(10.0, 610.0, 61) #print(Sigma2_sm_array) #const = 1.0 Sigma2_dd_array = np.array([ integrate.quad(Sigma2_dd_integrand, 1.05e-5, 100.0, args=(tck_Pk_linw, q_max, Sigma2_sm), epsabs=1.e-4, epsrel=1.e-4)[0] for Sigma2_sm in Sigma2_sm_array ]) Sigma2_sd_array = np.array([ integrate.quad(Sigma2_sd_integrand, 1.05e-5, 100.0, args=(tck_Pk_linw, q_max, Sigma2_sm), epsabs=1.e-4, epsrel=1.e-4)[0] for Sigma2_sm in Sigma2_sm_array ]) Sigma2_ss_array = np.array([ integrate.quad(Sigma2_ss_integrand, 1.05e-5, 100.0, args=(tck_Pk_linw, q_max, Sigma2_sm), epsabs=1.e-4, epsrel=1.e-4)[0] for Sigma2_sm in Sigma2_sm_array ]) tck_Sigma2_dd = interpolate.splrep(Sigma2_sm_array, Sigma2_dd_array, k=3) tck_Sigma2_sd = interpolate.splrep(Sigma2_sm_array, Sigma2_sd_array, k=3) tck_Sigma2_ss = interpolate.splrep(Sigma2_sm_array, Sigma2_ss_array, k=3) all_params = list(input_params['init_params'].values()) params_indices = input_params['params_indices'] params_name = list(input_params['init_params'].keys()) all_temperature = input_params['all_temperature'] N_params, theta, fix_params, params_T, params_name = set_params( all_params, params_indices, params_name, all_temperature) print(N_params, theta, fix_params, params_T, params_name) sim_z = input_params['sim_z'] # redshift of the simulated power spectrum N_walkers = input_params['N_walkers'] Omega_m = input_params['Omega_m'] G_0 = growth_factor(0.0, Omega_m) # G_0 at z=0, normalization factor norm_gf = growth_factor(sim_z, Omega_m) / G_0 const = 1.0 / (6.0 * np.pi**2.0) * norm_gf**2.0 pool = MPIPool(loadbalance=True) np.random.seed(1) # set random seed for random number generator params_mcmc = mcmc_routine(N_params, N_walkers, theta, params_T, params_indices, fix_params, k_p, mu_p, Pwnw_diff_obs, ivar_Pk_wnow, tck_Pk_linw, tck_Pk_sm, tck_Sigma2_dd, tck_Sigma2_sd, tck_Sigma2_ss, norm_gf, const, params_name, pool) print(params_mcmc) chi_square = chi2(params_mcmc[:, 0], params_indices, fix_params, k_p, mu_p, Pwnw_diff_obs, ivar_Pk_wnow, tck_Pk_linw, tck_Pk_sm, tck_Sigma2_dd, tck_Sigma2_sd, tck_Sigma2_ss, norm_gf, const) dof = N_fitbin - N_params reduced_chi2 = chi_square / dof odir = './output_files/' if not os.path.exists(odir): os.makedirs(odir) ofile_params = odir + input_params['ofile_name'].format( sim_z, ''.join(map(str, params_indices))) write_params(ofile_params, params_mcmc, params_name, reduced_chi2, fix_params, dof) pool.close()
def fit_subsamplefof_mean(): parser = argparse.ArgumentParser( description= 'This is the MCMC code to get the fitting parameters, made by Zhejie Ding.' ) parser.add_argument( '-rec_id', "--rec_id", help='The id of reconstruction, either 0 or 1.', required=True) #0: pre-reconstruct; 1: post-reconstruct parser.add_argument('-space_id', "--space_id", help='0 for real space, 1 for redshift space.', required=True) parser.add_argument( '-set_Sigma_xyz_theory', "--set_Sigma_xyz_theory", help= 'Determine whether the parameters \Sigma_xy and \Sigma_z are fixed or not, either True or False', required=True) parser.add_argument( '-set_Sigma_sm_theory', "--set_Sigma_sm_theory", help= 'Determine whether we use sigma_sm from theory in the fitting model. \ If False, sigma_sm=0 (be careful that sigma_sm=\inf in real space case)', required=True) args = parser.parse_args() print("args: ", args) rec_id = int(args.rec_id) space_id = int(args.space_id) set_Sigma_xyz_theory = args.set_Sigma_xyz_theory set_Sigma_sm_theory = args.set_Sigma_sm_theory print("rec_id: ", rec_id, "space_id: ", space_id) print("set_Sigma_xyz_theory: ", set_Sigma_xyz_theory, "set_Sigma_sm_theory: ", set_Sigma_sm_theory) N_walkers = 40 # increase N_walkers would decrease the minimum number of walk steps which make fitting parameters convergent, but running time increases. N_walkersteps = 5000 # simulation run name N_dataset = 20 N_mu_bin = 100 #N_skip_header = 11 #N_skip_footer = 31977 Omega_m = 0.3075 G_0 = growth_factor(0.0, Omega_m) # G_0 at z=0, normalization factor Volume = 1380.0**3.0 # the volume of simulation box sim_z = ['0', '0.6', '1.0'] sim_seed = [0, 9] sim_wig = ['NW', 'WG'] sim_a = ['1.0000', '0.6250', '0.5000'] sim_space = ['r', 's'] # r for real space; s for redshift space rec_dirs = [ 'DD', 'ALL' ] # "ALL" folder stores P(k, \mu) after reconstruction process, while DD is before reconstruction. rec_fprefix = ['', 'R'] mcut_Npar_list = [[37, 149, 516, 1524, 3830], [35, 123, 374, 962, 2105], [34, 103, 290, 681, 1390]] N_masscut = np.size(mcut_Npar_list, axis=1) # Sigma_sm = sqrt(2.* Sig_RR) in post-reconstruction case, for pre-reconstruction, we don't use sub_Sigma_RR. Sigma_RR_list = [[37, 48.5, 65.5, 84.2, 110], [33, 38, 48.5, 63.5, 91.5], [31, 38, 49, 65, 86]] sub_Sigma_RR = 50.0 # note from Hee-Jong's recording inputf = '../Zvonimir_data/planck_camb_56106182_matterpower_smooth_z0.dat' k_smooth, Pk_smooth = np.loadtxt(inputf, dtype='f8', comments='#', unpack=True) tck_Pk_sm = interpolate.splrep(k_smooth, Pk_smooth) inputf = '../Zvonimir_data/planck_camb_56106182_matterpower_z0.dat' k_wiggle, Pk_wiggle = np.loadtxt(inputf, dtype='f8', comments='#', unpack=True) tck_Pk_linw = interpolate.splrep(k_wiggle, Pk_wiggle) # firstly, read one file and get k bins we want for the fitting range dir0 = '/Users/ding/Documents/playground/WiggleNowiggle/subsample_FoF_data_HS/Pk_obs_2d_wnw_mean_DD_ksorted_mu_masscut/' inputf = dir0 + 'fof_kaver.wnw_diff_a_0.6250_mcut35_fraction0.126.dat' k_p, mu_p = np.loadtxt(inputf, dtype='f8', comments='#', delimiter=' ', usecols=(0, 1), unpack=True) #print(k_p, mu_p) N_fitbin = len(k_p) #print('# of (k, mu) bins: ', N_fitbin) # for output parameters fitted odir = './params_{}_wig-now_b_bscale_fitted_mean_dset/'.format( rec_dirs[rec_id]) if not os.path.exists(odir): os.makedirs(odir) print("N_walkers: ", N_walkers, "N_walkersteps: ", N_walkersteps, "\n") if rec_id == 0: ##Sigma_0 = 8.3364 # the approximated value of \Sigma_xy and \Sigma_z, unit Mpc/h, at z=0. Sigma_0 = 7.8364 # suggested by Zvonimir, at z=0 elif rec_id == 1: Sigma_0 = 2.84 ##space_id = 1 # in redshift space # 0: parameter fixed, 1: parameter free. #params_indices = [1, 1, 1, 1, 1, 1] # It doesn't fit \Sigma and bscale well. Yes, it dosen't work well (it's kind of overfitting). ##params_indices = [1, 1, 1, 1, 0, 1, 1, 0, 0] # b0 needs to be fitted. For this case, make sure \Sigma_xy and \Sigma_z positive, which should be set in mcmc_routine. ##params_indices = [0, 1, 0, 1, 1, 0] # make sure \alpha_xy = \alpha_z and \Sigma_xy = \Sigma_z params_indices = [ 1, 1, 1, 1, 0, 0, 0, 0, 0 ] # Set sigma_fog=0, f=0, b_scale=0, b_0=1.0 for subsamled DM case in real space. ##params_indices = [1, 1, 0, 0, 1, 1] # with fixed Sigma from theoretical value, then need to set sigma_xy, sigma_z equal to Sigma_z ##params_indices = [0, 1, 0, 0, 1, 1] # For this case, make sure \alpha_1 = \alpha_2 in the function lnlike(..) and set sigma_xy, sigma_z equal to Sigma_z. print("params_indices: ", params_indices) ##alpha_1, alpha_2, sigma_fog, f, b_0, b_scale = 1.0, 1.0, 2.0, 0.2, 1.0, 0.0 alpha_1, alpha_2, sigma_fog, f, b_0, b_scale = 1.0, 1.0, 0.0, 0.0, 1.0, 0.0 # ! only for real space, i.e., set sigma_fog, f equal to 0. all_names = "alpha_1", "alpha_2", "sigma_xy", "sigma_z", "sigma_sm", "sigma_fog", "f", "b_0", "b_scale" # the same order for params_indices all_temperature = 0.01, 0.01, 0.1, 0.1, 0.1, 0.1, 0.1, 0.01, 0.1 pool = MPIPool(loadbalance=True) for z_id in xrange(3): norm_gf = growth_factor(float(sim_z[z_id]), Omega_m) / G_0 Sigma_z = Sigma_0 * norm_gf / 2.0 # divided by 2.0 for estimated \Sigma of post-reconstruction ##Sigma_z = Sigma_0* norm_gf if set_Sigma_xyz_theory == "True": print("Sigma_z: ", Sigma_z) sigma_xy, sigma_z = Sigma_z, Sigma_z else: if params_indices[2] == 0: sigma_xy = 0.0 else: sigma_xy = 10.0 if params_indices[3] == 0: sigma_z = 0.0 else: sigma_z = 10.0 np.random.seed() # Set it for FoF fitting # for mcut_id in xrange(N_masscut): # if set_Sigma_sm_theory == "True": # sigma_sm = (float(Sigma_RR_list[z_id][mcut_id])*2.0)**0.5 # else: # sigma_sm = 0.0 # # all_params = alpha_1, alpha_2, sigma_xy, sigma_z, sigma_sm, sigma_fog, f, b_0, b_scale # N_params, theta, fix_params, params_T, params_name = set_params(all_params, params_indices, all_names, all_temperature) # # ifile_Pk = './run2_3_Pk_obs_2d_wnw_mean_{}_ksorted_mu_masscut/{}kave{}.wig_minus_now_mean_fof_a_{}_mcut{}.dat'.format(rec_dirs[rec_id], rec_fprefix[rec_id], sim_space[space_id], sim_a[z_id], mcut_Npar_list[z_id][mcut_id]) # Pk_wnw_diff_obs = np.loadtxt(ifile_Pk, dtype='f4', comments='#', usecols=(2,)) # be careful that there are k, \mu, P(k, \mu) columns. # # ifile_Cov_Pk = './run2_3_Cov_Pk_obs_2d_wnw_{}_ksorted_mu_masscut/{}kave{}.wig_minus_now_mean_fof_a_{}_mcut{}.dat'.format(rec_dirs[rec_id], rec_fprefix[rec_id], sim_space[space_id], sim_a[z_id], mcut_Npar_list[z_id][mcut_id]) # Cov_Pk_wnw = np.loadtxt(ifile_Cov_Pk, dtype='f4', comments='#') # ivar_Pk_wnow = N_dataset/np.diag(Cov_Pk_wnw) # the mean sigma error # # params_mcmc = mcmc_routine(N_params, N_walkers, N_walkersteps, theta, params_T, params_indices, fix_params, k_p, mu_p, Pk_wnw_diff_obs, ivar_Pk_wnow, tck_Pk_linw, tck_Pk_sm, norm_gf, params_name, pool) # # chi_square = chi2(params_mcmc[:, 0], params_indices, fix_params, k_p, mu_p, Pk_wnw_diff_obs, ivar_Pk_wnow, tck_Pk_linw, tck_Pk_sm, norm_gf) # reduced_chi2 = chi_square/(N_fitbin-N_params) # print("Reduced chi2: {}\n".format(reduced_chi2)) # # output parameters into a file # if set_Sigma_xyz_theory == "False": # ofile_params = odir + 'fof_{}kave{}.wnw_diff_a_{}_mcut{}_params{}_Sigma_sm{}.dat'.format(rec_fprefix[rec_id], sim_space[space_id], sim_a[z_id], mcut_Npar_list[z_id][mcut_id], ''.join(map(str, params_indices)), round(sigma_sm,3)) # else: # ofile_params = odir + 'fof_{}kave{}.wnw_diff_a_{}_mcut{}_params{}_isotropic_Sigmaz_{}_Sigma_sm{}.dat'.format(rec_fprefix[rec_id], sim_space[space_id], sim_a[z_id], mcut_Npar_list[z_id][mcut_id], ''.join(map(str, params_indices)), round(Sigma_z, 3), round(sigma_sm,3)) # print(ofile_params) # write_params(ofile_params, params_mcmc, params_name, reduced_chi2) # set it for DM subsample fitting sub_sigma_sm = (sub_Sigma_RR * 2.0)**0.5 print("sub_sigma_sm: ", sub_sigma_sm) all_params = alpha_1, alpha_2, sigma_xy, sigma_z, sub_sigma_sm, sigma_fog, f, b_0, b_scale # set \Sigma_sm = sqrt(50*2)=10, for post-rec N_params, theta, fix_params, params_T, params_name = set_params( all_params, params_indices, all_names, all_temperature) # Fit for DM power spectrum ifile_Pk = './run2_3_sub_Pk_2d_wnw_mean_{}_ksorted_mu/{}kave{}.wig_minus_now_mean_sub_a_{}.dat'.format( rec_dirs[rec_id], rec_fprefix[rec_id], sim_space[space_id], sim_a[z_id]) Pk_wnw_diff_true = np.loadtxt( ifile_Pk, dtype='f4', comments='#', usecols=( 2, )) # be careful that there are k, \mu, P(k, \mu) columns. ifile_Cov_Pk = './run2_3_sub_Cov_Pk_2d_wnw_{}_ksorted_mu/{}kave{}.wig_minus_now_mean_sub_a_{}.dat'.format( rec_dirs[rec_id], rec_fprefix[rec_id], sim_space[space_id], sim_a[z_id]) Cov_Pk_wnw = np.loadtxt(ifile_Cov_Pk, dtype='f4', comments='#') ivar_Pk_wnow = N_dataset / np.diag(Cov_Pk_wnw) # the mean sigma error params_mcmc = mcmc_routine(N_params, N_walkers, N_walkersteps, theta, params_T, params_indices, fix_params, k_p, mu_p, Pk_wnw_diff_true, ivar_Pk_wnow, tck_Pk_linw, tck_Pk_sm, norm_gf, params_name, pool) chi_square = chi2(params_mcmc[:, 0], params_indices, fix_params, k_p, mu_p, Pk_wnw_diff_true, ivar_Pk_wnow, tck_Pk_linw, tck_Pk_sm, norm_gf) reduced_chi2 = chi_square / (N_fitbin - N_params) print('Reduced chi2: {}\n'.format(reduced_chi2)) if rec_id == 1: if set_Sigma_xyz_theory == "False": ofile_params = odir + 'sub_{}kave{}.wnw_diff_a_{}_params{}_Sigma_sm{}.dat'.format( rec_fprefix[rec_id], sim_space[space_id], sim_a[z_id], ''.join(map(str, params_indices)), round(sub_sigma_sm, 3)) else: ofile_params = odir + 'sub_{}kave{}.wnw_diff_a_{}_params{}_isotropic_Sigmaz_{}_Sigma_sm{}.dat'.format( rec_fprefix[rec_id], sim_space[space_id], sim_a[z_id], ''.join(map(str, params_indices)), round(Sigma_z, 3), round(sub_sigma_sm, 3)) elif rec_id == 0: if set_Sigma_xyz_theory == "False": ofile_params = odir + 'sub_{}kave{}.wnw_diff_a_{}_params{}.dat'.format( rec_fprefix[rec_id], sim_space[space_id], sim_a[z_id], ''.join(map(str, params_indices))) else: ofile_params = odir + 'sub_{}kave{}.wnw_diff_a_{}_params{}_isotropic_Sigmaz_{}.dat'.format( rec_fprefix[rec_id], sim_space[space_id], sim_a[z_id], ''.join(map(str, params_indices)), round(Sigma_z, 3)) # write_params(ofile_params, params_mcmc, params_name, reduced_chi2) pool.close()
def mcmc_mpi(Nwalkers, Nchains, observables=['nbar', 'xi'], data_dict={ 'Mr': 21, 'b_normal': 0.25 }, prior_name='first_try', mcmcrun=None): ''' Standard MCMC implementaion Parameters ----------- - Nwalker : Number of walkers - Nchains : Number of MCMC chains - observables : list of observables. Options are: ['nbar','xi'],['nbar','gmf'],['xi'] - data_dict : dictionary that specifies the observation keywords ''' #Initializing the vector of observables and inverse covariance matrix if observables == ['xi']: fake_obs = Data.data_xi(**data_dict) #fake_obs_icov = Data.data_inv_cov('xi', **data_dict) fake_obs_icov = Data.data_cov(inference='mcmc', **data_dict)[1:16, 1:16] if observables == ['nbar', 'xi']: fake_obs = np.hstack( [Data.data_nbar(**data_dict), Data.data_xi(**data_dict)]) fake_obs_icov = Data.data_cov(inference='mcmc', **data_dict)[:16, :16] if observables == ['nbar', 'gmf']: ##### FIRST BIN OF GMF DROPPED ############### # CAUTION: hardcoded fake_obs = np.hstack( [Data.data_nbar(**data_dict), Data.data_gmf(**data_dict)[1:]]) fake_obs_icov = np.zeros((10, 10)) #print Data.data_cov(**data_dict)[17: , 17:].shape # Covariance matrix being adjusted accordingly fake_obs_icov[1:, 1:] = Data.data_cov(inference='mcmc', **data_dict)[17:, 17:] fake_obs_icov[0, 1:] = Data.data_cov(inference='mcmc', **data_dict)[0, 17:] fake_obs_icov[1:, 0] = Data.data_cov(inference='mcmc', **data_dict)[17:, 0] fake_obs_icov[0, 0] = Data.data_cov(inference='mcmc', **data_dict)[0, 0] # True HOD parameters data_hod_dict = Data.data_hod_param(Mr=data_dict['Mr']) data_hod = np.array([ data_hod_dict['logM0'], # log M0 np.log(data_hod_dict['sigma_logM']), # log(sigma) data_hod_dict['logMmin'], # log Mmin data_hod_dict['alpha'], # alpha data_hod_dict['logM1'] # log M1 ]) Ndim = len(data_hod) # Priors prior_min, prior_max = PriorRange(prior_name) prior_range = np.zeros((len(prior_min), 2)) prior_range[:, 0] = prior_min prior_range[:, 1] = prior_max # mcmc chain output file chain_file = ''.join([ util.mcmc_dir(), util.observable_id_flag(observables), '.', mcmcrun, '.mcmc_chain.dat' ]) #print chain_file if os.path.isfile(chain_file) and continue_chain: print 'Continuing previous MCMC chain!' sample = np.loadtxt(chain_file) Nchain = Niter - (len(sample) / Nwalkers ) # Number of chains left to finish if Nchain > 0: pass else: raise ValueError print Nchain, ' iterations left to finish' # Initializing Walkers from the end of the chain pos0 = sample[-Nwalkers:] else: # new chain f = open(chain_file, 'w') f.close() Nchain = Niter # Initializing Walkers random_guess = data_hod pos0 = np.repeat(random_guess, Nwalkers).reshape(Ndim, Nwalkers).T + \ 5.e-2 * np.random.randn(Ndim * Nwalkers).reshape(Nwalkers, Ndim) #print pos0.shape # Initializing MPIPool pool = MPIPool() if not pool.is_master(): pool.wait() sys.exit(0) # Initializing the emcee sampler hod_kwargs = { 'prior_range': prior_range, 'data': fake_obs, 'data_icov': fake_obs_icov, 'observables': observables, 'Mr': data_dict['Mr'] } sampler = emcee.EnsembleSampler(Nwalkers, Ndim, lnPost, pool=pool, kwargs=hod_kwargs) # Initializing Walkers for result in sampler.sample(pos0, iterations=Nchain, storechain=False): position = result[0] #print position f = open(chain_file, 'a') for k in range(position.shape[0]): output_str = '\t'.join(position[k].astype('str')) + '\n' f.write(output_str) f.close() pool.close()
def main(runmpi=False, nw=100, th=6, bi=10, fr=10): if runmpi: pool = MPIPool() if not pool.is_master(): pool.wait() sys.exit(0) else: pool=None time, flux, ferr = get_lc() toi = 175 cadence = 120 rho = 18 rho_unc = 1 nplanets = 3 ld1 = 0.1642 ld2 = 0.4259 dil=0.0 periods = [2.25321888449, 3.6906274382, 7.45131144274] impacts = [0.26, 0.21, 0.89] T0s = [1354.90455205, 1356.203624274, 1355.2866249] rprss = [0.02011, 0.038564, 0.0438550698] planet_guess = np.array([]) for i in range(nplanets): planet_guess = np.r_[planet_guess, T0s[i], periods[i], impacts[i], rprss[i], 0.0, 0.0 ] nwalkers = nw threads = th burnin = bi fullrun = fr thin = 1 M = tmod.transitmc2( nplanets, cadence) M.get_ld(ld1, ld2) M.already_open(time, flux, ferr) M.get_rho([rho, rho_unc]) M.get_zpt(0.0) M.get_sol(*planet_guess) outfile = 'koi{0}_np{1}.hdf5'.format( toi, nplanets) p0 = M.get_guess(nwalkers) l_var = np.shape(p0)[1] tom = tmod.logchi2 args = [M.nplanets, M.rho_0, M.rho_0_unc, M.ld1, M.ld1_unc, M.ld2, M.ld2_unc, M.flux, M.err, M.fixed_sol, M.time, M._itime, M._ntt, M._tobs, M._omc, M._datatype] N = len([indval for indval in range(fullrun) if indval%thin == 0]) with h5py.File(outfile, u"w") as f: f.create_dataset("time", data=M.time) f.create_dataset("flux", data=M.flux) f.create_dataset("err", data=M.err) f.attrs["rho_0"] = M.rho_0 f.attrs["rho_0_unc"] = M.rho_0_unc f.attrs["nplanets"] = M.nplanets f.attrs["ld1"] = M.ld1 f.attrs["ld2"] = M.ld2 g = f.create_group("mcmc") g.attrs["nwalkers"] = nwalkers g.attrs["burnin"] = burnin g.attrs["iterations"] = fullrun g.attrs["thin"] = thin g.create_dataset("fixed_sol", data= M.fixed_sol) g.create_dataset("fit_sol_0", data= M.fit_sol_0) c_ds = g.create_dataset("chain", (nwalkers, N, l_var), dtype=np.float64) lp_ds = g.create_dataset("lnprob", (nwalkers, N), dtype=np.float64) if runmpi: sampler = emcee.EnsembleSampler(nwalkers, l_var, tom, args=args,pool=pool) else: sampler = emcee.EnsembleSampler(nwalkers, l_var, tom, args=args,threads=th) time1 = thetime.time() p2, prob, state = sampler.run_mcmc(p0, burnin, storechain=False) sampler.reset() with h5py.File(outfile, u"a") as f: g = f["mcmc"] g.create_dataset("burnin_pos", data=p2) g.create_dataset("burnin_prob", data=prob) time2 = thetime.time() print('burn-in took ' + str((time2 - time1)/60.) + ' min') time1 = thetime.time() for i, (pos, lnprob, state) in enumerate(tqdm(sampler.sample(p2, iterations=fullrun, rstate0=state, storechain=False), total=fullrun)): #do the thinning in the loop here if i % thin == 0: ind = i / thin with h5py.File(outfile, u"a") as f: g = f["mcmc"] c_ds = g["chain"] lp_ds = g["lnprob"] c_ds[:, ind, :] = pos lp_ds[:, ind] = lnprob time2 = thetime.time() print('MCMC run took ' + str((time2 - time1)/60.) + ' min') print('') print("Mean acceptance: " + str(np.mean(sampler.acceptance_fraction))) print('') if runmpi: pool.close() else: sampler.pool.close() return sampler
def run_espei(run_settings): """Wrapper around the ESPEI fitting procedure, taking only a settings dictionary. Parameters ---------- run_settings : dict Dictionary of input settings Returns ------- Either a Database (for generate parameters only) or a tuple of (Database, sampler) """ run_settings = get_run_settings(run_settings) system_settings = run_settings['system'] output_settings = run_settings['output'] generate_parameters_settings = run_settings.get('generate_parameters') mcmc_settings = run_settings.get('mcmc') # handle verbosity verbosity = {0: logging.WARNING, 1: logging.INFO, 2: logging.DEBUG} logging.basicConfig(level=verbosity[output_settings['verbosity']]) # load datasets and handle i/o logging.debug('Loading and checking datasets.') dataset_path = system_settings['datasets'] datasets = load_datasets(sorted(recursive_glob(dataset_path, '*.json'))) if len(datasets.all()) == 0: logging.warning( 'No datasets were found in the path {}. This should be a directory containing dataset files ending in `.json`.' .format(dataset_path)) logging.debug('Finished checking datasets') with open(system_settings['phase_models']) as fp: phase_models = json.load(fp) if generate_parameters_settings is not None: refdata = generate_parameters_settings['ref_state'] excess_model = generate_parameters_settings['excess_model'] dbf = generate_parameters(phase_models, datasets, refdata, excess_model) dbf.to_file(output_settings['output_db'], if_exists='overwrite') if mcmc_settings is not None: tracefile = output_settings['tracefile'] probfile = output_settings['probfile'] # check that the MCMC output files do not already exist # only matters if we are actually running MCMC if os.path.exists(tracefile): raise OSError( 'Tracefile "{}" exists and would be overwritten by a new run. Use the ``output.tracefile`` setting to set a different name.' .format(tracefile)) if os.path.exists(probfile): raise OSError( 'Probfile "{}" exists and would be overwritten by a new run. Use the ``output.probfile`` setting to set a different name.' .format(probfile)) # scheduler setup if mcmc_settings['scheduler'] == 'MPIPool': # check that cores is not an input setting if mcmc_settings.get('cores') != None: logging.warning("MPI does not take the cores input setting.") from emcee.utils import MPIPool # code recommended by emcee: if not master, wait for instructions then exit client = MPIPool() if not client.is_master(): logging.debug( 'MPIPool is not master. Waiting for instructions...') client.wait() sys.exit(0) logging.info("Using MPIPool on {} MPI ranks".format(client.size)) elif mcmc_settings['scheduler'] == 'dask': from distributed import LocalCluster cores = mcmc_settings.get('cores', multiprocessing.cpu_count()) if (cores > multiprocessing.cpu_count()): cores = multiprocessing.cpu_count() logging.warning( "The number of cores chosen is larger than available. " "Defaulting to run on the {} available cores.".format( cores)) scheduler = LocalCluster(n_workers=cores, threads_per_worker=1, processes=True) client = ImmediateClient(scheduler) client.run(logging.basicConfig, level=verbosity[output_settings['verbosity']]) logging.info("Running with dask scheduler: %s [%s cores]" % (scheduler, sum(client.ncores().values()))) try: logging.info( "bokeh server for dask scheduler at localhost:{}".format( client.scheduler_info()['services']['bokeh'])) except KeyError: logging.info("Install bokeh to use the dask bokeh server.") elif mcmc_settings['scheduler'] == 'emcee': from emcee.interruptible_pool import InterruptiblePool cores = mcmc_settings.get('cores', multiprocessing.cpu_count()) if (cores > multiprocessing.cpu_count()): cores = multiprocessing.cpu_count() logging.warning( "The number of cores chosen is larger than available. " "Defaulting to run on the {} available cores.".format( cores)) client = InterruptiblePool(processes=cores) logging.info("Using multiprocessing on {} cores".format(cores)) elif mcmc_settings['scheduler'] == 'None': client = None logging.info( "Not using a parallel scheduler. ESPEI is running MCMC on a single core." ) # get a Database if mcmc_settings.get('input_db'): dbf = Database(mcmc_settings.get('input_db')) # load the restart chain if needed if mcmc_settings.get('restart_chain'): restart_chain = np.load(mcmc_settings.get('restart_chain')) else: restart_chain = None # load the remaning mcmc fitting parameters mcmc_steps = mcmc_settings.get('mcmc_steps') save_interval = mcmc_settings.get('mcmc_save_interval') chains_per_parameter = mcmc_settings.get('chains_per_parameter') chain_std_deviation = mcmc_settings.get('chain_std_deviation') deterministic = mcmc_settings.get('deterministic') dbf, sampler = mcmc_fit( dbf, datasets, scheduler=client, mcmc_steps=mcmc_steps, chains_per_parameter=chains_per_parameter, chain_std_deviation=chain_std_deviation, save_interval=save_interval, tracefile=tracefile, probfile=probfile, restart_chain=restart_chain, deterministic=deterministic, ) dbf.to_file(output_settings['output_db'], if_exists='overwrite') # close the scheduler, if possible if hasattr(client, 'close'): client.close() return dbf, sampler return dbf
lnp_phot = lnlike_phot(phot, obs=obs, phot_noise=phot_noise, **vectors) d2 = time.time() - t2 if verbose: write_log(theta, lnp_prior, lnp_spec, lnp_phot, d1, d2) return lnp_prior + lnp_phot + lnp_spec # ----------------- # MPI pool. This must be done *after* lnprob and # chi2 are defined since slaves will only see up to # sys.exit() # ------------------ try: from emcee.utils import MPIPool pool = MPIPool(debug=False, loadbalance=True) if not pool.is_master(): # Wait for instructions from the master process. pool.wait() sys.exit(0) except (ImportError, ValueError): pool = None print('Not using MPI') def halt(message): """Exit, closing pool safely. """ print(message) try: pool.close()
def main(argv): ################## #These change a lot numWaveforms = 16 numThreads = 12 ndim = 6 * numWaveforms + 8 nwalkers = 2 * ndim iter = 50 burnIn = 40 wfPlotNumber = 10 ###################### # plt.ion() fitSamples = 200 #Prepare detector zero_1 = -5.56351644e+07 pole_1 = -1.38796386e+04 pole_real = -2.02559385e+07 pole_imag = 9885315.37450211 zeros = [zero_1, 0] poles = [pole_real + pole_imag * 1j, pole_real - pole_imag * 1j, pole_1] system = signal.lti(zeros, poles, 1E7) tempGuess = 77.89 gradGuess = 0.0483 pcRadGuess = 2.591182 pcLenGuess = 1.613357 #Create a detector model detName = "conf/P42574A_grad%0.2f_pcrad%0.2f_pclen%0.2f.conf" % (0.05, 2.5, 1.65) det = Detector(detName, temperature=tempGuess, timeStep=1., numSteps=fitSamples * 10, tfSystem=system) det.LoadFields("P42574A_fields_v3.npz") det.SetFields(pcRadGuess, pcLenGuess, gradGuess) tempIdx = -8 gradIdx = -7 pcRadIdx = -6 pcLenIdx = -5 #and the remaining 4 are for the transfer function fig_size = (20, 10) #Create a decent start guess by fitting waveform-by-waveform wfFileName = "P42574A_512waveforms_%drisetimeculled.npz" % numWaveforms if os.path.isfile(wfFileName): data = np.load(wfFileName) results = data['results'] wfs = data['wfs'] numWaveforms = wfs.size else: print "No saved waveforms available. Loading from Data" exit(0) #prep holders for each wf-specific param r_arr = np.empty(numWaveforms) phi_arr = np.empty(numWaveforms) z_arr = np.empty(numWaveforms) scale_arr = np.empty(numWaveforms) t0_arr = np.empty(numWaveforms) smooth_arr = np.ones(numWaveforms) * 7. simWfArr = np.empty((1, numWaveforms, fitSamples)) #Prepare the initial value arrays for (idx, wf) in enumerate(wfs): wf.WindowWaveformTimepoint(fallPercentage=.99) r_arr[idx], phi_arr[idx], z_arr[idx], scale_arr[idx], t0_arr[ idx], smooth_arr[idx] = results[idx]['x'] t0_arr[ idx] += 10 #because i had a different windowing offset back in the day #Plot the waveforms to take a look at the initial guesses if False: fig = plt.figure() for (idx, wf) in enumerate(wfs): print "WF number %d:" % idx print " >>r: %f\n >>phi %f\n >>z %f\n >>e %f\n >>t0 %f\n >>smooth %f" % ( r_arr[idx], phi_arr[idx], z_arr[idx], scale_arr[idx], t0_arr[idx], smooth_arr[idx]) ml_wf = det.GetSimWaveform(r_arr[idx], phi_arr[idx], z_arr[idx], scale_arr[idx] * 100, t0_arr[idx], fitSamples, smoothing=smooth_arr[idx]) plt.plot(ml_wf, color="b") plt.plot(wf.windowedWf, color="r") value = raw_input(' --> Press q to quit, any other key to continue\n') if value == 'q': exit(0) #Initialize this thread's globals initializeDetectorAndWaveforms(det, wfs) #Initialize the multithreading pool = MPIPool() if not pool.is_master(): pool.wait() sys.exit(0) #Do the MCMC mcmc_startguess = np.hstack(( r_arr[:], phi_arr[:], z_arr[:], scale_arr[:] * 100., t0_arr[:], smooth_arr[:], # waveform-specific params tempGuess, gradGuess, pcRadGuess, pcLenGuess, zero_1, pole_1, pole_real, pole_imag)) # detector-specific #number of walkers _must_ be even if nwalkers % 2: nwalkers += 1 #Initialize walkers with a random, narrow ball around the start guess pos0 = [ mcmc_startguess + 1e-2 * np.random.randn(ndim) * mcmc_startguess for i in range(nwalkers) ] #Make sure everything in the initial guess is within bounds for pos in pos0: pos[:numWaveforms] = np.clip(pos[:numWaveforms], 0, np.floor(det.detector_radius * 10.) / 10.) pos[numWaveforms:2 * numWaveforms] = np.clip( pos[numWaveforms:2 * numWaveforms], 0, np.pi / 4) pos[2 * numWaveforms:3 * numWaveforms] = np.clip( pos[2 * numWaveforms:3 * numWaveforms], 0, np.floor(det.detector_length * 10.) / 10.) pos[4 * numWaveforms:5 * numWaveforms] = np.clip( pos[4 * numWaveforms:5 * numWaveforms], 0, fitSamples) pos[5 * numWaveforms:6 * numWaveforms] = np.clip( pos[5 * numWaveforms:6 * numWaveforms], 0, 20.) pos[tempIdx] = np.clip(pos[tempIdx], 40, 120) pos[gradIdx] = np.clip(pos[gradIdx], det.gradList[0], det.gradList[-1]) pos[pcRadIdx] = np.clip(pos[pcRadIdx], det.pcRadList[0], det.pcRadList[-1]) pos[pcLenIdx] = np.clip(pos[pcLenIdx], det.pcLenList[0], det.pcLenList[-1]) prior = lnprior(pos, ) if not np.isfinite(prior): print "BAD PRIOR WITH START GUESS YOURE KILLING ME SMALLS" print pos exit(0) #Initialize, run the MCMC sampler = emcee.EnsembleSampler(nwalkers, ndim, lnprob, pool=p) #w/ progress bar, & time the thing start = timer() for (idx, result) in enumerate( sampler.sample(pos0, iterations=iter, storechain=True)): continue end = timer() pool.close() print "Elapsed time: " + str(end - start) print "Dumping chain to file..." np.save("sampler_mpi_%dwfs.npy" % numWaveforms, sampler.chain) print "Making MCMC steps figure..." ######### Plots for Waveform params stepsFig = plt.figure(2, figsize=fig_size) plt.clf() ax0 = stepsFig.add_subplot(611) ax1 = stepsFig.add_subplot(612, sharex=ax0) ax2 = stepsFig.add_subplot(613, sharex=ax0) ax3 = stepsFig.add_subplot(614, sharex=ax0) ax4 = stepsFig.add_subplot(615, sharex=ax0) ax5 = stepsFig.add_subplot(616, sharex=ax0) ax0.set_ylabel('r') ax1.set_ylabel('phi') ax2.set_ylabel('z') ax3.set_ylabel('scale') ax4.set_ylabel('t0') ax5.set_ylabel('smoothing') for i in range(nwalkers): for j in range(wfs.size): ax0.plot(sampler.chain[i, :, 0 + j], alpha=0.3) # r ax1.plot(sampler.chain[i, :, numWaveforms + j], alpha=0.3) # phi ax2.plot(sampler.chain[i, :, 2 * numWaveforms + j], alpha=0.3) #z ax3.plot(sampler.chain[i, :, 3 * numWaveforms + j], alpha=0.3) #energy ax4.plot(sampler.chain[i, :, 4 * numWaveforms + j], alpha=0.3) #t0 ax5.plot(sampler.chain[i, :, 5 * numWaveforms + j], alpha=0.3) #smoothing plt.savefig("emcee_mpi_wfchain_%dwfs.png" % numWaveforms) ######### Plots for Detector params stepsFigDet = plt.figure(3, figsize=fig_size) plt.clf() ax0 = stepsFigDet.add_subplot(411) ax1 = stepsFigDet.add_subplot(412, sharex=ax0) ax2 = stepsFigDet.add_subplot(413, sharex=ax0) ax3 = stepsFigDet.add_subplot(414, sharex=ax0) ax0.set_ylabel('temp') ax1.set_ylabel('grad') ax2.set_ylabel('pcRad') ax3.set_ylabel('pcLen') for i in range(nwalkers): ax0.plot(sampler.chain[i, :, tempIdx], "b", alpha=0.3) #temp ax1.plot(sampler.chain[i, :, gradIdx], "b", alpha=0.3) #grad ax2.plot(sampler.chain[i, :, pcRadIdx], "b", alpha=0.3) #pcrad ax3.plot(sampler.chain[i, :, pcLenIdx], "b", alpha=0.3) #pclen plt.savefig("emcee_mpi_detchain_%dwfs.png" % numWaveforms) #and for the transfer function stepsFigTF = plt.figure(4, figsize=fig_size) plt.clf() tf0 = stepsFigTF.add_subplot(411) tf1 = stepsFigTF.add_subplot(412, sharex=ax0) tf2 = stepsFigTF.add_subplot(413, sharex=ax0) tf3 = stepsFigTF.add_subplot(414, sharex=ax0) tf0.set_ylabel('zero_1') tf1.set_ylabel('pole_1') tf2.set_ylabel('pole_real') tf3.set_ylabel('pole_imag') for i in range(nwalkers): tf0.plot(sampler.chain[i, :, -4], "b", alpha=0.3) #2 tf1.plot(sampler.chain[i, :, -3], "b", alpha=0.3) #den1 tf2.plot(sampler.chain[i, :, -2], "b", alpha=0.3) #2 tf3.plot(sampler.chain[i, :, -1], "b", alpha=0.3) #3 plt.savefig("emcee_mpi_tfchain_%dwfs.png" % numWaveforms) samples = sampler.chain[:, burnIn:, :].reshape((-1, ndim)) print "temp is %f" % np.median(samples[:, tempIdx]) print "grad is %f" % np.median(samples[:, gradIdx]) print "pcrad is %f" % np.median(samples[:, pcRadIdx]) print "pclen is %f" % np.median(samples[:, pcLenIdx]) print "zero_1 is %f" % np.median(samples[:, -4]) print "pole_1 is %f" % np.median(samples[:, -3]) print "pole_real is %f" % np.median(samples[:, -2]) print "pole_imag is %f" % np.median(samples[:, -1]) #TODO: Aaaaaaand plot some waveforms.. simWfs = np.empty((wfPlotNumber, numWaveforms, fitSamples)) for idx, (theta) in enumerate(samples[np.random.randint( len(samples), size=wfPlotNumber)]): temp, impGrad, pcRad, pcLen = theta[tempIdx], theta[gradIdx], theta[ pcRadIdx], theta[pcLenIdx] zero_1, pole_1, pole_real, pole_imag = theta[-4:] r_arr, phi_arr, z_arr, scale_arr, t0_arr, smooth_arr = theta[:-8].reshape( (6, numWaveforms)) det.SetTemperature(temp) det.SetFields(pcRad, pcLen, impGrad) zeros = [zero_1, 0] poles = [ pole_real + pole_imag * 1j, pole_real - pole_imag * 1j, pole_1 ] det.SetTransferFunction(zeros, poles, 1E7) for wf_idx in range(wfs.size): wf_i = det.GetSimWaveform(r_arr[wf_idx], phi_arr[wf_idx], z_arr[wf_idx], scale_arr[wf_idx], t0_arr[wf_idx], fitSamples) simWfs[idx, wf_idx, :] = wf_i if wf_i is None: print "Waveform %d, %d is None" % (idx, wf_idx) residFig = plt.figure(4, figsize=(20, 15)) helpers.plotManyResidual(simWfs, wfs, figure=residFig) plt.savefig("emcee_mpi_waveforms_%dwfs.png" % numWaveforms)
def rerunPosteriorPredictive(): ''' Rerun the posterior predictive distribution. This can be used to e.g. increase the resolution spatially or in terms of the age of stellar populations, or vary some parameter systematically. The mandatory argument func is a user-provided function that specifies how a model with known parameters should be modified and (re) run.''' pool = MPIPool(comm=comm, loadbalance=True) if not pool.is_master(): pool.wait() sys.exit(0) output = readoutput.Experiment(chainDirRel+'-ppd') # read in the posterior predictive distribution. output.read(paramsOnly=True,keepStars=False) emcee_params = [] print "output.models: ",len(output.models) # For each model, take the parameters we have read in and construct the corresponding emcee parameters. for model in output.models: #eta, epsff, fg0, muNorm, muMhScaling, fixedQ, accScaleLength, fcool, Mh0, fscatter, x0, x1, x2, x3, obsScale, conRF, muHgScaling = emceeParams eta = model.p['eta'] epsff = model.p['epsff'] fg0 = model.p['fg0'] muNorm = model.p['muNorm'] muMhScaling = model.p['muMhScaling'] fixedQ = model.p['fixedQ'] accScaleLength = model.p['accScaleLength'] fcool = model.p['fcool'] Mh0 = model.p['Mh0'] fscatter = model.p['fscatter'] x0 = model.p['x0'] x1 = model.p['x1'] x2 = model.p['x2'] x3 = model.p['x3'] obsScale = 1.0 # doesn't matter.. see below conRF = model.p['concentrationRandomFactor'] muHgScaling = model.p['muHgScaling'] # We have everything except obsScale, but actually that doesn't matter, # since it only affects the model in post-processing, i.e. in comparing to the data, # not the running of the model itself. So..... we good! theList = [ eta, epsff, fg0, muNorm, muMhScaling, fixedQ, accScaleLength, fcool, Mh0, fscatter, x0, x1, x2, x3, obsScale, conRF, muHgScaling] try: assert eta>0 and epsff>0 and fg0>0 and fg0<=1 and fixedQ>0 and muNorm>=0 and fcool>=0 and fcool<=1 and Mh0>0 except: print 'Unexpected ppd params: ',theList emcee_params.append( copy.deepcopy(theList) ) # OK, from here on out, we just need to emulate parts of the run() function to trick emcee into running a single iteration of the algorithm with this IC. ndim = 17 restart = {} restart['currentPosition'] = emcee_params restart['chain'] = None restart['state'] = None restart['prob'] = None restart['iterationCounter'] = 0 restart['mcmcRunCounter'] = 0 nwalkers = len(emcee_params) # Need one walker per sample from posterior predictive distribution print "Starting up the ensemble sampler!" sampler = emcee.EnsembleSampler(nwalkers, ndim, fakeProb, pool=pool) #pos, prob, state = sampler.run_mcmc(restart['currentPosition'], N, rstate0=restart['state'], lnprob0=restart['prob']) print "Take a step with the ensemble sampler" # Take a single step with the ensemble sampler. print np.shape(restart['currentPosition']), np.shape(np.random.uniform(0,1,nwalkers)) sampler._get_lnprob(pos = restart['currentPosition']) #result = sampler.sample(restart['currentPosition'], iterations=1, lnprob0=None, rstate0=None) #pos, prob, state = result print "Close the pool" pool.close()
def fit_BAO(): parser = argparse.ArgumentParser(description='Use mcmc routine to get the BAO peak stretching parameter alpha and damping parameter, made by Zhejie.',\ formatter_class=argparse.ArgumentDefaultsHelpFormatter) parser.add_argument( "--lt", help= '*The type of weak lensing survey. (TF: Tully-Fisher; TW: traditional (photo-z) weak lensing.)', required=True) parser.add_argument("--nrbin", help='*Number of tomographic bins.', type=int, required=True) parser.add_argument("--nkbin", help='*Number of output k bins.', type=int, required=True) parser.add_argument("--shapenf", help='*Shape noise factor.', required=True) parser.add_argument("--kmin", help='*kmin fit boundary.', required=True) parser.add_argument("--kmax", help='*kmax fit boundary.', required=True) parser.add_argument("--params_str", help='Set fitting parameters. 1: free; 0: fixed.', required=True) parser.add_argument("--Sigma2_inf", help='Whether setting Sigma2_xy as infinity or not.', default='False') parser.add_argument("--alpha", help='Fix the parameter alpha value.', default=1.0, type=np.float) parser.add_argument( "--Pwig_type", help= '*The spatial P(k)_wig whether is linear or not. Type Pwig_linear or Pwig_nonlinear; in nonlinear, BAO is damped.', required=True) parser.add_argument( "--Psm_type", help= 'The expression of Pnorm. The default case, Pnorm from Eisenstein & Zaldarriaga 1999. \ Test Pnorm=Pnow, which is derived from transfer function.' ) #parser.add_argument("--Sigma", help = 'BAO damping parameter Sigma value. (Either 0.0 or 100.0.)', type=float, required=True) parser.add_argument( "--survey_stage", help= 'Optional parameter. KW_stage_IV (kinematic weak lensing) or PW_stage_IV (photo-z). It could also be considered as the directory of data\ files to be fitted.') #parser.add_argument("--mpi_used", help = 'Whether MPI is implemented in the calculation of Cijl_prime, Gm_prime. Either True or False.') parser.add_argument( "--f_sky", help= 'This addtional argument is for data files in TF_cross-ps. Distinguish cases with different f_sky value.' ) parser.add_argument( "--set_SVc_on_CovP", help= '*Whether we replace smaller SV to be SVc in W matrix for the output inverse covariance matrix of Pk. Either True or False.', required=True) parser.add_argument( "--start_lmin", help= "*The minimum ell value considered in the analysis. Default is 1. For Stage III, it's 10 for Stage III and 4 for Stage IV", default=1, type=int) parser.add_argument( "--nSV_min", help= "Fitting the extracted power spectrum with nSV used. nSV_min is the strating point.", type=int, default=1) parser.add_argument( "--save_sampler", help= "Whether we save samplers of MCMC into a data file or not. True or False.", default='False') parser.add_argument( "--modify_Cov_cij_cpq", help= "Whether we have modified the Cov_cij_cpq while outputing Cijl^prime and Gm^prime.", default='False') args = parser.parse_args() #print("args: ", args.lt) lt = args.lt num_rbin = args.nrbin num_kbin = args.nkbin shapenf = args.shapenf kmin = float(args.kmin) kmax = float(args.kmax) params_str = args.params_str params_indices = [int(i) for i in params_str] Pwig_type = args.Pwig_type Psm_type = args.Psm_type survey_stage = args.survey_stage #mpi_used = args.mpi_used f_sky = args.f_sky set_SVc_on_CovP = args.set_SVc_on_CovP start_lmin = args.start_lmin comm = MPI.COMM_WORLD rank = comm.Get_rank() size = comm.Get_size() lt_prefix = {'TF': 'TF', 'TW': 'TW_zext'} old_stdout = sys.stdout if survey_stage: odir = './{}/fit_kmin{}_kmax{}_Pwig_over_Pnow/{}/'.format( survey_stage, kmin, kmax, Pwig_type) else: odir = './fit_kmin{}_kmax{}_Pwig_over_Pnow_fsky{}/{}/'.format( kmin, kmax, f_sky, Pwig_type) # if mpi_used == 'True': # odir = './{}/fit_kmin{}_kmax{}_Pwig_over_Pnow/mpi_{}/'.format(survey_stage, kmin, kmax, Pwig_type) if Psm_type == 'Pnow': odir = odir + 'set_Pnorm_Pnow/' if start_lmin != 1: odir = odir + 'start_ell_{}/'.format(start_lmin) if rank == 0: if not os.path.exists(odir): os.makedirs(odir) comm.Barrier() Sigma2_xy_dict = {'stage_III': 31.176, 'stage_IV': 22.578} for ss_name in Sigma2_xy_dict: if ss_name in survey_stage: stage_name = ss_name if Pwig_type == 'Pwig_nonlinear': Sigma2_xy = Sigma2_xy_dict[stage_name] elif Pwig_type == 'Pwig_linear': Sigma2_xy = 0.0 if args.Sigma2_inf == 'True': #Sigma2_xy = np.inf Sigma2_xy = 100 * 100 # test whether it influences results or not ofile = odir + "mcmc_fit_{}_{}rbin_{}kbin_snf{}_params{}_Sigma2_{}.log".format( lt, num_rbin, num_kbin, shapenf, params_str, Sigma2_xy) if params_str == '001': ofile = odir + "mcmc_fit_{0}_{1}rbin_{2}kbin_snf{3}_params{4}_alpha{5:.2f}_Sigma2_{6}.log".format( lt, num_rbin, num_kbin, shapenf, params_str, args.alpha, Sigma2_xy) log_file = open(ofile, "w") sys.stdout = log_file print('Arguments of fitting: ', args) ifile = '../Input_files/CAMB_Planck2015_matterpower.dat' kcamb, Pkcamb = np.loadtxt(ifile, dtype='f8', comments='#', unpack=True) Pwig_spl = InterpolatedUnivariateSpline(kcamb, Pkcamb) k_0 = 0.001 # unit h*Mpc^-1 Pk_0 = Pwig_spl(k_0) ifile = '../Input_files/transfer_fun_Planck2015.dat' kk, Tf = np.loadtxt(ifile, dtype='f8', comments='#', usecols=(0, 1), unpack=True) #print kk==kcamb Tf_spl = InterpolatedUnivariateSpline(kk, Tf) Tf_0 = Tf_spl(k_0) P0_a = Pk_0 / (pow(k_0, cosmic_params.ns) * Tf_0**2.0) Psm = P0_a * pow( kcamb, cosmic_params.ns ) * Tf**2.0 # Get primordial (smooth) power spectrum from the transfer function Psm_spl = InterpolatedUnivariateSpline(kcamb, Psm) norm_gf = 1.0 N_walkers = 40 all_param_names = 'alpha', 'Sigma2_xy', 'A' all_temperature = 0.01, 1.0, 0.1 alpha, A = args.alpha, 1.0 # initial guess for fitting, the value of Sigma2_xy at z=0.65 is from theory prediction (see code TF_cross_convergence_ps_bin.py) all_params = alpha, Sigma2_xy, A N_params, theta, fix_params, params_T, params_name = set_params( all_params, params_indices, all_param_names, all_temperature) # Fit for DM power spectrum if lt == 'TF': idir0 = '../{}_cross-ps/'.format(lt) if survey_stage: idir0 = '../{}/'.format(survey_stage) idir1 = '{}_Pk_output_dset_{}/' # if mpi_used == 'True': # idir1 = 'mpi_{}_Pk_output_dset_{}/' if f_sky: idir1 = '{}_Pk_output_dset_{}_fsky{}/' elif lt == 'TW': idir0 = '../{}_f2py_SVD/'.format(lt) if survey_stage: idir0 = '../{}/'.format(survey_stage) idir1 = '{}_Pk_output_dset_{}/' idir2 = '{}rbins_{}kbins_snf{}/'.format(num_rbin, num_kbin, shapenf) odir1 = 'mcmc_fit_params_Pwig_over_Pnow/{}/kmin{}_kmax{}/{}rbins_{}kbins_snf{}/'.format( Pwig_type, kmin, kmax, num_rbin, num_kbin, shapenf) if f_sky: odir1 = 'mcmc_fit_params_Pwig_over_Pnow/{}/kmin{}_kmax{}/{}rbins_{}kbins_snf{}_fsky{}/'.format( Pwig_type, kmin, kmax, num_rbin, num_kbin, shapenf, f_sky) # if mpi_used == 'True': # odir1 = 'mcmc_fit_params_Pwig_over_Pnow/{}/kmin{}_kmax{}/mpi_{}rbins_{}kbins_snf{}/'.format(Pwig_type, kmin, kmax, num_rbin, num_kbin, shapenf) if Psm_type == 'Pnow': idir2 = idir2 + 'set_Pnorm_Pnow/' odir1 = odir1 + 'set_Pnorm_Pnow/' if start_lmin != 1: idir2 = idir2 + 'start_ell_{}/'.format(start_lmin) odir1 = odir1 + 'start_ell_{}/'.format(start_lmin) if args.modify_Cov_cij_cpq == 'True': idir2 = idir2 + 'modify_Cov_cij_cpq/' odir1 = odir1 + 'modify_Cov_cij_cpq/' if params_str == '001': odir1 = odir1 + 'alpha_{0:.2f}/'.format(args.alpha) if f_sky: idir = idir0 + idir1.format( lt_prefix[lt], Pwig_type, f_sky ) + idir2 # not exactly matching format, but it's ok if there is no f_sky parameter. else: idir = idir0 + idir1.format(lt_prefix[lt], Pwig_type) + idir2 odir = idir0 + odir1 if rank == 0: if not os.path.exists(odir): os.makedirs(odir) if args.save_sampler == 'True': os.makedirs( odir + 'chain_samplers/') # store positions of samplers in chains comm.Barrier() ifile = idir + 'Pk_wnw_{}rbin_{}kbin_withshapenoisefactor{}_{}eigenvW.out'.format( num_rbin, num_kbin, shapenf, num_kbin) k_all = np.loadtxt(ifile, dtype='f8', comments='#', usecols=(0, )) indices = filter_krange( k_all, kmin, kmax) # Here we don't need to use sigma_Pk_wnw anymore. N_fitbin = len(indices) print('# of fit k bins: ', N_fitbin) print('The indices of fitting k bins: ', indices) if set_SVc_on_CovP != 'True': ifile = idir + 'Cov_Pwnw_inv_{}rbin_{}kbin_withshapenoisefactor{}.npz'.format( num_rbin, num_kbin, shapenf) npzfile = np.load(ifile) icov_Pk_wnw = npzfile['arr_0'] print('icov_Pk_wnw: ', icov_Pk_wnw) cov_Pwnw = linalg.inv(icov_Pk_wnw) print('cov_Pwnw: ', cov_Pwnw) identity = np.dot(cov_Pwnw, icov_Pk_wnw) print('cov * icov: ', identity) print('Inverse process is good?', np.allclose(identity, np.eye(num_kbin)) ) # test whether the inverse process is very successful part_cov_Pwnw = cov_Pwnw[np.ix_(indices, indices)] part_icov_Pk_wnw = linalg.inv(part_cov_Pwnw) print( 'Inverse process for marginalized matrix is good?', np.allclose(np.dot(part_cov_Pwnw, part_icov_Pk_wnw), np.eye(N_fitbin))) print('The inverse covariance matrix for fitting: ', part_icov_Pk_wnw) pool = MPIPool(loadbalance=True) for num_eigv in range(args.nSV_min, num_kbin + 1): np.random.seed(1) if f_sky: idir_Pwig = idir0 + idir1.format(lt_prefix[lt], Pwig_type, f_sky) + idir2 idir_Pnow = idir0 + idir1.format(lt_prefix[lt], 'Pnow', f_sky) + idir2 else: idir_Pwig = idir0 + idir1.format(lt_prefix[lt], Pwig_type) + idir2 idir_Pnow = idir0 + idir1.format(lt_prefix[lt], 'Pnow') + idir2 ifile = idir_Pwig + 'Pk_wig_{}rbin_{}kbin_withshapenoisefactor{}_{}eigenvW.out'.format( num_rbin, num_kbin, shapenf, num_eigv) print(ifile) data_m = np.loadtxt(ifile, dtype='f8', comments='#') # k, P(k), sigma_Pk k_obs, Pk_wig_obs = data_m[indices, 0], data_m[indices, 1] ifile = idir_Pnow + 'Pk_now_{}rbin_{}kbin_withshapenoisefactor{}_{}eigenvW.out'.format( num_rbin, num_kbin, shapenf, num_eigv) print(ifile) data_m = np.loadtxt(ifile, dtype='f8', comments='#') k_obs, Pk_now_obs = data_m[indices, 0], data_m[indices, 1] Pk_wnw_obs = Pk_wig_obs / Pk_now_obs if set_SVc_on_CovP == 'True': ifile = idir_Pwig + 'Cov_Pwnw_inv_{}rbin_{}kbin_withshapenoisefactor{}_{}eigenvW_SVc.npz'.format( num_rbin, num_kbin, shapenf, num_eigv) ####ifile = idir_Pwig + 'Cov_Pwnw_inv_{}rbin_{}kbin_withshapenoisefactor1.0_{}eigenvW_SVc.npz'.format(num_rbin, num_kbin, num_eigv) # only for some special case, it's temporary part_icov_Pk_wnw = filter_invCov_on_k(ifile, indices) mergedsamples, params_mcmc = mcmc_routine(N_params, N_walkers, theta, params_T, params_indices, fix_params, k_obs, Pk_wnw_obs, part_icov_Pk_wnw, Pwig_spl, Psm_spl, norm_gf, params_name, pool) print(params_mcmc) chi_square = chi2(params_mcmc[:, 0], params_indices, fix_params, k_obs, Pk_wnw_obs, part_icov_Pk_wnw, Pwig_spl, Psm_spl, norm_gf) dof = N_fitbin - N_params reduced_chi2 = chi_square / dof print("chi^2/dof: ", reduced_chi2, "\n") if args.Sigma2_inf == 'True': # in order to distinguish the Sigma2_xy value, we need to specify it. filename = 'Pk_wnw_{}rbin_{}kbin_snf{}_{}eigenvW_params{}_Sigma2_{}.dat'.format( num_rbin, num_kbin, shapenf, num_eigv, params_str, Sigma2_xy) else: filename = 'Pk_wnw_{}rbin_{}kbin_snf{}_{}eigenvW_params{}.dat'.format( num_rbin, num_kbin, shapenf, num_eigv, params_str) # Sigma2_xy is fixed in the default case ofile_params = odir + filename print('ofile_params:', ofile_params) write_params(ofile_params, params_mcmc, params_name, reduced_chi2, fix_params, dof) if (args.save_sampler == 'True') and (rank == 0): odir_merged_sample = odir + 'chain_samplers/' if not os.path.exists(odir_merged_sample): os.makedirs(odir_merged_sample) ofile = odir_merged_sample + 'mergedsamples_' + filename[ 0:len(filename) - 4] + '.npz' np.savez_compressed(ofile, np.array(mergedsamples)) pool.close() sys.stdout = old_stdout log_file.close()
def LensModelMCMC(data,lens,source, xmax=30.,highresbox=[-3.,3.,-3.,3.],emitres=None,fieldres=None, sourcedatamap=None, scaleamp=False, shiftphase=False, modelcal=True,cosmo=Planck15, nwalkers=1e3,nburn=1e3,nstep=1e3,pool=None,nthreads=1,mpirun=False): """ Wrapper function which basically takes what the user wants and turns it into the format needed for the acutal MCMC lens modeling. Inputs: data: One or more visdata objects; if multiple datasets are being fit to, should be a list of visdata objects. lens: Any of the currently implemented lens objects or ExternalShear. source: One or more of the currently implemented source objects; if more than one source to be fit, should be a list of multiple sources. xmax: (Half-)Grid size, in arcseconds; the grid will span +/-xmax in x&y highresbox: The region to model at higher resolution (to account for high-magnification and differential lensing effects), as [xmin, xmax, ymin, ymax]. Note the sign convention is: +x = West, +y = North, like the lens positions. sourcedatamap: A list of length the number of datasets which tells which source(s) are to be fit to which dataset(s). Eg, if two sources are to be fit to two datasets jointly, should be [[0,1],[0,1]]. If we have four sources and three datasets, could be [[0,1],[0,1],[2,3]] to say that the first two sources should both be fit to the first two datasets, while the second two should be fit to the third dataset. If None, will assume all sources should be fit to all datasets. scaleamp: A list of length the number of datasets which tells whether a flux rescaling is allowed and which dataset the scaling should be relative to. False indicates no scaling should be done, while True indicates that amplitude scaling should be allowed. shiftphase: Similar to scaleamp above, but allowing for positional/astrometric offsets. modelcal: Whether or not to perform the pseudo-selfcal procedure of H+13 cosmo: The cosmology to use, as an astropy object, e.g., from astropy.cosmology import WMAP9; cosmo=WMAP9 Default is Planck15. nwalkers: Number of walkers to use in the mcmc process; see dan.iel.fm/emcee/current for more details. nburn: Number of burn-in steps to take with the chain. nstep: Number of actual steps to take in the mcmc chains after the burn-in nthreads: Number of threads (read: cores) to use during the fitting, default 1. mpirun: Whether to parallelize using MPI instead of multiprocessing. If True, nthreads has no effect, and your script should be run with, eg, mpirun -np 16 python lensmodel.py. Returns: mcmcresult: A nested dict containing the chains requested. Will have all the MCMC chain results, plus metadata about the run (initial params, data used, etc.). Formatting still a work in progress (esp. for modelcal phases). chains: The raw chain data, for testing. blobs: Everything else returned by the likelihood function; will have magnifications and any modelcal phase offsets at each step; eventually will remove this once get everything packaged up for mcmcresult nicely. colnames: Basically all the keys to the mcmcresult dict; eventually won't need to return this once mcmcresult is packaged up nicely. """ if pool: nthreads = 1 elif mpirun: nthreads = 1 from emcee.utils import MPIPool pool = MPIPool(debug=False,loadbalance=True) if not pool.is_master(): pool.wait() sys.exit(0) else: pool = None # Making these lists just makes later stuff easier since we now know the dtype lens = list(np.array([lens]).flatten()) source = list(np.array([source]).flatten()) # Ensure source(s) are a list data = list(np.array([data]).flatten()) # Same for dataset(s) scaleamp = list(np.array([scaleamp]).flatten()) shiftphase = list(np.array([shiftphase]).flatten()) modelcal = list(np.array([modelcal]).flatten()) if len(scaleamp)==1 and len(scaleamp)<len(data): scaleamp *= len(data) if len(shiftphase)==1 and len(shiftphase)<len(data): shiftphase *= len(data) if len(modelcal)==1 and len(modelcal)<len(data): modelcal *= len(data) if sourcedatamap is None: sourcedatamap = [None]*len(data) # emcee isn't very flexible in terms of how it gets initialized; start by # assembling the user-provided info into a form it likes ndim, p0, colnames = 0, [], [] # Lens(es) first for i,ilens in enumerate(lens): if ilens.__class__.__name__=='SIELens': for key in ['x','y','M','e','PA']: if not vars(ilens)[key]['fixed']: ndim += 1 p0.append(vars(ilens)[key]['value']) colnames.append(key+'L'+str(i)) elif ilens.__class__.__name__=='ExternalShear': for key in ['shear','shearangle']: if not vars(ilens)[key]['fixed']: ndim += 1 p0.append(vars(ilens)[key]['value']) colnames.append(key) # Then source(s) for i,src in enumerate(source): if src.__class__.__name__=='GaussSource': for key in ['xoff','yoff','flux','width']: if not vars(src)[key]['fixed']: ndim += 1 p0.append(vars(src)[key]['value']) colnames.append(key+'S'+str(i)) elif src.__class__.__name__=='SersicSource': for key in ['xoff','yoff','flux','majax','index','axisratio','PA']: if not vars(src)[key]['fixed']: ndim += 1 p0.append(vars(src)[key]['value']) colnames.append(key+'S'+str(i)) elif src.__class__.__name__=='PointSource': for key in ['xoff','yoff','flux']: if not vars(src)[key]['fixed']: ndim += 1 p0.append(vars(src)[key]['value']) colnames.append(key+'S'+str(i)) # Then flux rescaling; only matters if >1 dataset for i,t in enumerate(scaleamp[1:]): if t: ndim += 1 p0.append(1.) # Assume 1.0 scale factor to start colnames.append('ampscale_dset'+str(i+1)) # Then phase/astrometric shift; each has two vals for a shift in x&y for i,t in enumerate(shiftphase[1:]): if t: ndim += 2 p0.append(0.); p0.append(0.) # Assume zero initial offset colnames.append('astromshift_x_dset'+str(i+1)) colnames.append('astromshift_y_dset'+str(i+1)) # Get any model-cal parameters set up. The process involves some expensive # matrix inversions, but these only need to be done once, so we'll do them # now and pass the results as arguments to the likelihood function. See docs # in calc_likelihood.model_cal for more info. for i,dset in enumerate(data): if modelcal[i]: uniqant = np.unique(np.asarray([dset.ant1,dset.ant2]).flatten()) dPhi_dphi = np.zeros((uniqant.size-1,dset.u.size)) for j in range(1,uniqant.size): dPhi_dphi[j-1,:]=(dset.ant1==uniqant[j])-1*(dset.ant2==uniqant[j]) C = scipy.sparse.diags((dset.sigma/dset.amp)**-2.,0) F = np.dot(dPhi_dphi,C*dPhi_dphi.T) Finv = np.linalg.inv(F) FdPC = np.dot(-Finv,dPhi_dphi*C) modelcal[i] = [dPhi_dphi,FdPC] # Create our lensing grid coordinates now, since those shouldn't be # recalculated with every call to the likelihood function xmap,ymap,xemit,yemit,indices = GenerateLensingGrid(data,xmax,highresbox, fieldres,emitres) # Calculate the uv coordinates we'll interpolate onto; only need to calculate # this once, so do it here. kmax = 0.5/((xmap[0,1]-xmap[0,0])*arcsec2rad) ug = np.linspace(-kmax,kmax,xmap.shape[0]) # Calculate some distances; we only need to calculate these once. # This assumes multiple sources are all at same z; should be this # way anyway or else we'd have to deal with multiple lensing planes if cosmo is None: cosmo = Planck15 Dd = cosmo.angular_diameter_distance(lens[0].z).value Ds = cosmo.angular_diameter_distance(source[0].z).value Dds= cosmo.angular_diameter_distance_z1z2(lens[0].z,source[0].z).value p0 = np.array(p0) # Create a ball of starting points for the walkers, gaussian ball of # 10% width; if initial value is 0 (eg, astrometric shift), give a small sigma # for angles, generally need more spread than 10% to sample well, do 30% for those cases [~0.5% >180deg for p0=100deg] isangle = np.array([0.30 if 'PA' in s or 'angle' in s else 0.1 for s in colnames]) initials = emcee.utils.sample_ball(p0,np.asarray([isangle[i]*x if x else 0.05 for i,x in enumerate(p0)]),int(nwalkers)) # All the lens objects know if their parameters have been altered since the last time # we calculated the deflections. If all the lens pars are fixed, we only need to do the # deflections once. This step ensures that the lens object we create the sampler with # has these initial deflections. for i,ilens in enumerate(lens): if ilens.__class__.__name__ == 'SIELens': ilens.deflect(xemit,yemit,Dd,Ds,Dds) elif ilens.__class__.__name__ == 'ExternalShear': ilens.deflect(xemit,yemit,lens[0]) # Create the sampler object; uses calc_likelihood function defined elsewhere lenssampler = emcee.EnsembleSampler(nwalkers,ndim,calc_vis_lnlike, args = [data,lens,source,Dd,Ds,Dds,ug, xmap,ymap,xemit,yemit,indices, sourcedatamap,scaleamp,shiftphase,modelcal], threads=nthreads,pool=pool) # Run burn-in phase print("Running burn-in... ") #pos,prob,rstate,mus = lenssampler.run_mcmc(initials,nburn,storechain=False) for i,result in enumerate(lenssampler.sample(initials,iterations=nburn,storechain=False)): if i%20==0: print('Burn-in step ',i,'/',nburn) pos,prob,rstate,blob = result lenssampler.reset() # Run actual chains print("Done. Running chains... ") for i,result in enumerate(lenssampler.sample(pos,rstate0=rstate,iterations=nstep,storechain=True)): if i%20==0: print('Chain step ',i,'/',nstep) #lenssampler.run_mcmc(pos,nstep,rstate0=rstate) if mpirun: pool.close() print("Mean acceptance fraction: ",np.mean(lenssampler.acceptance_fraction)) #return lenssampler.flatchain,lenssampler.blobs,colnames # Package up the magnifications and modelcal phases; disregards nan points (where # we failed the prior, usu. because a periodic angle wrapped). blobs = lenssampler.blobs mus = np.asarray([[a[0] for a in l] for l in blobs]).flatten(order='F') bad = np.where(np.asarray([np.any(np.isnan(m)) for m in mus],dtype=bool))[0] for k in bad: mus[k] = np.array([np.nan]*len(source)) mus = np.asarray(list(mus),dtype=float).reshape((-1,len(source)),order='F') # stupid-ass hack bad = np.isnan(mus)[:,0] #bad = bad.reshape((-1,len(source)),order='F')[:,0] #mus = np.atleast_2d(np.asarray([mus[i] if not bad[i] else [np.nan]*len(source) for i in range(mus.size)])).T colnames.extend(['mu{0:.0f}'.format(i) for i in range(len(source))]) # Assemble the output. Want to return something that contains both the MCMC chains # themselves, but also metadata about the run. mcmcresult = {} # keep track of git revision, for reproducibility's sake # if run under mpi, this will spew some scaremongering warning text, # but it's fine. use --mca mpi_warn_on_fork 0 in the mpirun statement to disable try: import subprocess gitd = os.path.abspath(os.path.join(os.path.dirname(__file__),os.pardir)) mcmcresult['githash'] = subprocess.check_output('git --git-dir={0:s} --work-tree={1:s} '\ 'rev-parse HEAD'.format(gitd+'/.git',gitd),shell=True).rstrip() except: mcmcresult['githash'] = 'No repo found' mcmcresult['datasets'] = [dset.filename for dset in data] # Data files used mcmcresult['lens_p0'] = lens # Initial params for lens,src(s),shear; also tells if fixed, priors, etc. mcmcresult['source_p0'] = source if sourcedatamap: mcmcresult['sourcedatamap'] = sourcedatamap mcmcresult['xmax'] = xmax mcmcresult['highresbox'] = highresbox mcmcresult['fieldres'] = fieldres mcmcresult['emitres'] = emitres if any(scaleamp): mcmcresult['scaleamp'] = scaleamp if any(shiftphase): mcmcresult['shiftphase'] = shiftphase mcmcresult['chains'] = np.core.records.fromarrays(np.hstack((lenssampler.flatchain[~bad],mus[~bad])).T,names=colnames) mcmcresult['lnlike'] = lenssampler.flatlnprobability[~bad] # Keep track of best-fit params, derived from chains. c = copy.deepcopy(mcmcresult['chains']) mcmcresult['best-fit'] = {} pbest = [] # Calculate the best fit values as medians of each param lens,source = copy.deepcopy(mcmcresult['lens_p0']), copy.deepcopy(mcmcresult['source_p0']) for i,ilens in enumerate(lens): if ilens.__class__.__name__ == 'SIELens': ilens.__dict__['_altered'] = True for key in ['x','y','M','e','PA']: if not vars(ilens)[key]['fixed']: ilens.__dict__[key]['value'] = np.median(c[key+'L'+str(i)]) pbest.append(np.median(c[key+'L'+str(i)])) elif ilens.__class__.__name__ == 'ExternalShear': for key in ['shear','shearangle']: if not vars(ilens)[key]['fixed']: ilens.__dict__[key]['value'] = np.median(c[key]) pbest.append(np.median(c[key])) mcmcresult['best-fit']['lens'] = lens # now do the source(s) for i,src in enumerate(source): # Source is a list of source objects if src.__class__.__name__ == 'GaussSource': for key in ['xoff','yoff','flux','width']: if not vars(src)[key]['fixed']: src.__dict__[key]['value'] = np.median(c[key+'S'+str(i)]) pbest.append(np.median(c[key+'S'+str(i)])) elif src.__class__.__name__ == 'SersicSource': for key in ['xoff','yoff','flux','majax','index','axisratio','PA']: if not vars(src)[key]['fixed']: src.__dict__[key]['value'] = np.median(c[key+'S'+str(i)]) pbest.append(np.median(c[key+'S'+str(i)])) elif src.__class__.__name__ == 'PointSource': for key in ['xoff','yoff','flux']: if not vars(src)[key]['fixed']: src.__dict__[key]['value'] = np.median(c[key+'S'+str(i)]) pbest.append(np.median(c[key+'S'+str(i)])) mcmcresult['best-fit']['source'] = source mcmcresult['best-fit']['magnification'] = np.median(mus[~bad],axis=0) # Any amplitude scaling or astrometric shifts bfscaleamp = np.ones(len(data)) if 'scaleamp' in mcmcresult.keys(): for i,t in enumerate(mcmcresult['scaleamp']): # only matters if >1 datasets if i==0: pass elif t: bfscaleamp[i] = np.median(c['ampscale_dset'+str(i)]) pbest.append(np.median(c['ampscale_dset'+str(i)])) else: pass mcmcresult['best-fit']['scaleamp'] = bfscaleamp bfshiftphase = np.zeros((len(data),2)) if 'shiftphase' in mcmcresult.keys(): for i,t in enumerate(mcmcresult['shiftphase']): if i==0: pass # only matters if >1 datasets elif t: bfshiftphase[i][0] = np.median(c['astromshift_x_dset'+str(i)]) bfshiftphase[i][1] = np.median(c['astromshift_y_dset'+str(i)]) pbest.append(np.median(c['astromshift_x_dset'+str(i)])) pbest.append(np.median(c['astromshift_y_dset'+str(i)])) else: pass # no shifting mcmcresult['best-fit']['shiftphase'] = bfshiftphase mcmcresult['best-fit']['lnlike'] = calc_vis_lnlike(pbest,data,mcmcresult['best-fit']['lens'], mcmcresult['best-fit']['source'], Dd,Ds,Dds,ug,xmap,ymap,xemit,yemit,indices, sourcedatamap,scaleamp,shiftphase,modelcal)[0] # Calculate the deviance information criterion, using the Spiegelhalter+02 definition (cf Gelman+04) mcmcresult['best-fit']['DIC'] = -4*np.mean(mcmcresult['lnlike']) + 2*mcmcresult['best-fit']['lnlike'] # If we did any modelcal stuff, keep the antenna phase offsets here if any(modelcal): mcmcresult['modelcal'] = [True if j else False for j in modelcal] dp = np.squeeze(np.asarray([[a[1] for a in l if ~np.any(np.isnan(a[0]))] for l in blobs])) a = [x for l in dp for x in l] # Have to dick around with this if we had any nan's dphases = np.squeeze(np.reshape(a,(nwalkers*nstep-bad.sum(),len(data),-1),order='F')) if len(data) > 1: for i in range(len(data)): if modelcal[i]: mcmcresult['calphases_dset'+str(i)] = np.vstack(dphases[:,i]) else: if any(modelcal): mcmcresult['calphases_dset0'] = dphases return mcmcresult
def run(N, p00=None, nwalkers=500): fn = chainDirRel+'.pickle' ndim = 17 if p00 is not None: p0 = [p00*(1.0+0.001*np.random.randn( ndim )) for i in range(nwalkers)] else: p0 = [sampleFromPrior() for i in range(nwalkers)] restart = {} restart['currentPosition'] = p0 restart['chain'] = None restart['state'] = None restart['prob'] = None restart['iterationCounter'] = 0 restart['mcmcRunCounter'] = 0 # Read in our past progress UNLESS we've been given a new starting location. if p00 is None: updateRestart(fn,restart) if restart['chain'] is not None: # This may save some time if you change something and forget to delete the .pickle file. restartedShape = np.shape(restart['chain']) print restartedShape, nwalkers, ndim assert restartedShape[0] == nwalkers assert restartedShape[2] == ndim global runNumber runNumber = restart['mcmcRunCounter'] restart['iterationCounter'] += N restart['mcmcRunCounter'] += 1 pool = MPIPool(comm=comm, loadbalance=True) if not pool.is_master(): pool.wait() sys.exit(0) sampler = emcee.EnsembleSampler(nwalkers, ndim, lnProb, pool=pool) #pos, prob, state = sampler.run_mcmc(restart['currentPosition'], N, rstate0=restart['state'], lnprob0=restart['prob']) for result in sampler.sample(restart['currentPosition'], iterations=N, lnprob0=restart['prob'], rstate0=restart['state']): pos, prob, state = result restart['acor'] = sampler.acor[:] # autocorr length for each param (ndim) restart['accept'] = sampler.acceptance_fraction[:] # acceptance frac for each walker. restart['currentPosition'] = pos # same shape as p0: nwalkers x ndim restart['state'] = state # random number generator state restart['prob'] = prob # nwalkers x __ if restart['chain'] is None: restart['chain'] = np.expand_dims(sampler.chain[:,0,:],1) # nwalkers x niterations x ndim restart['allProbs'] = np.expand_dims(prob,1) # nwalkers x niterations else: print np.shape(restart['chain']), np.shape(sampler.chain[:,-1,:]), np.shape(sampler.chain) print restart['mcmcRunCounter'], restart['iterationCounter'] #restart['chain'] = np.concatenate((restart['chain'], sampler.chain[:,-1,:]), axis=1) print "dbg1: ",np.shape(restart['chain']), np.shape(np.zeros((nwalkers, 1, ndim))), np.shape(np.expand_dims(pos,1)) restart['chain'] = np.concatenate((restart['chain'], np.expand_dims(pos, 1)),axis=1) restart['allProbs'] = np.concatenate((restart['allProbs'], np.expand_dims(prob, 1)),axis=1) saveRestart(fn,restart) pool.close()
def mcmc(result_path, result_name, param_dict, nwalkers=10, steps=2, output='./mcmc_chains.h5', overwrite=False, covmat=None, redshift=0.05, logmass_min=11., logmass_max=16.): pool = MPIPool(loadbalance=True, debug=False) rank = pool.rank comm = pool.comm # initializign CAMB power spectrum _camb.set_matter_power(redshifts=[redshift, ], kmax=3.e1) if rank == 0: result = h5py.File(result_path + result_name, 'r') pksz = result['pkSZ'][:] jk_sample = False if 'pkSZ_random' in result.keys(): pksz_random = result['pkSZ_random'][:] jk_sample = False elif 'pkSZ_jk' in result.keys(): pksz_random = result['pkSZ_jk'][:] jk_sample = True else: print "Need random samples" exit() pksz_bins = result['pkSZ_bins'][:] d_bins = pksz_bins[1:] - pksz_bins[:-1] pksz_bins = pksz_bins[:-1] + 0.5 * d_bins result.close() lin_scale = pksz_bins > 25. pksz_obs = pksz[lin_scale] pksz_err = None pksz_cov = np.cov(pksz_random, rowvar=False, bias=jk_sample) if jk_sample: spl_n = float(pksz_random.shape[0] - 1) bin_n = pksz_cov.shape[0] pksz_cov *= spl_n pksz_covm = np.linalg.inv(pksz_cov[:, lin_scale][lin_scale, :]) pksz_covm *= (spl_n - bin_n) / spl_n else: pksz_covm = np.linalg.inv(pksz_cov[:, lin_scale][lin_scale, :]) pksz_bins = pksz_bins[lin_scale] else: pksz_obs = None pksz_err = None pksz_bins = None pksz_covm = None pksz_obs = comm.bcast(pksz_obs, root=0) #pksz_err = comm.bcast(pksz_err, root=0) pksz_bins = comm.bcast(pksz_bins, root=0) pksz_covm = comm.bcast(pksz_covm, root=0) comm.barrier() if rank != 0: pool.wait() sys.exit(0) param, theta, theta_min, theta_max, param_tex, cosm_param_idx, camb_run\ = read_param_dict(param_dict, camb_param=_camb.params.keys()) paramnames = open(output.replace('.h5', '.paramnames'), 'w') for i in range(param.shape[0]): paramnames.write('%s\t\t'%param[i] + param_tex[i] + '\n') paramnames.close() # param = [cent, min, max] #tau_bar = [1., -1., 3.] #mnu = [0.4, 0.2, 0.6] #ombh2 = [0.0221, 0.005, 0.1] #omch2 = [0.12, 0.001, 0.99] #w = [-1, -3., 1.] #wa = [0., -3., 3.] #theta = param #theta_min = np.array([tau_bar[1], w[1]]) #theta_max = np.array([tau_bar[2], w[2]]) ndim = theta.shape[0] #threads = nwalkers threads = 1 pos = [theta + 1.e-4 * np.random.randn(ndim) for i in range(nwalkers)] sampler = emcee.EnsembleSampler(nwalkers, ndim, lnprob, threads=threads, pool=pool, args=(pksz_bins, pksz_obs, pksz_err, pksz_covm, param), kwargs={ 'theta_min' :theta_min, 'theta_max' :theta_max, 'cosm_param_idx' :cosm_param_idx, 'camb_run' :camb_run, 'T_CMB' :2.7255, 'logmass_min' :logmass_min, 'logmass_max' :logmass_max}) # Run 100 steps as a burn-in. #pos, prob, state = sampler.run_mcmc(pos, 100) # Reset the chain to remove the burn-in samples. #sampler.reset() # Starting from the final position in the burn-in chain #sampler.run_mcmc(pos, steps, rstate0=state) #chain = np.zeros((size, ) + chain_local.shape ) #comm.Gather(chain_local, chain, root=0) step_group = 100 if step_group > steps: step_group = steps n_steps = steps / step_group state = None if rank == 0: if overwrite or not os.path.exists(output): mcmc_chains = h5py.File(output, 'w') #mcmc_chains['n_steps'] = n_steps mcmc_chains['n_steps'] = 0 mcmc_chains['pos'] = pos if covmat is not None: mcmc_chains['covmat'] = np.zeros((ndim, ndim)) #mcmc_chains['state'] = 0 n_steps0 = 0 mcmc_chains.close() else: mcmc_chains = h5py.File(output, 'a') n_steps0 = mcmc_chains['n_steps'][...] pos = mcmc_chains['pos'][...] if covmat is not None and n_steps0 != 0: covmat = mcmc_chains['covmat'] #state = mcmc_chains['state'][...] #if state == 0: state = None mcmc_chains.close() for i in range(n_steps0, n_steps0 + n_steps): if rank == 0: t1 = time.time() if covmat is not None: mh_proposal = MH_proposal(covmat, sampler._random) else: mh_proposal = None #mh_proposal = None pos, prob, state = sampler.run_mcmc(pos, step_group, state, mh_proposal=mh_proposal) if rank == 0: chain = sampler.chain chisq = sampler.lnprobability * (-2.) mcmc_chains = h5py.File(output, 'a') mcmc_chains['chains_%02d'%i] = chain mcmc_chains['chisqs_%02d'%i] = chisq mcmc_chains['n_steps'][...] = i + 1 mcmc_chains['pos'][...] = pos #mcmc_chains['state'][...] = state #pksz_covm = np.cov(pksz_random[:, lin_scale], rowvar=False) if covmat is not None: covmat = np.cov(chain.reshape(-1, chain.shape[-1]), rowvar=False) mcmc_chains['covmat'][...] = covmat mcmc_chains.close() print "[TIMING] %3d x %4d Steps: %8.4f [s]"\ %(nwalkers, step_group, time.time() - t1) print "Mean acceptance fraction: ", np.mean(sampler.acceptance_fraction) sampler.reset() pool.close()
def fit_BAO(args): kmin = float(args.kmin) kmax = float(args.kmax) params_str = args.params_str Pk_type = args.Pk_type params_indices = [int(i) for i in params_str] old_stdout = sys.stdout odir = './fit_kmin{}_kmax{}_{}/'.format(kmin, kmax, Pk_type) if not os.path.exists(odir): os.makedirs(odir) ofile = odir + "mcmc_fit_params{}.log".format(params_str) log_file = open(ofile, "w") sys.stdout = log_file print('Arguments for the fitting: ', args) ifile = '/Users/mehdi/work/quicksurvey/ELG/run8/planck_camb_56106182_matterpower_z0.dat' klin, Pk_linw = np.loadtxt(ifile, dtype='f8', comments='#', unpack=True) Pwig_spl = InterpolatedUnivariateSpline(klin, Pk_linw) ifile = '/Users/mehdi/work/quicksurvey/ELG/run8/planck_camb_56106182_matterpower_smooth_z0.dat' klin, Pk_sm = np.loadtxt(ifile, dtype='f8', comments='#', unpack=True) Psm_spl = InterpolatedUnivariateSpline(klin, Pk_sm) norm_gf = 1.0 N_walkers = 40 ##params_indices = [1, 0, 0] # 1: free parameter; 0: fixed parameter all_param_names = 'alpha', 'Sigma2_xy', 'A', 'B' all_temperature = 0.01, 1.0, 0.1, 0.1 Omega_m = 0.3075 # matter density G_0 = growth_factor(0.0, Omega_m) Sigma_0 = 7.7840 # This is exactly calculated from theoretical prediction with q_{BAO}=110 Mpc/h. z_list = [0.625] #, 0.875, 1.125, 1.375] # z list for the data files cut_list = ['F'] #, 'T'] # initial guess for fitting, Sigma2_xy=31.176 at z=0.65 is from theory prediction alpha, A, B = 1.0, 1.0, 10.0 idir = '../kp0kp2knmodes/surveyscaled-nmodes/' odir = './mcmc_fit_params_{}/kmin{}_kmax{}/'.format(Pk_type, kmin, kmax) comm = MPI.COMM_WORLD rank = comm.Get_rank() size = comm.Get_size() if rank == 0: if not os.path.exists(odir): os.makedirs(odir) pool = MPIPool(loadbalance=True) for z_value in z_list: norm_gf = growth_factor(z_value, Omega_m) / G_0 Sigma2_xy = (Sigma_0 * norm_gf)**2.0 print('z, Sigma2_xy: ', z_value, Sigma2_xy) all_params = alpha, Sigma2_xy, A, B N_params, theta, fix_params, params_T, params_name = set_params( all_params, params_indices, all_param_names, all_temperature) for cut_type in cut_list: ifile = idir + 'kp0kp2knmodes_z{}RADECcut{}.dat'.format( z_value, cut_type) print(ifile) data_m = np.loadtxt(ifile, dtype='f8', comments='#') # k, P0(k), P2(k), N_modes indices = np.argwhere((data_m[:, 0] >= kmin) & (data_m[:, 0] <= kmax)).flatten() N_fitbin = len(indices) k_obs, Pk_wig_obs, N_modes = data_m[indices, 0], data_m[indices, 1], data_m[indices, 3] ivar_Pk_wig = N_modes / (2.0 * Pk_wig_obs**2.0) #print('ivar_Pk_wig', ivar_Pk_wig) params_mcmc = mcmc_routine(N_params, N_walkers, theta, params_T, params_indices, fix_params, k_obs, Pk_wig_obs, ivar_Pk_wig, Pwig_spl, Psm_spl, norm_gf, params_name, pool) print(params_mcmc) chi_square = chi2(params_mcmc[:, 0], params_indices, fix_params, k_obs, Pk_wig_obs, ivar_Pk_wig, Pwig_spl, Psm_spl, norm_gf) reduced_chi2 = chi_square / (N_fitbin - N_params) print("chi^2/dof: ", reduced_chi2, "\n") ofile_params = odir + 'fit_p0_z{}RADECcut{}_params{}.dat'.format( z_value, cut_type, params_str) write_params(ofile_params, params_mcmc, params_name, reduced_chi2) pool.close() sys.stdout = old_stdout log_file.close()
mu_max[i] = dspec2.max() plt.figure() plt.plot(beta_dspec, dspec2[om.value == om_max[i], :][0, :]) plt.xlabel('Pulsar Position (mas)') plt.ylabel('Magnification') plt.savefig('%sDspec_Slice_%s_%s_%sdense.png' % (dr, rat, widths[i], dens[0])) plt.close('all') np.savez('%sEvo.npz' % dr, mu_max=mu_max, om_max=om_max, beta_max=beta_max) return (mu_max, om_max, beta_max, dens, rat) pool = MPIPool(loadbalance=True) if not pool.is_master(): pool.wait() sys.exit(0) tasks = list() for i in range(len(filelist)): tasks.append((direct, float(filelist[i][2:-4]), 'Under')) vals = pool.map(dspec_find, tasks) pool.close() mu_max = np.zeros((len(vals), widths.shape[0])) om_max = np.zeros((len(vals), widths.shape[0])) beta_max = np.zeros((len(vals), widths.shape[0])) rats = np.zeros(len(vals))
def fit_subsamplefof_mean(): parser = argparse.ArgumentParser( description= 'This is the MCMC code to get the fitting parameters, made by Zhejie Ding.' ) parser.add_argument( '-rec_id', "--rec_id", help='The id of reconstruction, either 0 or 1.', required=True) #0: pre-reconstruct; 1: post-reconstruct parser.add_argument( '-space_id', "--space_id", help= 'The type of space for fitting, 0 for real-space and 1 for redshift space.', required=True) args = parser.parse_args() print("args: ", args) rec_id = int(args.rec_id) print("rec_id: ", rec_id) space_id = int(args.space_id) print("space_id: ", space_id) N_walkers = 40 # simulation run name N_dataset = 20 N_mu_bin = 100 #N_skip_header = 11 #N_skip_footer = 31977 Omega_m = 0.3075 G_0 = growth_factor(0.0, Omega_m) # G_0 at z=0, normalization factor Volume = 1380.0**3.0 # the volume of simulation box sim_z = ['0', '0.6', '1.0'] sim_seed = [0, 9] sim_wig = ['NW', 'WG'] sim_a = ['1.0000', '0.6250', '0.5000'] sim_space = ['r', 's'] # r for real space; s for redshift space rec_dirs = [ 'DD', 'ALL' ] # "ALL" folder stores P(k, \mu) after reconstruction process, while DD is before reconstruction. rec_fprefix = ['', 'R'] mcut_Npar_list = [[37, 149, 516, 1524, 3830], [35, 123, 374, 962, 2105], [34, 103, 290, 681, 1390]] N_masscut = np.size(mcut_Npar_list, axis=1) inputf = '../Zvonimir_data/planck_camb_56106182_matterpower_smooth_z0.dat' k_smooth, Pk_smooth = np.loadtxt(inputf, dtype='f8', comments='#', unpack=True) tck_Pk_sm = interpolate.splrep(k_smooth, Pk_smooth) inputf = '../Zvonimir_data/planck_camb_56106182_matterpower_z0.dat' k_wiggle, Pk_wiggle = np.loadtxt(inputf, dtype='f8', comments='#', unpack=True) tck_Pk_linw = interpolate.splrep(k_wiggle, Pk_wiggle) # firstly, read one file and get k bins we want for the fitting range dir0 = '/Users/ding/Documents/playground/WiggleNowiggle/subsample_FoF_data_HS/Pk_obs_2d_wnw_mean_DD_ksorted_mu_masscut/' inputf = dir0 + 'fof_kaver.wnw_diff_a_0.6250_mcut35_fraction0.126.dat' k_p, mu_p = np.loadtxt(inputf, dtype='f8', comments='#', delimiter=' ', usecols=(0, 1), unpack=True) #print(k_p, mu_p) N_fitbin = len(k_p) #print('# of (k, mu) bins: ', N_fitbin) # for output parameters fitted odir = './ZV_lagrange_params_{}_wig-now_mean_dset/'.format( rec_dirs[rec_id]) if not os.path.exists(odir): os.makedirs(odir) #space_id = 1 # in redshift space if rec_id == 0: ##Sigma_0 = 8.3364 # the approximated value of \Sigma_xy and \Sigma_z, unit Mpc/h, at z=0. Sigma_0 = 7.8364 # suggested by Zvonimir, at z=0 elif rec_id == 1: Sigma_0 = 2.84 # 0: parameter fixed, 1: parameter free. #params_indices = [1, 1, 1, 1, 1, 1, 0, 0] # Only for DM case. b0 needs to be fixed for DM subsample power spectrum. For this case, make sure \Sigma_xy and \Sigma_z positive, which should be set in mcmc_routine. if space_id == 0: # It may not work in real-space just setting f=0. params_indices = [1, 1, 1, 0, 1, 1] # Set f=0 for real-space. f = 0.0 elif space_id == 1: params_indices = [1, 1, 1, 1, 1, 1] # for redshift space f = 1.0 print("params_indices: ", params_indices) alpha_1, alpha_2, sigma, b_0, b_scale = 1.0, 1.0, Sigma_0, 1.0, 0.0 all_names = "alpha_1", "alpha_2", "Sigma_qmax", "f", "b_1", "b_partial" # the same order for params_indices all_temperature = 0.01, 0.01, 0.1, 0.1, 0.01, 0.1 pool = MPIPool(loadbalance=True) for z_id in xrange(3): norm_gf = growth_factor(float(sim_z[z_id]), Omega_m) / G_0 # ##Sigma_z = Sigma_0* norm_gf # # if set_Sigma_xyz_theory == "True": # print("Sigma_z: ", Sigma_z) # sigma_xy, sigma_z = Sigma_z, Sigma_z # else: # if params_indices[2] == 0: # sigma_xy = 0.0 # else: # sigma_xy = 10.0 # if params_indices[3] == 0: # sigma_z = 0.0 # else: # sigma_z = 10.0 all_params = alpha_1, alpha_2, sigma, f, b_0, b_scale np.random.seed() for mcut_id in xrange(N_masscut): N_params, theta, fix_params, params_T, params_name = set_params( all_params, params_indices, all_names, all_temperature) ifile_Pk = './run2_3_Pk_obs_2d_wnw_mean_{}_ksorted_mu_masscut/{}kave{}.wig_minus_now_mean_fof_a_{}_mcut{}.dat'.format( rec_dirs[rec_id], rec_fprefix[rec_id], sim_space[space_id], sim_a[z_id], mcut_Npar_list[z_id][mcut_id]) Pk_wnw_diff_obs = np.loadtxt( ifile_Pk, dtype='f4', comments='#', usecols=(2, ) ) # be careful that there are k, \mu, P(k, \mu) columns. ifile_Cov_Pk = './run2_3_Cov_Pk_obs_2d_wnw_{}_ksorted_mu_masscut/{}kave{}.wig_minus_now_mean_fof_a_{}_mcut{}.dat'.format( rec_dirs[rec_id], rec_fprefix[rec_id], sim_space[space_id], sim_a[z_id], mcut_Npar_list[z_id][mcut_id]) Cov_Pk_wnw = np.loadtxt(ifile_Cov_Pk, dtype='f4', comments='#') ivar_Pk_wnow = N_dataset / np.diag( Cov_Pk_wnw) # the mean sigma error params_mcmc = mcmc_routine(N_params, N_walkers, theta, params_T, params_indices, fix_params, k_p, mu_p, Pk_wnw_diff_obs, ivar_Pk_wnow, tck_Pk_linw, tck_Pk_sm, norm_gf, params_name, pool) chi_square = chi2(params_mcmc[:, 0], params_indices, fix_params, k_p, mu_p, Pk_wnw_diff_obs, ivar_Pk_wnow, tck_Pk_linw, tck_Pk_sm, norm_gf) reduced_chi2 = chi_square / (N_fitbin - N_params) print("Reduced chi2: {}\n".format(reduced_chi2)) # output parameters into a file ofile_params = odir + 'fof_{}kave{}.wnw_diff_a_{}_mcut{}_params{}.dat'.format( rec_fprefix[rec_id], sim_space[space_id], sim_a[z_id], mcut_Npar_list[z_id][mcut_id], ''.join(map( str, params_indices))) print(ofile_params) write_params(ofile_params, params_mcmc, params_name, reduced_chi2) np.random.seed() N_params, theta, fix_params, params_T, params_name = set_params( all_params, params_indices, all_names, all_temperature) # Fit for DM power spectrum ifile_Pk = './run2_3_sub_Pk_2d_wnw_mean_{}_ksorted_mu/{}kave{}.wig_minus_now_mean_sub_a_{}.dat'.format( rec_dirs[rec_id], rec_fprefix[rec_id], sim_space[space_id], sim_a[z_id]) Pk_wnw_diff_true = np.loadtxt( ifile_Pk, dtype='f4', comments='#', usecols=( 2, )) # be careful that there are k, \mu, P(k, \mu) columns. print(ifile_Pk) ifile_Cov_Pk = './run2_3_sub_Cov_Pk_2d_wnw_{}_ksorted_mu/{}kave{}.wig_minus_now_mean_sub_a_{}.dat'.format( rec_dirs[rec_id], rec_fprefix[rec_id], sim_space[space_id], sim_a[z_id]) Cov_Pk_wnw = np.loadtxt(ifile_Cov_Pk, dtype='f4', comments='#') ivar_Pk_wnow = N_dataset / np.diag(Cov_Pk_wnw) # the mean sigma error params_mcmc = mcmc_routine(N_params, N_walkers, theta, params_T, params_indices, fix_params, k_p, mu_p, Pk_wnw_diff_true, ivar_Pk_wnow, tck_Pk_linw, tck_Pk_sm, norm_gf, params_name, pool) chi_square = chi2(params_mcmc[:, 0], params_indices, fix_params, k_p, mu_p, Pk_wnw_diff_true, ivar_Pk_wnow, tck_Pk_linw, tck_Pk_sm, norm_gf) reduced_chi2 = chi_square / (N_fitbin - N_params) print('Reduced chi2: {}\n'.format(reduced_chi2)) ofile_params = odir + 'sub_{}kave{}.wnw_diff_a_{}_params{}.dat'.format( rec_fprefix[rec_id], sim_space[space_id], sim_a[z_id], ''.join( map(str, params_indices))) write_params(ofile_params, params_mcmc, params_name, reduced_chi2) pool.close()
### skips if files do not exist for some reason print 'Warning: file does not exist, cosmo, snap' return if os.path.isfile(out_fn): ###### in case the code breaks return ######### read rockstar files print 'Opening rockstar files:', rockstar_fn reader = sim_manager.TabularAsciiReader(rockstar_fn, columns_to_keep_dict) rock_arr = reader.read_ascii() logM = log10(rock_arr['halo_mvir']) rock_arr = 0 ## release memory hmf = histogram(logM, bins=hist_bins)[0] save(out_fn, hmf) all_snaps = [] for i in range(len(cosmo_arr)): for isnap in arange(30, nsnaps_arr[i]): all_snaps.append([cosmo_arr[i], int(isnap)]) pool = MPIPool() if not pool.is_master(): pool.wait() sys.exit(0) pool.map(Phm_gen, all_snaps) #pool.map(Phh_gen, all_snaps) #pool.map(hmf_gen, all_snaps) pool.close()
f.write( "## loadbalance runtime_sorting iteration mean_time variance ideal actual \n" ) f.write( "#####################################################################################################\n" ) for mean_time in mean_times: for variance_fac in xrange(len(variances)): first = 0 variance = variances[variance_fac] * mean_time for loadbalance in loadbalancing_options: for runtime_sorting_option in runtime_sorting_options: # Initialize the MPI-based pool used for parallelization. pool = MPIPool(loadbalance=loadbalance) if not pool.is_master(): # Wait for instructions from the master process. pool.wait() sys.exit(0) # Initialize the sampler with the chosen specs. sampler = emcee.EnsembleSampler( nwalkers, ndim, lnprob, pool=pool, runtime_sortingfn=runtime_sorting_option) tstart = time.time()