diff --git a/PyROA/PyROA.py b/PyROA/PyROA.py index 6c989f8..53b6477 100644 --- a/PyROA/PyROA.py +++ b/PyROA/PyROA.py @@ -1021,9 +1021,12 @@ def log_prior(params, priors, add_var, data, delay_dist, AccDisc, wavelengths, i if (add_var == True): V = params_chunks[i][-1]/(init_params_chunks[i][-1]*5) - - if (delay_dist == True and i>0): - if (params_chunks[i][3]>=0.0): + + if (delay_dist == True and i>0): #Upper limit on delta_i stops it from blowing up in MCMC + baseline_0 = np.nanmax(data[0][:,0]) - np.nanmin(data[0][:,0]) + baseline_i = np.nanmax(data[i][:,0]) - np.nanmin(data[i][:,0]) + maxval = max(baseline_0, baseline_i)*5 + if (params_chunks[i][3]>=0.0) and (params_chunks[i][3] <= maxval): tau_rms = params_chunks[i][3] #pr.append(2.0*np.log((1.0/np.sqrt(2.0*np.pi*(rms_prior_width**2)))*np.exp(-0.5*(tau_rms/rms_prior_width)**2))) check.append(0.0) @@ -1138,7 +1141,7 @@ def Slow(t, S0, dS, t0): def FullFit(data, priors, init_tau, init_delta, add_var, sig_level, Nsamples, Nburnin, include_slow_comp, slow_comp_delta, calc_P, delay_dist, psi_types, pos_ref, AccDisc, wavelengths, filters, use_backend, - resume_progress, plot_corner,memfunction, gridsize): + resume_progress, plot_corner,memfunction, gridsize, nthread=1): Nchunk = 2 @@ -1400,7 +1403,8 @@ def FullFit(data, priors, init_tau, init_delta, add_var, sig_level, Nsamples, else: backend = None - with Pool() as pool: + + with Pool(nthread) as pool: sampler = emcee.EnsembleSampler(nwalkers, ndim, log_probability, args=[data, priors, add_var, size,sig_level, include_slow_comp, slow_comp_delta, P_func, slow_comps, P_slow, init_delta, delay_dist, psi_types, @@ -1667,9 +1671,12 @@ def FullFit(data, priors, init_tau, init_delta, add_var, sig_level, Nsamples, mx=max(merged_mjd) mn=min(merged_mjd) length = abs(mx-mn) - t = np.arange(mn, mx, length/(gridsize)) - + if gridsize is not None: + t = np.arange(mn, mx, length/(gridsize)) + else: + t = np.arange(mn, mx, length/1000) + ts, Xs, errss = RunningOptimalAverageOutConv(t, merged_mjd, merged_flux, merged_err, factors, conv, prev, x, delta_new) @@ -1814,13 +1821,15 @@ def __init__(self, datadir, objName, filters, priors, delay_ref = None, init_tau delay_dist=False , psi_types = None, add_var=True, sig_level = 4.0, Nsamples=10000, Nburnin=0, include_slow_comp=False, slow_comp_delta=30.0, calc_P=False, AccDisc=False, wavelengths=None, - use_backend = False, resume_progress = False, plot_corner=False,memfunction='gaussian', gridsize = None): + use_backend = False, resume_progress = False, plot_corner=False,memfunction='gaussian', gridsize = None, + nthread=1): if datadir[-1] != '/': datadir += '/' #Add forward slash in case it isn't there self.datadir=datadir self.objName=objName self.filters=filters self.gridsize = gridsize + self.nthread = nthread data=[] for i in range(len(filters)): file = datadir + str(self.objName) +"_"+ str(self.filters[i]) + ".dat" @@ -1888,7 +1897,8 @@ def __init__(self, datadir, objName, filters, priors, delay_ref = None, init_tau self.sig_level, self.Nsamples, self.Nburnin, self.include_slow_comp, self.slow_comp_delta, self.calc_P, self.delay_dist, self.psi_types, self.delay_ref_pos, self.AccDisc, self.wavelengths, self.filters, - self.use_backend, self.resume_progress,plot_corner,memfunction, self.gridsize) + self.use_backend, self.resume_progress,plot_corner,memfunction, self.gridsize, + nthread=self.nthread) self.samples = run[0] self.samples_flat = run[1] @@ -2356,7 +2366,7 @@ def log_probability2(params, data, priors, sig_level, init_params_chunks,memfunc -def InterCalib(data, priors, init_delta, sig_level, Nsamples, Nburnin, filter,plot_corner,memfunction, gridsize): +def InterCalib(data, priors, init_delta, sig_level, Nsamples, Nburnin, filter,plot_corner,memfunction, gridsize, nthread=1): ######################################################################################## #Run MCMC to fit to data @@ -2426,7 +2436,7 @@ def InterCalib(data, priors, init_delta, sig_level, Nsamples, Nburnin, filter,pl np.savetxt('test_initial_points.txt',pos.T) nwalkers, ndim = pos.shape - with Pool() as pool: + with Pool(nthread) as pool: sampler = emcee.EnsembleSampler(nwalkers, ndim, log_probability2, args=(data, priors, sig_level, init_params_chunks,memfunction, gridsize), pool=pool) @@ -2564,11 +2574,12 @@ def InterCalib(data, priors, init_delta, sig_level, Nsamples, Nburnin, filter,pl class InterCalibrate(): def __init__(self, datadir, objName, filter, scopes, priors, init_delta=1.0, sig_level = 3.0, - Nsamples=15000, Nburnin=10000,plot_corner=False,memfunction='gaussian'): + Nsamples=15000, Nburnin=10000,plot_corner=False,memfunction='gaussian', nthread=1): self.datadir=datadir self.objName=objName self.filter=filter self.scopes=scopes + self.nthread = nthread scopes_array = [] data=[] for i in range(len(scopes)): @@ -2590,7 +2601,7 @@ def __init__(self, datadir, objName, filter, scopes, priors, init_delta=1.0, sig run = InterCalib(data, self.priors, self.init_delta, self.sig_level, self.Nsamples, - self.Nburnin, self.filter,plot_corner,memfunction) + self.Nburnin, self.filter,plot_corner,memfunction, nthread=self.nthread) self.samples = run[0] self.samples_flat = run[1] @@ -2958,7 +2969,7 @@ def log_probability3(params, data, priors, add_var, size, sig_level): -def LensFit(data, priors, init_tau, init_delta, add_var, sig_level, Nsamples, Nburnin, image, file): +def LensFit(data, priors, init_tau, init_delta, add_var, sig_level, Nsamples, Nburnin, image, file, nthread=1): if (add_var == True): Npar = 7*len(data) + 3 @@ -3085,7 +3096,7 @@ def LensFit(data, priors, init_tau, init_delta, add_var, sig_level, Nsamples, Nb nwalkers, ndim = pos.shape print("Nwalkers = ", nwalkers) - with Pool() as pool: + with Pool(nthread) as pool: sampler = emcee.EnsembleSampler(nwalkers, ndim, log_probability3, args=(data, priors, add_var, size, sig_level), pool=pool) sampler.run_mcmc(pos, Nsamples, progress=True);