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 get_random_positions(N, t_b, ra_in=None, dec_in=None):
""" Use the star formation history to generate a population of new binaries
Parameters
----------
N : integer
Number of positions to calculate
t_b : float
Birth time to calculate star formation history (Myr)
ra_in : float
RA of system (optional)
dec_in : float
Dec of system (optional)
Returns
-------
ra_out : ndarray
Array of output RA's (degrees)
dec_out : ndarray
Array of output Dec's (degrees)
N_stars : int
Normalization constant calculated from number of stars formed at time t_b
"""
if sf_history.smc_sfh is None or sf_history.smc_coor is None:
sf_history.load_sf_history()
N_regions = len(sf_history.smc_sfh)
# If provided with an ra and dec, only generate stars within 3 degrees of input position
SF_regions = np.zeros((2,N_regions))
for i in np.arange(N_regions):
SF_regions[0,i] = i
if ra_in is None or dec_in is None:
SF_regions[1,i] = sf_history.smc_sfh[i](np.log10(t_b*1.0e6))
elif sf_history.get_theta_proj_degree(sf_history.smc_coor["ra"][i], sf_history.smc_coor["dec"][i], ra_in, dec_in) < c.deg_to_rad * 3.0:
SF_regions[1,i] = sf_history.smc_sfh[i](np.log10(t_b*1.0e6))
else:
SF_regions[1,i] = 0.0
N_stars = np.sum(SF_regions, axis=1)[1]
# Normalize
SF_regions[1] = SF_regions[1] / N_stars
# Sort
SF_sort = SF_regions[:,SF_regions[1].argsort()]
# Move from normed PDF to CDF
SF_sort[1] = np.cumsum(SF_sort[1])
# TEST #
# ra_out = lmc_coor["ra"][SF_sort[0][-100:].astype(int)]
# dec_out = lmc_coor["dec"][SF_sort[0][-100:].astype(int)]
# return ra_out, dec_out
# TEST #
# Random numbers
y = uniform.rvs(size=N)
# Create a 2D grid of CDFs, and random numbers
SF_out, y_out = np.meshgrid(SF_sort[1], y)
# Get index of closest region in sorted array
indices = np.argmin((SF_out - y_out)**2,axis=1)
# Move to indices of stored LMC SFH data array
indices = SF_sort[0][indices].astype(int)
# Get random ra's and dec's of each region
ra_out = sf_history.smc_coor["ra"][indices]
dec_out = sf_history.smc_coor["dec"][indices]
# Width is 12 arcmin or 12/60 degrees for outermost regions
# Width is 6 arcmin or 6/60 degrees for inner regions
# width = 12.0 / 60.0 * np.ones(len(indices))
width = 6.0 / 60.0 * np.ones(len(indices))
# for i in np.arange(len(indices)):
# if str(smc_coor["region"][indices[i]]).find("_") != -1:
# width[i] = 6.0 / 60.0
tmp_delta_ra = width * (2.0 * uniform.rvs(size=len(indices)) - 1.0) / np.cos(c.deg_to_rad * dec_out) * 2.0
tmp_delta_dec = width * (2.0 * uniform.rvs(size=len(indices)) - 1.0)
ra_out = ra_out + tmp_delta_ra
dec_out = dec_out + tmp_delta_dec
return ra_out, dec_out, N_stars
def ln_priors(y):
""" Priors on the model parameters
Parameters
----------
y : ra, dec, M1, M2, A, ecc, v_k, theta, phi, ra_b, dec_b, t_b
Current HMXB location (ra, dec) and 10 model parameters
Returns
-------
lp : float
Natural log of the prior
"""
# M1, M2, A, v_k, theta, phi, ra_b, dec_b, t_b = y
ra, dec, M1, M2, A, ecc, v_k, theta, phi, ra_b, dec_b, t_b = y
lp = 0.0
# P(M1)
if M1 < c.min_mass or M1 > c.max_mass: return -np.inf
norm_const = (c.alpha+1.0) / (np.power(c.max_mass, c.alpha+1.0) - np.power(c.min_mass, c.alpha+1.0))
lp += np.log( norm_const * np.power(M1, c.alpha) )
# M1 must be massive enough to evolve off the MS by t_obs
if load_sse.func_sse_tmax(M1) > t_b: return -np.inf
# P(M2)
# Normalization is over full q in (0,1.0)
q = M2 / M1
if q < 0.3 or q > 1.0: return -np.inf
lp += np.log( (1.0 / M1 ) )
# P(ecc)
if ecc < 0.0 or ecc > 1.0: return -np.inf
lp += np.log(2.0 * ecc)
# P(A)
if A*(1.0-ecc) < c.min_A or A*(1.0+ecc) > c.max_A: return -np.inf
norm_const = np.log(c.max_A) - np.log(c.min_A)
lp += np.log( norm_const / A )
# A must avoid RLOF at ZAMS, by a factor of 2
r_1_roche = binary_evolve.func_Roche_radius(M1, M2, A*(1.0-ecc))
if 2.0 * load_sse.func_sse_r_ZAMS(M1) > r_1_roche: return -np.inf
# P(v_k)
if v_k < 0.0: return -np.inf
lp += np.log( maxwell.pdf(v_k, scale=c.v_k_sigma) )
# P(theta)
if theta <= 0.0 or theta >= np.pi: return -np.inf
lp += np.log(np.sin(theta) / 2.0)
# P(phi)
if phi < 0.0 or phi > np.pi: return -np.inf
lp += -np.log( np.pi )
# Get star formation history
sf_history.load_sf_history()
sfh = sf_history.get_SFH(ra_b, dec_b, t_b, sf_history.smc_coor, sf_history.smc_sfh)
if sfh <= 0.0: return -np.inf
# P(alpha, delta)
# From spherical geometric effect, we need to care about cos(declination)
lp += np.log(np.cos(c.deg_to_rad * dec_b) / 2.0)
##################################################################
# We add an additional prior that scales the RA and Dec by the
# area available to it, i.e. pi theta^2, where theta is the angle
# of the maximum projected separation over the distance.
#
# Still under construction
##################################################################
M1_b, M2_b, A_b = binary_evolve.func_MT_forward(M1, M2, A, ecc)
A_c, v_sys, ecc = binary_evolve.func_SN_forward(M1_b, M2_b, A_b, v_k, theta, phi)
if ecc < 0.0 or ecc > 1.0 or np.isnan(ecc): return -np.inf
# Ensure that we only get non-compact object companions
tobs_eff = binary_evolve.func_get_time(M1, M2, t_b)
M_tmp, M_dot_tmp, R_tmp, k_type = load_sse.func_get_sse_star(M2_b, tobs_eff)
if int(k_type) > 9: return -np.inf
# t_sn = (t_b - func_sse_tmax(M1)) * 1.0e6 * yr_to_sec # The time since the primary's core collapse
# theta_max = (v_sys * t_sn) / dist_LMC # Unitless
# area = np.pi * rad_to_dec(theta_max)**2
# lp += np.log(1.0 / area)
##################################################################
# Instead, let's estimate the number of stars formed within a cone
# around the observed position, over solid angle and time.
# Prior is in Msun/Myr/steradian
##################################################################
t_min = load_sse.func_sse_tmax(M1) * 1.0e6 * c.yr_to_sec
t_max = (load_sse.func_sse_tmax(M2_b) - binary_evolve.func_get_time(M1, M2, 0.0)) * 1.0e6 * c.yr_to_sec
if t_max-t_min < 0.0: return -np.inf
theta_C = (v_sys * (t_max - t_min)) / c.dist_SMC
stars_formed = get_stars_formed(ra, dec, t_min, t_max, v_sys, c.dist_SMC)
if stars_formed == 0.0: return -np.inf
volume_cone = (np.pi/3.0 * theta_C**2 * (t_max - t_min) / c.yr_to_sec / 1.0e6)
lp += np.log(sfh / stars_formed / volume_cone)
##################################################################
# # P(t_b | alpha, delta)
# sfh_normalization = 1.0e-6
# lp += np.log(sfh_normalization * sfh)
# Add a prior so that the post-MT secondary is within the correct bounds
M2_c = M1 + M2 - load_sse.func_sse_he_mass(M1)
if M2_c > c.max_mass or M2_c < c.min_mass: return -np.inf
# Add a prior so the effective time remains bounded
t_eff_obs = binary_evolve.func_get_time(M1, M2, t_b)
if t_eff_obs < 0.0: return -np.inf
if t_b * 1.0e6 * c.yr_to_sec < t_min: return -np.inf
if t_b * 1.0e6 * c.yr_to_sec > t_max: return -np.inf
return lp
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
