def run_emcee(M2_d, P_orb_obs, ecc_obs, ra, dec, M2_d_err=1.0, P_orb_obs_err=1.0, ecc_obs_err=0.05, nburn=1000, nsteps=1000): """ Run the emcee function Parameters ---------- M2_d : float Observed secondary mass P_orb_obs : float Observed orbital period ecc_obs : float Observed orbital eccentricity ra : float Observed right ascension dec : float Observed declination Returns ------- sampler : emcee object """ # First thing is to load the sse data and SF_history data load_sse.load_sse() sf_history.load_sf_history() # Get initial values initial_vals = get_initial_values(M2_d, nwalkers=nwalkers) # Define sampler args = [[M2_d, M2_d_err, P_orb_obs, P_orb_obs_err, ecc_obs, ecc_obs_err, ra, dec]] sampler = emcee.EnsembleSampler(nwalkers=nwalkers, dim=10, lnpostfn=ln_posterior, args=args) # Assign initial values p0 = np.zeros((nwalkers,10)) p0 = set_walkers(initial_vals, args[0], nwalkers=nwalkers) # Burn-in pos,prob,state = sampler.run_mcmc(p0, N=nburn) # Full run sampler.reset() pos,prob,state = sampler.run_mcmc(pos, N=nsteps) return sampler
def run_emcee_population(nburn=10000, nsteps=100000, nwalkers=80): """ Run emcee on the entire X-ray binary population Parameters ---------- nburn : float (optional) number of steps for the Burn-in (default=10000) nsteps : float (optional) number of steps for the simulation (default=100000) Returns ------- sampler : emcee object """ # First thing is to load the sse data and SF_history data load_sse.load_sse() sf_history.load_sf_history() # Get initial values - choose 12 Msun as a seed for initial position initial_vals = get_initial_values(12.0, nwalkers=nwalkers) # Define sampler sampler = emcee.EnsembleSampler(nwalkers=nwalkers, dim=10, lnpostfn=ln_posterior_population) # Assign initial values based on a random binary args = [12.0, 2.0, 500.0, 20.0, 0.50, 0.2, 13.5, -72.7] # SMC p0 = np.zeros((nwalkers,10)) p0 = set_walkers(initial_vals, args, nwalkers=nwalkers) # Burn-in pos,prob,state = sampler.run_mcmc(p0, N=nburn) # Full run sampler.reset() pos,prob,state = sampler.run_mcmc(pos, N=nsteps) return sampler
def run_pop_synth(input_sys, N_sys=10000, t_low=15.0, t_high=60.0, delta_t=1): """ Run a forward population synthesis Parameters ---------- input_sys : ra_sys, dec_sys, P_orb_sys, ecc_sys, M2_d_sys Observed values of individual system N_sys : int Number of systems per age calculated (10000) t_low : float Lower limit to age range tested (15 Myr) t_high : float Upper limit to age range tested (60 Myr) delta_t : float Age resolution (1 Myr) Returns ------- HMXB_sys : numpy recarray Systems that evolve into the observed system names = ["ra", "dec", "ra_b", "dec_b", "P_orb", "ecc", "M_2_d", "theta_proj", "age", "norm"] init_params_sys : numpy recarray Initial conditions for the observed Systems names = ["M1","M2","A","ecc","v_k","theta","phi","ra_b","dec_b","t_b"] """ ra_sys, dec_sys, P_orb_sys, ecc_sys, M2_d_sys = input_sys # First thing is to load the sse data and SF_history data load_sse.load_sse() sf_history.load_sf_history() names = ["ra", "dec", "ra_b", "dec_b", "P_orb", "ecc", "M_2_d", "theta_proj", "age", "norm"] HMXB = np.recarray(0, names=names, formats=['float64,float64,float64,float64,float64,float64,float64,float64,float64,float64']) HMXB_sys = np.recarray(0, names=names, formats=['float64,float64,float64,float64,float64,float64,float64,float64,float64,float64']) names = ["M1","M2","A","ecc","v_k","theta","phi","ra_b","dec_b","t_b"] init_params_sys = np.recarray(0, names=names, formats=['float64,float64,float64,float64,float64,float64,float64,float64,float64,float64']) theta_sep = np.array([]) batch_size = 1000 n_batch = int(np.ceil(float(N_sys)/float(batch_size))) for t_b in np.linspace(14.0, 56.0, 43): for batch in np.arange(n_batch): n_run = min(batch_size, N_sys - (batch)*batch_size) HMXB_t, init_params_t = create_HMXBs(t_b, N_sys=n_run, ra_in=ra_sys, dec_in=dec_sys) HMXB = np.concatenate((HMXB, HMXB_t)) for i in np.arange(len(HMXB_t)): h = HMXB_t[i] p = init_params_t[i] angle = c.rad_to_deg * sf_history.get_theta_proj_degree(ra_sys, dec_sys, h["ra"], h["dec"]) theta_sep = np.append(theta_sep, angle) if angle < 0.2 \ and np.abs(h["P_orb"] - P_orb_sys) < 5.0 \ and np.abs(h["ecc"] - ecc_sys) < 0.1 \ and np.abs(h["M_2_d"] - M2_d_sys) < 1.0: HMXB_sys = np.append(HMXB_sys, h) init_params_sys = np.append(init_params_sys, p) return HMXB_sys, init_params_sys
def full_forward(M1, M2, A, ecc, v_k, theta, phi, t_obs): """ Evolve a binary forward from its initial conditions Parameters ---------- M1 : float Initial primary mass (Msun) M2 : float Initial secondary mass (Msun) A : float Initial orbital separation (Rsun) ecc : float Initial orbital eccentricity (unitless) v_k : float SN kick velocity theta : float SN kick polar angle phi : float SN kick azimuthal angle t_obs : float observation time Returns ------- M_NS : float or ndarray Array of final primary masses (Currently set to the NS mass, c.M_NS) M_2 : float or ndarray Array of final secondary masses (Msun) L_x : float or ndarray X-ray luminosity (erg/s) v_sys : float or ndarray Systemic velocity (km/s) M2_dot : float or ndarray Mass accretion rate (Msun/yr) A : float or ndarray Orbital separation (Rsun) ecc : float or ndarray Orbital eccentricity (unitless) theta : float or ndarray Projected angular distance traveled from birth location (radians) k_type : int k-type of HMXB donor """ if load_sse.func_sse_mass is None: load_sse.load_sse() if isinstance(M1, np.ndarray): dtypes = [('M_NS','<f8'), \ ('M_2','<f8'), \ ('L_x','<f8'), \ ('v_sys','<f8'), \ ('M2_dot','<f8'), \ ('A','<f8'), \ ('ecc','<f8'), \ ('theta','<f8'), \ ('k_type','<i8')] HMXB = np.recarray(len(M1), dtype=dtypes) for i in np.arange(len(M1)): if isinstance(t_obs, np.ndarray): if t_obs[i] < load_sse.func_sse_ms_time(M1[i]): HMXB["M_NS"][i] = M1[i] HMXB["M_2"][i] = M2[i] HMXB["A"][i] = A[i] continue else: if t_obs < load_sse.func_sse_ms_time(M1[i]): HMXB["M_NS"][i] = M1[i] HMXB["M_2"][i] = M2[i] HMXB["A"][i] = A[i] continue # First MT phase M_1_b, M_2_b, A_b = binary_evolve.func_MT_forward(M1[i], M2[i], A[i], ecc[i]) if isinstance(t_obs, np.ndarray): if t_obs[i] < load_sse.func_sse_tmax(M1[i]): HMXB["M_NS"][i] = M_1_b HMXB["M_2"][i] = M_2_b HMXB["A"][i] = A_b continue else: if t_obs < load_sse.func_sse_tmax(M1[i]): HMXB["M_NS"][i] = M_1_b HMXB["M_2"][i] = M_2_b HMXB["A"][i] = A_b continue # SN A_tmp, v_sys_tmp, e_tmp = binary_evolve.func_SN_forward(M_1_b, M_2_b, A_b, v_k[i], theta[i], phi[i]) # XRB if isinstance(t_obs, np.ndarray): M_2_tmp, L_x_tmp, M2_dot_out, A_out = binary_evolve.func_Lx_forward(M1[i], M2[i], M_2_b, A_tmp, e_tmp, t_obs[i]) theta_out = (t_obs[i] - load_sse.func_sse_tmax(M1[i])) * v_sys_tmp / c.dist_SMC * c.yr_to_sec * 1.0e6 * np.sin(get_theta(1)) tobs_eff = binary_evolve.func_get_time(M1[i], M2[i], t_obs[i]) else: M_2_tmp, L_x_tmp, M2_dot_out, A_out = binary_evolve.func_Lx_forward(M1[i], M2[i], M_2_b, A_tmp, e_tmp, t_obs) theta_out = (t_obs - load_sse.func_sse_tmax(M1[i])) * v_sys_tmp / c.dist_SMC * c.yr_to_sec * 1.0e6 * np.sin(get_theta(1)) tobs_eff = binary_evolve.func_get_time(M1[i], M2[i], t_obs) # To get k-type of HMXB donor if M_2_b > c.max_mass: k_type = -999 else: M_tmp, M_dot_tmp, R_tmp, k_type = load_sse.func_get_sse_star(M_2_b, tobs_eff) HMXB["M_NS"][i] = c.M_NS HMXB["M_2"][i] = M_2_tmp HMXB["L_x"][i] = L_x_tmp HMXB["v_sys"][i] = v_sys_tmp HMXB["M2_dot"][i] = M2_dot_out HMXB["A"][i] = A_out HMXB["ecc"][i] = e_tmp HMXB["theta"][i] = theta_out HMXB["k_type"][i] = int(k_type) return HMXB["M_NS"], HMXB["M_2"], HMXB["L_x"], HMXB["v_sys"], HMXB["M2_dot"], HMXB["A"], HMXB["ecc"], HMXB["theta"], HMXB["k_type"] else: # Star does not make it to MT phase if t_obs < load_sse.func_sse_ms_time(M1): return M1, M2, 0.0, 0.0, 0.0, A, ecc, 0.0 # MT phase M_1_b, M_2_b, A_b = binary_evolve.func_MT_forward(M1, M2, A, ecc) # Star does not make it to SN if t_obs < load_sse.func_sse_tmax(M1): return M_1_b, M_2_b, 0.0, 0.0, 0.0, A_b, ecc, 0.0 # SN A_tmp, v_sys_tmp, e_tmp = binary_evolve.func_SN_forward(M_1_b, M_2_b, A_b, v_k, theta, phi) # XRB M_2_tmp, L_x_tmp, M2_dot_out, A_out = binary_evolve.func_Lx_forward(M1, M2, M_2_b, A_tmp, e_tmp, t_obs) theta_out = (t_obs - load_sse.func_sse_tmax(M1)) * v_sys_tmp / c.dist_SMC * c.yr_to_sec * 1.0e6 * np.sin(get_theta(1)) # To get k-type of HMXB donor tobs_eff = binary_evolve.func_get_time(M1, M2, t_obs) M_tmp, M_dot_tmp, R_tmp, k_type = load_sse.func_get_sse_star(M_2_b, tobs_eff) return c.M_NS, M_2_tmp, L_x_tmp, v_sys_tmp, M2_dot_out, A_out, e_tmp, theta_out, int(k_type)
corner.hist2d(np.log10(HMXB[2]), HMXB[3], ax=ax[0], bins=40, range=plt_range, plot_datapoints=False) ax[0].set_xlabel(r"${\rm log}\ P_{\rm orb}\ {\rm (days)}$", size=16) ax[0].set_ylabel(r"$e$", size=16) ax[0].set_xticks([0,1,2,3,4]) ax[0].set_yticks([0,0.25,0.5,0.75,1.0]) plt_range = ([8, 24], [0,80]) corner.hist2d(HMXB[4], HMXB[5], ax=ax[1], bins=40, range=plt_range, plot_datapoints=False) ax[1].set_xlabel(r"$M_2\ ({\rm M}_{\odot})$", size=16) ax[1].set_ylabel(r"$v_{\rm sys}\ {\rm (km\ s}^{-1})$", size=16) ax[1].set_xticks([8,12,16,20,24]) ax[1].set_yticks([0,20, 40, 60, 80]) # Get flight time from birth time and M1 lifetime load_sse.load_sse() t_flight = sampler.flatchain.T[9] - load_sse.func_sse_tmax(sampler.flatchain.T[0]) plt_range = ([0,55], [0,25]) contour_kwargs = {'colors':'r', 'linestyles':'dashed'} corner.hist2d(t_flight, HMXB[7], ax=ax[2], bins=40, range=plt_range, plot_density=False, plot_datapoints=False, contour_kwargs=contour_kwargs) corner.hist2d(sampler.flatchain.T[9], HMXB[7], ax=ax[2], bins=40, range=plt_range, plot_density=False, plot_datapoints=False) ax[2].set_xlabel(r"$t_{\rm i}\ {\rm (Myr)}$", size=16) ax[2].set_ylabel(r"$\theta\ {\rm (amin)}$", size=16) plt.tight_layout() plt.savefig('../figures/smc_population_HMXB.pdf')
def run_emcee_2(M2_d, P_orb_obs, ecc_obs, ra, dec, M2_d_err=1.0, P_orb_obs_err=1.0, ecc_obs_err=0.05, nwalkers=80, nburn=1000, nsteps=1000, threads=1, mpi=False): """ Run the emcee function Parameters ---------- M2_d : float Observed secondary mass P_orb_obs : float Observed orbital period ecc_obs : float Observed orbital eccentricity ra : float Observed right ascension dec : float Observed declination threads : int Number of threads to use for parallelization mpi : bool If true, use MPIPool for parallelization Returns ------- sampler : emcee object """ # First thing is to load the sse data and SF_history data load_sse.load_sse() sf_history.load_sf_history() # Get initial values initial_vals = get_initial_values(M2_d, nwalkers=nwalkers) # Define sampler args = [[M2_d, M2_d_err, P_orb_obs, P_orb_obs_err, ecc_obs, ecc_obs_err, ra, dec]] if mpi == True: pool = MPIPool() if not pool.is_master(): pool.wait() sys.exit(0) sampler = emcee.EnsembleSampler(nwalkers=nwalkers, dim=10, lnpostfn=ln_posterior, args=args, pool=pool) elif threads != 1: sampler = emcee.EnsembleSampler(nwalkers=nwalkers, dim=10, lnpostfn=ln_posterior, args=args, threads=threads) else: sampler = emcee.EnsembleSampler(nwalkers=nwalkers, dim=10, lnpostfn=ln_posterior, args=args) # Assign initial values p0 = np.zeros((nwalkers,10)) p0 = set_walkers(initial_vals, args[0], nwalkers=nwalkers) # Burn-in 1 pos,prob,state = sampler.run_mcmc(p0, N=nburn) sampler1 = copy.copy(sampler) # TESTING BEGIN - Get limiting ln_prob for worst 10 chains prob_lim = (np.sort(prob)[9] + np.sort(prob)[10])/2.0 index_best = np.argmax(prob) for i in np.arange(len(prob)): # if sampler1.acceptance_fraction[i] == 0.0: pos[i] = np.copy(pos[index_best]) + np.random.normal(0.0, 0.005, size=10) if prob[i] < prob_lim: pos[i] = np.copy(pos[index_best]) + np.random.normal(0.0, 0.005, size=10) # TESTING END print "Burn-in 1 finished." print "Starting burn-in 2..." # Burn-in 2 sampler.reset() pos,prob,state = sampler.run_mcmc(pos, N=nburn) sampler2 = copy.copy(sampler) # TESTING BEGIN - Get limiting ln_prob for worst 10 chains prob_lim = (np.sort(prob)[9] + np.sort(prob)[10])/2.0 index_best = np.argmax(prob) for i in np.arange(len(prob)): # if sampler2.acceptance_fraction[i] == 0.0: pos[i] = np.copy(pos[index_best]) + np.random.normal(0.0, 0.005, size=10) if prob[i] < prob_lim: pos[i] = np.copy(pos[index_best]) + np.random.normal(0.0, 0.005, size=10) # TESTING END print "Burn-in 2 finished." print "Starting burn-in 3..." # Burn-in 3 sampler.reset() pos,prob,state = sampler.run_mcmc(pos, N=nburn) sampler3 = copy.copy(sampler) # TESTING BEGIN - Get limiting ln_prob for worst 10 chains prob_lim = (np.sort(prob)[9] + np.sort(prob)[10])/2.0 index_best = np.argmax(prob) for i in np.arange(len(prob)): # if sampler3.acceptance_fraction[i] == 0.0: pos[i] = np.copy(pos[index_best]) + np.random.normal(0.0, 0.005, size=10) if prob[i] < prob_lim: pos[i] = np.copy(pos[index_best]) + np.random.normal(0.0, 0.005, size=10) # TESTING END print "Burn-in 3 finished." print "Starting burn-in 4..." # Burn-in 4 sampler.reset() pos,prob,state = sampler.run_mcmc(pos, N=nburn) sampler4 = copy.copy(sampler) print "Burn-in 4 finished." print "Starting production run..." # Full run sampler.reset() pos,prob,state = sampler.run_mcmc(pos, N=nsteps) print "Finished production run" if mpi is True: pool.close() return sampler1, sampler2, sampler3, sampler4, sampler