Polytrope EOS emcee workflow
Import the libraries
[ ]:
import numpy as np
import pandas as pd
import emcee
import corner
import matplotlib.pyplot as plt
from scipy.interpolate import interp1d
# Import your existing modules
import InferenceWorkflow.prior as prior
# import ultranest # Not used anymore
# import TOVsolver.unit as unit
# from TOVsolver.unit import g_cm_3, dyn_cm_2, km, Msun, e0
import TOVsolver.maximum_central_density as maximum_central_density
# from TOVsolver.maxium_central_density import maxium_central_density
from TOVsolver.solver_code import solveTOV2
import Polytrope_EOS as Polytrope
Define the Constants
[ ]:
fm = 1
hbarc = 0.197327053
c = 1 # speed of light
hbar = 1 # reduced Planck constant
GeV = 1 / hbarc # Giga-electronvolt
MeV = 1e-3 * GeV # Mega-electronvolt
g = 5.625e26 * MeV # gram
kg = 1e3 * g # kilogram
cm = 1e13 * fm # centimeter
m = 100 * cm # meter
km = 1e5 * cm # kilometer
s = 3e10 * cm # second
dyn = g * cm / s**2 # dyne
dyn_cm_2 = dyn / cm**2 # dyne / cm^2
g_cm_3 = g / cm**3 # gram / cm^3
erg = dyn * cm # ἐργον energy unit
m_n = 939.565 * MeV # mass of neutron
n0 = 0.16 / fm**3 # saturation density
e0 = m_n * n0 # saturation energy density
G = 6.6743e-8 * dyn * cm**2 / g**2 # gravitational constant
Msun = 1.989e33 * g # mass of sun
import numpy as np
from scipy.interpolate import interp1d
import TOVsolver.unit as unit
# from TOVsolver.unit import unit
from TOVsolver.main import OutputMR
Defining the Maximum Central Density
Outputs the maxium central density of a stable EoS Args: energy_density (numpy 1Darray): Density of EoS pressure (numpy 1Darray): Pressure of EoS central_densitys (numpy 1Darray): The range of central density num2 (int): The number of segments in the density interval of second search. At least 5 points Returns: density (float): The maxium central density, in unit.g_cm_3
[ ]:
def maxium_central_density(energy_density, pressure, central_densitys=np.logspace(14.3, 15.6, 5) * unit.g_cm_3, num2=30):
############## Below is the main part of two searches for peaks
######## First search
Ms = [-1,-2] # Meaningless initialization, only to ensure that the following 'if (i>0) and (Ms [-1]<=Ms [-2])' statements can run properly
store_d_range = [central_densitys[-2], central_densitys[-1]] # When the following loop does not output a result, initialization here will have its meaning
# Find the extremum point within the predetermined range of central density and return the central density near the extremum point
for i, rho in enumerate(central_densitys):
M, R = OutputMR('', energy_density, pressure, [rho])[0]
Ms.append(M)
if (i>0) and (Ms[-1] <= Ms[-2]):
store_d_range = [central_densitys[i-2], central_densitys[i]]
break
Ms_larg = Ms[-1] # Used to store the mass corresponding to the maximum density during the first peak search
######## Second search
Ms = [-1,-2] # Reinitialize
# Update and refine the central density range, and ultimately return the central density of the extremum points
store_d_range = np.geomspace(store_d_range[0], store_d_range[1], num2) # At least 5 points
# Note that the first and last points have already been calculated, so there is no need to traverse them
for i, rho in enumerate(store_d_range[1:-1]):
M, R = OutputMR('', energy_density, pressure, [rho])[0]
Ms.append(M)
if Ms[-1] <= Ms[-2]: # Note that due to the second peak search refining the interval, the result is generally not obtained at the first point.
# Therefore, initializing Ms=[-1, -2] is acceptable
return store_d_range[1:][i-1]
# When the above peak search fails, continue comparing the last point
if Ms_larg < Ms[-1]:
return store_d_range[-2]
else:
return store_d_range[-1]
Load the data
[ ]:
# Load crust data
Tolos_crust_out = np.loadtxt('', delimiter=None, comments='#', usecols=(0, 1, 2, 3, 4))
eps_crust = Tolos_crust_out[:,3] * g_cm_3
pres_crust = Tolos_crust_out[:,4] * dyn_cm_2
rho_crust_end = eps_crust[-1]
P_crust_end = pres_crust[-1]
eps_core = np.logspace(11.7, 15.6, 1000, base=10) * g_cm_3
cent_densitys = np.logspace(14.3, 15.6, 30) * g_cm_3
parameters = ['rho1', 'rho2', 'rho3', 'gamma1', 'gamma2', 'gamma3', 'gamma4', 'd1', 'd2', 'd3']
derived_param_names = ['d_max']
gamma_lo_ranges = [1, 0.1, 0.1, 0.1]
gamma_up_ranges = [1.5, 4, 8, 4]
Define the Prior
[ ]:
# Define log_prior as shown above
def log_prior(theta):
para = np.zeros(len(theta))
try:
# Assign rho parameters
para[0] = prior.flat_prior(0.04/0.16*e0, 0.065/0.16*e0, theta[0])
para[1] = prior.flat_prior(1.5*e0, 8.3*e0, theta[1])
para[2] = prior.flat_prior(para[1], 8.3*e0, theta[2])
rho_s = [para[0], para[1], para[2], eps_core[-1]]
gamma_s = []
k = 0
pt = 0
for i in range(4):
if i == 0:
gamma_max = Polytrope.fun_gamma_max(rho_s[i], rho_crust_end, P_crust_end)
upper = min(gamma_up_ranges[i], gamma_max)
if upper < gamma_lo_ranges[i]:
return -np.inf
para[3+i] = prior.flat_prior(gamma_lo_ranges[i], upper, theta[3+i])
else:
gamma_max = Polytrope.fun_gamma_max(rho_s[i], rho_s[i-1], pt)
upper = min(gamma_up_ranges[i], gamma_max)
if upper < gamma_lo_ranges[i]:
return -np.inf
para[3+i] = prior.flat_prior(gamma_lo_ranges[i], upper, theta[3+i])
gamma_s.append(para[3+i])
if i == 0:
k = P_crust_end / (rho_s[0]**gamma_s[-1])
else:
k = pt / (rho_s[i-1]**gamma_s[-1])
pt = k * rho_s[i]**gamma_s[-1]
# Check derived parameters
pres = Polytrope.compute_EOS(eps_core, para[3:7])[0]
eps = np.hstack((eps_crust, eps_core))
pres = np.hstack((pres_crust, pres))
d_max = maxium_central_density(eps, pres, cent_densitys)
log_d_max = 14.3 + (np.log10(d_max / g_cm_3) - 14.3) * theta[7:]
if not np.all((14.3 <= log_d_max) & (log_d_max <= np.log10(cent_densitys[-1]/g_cm_3))):
return -1e101
return 0.0
except:
return -1e101
Define the Log likelihood and the log posterior
[ ]:
# Define log_likelihood as shown above
def log_likelihood(theta):
gammas = np.array([theta[3], theta[4], theta[5], theta[6]])
rho_ts = np.array([theta[0], theta[1], theta[2]])
d_s = 10 ** np.array(theta[7:]) * g_cm_3
theta_full = np.append(gammas, rho_ts)
Ms, Rs = from_para_to_MR_points(theta_full, d_s).T
# Gaussian likelihood for M and R
fm = np.where(Ms > 3 * Msun, -(Ms - 2 * Msun)**20, -(Ms - Mvalue)**2 / (2 * M_sigma**2))
fr = np.where(Rs > 20 * km, -(Rs - 16 * km)**20, -(Rs - Rvalue)**2 / (2 * R_sigma**2))
f2 = (fm + fr).sum()
likelihood = f1 + f2
return likelihood
# Define log_posterior
def log_posterior(theta):
lp = log_prior(theta)
if not np.isfinite(lp):
return -np.inf
ll = log_likelihood(theta)
return lp + ll
Define the parameters to the Mass and Radius points
[2]:
# # Prepare data
# datas = [
# [10, 1.97, 0.5, 0.0985],
# [13, 1.44, 0.65, 0.072],
# [12, 2.07, 0.6, 0.1035]
# ]
# datas = pd.DataFrame(datas, index=[1,2,3], columns=['R', 'M', 'Rsigma', 'Msigma'])
# Rvalue = datas['R'].values * km
# Mvalue = datas['M'].values * Msun
# R_sigma = datas['Rsigma'].values * km
# M_sigma = datas['Msigma'].values * Msun
# f1 = np.log(1/(R_sigma*M_sigma*(np.sqrt(2*np.pi))**2)).sum()
# Define from_para_to_MR_points as in your original code
def MR_likelihood_Gaussian(theta, x, d1):
"""Compute the likelihood from MR measurement with Gaussian likelihood.
Args:
theta: Parameter set for the equation of state (EoS).
x: [Mvalue, Rvalue, Mwidth, Rwidth] - parameters of the Gaussian distribution.
d1: Sampled logarithmic density (log10 scale) for the measurement.
Returns:
likelihood (float): Log-likelihood value for the given parameter setup.
"""
Mvalue, Rvalue, Mwidth, Rwidth = x
sigma_M = Mwidth
sigma_R = Rwidth
# Compute pressure and energy density from the EoS using theta
pres_core = Polytrope.compute_EOS(eps_core, theta)[0]
eps_total = np.hstack((eps_crust, eps_core))
pres_total = np.hstack((pres_crust, pres_core))
# Ensure unique pressure values for interpolation
unique_pressure_indices = np.unique(pres_total, return_index=True)[1]
unique_pressure = pres_total[np.sort(unique_pressure_indices)]
# Create interpolating functions for EoS and its inverse
eos = interp1d(eps_total, pres_total, kind='cubic', fill_value='extrapolate')
inveos = interp1d(unique_pressure, eps_total[unique_pressure_indices], kind='cubic', fill_value='extrapolate')
if d1 == 0:
return -1e101 # Return a large negative value if density is zero
else:
density = 10**d1 # Convert log density to linear scale
try:
# Solve the TOV equations to get mass and radius for the given density
MR = solveTOV2(density, pres_total[20], eos, inveos) # Returns [mass, radius]
if MR[0] == 0 or MR[1] == 0:
return -1e101 # Return a large negative value if mass or radius is zero
else:
# Compute the Gaussian likelihood
fx = (1 / (sigma_M * sigma_R * (np.sqrt(2 * np.pi)) ** 2)) * np.exp(
-((MR[1] / km - Rvalue) ** 2) / (2 * sigma_R ** 2)
- ((MR[0] / Msun - Mvalue) ** 2) / (2 * sigma_M ** 2)
)
likelihood = np.log(fx)
except OverflowError:
return -1e101 # Return a large negative value if an overflow occurs
return likelihood if likelihood > -1e101 else -1e101
Run the EMCEE for MCMC Analysis
[ ]:
# Initialize emcee
ndim = len(parameters) # 10
nwalkers = 2 * ndim # 20
nsteps = 25000 # Adjust based on convergence
# Initialize walkers
initial_positions = []
for i in range(nwalkers):
initial_position = [
prior.flat_prior(0.04/0.16*e0, 0.065/0.16*e0, np.random.uniform(0, 1)),
prior.flat_prior(1.5*e0, 8.3*e0, np.random.uniform(0, 1)),
prior.flat_prior(1.5*e0, 8.3*e0, np.random.uniform(0, 1)),
prior.flat_prior(gamma_lo_ranges[0], gamma_up_ranges[0], np.random.uniform(0, 1)),
prior.flat_prior(gamma_lo_ranges[1], gamma_up_ranges[1], np.random.uniform(0, 1)),
prior.flat_prior(gamma_lo_ranges[2], gamma_up_ranges[2], np.random.uniform(0, 1)),
prior.flat_prior(gamma_lo_ranges[3], gamma_up_ranges[3], np.random.uniform(0, 1)),
np.random.uniform(0, 1),
np.random.uniform(0, 1),
np.random.uniform(0, 1)
]
initial_positions.append(initial_position)
initial_positions = np.array(initial_positions)
# Create the sampler
sampler = emcee.EnsembleSampler(nwalkers, ndim, log_posterior)
# Run the sampler
print("Running MCMC...")
sampler.run_mcmc(initial_positions, nsteps, progress=True)
print("MCMC completed.")
# Flatten the chain
flat_samples = sampler.get_chain(discard=1000, thin=15, flat=True)
# Save the samples
np.savetxt("emcee_samples.txt", flat_samples)
# Plot corner plot
fig = corner.corner(flat_samples, labels=parameters, show_titles=True, title_fmt=".2f")
plt.show()