{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Strangeon EOS emcee workflow"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Import Libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import numpy as np\n",
    "import emcee\n",
    "import InferenceWorkflow.Likelihood as likelihood\n",
    "import InferenceWorkflow.prior as prior\n",
    "import EOSgenerators.Strangeon_EOS as Strangeon\n",
    "from TOVsolver.constant import oneoverfm_MeV, m_rho, m_w, G, c\n",
    "import TOVsolver.main as main\n",
    "import pandas as pd\n",
    "import corner\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define parameter dimensions\n",
    "\n",
    "\n",
    "1. Here, we define the parameters of interest that will be sampled: epsilon, ns, and d1.\n",
    "\n",
    "2. ndim is the number of parameters, and Nq is defined for later use in the EOS calculations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the number of dimensions based on parameters\n",
    "parameters = ['epsilon', 'ns', 'd1']\n",
    "ndim = len(parameters)\n",
    "Nq = 18"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define the Prior Transform Function\n",
    "\n",
    "1. 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.\n",
    "\n",
    "2. It computes the energy density and pressure from the EOS using the Strangeon.compute_EOS function based on the parameters, validating the EOS.\n",
    "\n",
    "3. The function returns the parameters with d1 scaled based on the maximum density computed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Existing prior_transform function (unchanged)\n",
    "def prior_transform(cube):\n",
    "    params = cube.copy()\n",
    "    params[0] = prior.flat_prior(10, 170, cube[0])   # epsilon=10-170MeV\n",
    "    params[1] = prior.flat_prior(0.17, 0.36, cube[1])  # ns=0.17-0.36 fm^-3    \n",
    "    \n",
    "    epsilon = params[0]\n",
    "    ns = params[1]\n",
    "\n",
    "    theta1 = np.array([Nq, epsilon, ns])\n",
    "    n_min = 3 * theta1[2] / theta1[0]  \n",
    "    n_max = 0.16 * 8 * 3 / theta1[0]   \n",
    "    n_values = np.linspace(n_min, n_max, 2000) \n",
    "    energy_density, pressure = Strangeon.compute_EOS(n_values, theta1)\n",
    "\n",
    "    eps_total = energy_density\n",
    "    pres_total = pressure\n",
    "\n",
    "    RFSU2R = [] \n",
    "    MFSU2R = []\n",
    "    density = np.logspace(14.3, 15.6, 50) \n",
    "    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:])):\n",
    "        MR = main.OutputMR('', energy_density, pressure).T  \n",
    "    else:\n",
    "        MR = []\n",
    "    if len(MR) == 0: \n",
    "        params[2] = 0\n",
    "    else:\n",
    "        for i in range(len(MR[1])):\n",
    "            RFSU2R.append(MR[0][i])\n",
    "            MFSU2R.append(MR[1][i])   \n",
    "            if i > 20 and MR[1][i] - MR[1][i-1] < 0:    \n",
    "                break\n",
    "    if len(MFSU2R) == 0:\n",
    "        params[2] = 0\n",
    "    else:\n",
    "        max_index = len(MFSU2R)\n",
    "        max_d = np.log10(density[max_index - 1])\n",
    "        params[2] = 14.3 + (max_d - 14.3) * (cube[2])  # cube[2] is the third parameter in the transformed space\n",
    "    return params"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define likelihood transform \n",
    "\n",
    "1. This function computes the likelihood of observing data given the parameters.\n",
    "\n",
    "2. 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.\n",
    "\n",
    "3. The likelihood value is returned."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Existing likelihood_transform function (unchanged)\n",
    "def likelihood_transform(theta):\n",
    "    epsilon, ns, d1 = theta \n",
    "    \n",
    "    theta1 = np.array([Nq, epsilon, ns])\n",
    "    n_min = 3 * theta1[2] / theta1[0]  \n",
    "    n_max = 0.16 * 8 * 3 / theta1[0]   \n",
    "    n_values = np.linspace(n_min, n_max, 2000) \n",
    "    energy_density, pressure = Strangeon.compute_EOS(n_values, theta1)\n",
    "    \n",
    "    probMRgaussian = likelihood.MRlikihood_Gaussian(energy_density, pressure, (1.4, 13, 0.07, 0.65), d1)\n",
    "    \n",
    "    return probMRgaussian"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define Log-Prior and Log-Likelihood Functions\n",
    "\n",
    "1. 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.\n",
    "\n",
    "2. log_likelihood computes the log likelihood based on the likelihood_transform function, returning -inf if the likelihood is non-positive."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define log_prior based on prior_transform\n",
    "def log_prior(theta):\n",
    "    epsilon, ns, d1 = theta\n",
    "\n",
    "    # Check epsilon and ns within their flat prior ranges\n",
    "    if not (10 <= epsilon <= 170):\n",
    "        return -np.inf\n",
    "    if not (0.17 <= ns <= 0.36):\n",
    "        return -1e101\n",
    "\n",
    "    # Compute max_d based on epsilon and ns\n",
    "    theta1 = np.array([Nq, epsilon, ns])\n",
    "    n_min = 3 * theta1[2] / theta1[0]  \n",
    "    n_max = 0.16 * 8 * 3 / theta1[0]   \n",
    "    density = np.logspace(14.3, 15.6, 50)\n",
    "    n_values = np.linspace(n_min, n_max, 2000)\n",
    "    energy_density, pressure = Strangeon.compute_EOS(n_values, theta1)\n",
    "\n",
    "    if not (np.all(np.diff(energy_density) > 0) and np.all(np.diff(pressure) > 0)):\n",
    "        return -1e101  # Invalid EOS\n",
    "\n",
    "    MR = main.OutputMR('', energy_density, pressure).T\n",
    "\n",
    "    if len(MR) == 0:\n",
    "        return -1e101  # Invalid MR\n",
    "\n",
    "    # Extract MFSU2R to determine max_d\n",
    "    MFSU2R = []\n",
    "    for i in range(len(MR[1])):\n",
    "        MFSU2R.append(MR[1][i])\n",
    "        if i > 20 and MR[1][i] - MR[1][i-1] < 0:\n",
    "            break\n",
    "\n",
    "    if len(MFSU2R) == 0:\n",
    "        return -1e101  # Invalid MFSU2R\n",
    "\n",
    "    max_index = len(MFSU2R)\n",
    "    max_d = np.log10(density[max_index - 1])\n",
    "\n",
    "    # Validate d1\n",
    "    if not (14.3 <= d1 <= max_d):\n",
    "        return -1e101\n",
    "\n",
    "    return 0.0  # Uniform prior log(1)\n",
    "\n",
    "# Define log_likelihood based on likelihood_transform\n",
    "def log_likelihood(theta):\n",
    "    prob = likelihood_transform(theta)\n",
    "    if prob <= 0:\n",
    "        return -1e101\n",
    "    return np.log(prob)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define Combined Log Probability Function\n",
    "\n",
    "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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the combined log probability\n",
    "def log_prob(theta):\n",
    "    lp = log_prior(theta)\n",
    "    if not np.isfinite(lp):\n",
    "        return -np.inf\n",
    "    ll = log_likelihood(theta)\n",
    "    if not np.isfinite(ll):\n",
    "        return -np.inf\n",
    "    return lp + ll"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Initialize Walkers for MCMC Sampling\n",
    "\n",
    "1. This block initializes the MCMC walkers, where nwalkers is set to 32 (generally 2-4 times the number of dimensions).\n",
    "\n",
    "2. 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.\n",
    "\n",
    "3. If not enough valid walkers are found, a ValueError is raised."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize the walkers\n",
    "nwalkers = 32  # Typically 2-4 times the number of dimensions\n",
    "initial = []\n",
    "np.random.seed(42)  # For reproducibility\n",
    "\n",
    "# Sample initial positions uniformly within the prior bounds\n",
    "while len(initial) < nwalkers:\n",
    "    # Sample from the prior\n",
    "    cube = np.random.uniform(0, 1, ndim)\n",
    "    params = prior_transform(cube)\n",
    "    epsilon, ns, d1 = params\n",
    "\n",
    "    # Check if the d1 is within the allowed range\n",
    "    theta1 = np.array([Nq, epsilon, ns])\n",
    "    n_min = 3 * theta1[2] / theta1[0]  \n",
    "    n_max = 0.16 * 8 * 3 / theta1[0]   \n",
    "    density = np.logspace(14.3, 15.6, 50)\n",
    "    n_values = np.linspace(n_min, n_max, 2000)\n",
    "    energy_density, pressure = Strangeon.compute_EOS(n_values, theta1)\n",
    "\n",
    "    if not (np.all(np.diff(energy_density) > 0) and np.all(np.diff(pressure) > 0)):\n",
    "        continue  # Skip invalid initial positions\n",
    "\n",
    "    MR = main.OutputMR('', energy_density, pressure).T\n",
    "    if len(MR) == 0:\n",
    "        continue\n",
    "\n",
    "    MFSU2R = []\n",
    "    for i in range(len(MR[1])):\n",
    "        MFSU2R.append(MR[1][i])\n",
    "        if i > 20 and MR[1][i] - MR[1][i-1] < 0:\n",
    "            break\n",
    "    if len(MFSU2R) == 0:\n",
    "        continue\n",
    "\n",
    "    max_index = len(MFSU2R)\n",
    "    max_d = np.log10(density[max_index - 1])\n",
    "\n",
    "    # Validate d1 and prepare for initial position\n",
    "    if max_d == 14.3:\n",
    "        continue  # Avoid division by zero\n",
    "\n",
    "    cube2 = (d1 - 14.3) / (max_d - 14.3)\n",
    "    if not (0 <= cube2 <= 1):\n",
    "        continue  # d1 not within the allowed transformed range\n",
    "\n",
    "    initial.append([epsilon, ns, d1])\n",
    "\n",
    "# Ensure we have enough initial walkers\n",
    "if len(initial) < nwalkers:\n",
    "    raise ValueError(\"Not enough valid initial walkers. Please check your prior and EOS computations.\")\n",
    "\n",
    "initial = np.array(initial[:nwalkers])\n",
    "\n",
    "# Set up the sampler\n",
    "sampler_emcee = emcee.EnsembleSampler(nwalkers, ndim, log_prob)\n",
    "\n",
    "# Run the sampler for 50 steps\n",
    "print(\"Running emcee sampler for 50 steps...\")\n",
    "sampler_emcee.run_mcmc(initial, 50, progress=True)\n",
    "\n",
    "# Extract the samples\n",
    "samples = sampler_emcee.get_chain(flat=True)\n",
    "\n",
    "# Optional: Save the samples to a file\n",
    "df_samples = pd.DataFrame(samples, columns=parameters)\n",
    "df_samples.to_csv(\"emcee_samples_50_steps.csv\", index=False)\n",
    "\n",
    "print(\"Sampling complete. Samples saved to 'emcee_samples_50_steps.csv'.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plotting use Corner\n",
    "\n",
    "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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a corner plot of the samples\n",
    "fig = corner.corner(samples, labels=parameters, show_titles=True, title_kwargs={\"fontsize\": 12})\n",
    "\n",
    "# Save the corner plot to a file\n",
    "plt.savefig(\"corner_plot.png\")\n",
    "plt.show()  # Display the plot"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}