Strangeon EOS emcee workflow
Import Libraries
[ ]:
import math
import numpy as np
import emcee
import InferenceWorkflow.Likelihood as likelihood
import InferenceWorkflow.prior as prior
import EOSgenerators.Strangeon_EOS as Strangeon
from TOVsolver.constant import oneoverfm_MeV, m_rho, m_w, G, c
import TOVsolver.main as main
import pandas as pd
import corner
import matplotlib.pyplot as plt
Define parameter dimensions
Here, we define the parameters of interest that will be sampled: epsilon, ns, and d1.
ndim is the number of parameters, and Nq is defined for later use in the EOS calculations.
[2]:
# Define the number of dimensions based on parameters
parameters = ['epsilon', 'ns', 'd1']
ndim = len(parameters)
Nq = 18
Define the Prior Transform Function
The prior_transform function takes a unit cube sample (from the MCMC) and transforms it into the parameter space defined by epsilon, ns, and d1.
It computes the energy density and pressure from the EOS using the Strangeon.compute_EOS function based on the parameters, validating the EOS.
The function returns the parameters with d1 scaled based on the maximum density computed.
[ ]:
# Existing prior_transform function (unchanged)
def prior_transform(cube):
params = cube.copy()
params[0] = prior.flat_prior(10, 170, cube[0]) # epsilon=10-170MeV
params[1] = prior.flat_prior(0.17, 0.36, cube[1]) # ns=0.17-0.36 fm^-3
epsilon = params[0]
ns = params[1]
theta1 = np.array([Nq, epsilon, ns])
n_min = 3 * theta1[2] / theta1[0]
n_max = 0.16 * 8 * 3 / theta1[0]
n_values = np.linspace(n_min, n_max, 2000)
energy_density, pressure = Strangeon.compute_EOS(n_values, theta1)
eps_total = energy_density
pres_total = pressure
RFSU2R = []
MFSU2R = []
density = np.logspace(14.3, 15.6, 50)
if all(x < y for x, y in zip(eps_total[:], eps_total[1:])) and all(x < y for x, y in zip(pres_total[:], pres_total[1:])):
MR = main.OutputMR('', energy_density, pressure).T
else:
MR = []
if len(MR) == 0:
params[2] = 0
else:
for i in range(len(MR[1])):
RFSU2R.append(MR[0][i])
MFSU2R.append(MR[1][i])
if i > 20 and MR[1][i] - MR[1][i-1] < 0:
break
if len(MFSU2R) == 0:
params[2] = 0
else:
max_index = len(MFSU2R)
max_d = np.log10(density[max_index - 1])
params[2] = 14.3 + (max_d - 14.3) * (cube[2]) # cube[2] is the third parameter in the transformed space
return params
Define likelihood transform
This function computes the likelihood of observing data given the parameters.
It uses the EOS to compute energy density and pressure, and then evaluates the likelihood using a Gaussian likelihood function based on mass-radius observations.
The likelihood value is returned.
[3]:
# Existing likelihood_transform function (unchanged)
def likelihood_transform(theta):
epsilon, ns, d1 = theta
theta1 = np.array([Nq, epsilon, ns])
n_min = 3 * theta1[2] / theta1[0]
n_max = 0.16 * 8 * 3 / theta1[0]
n_values = np.linspace(n_min, n_max, 2000)
energy_density, pressure = Strangeon.compute_EOS(n_values, theta1)
probMRgaussian = likelihood.MRlikihood_Gaussian(energy_density, pressure, (1.4, 13, 0.07, 0.65), d1)
return probMRgaussian
Define Log-Prior and Log-Likelihood Functions
log_prior evaluates the prior probability of the parameters. It checks if the parameters fall within their defined ranges. If they do, it computes max_d to validate d1, returning 0.0 for valid parameters or -1e101 for invalid ones.
log_likelihood computes the log likelihood based on the likelihood_transform function, returning -inf if the likelihood is non-positive.
[4]:
# Define log_prior based on prior_transform
def log_prior(theta):
epsilon, ns, d1 = theta
# Check epsilon and ns within their flat prior ranges
if not (10 <= epsilon <= 170):
return -np.inf
if not (0.17 <= ns <= 0.36):
return -1e101
# Compute max_d based on epsilon and ns
theta1 = np.array([Nq, epsilon, ns])
n_min = 3 * theta1[2] / theta1[0]
n_max = 0.16 * 8 * 3 / theta1[0]
density = np.logspace(14.3, 15.6, 50)
n_values = np.linspace(n_min, n_max, 2000)
energy_density, pressure = Strangeon.compute_EOS(n_values, theta1)
if not (np.all(np.diff(energy_density) > 0) and np.all(np.diff(pressure) > 0)):
return -1e101 # Invalid EOS
MR = main.OutputMR('', energy_density, pressure).T
if len(MR) == 0:
return -1e101 # Invalid MR
# Extract MFSU2R to determine max_d
MFSU2R = []
for i in range(len(MR[1])):
MFSU2R.append(MR[1][i])
if i > 20 and MR[1][i] - MR[1][i-1] < 0:
break
if len(MFSU2R) == 0:
return -1e101 # Invalid MFSU2R
max_index = len(MFSU2R)
max_d = np.log10(density[max_index - 1])
# Validate d1
if not (14.3 <= d1 <= max_d):
return -1e101
return 0.0 # Uniform prior log(1)
# Define log_likelihood based on likelihood_transform
def log_likelihood(theta):
prob = likelihood_transform(theta)
if prob <= 0:
return -1e101
return np.log(prob)
Define Combined Log Probability Function
The log_prob function combines the log prior and log likelihood, returning the total log probability for a set of parameters. If either the prior or likelihood is invalid, it returns -1e101.
[ ]:
# Define the combined log probability
def log_prob(theta):
lp = log_prior(theta)
if not np.isfinite(lp):
return -np.inf
ll = log_likelihood(theta)
if not np.isfinite(ll):
return -np.inf
return lp + ll
Initialize Walkers for MCMC Sampling
This block initializes the MCMC walkers, where nwalkers is set to 32 (generally 2-4 times the number of dimensions).
The while loop samples the initial positions using the prior transform until enough valid walkers are collected. It ensures each sample corresponds to a valid EOS and MR configuration.
If not enough valid walkers are found, a ValueError is raised.
[ ]:
# Initialize the walkers
nwalkers = 32 # Typically 2-4 times the number of dimensions
initial = []
np.random.seed(42) # For reproducibility
# Sample initial positions uniformly within the prior bounds
while len(initial) < nwalkers:
# Sample from the prior
cube = np.random.uniform(0, 1, ndim)
params = prior_transform(cube)
epsilon, ns, d1 = params
# Check if the d1 is within the allowed range
theta1 = np.array([Nq, epsilon, ns])
n_min = 3 * theta1[2] / theta1[0]
n_max = 0.16 * 8 * 3 / theta1[0]
density = np.logspace(14.3, 15.6, 50)
n_values = np.linspace(n_min, n_max, 2000)
energy_density, pressure = Strangeon.compute_EOS(n_values, theta1)
if not (np.all(np.diff(energy_density) > 0) and np.all(np.diff(pressure) > 0)):
continue # Skip invalid initial positions
MR = main.OutputMR('', energy_density, pressure).T
if len(MR) == 0:
continue
MFSU2R = []
for i in range(len(MR[1])):
MFSU2R.append(MR[1][i])
if i > 20 and MR[1][i] - MR[1][i-1] < 0:
break
if len(MFSU2R) == 0:
continue
max_index = len(MFSU2R)
max_d = np.log10(density[max_index - 1])
# Validate d1 and prepare for initial position
if max_d == 14.3:
continue # Avoid division by zero
cube2 = (d1 - 14.3) / (max_d - 14.3)
if not (0 <= cube2 <= 1):
continue # d1 not within the allowed transformed range
initial.append([epsilon, ns, d1])
# Ensure we have enough initial walkers
if len(initial) < nwalkers:
raise ValueError("Not enough valid initial walkers. Please check your prior and EOS computations.")
initial = np.array(initial[:nwalkers])
# Set up the sampler
sampler_emcee = emcee.EnsembleSampler(nwalkers, ndim, log_prob)
# Run the sampler for 50 steps
print("Running emcee sampler for 50 steps...")
sampler_emcee.run_mcmc(initial, 50, progress=True)
# Extract the samples
samples = sampler_emcee.get_chain(flat=True)
# Optional: Save the samples to a file
df_samples = pd.DataFrame(samples, columns=parameters)
df_samples.to_csv("emcee_samples_50_steps.csv", index=False)
print("Sampling complete. Samples saved to 'emcee_samples_50_steps.csv'.")
Plotting use Corner
This block creates a corner plot to visualize the posterior distributions of the sampled parameters using the corner library. The labels for the axes are specified, and the plot is saved as corner_plot.png before being displayed.
[ ]:
# Create a corner plot of the samples
fig = corner.corner(samples, labels=parameters, show_titles=True, title_kwargs={"fontsize": 12})
# Save the corner plot to a file
plt.savefig("corner_plot.png")
plt.show() # Display the plot