pegs_tutorial.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "PEGS @ MMCW\n",
    "----\n",
    "Tom Killestein (t.killestein@warwick.ac.uk)\n",
    "\n",
    "This short tutorial is a step-by-step guide to producing your own ephemeris for Sco X-1,\n",
    "using a small subset of the VLT/UVES spectra from the real PEGS IV analysis.\n",
    "\n",
    "Some links:\n",
    "[PEGS IV on arXiv](https://arxiv.org/abs/2302.00018)\n",
    "[VLT/UVES instrument page](https://www.eso.org/sci/facilities/paranal/instruments/uves.html)"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "!pip install -r requirements.txt\n",
    "import os, glob\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import pandas as pd\n",
    "from scipy.optimize import least_squares\n",
    "\n",
    "# configure JAX for use - 64-bit precision is important\n",
    "from jax.config import config\n",
    "import corner\n",
    "\n",
    "config.update(\"jax_enable_x64\", True)\n",
    "import jax.numpy as jnp\n",
    "from jax import jacfwd, random\n",
    "import numpyro\n",
    "import numpyro.distributions as dist\n",
    "\n",
    "# astropy to handle necessary time/velocity conversions\n",
    "from astropy.io import fits\n",
    "from astropy.time import Time\n",
    "from astropy.coordinates import SkyCoord, EarthLocation\n",
    "import astropy.units as u\n",
    "from astropy.coordinates import solar_system_ephemeris\n",
    "\n",
    "solar_system_ephemeris.set(\"jpl\")\n",
    "\n",
    "# for progress bars\n",
    "from tqdm.notebook import tqdm\n",
    "\n",
    "# tell numpyro to use 64-bit mode\n",
    "numpyro.set_host_device_count(\n",
    "    8\n",
    ")  # set to n_cpus - 1 ideally, although beware IO-bottlenecked at high cores.\n",
    "numpyro.enable_x64()"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Configure plots to have nice defaults\n",
    "plt.rcParams.update({\n",
    "    # \"text.usetex\": True,\n",
    "    # \"font.family\": \"serif\",\n",
    "    # \"font.serif\": [\"Times\"],\n",
    "    \"xtick.direction\": \"in\",\n",
    "    \"xtick.top\": True,\n",
    "    \"xtick.minor.visible\": True,\n",
    "    \"ytick.right\": True,\n",
    "    \"ytick.direction\": \"in\",\n",
    "    \"ytick.minor.visible\": True\n",
    "})"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Utility functions"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# abstracted out of scipy.optimize.minpack - prefer to use least_squares over curve_fit hence need this\n",
    "# https://github.com/scipy/scipy/blob/47bb6febaa10658c72962b9615d5d5aa2513fa3a/scipy/optimize/minpack.py#L810-L816\n",
    "def cov_from_jacobian(jac: jnp.array):\n",
    "    _, s, VT = jnp.linalg.svd(jac, full_matrices=False)\n",
    "    threshold = jnp.finfo(float).eps * max(jac.shape) * s[0]\n",
    "    s = s[s > threshold]\n",
    "    VT = VT[: s.size]\n",
    "    pcov = jnp.dot(VT.T / s**2, VT)\n",
    "\n",
    "    return pcov\n",
    "\n",
    "\n",
    "def doppler_corr(wav: float, vel: float):\n",
    "    return wav * (1 + vel / 3e8)"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Define the Bowen line model"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "def circular_orbit(t: float, pars: jnp.array):\n",
    "    t0, P, gamma, K = pars\n",
    "    return K * jnp.sin(2 * jnp.pi * (t - t0) / P) + gamma\n",
    "\n",
    "\n",
    "def circ_orbit_resids(pars, x, y, yerr):\n",
    "    return (y - circular_orbit(x, pars)) / yerr\n",
    "\n",
    "\n",
    "jac_circ_orbit_resids = jacfwd(circ_orbit_resids)\n",
    "\n",
    "\n",
    "def circ_orbit_chisq(pars, x, y, yerr):\n",
    "    return jnp.sum(jnp.square(circ_orbit_resids(pars, x, y, yerr)))\n",
    "\n",
    "\n",
    "# Bowen model - line centres taken from NIST line list.\n",
    "line_centres = jnp.array(\n",
    "    [\n",
    "        4640.0,  # broad\n",
    "        4634.13,  # NIII\n",
    "        4640.64,  # NIII\n",
    "        4647.42,  # CIII\n",
    "        4650.25,  # CIII\n",
    "        4643.386,  # OII\n",
    "    ]\n",
    ")\n",
    "\n",
    "# FWHMs of broad components and emission lines (m/s)\n",
    "line_fwhms = jnp.array([1250e3, 50.0e3, 50.0e3, 50.0e3, 50.0e3, 50.0e3])\n",
    "\n",
    "# immediately convert into Gaussian sigmas\n",
    "line_sigmas = line_fwhms * line_centres / 3e8 / (2 * jnp.sqrt(2 * jnp.log(2)))\n",
    "\n",
    "\n",
    "def gaussian(x, a, mean, sigma_g):\n",
    "    return a * jnp.exp(-((x - mean) ** 2) / (2 * sigma_g**2))\n",
    "\n",
    "\n",
    "def bowen_model(\n",
    "    x, params, line_centres_list=line_centres, line_sigmas_list=line_sigmas\n",
    "):\n",
    "    vel, bl_amp, bl_cen, nl_amp1, nl_amp2, nl_amp3, nl_amp4, nl_amp5 = params\n",
    "    model = jnp.zeros_like(x)\n",
    "\n",
    "    # manually specify, can write a nice loop later\n",
    "    model = (\n",
    "        gaussian(x, bl_amp, bl_cen, line_sigmas_list[0])\n",
    "        + gaussian(\n",
    "            x, nl_amp1, doppler_corr(line_centres_list[1], vel), line_sigmas_list[1]\n",
    "        )\n",
    "        + gaussian(\n",
    "            x, nl_amp2, doppler_corr(line_centres_list[2], vel), line_sigmas_list[2]\n",
    "        )\n",
    "        + gaussian(\n",
    "            x, nl_amp3, doppler_corr(line_centres_list[3], vel), line_sigmas_list[3]\n",
    "        )\n",
    "        + gaussian(\n",
    "            x, nl_amp4, doppler_corr(line_centres_list[4], vel), line_sigmas_list[4]\n",
    "        )\n",
    "        + gaussian(\n",
    "            x, nl_amp5, doppler_corr(line_centres_list[5], vel), line_sigmas_list[5]\n",
    "        )\n",
    "    )\n",
    "\n",
    "    return model\n",
    "\n",
    "\n",
    "def bowen_residual_vector(params, x, y, yerr):\n",
    "    return (y - bowen_model(x, params)) / yerr\n",
    "\n",
    "\n",
    "# Housekeeping to prepare the jacobians for our residual function\n",
    "jac_bowen_residual_vector = jacfwd(bowen_residual_vector)"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# data ingest and parse\n",
    "storage_loc = \"data/\"\n",
    "framelist = glob.glob(os.path.join(storage_loc, \"*.fits\"))\n",
    "\n",
    "# UVES has two arms - one red and one blue. We only want the BLUE arm for this analysis\n",
    "blueframes = [\n",
    "    f\n",
    "    for f in framelist\n",
    "    if (fits.getval(f, \"HIERARCH ESO INS PATH\") == \"BLUE\")\n",
    "    & (fits.getval(f, \"NCOMBINE\") == 1)\n",
    "]\n",
    "print(f\"{len(blueframes)} of {len(framelist)} frames valid - ingesting\")\n",
    "\n",
    "spectra_table = []\n",
    "\n",
    "# The observatory and Sco X-1 don't  move (we hope!), so let's pre-compute some properties.\n",
    "summary_header = fits.getheader(blueframes[0])\n",
    "obsloc = (\n",
    "    summary_header[\"HIERARCH ESO TEL GEOLON\"],\n",
    "    summary_header[\"HIERARCH ESO TEL GEOLAT\"],\n",
    "    summary_header[\"HIERARCH ESO TEL GEOELEV\"],\n",
    ")\n",
    "obsloc = EarthLocation(*obsloc)\n",
    "skyloc = SkyCoord(summary_header[\"RA\"], summary_header[\"DEC\"], unit=\"deg\", frame=\"fk5\")\n",
    "\n",
    "for f in tqdm(blueframes, total=len(blueframes)):\n",
    "    hdul = fits.open(f)\n",
    "    obs_epoch = Time(hdul[0].header[\"DATE-OBS\"])\n",
    "    jd = obs_epoch.jd\n",
    "\n",
    "    exptime = hdul[0].header[\"EXPTIME\"]\n",
    "    jd += 0.5 * exptime / 86400  # move to midpoint of exposure\n",
    "    barycor = skyloc.radial_velocity_correction(\n",
    "        kind=\"barycentric\", location=obsloc, obstime=obs_epoch\n",
    "    )\n",
    "    snr = hdul[0].header[\"SNR\"]\n",
    "\n",
    "    # explicit coercion to 64-bit floats here\n",
    "    wav = hdul[1].data[\"WAVE\"].flatten().astype(float)\n",
    "    flux = hdul[1].data[\"FLUX\"].flatten().astype(float)\n",
    "    flux_err = hdul[1].data[\"ERR\"].flatten().astype(float)\n",
    "\n",
    "    hdul.close()\n",
    "\n",
    "    spectra_table.append(\n",
    "        (f, jd, wav, flux, flux_err, barycor.to(u.m / u.s).value, snr, exptime)\n",
    "    )\n",
    "\n",
    "spectra_table = pd.DataFrame(\n",
    "    spectra_table,\n",
    "    columns=[\"filename\", \"jd\", \"wav\", \"flux\", \"flux_err\", \"barycor\", \"snr\", \"exptime\"],\n",
    ")\n",
    "\n",
    "# Apply heliocentric/barycentric correction, but save for later so we can undo if needed\n",
    "obs_epochs = Time(spectra_table.jd, format=\"jd\", location=obsloc)\n",
    "bary_tt = obs_epochs.light_travel_time(skyloc, \"barycentric\", obsloc)\n",
    "spectra_table[\"barycor_dt\"] = bary_tt.jd\n",
    "obs_epochs += bary_tt\n",
    "\n",
    "# overwrite times\n",
    "spectra_table[\"jd\"] = obs_epochs.utc.jd\n",
    "\n",
    "# Coarse ephemeris estimate\n",
    "init_ephemeris = [2455444.7284, 0.7873143, -113.7e3, 74.6e3]\n",
    "spectra_table[\"expected_vel\"] = (\n",
    "    circular_orbit(spectra_table.jd.values, init_ephemeris)\n",
    "    - spectra_table.barycor.values\n",
    ")\n",
    "\n",
    "spectra_table = spectra_table.sort_values(\"jd\")"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "bowen_roi = (4605, 4675)\n",
    "continuum_regions = [(4605, 4615), (4666, 4675)]\n",
    "plot_wav = np.linspace(*bowen_roi, 1000)\n",
    "\n",
    "for i, spectrum in tqdm(spectra_table.iterrows(), total=len(spectra_table)):\n",
    "    # build fitting masks\n",
    "    wavmask = np.logical_and.reduce(\n",
    "        [spectrum.wav > bowen_roi[0], spectrum.wav < bowen_roi[1]]\n",
    "    )\n",
    "    cont1mask = np.logical_and.reduce(\n",
    "        [spectrum.wav > continuum_regions[0][0], spectrum.wav < continuum_regions[0][1]]\n",
    "    )\n",
    "    cont2mask = np.logical_and.reduce(\n",
    "        [spectrum.wav > continuum_regions[1][0], spectrum.wav < continuum_regions[1][1]]\n",
    "    )\n",
    "\n",
    "    contmask = cont1mask | cont2mask\n",
    "\n",
    "    # Fit a low-order Chebyshev polynomial to the continuum to remove background flux\n",
    "    bg_poly = np.polynomial.Chebyshev.fit(\n",
    "        spectrum.wav[contmask],\n",
    "        spectrum.flux[contmask],\n",
    "        w=1 / spectrum.flux_err[contmask],\n",
    "        deg=3,\n",
    "    )\n",
    "\n",
    "    # ideally don't share pixels between line and continuum fitting\n",
    "    fitregion_mask = ~contmask & wavmask\n",
    "\n",
    "    detrend_wav = spectrum.wav[fitregion_mask]\n",
    "    detrend_flux = spectrum.flux[fitregion_mask] - bg_poly(spectrum.wav[fitregion_mask])\n",
    "    detrend_fluxerr = spectrum.flux_err[fitregion_mask]\n",
    "\n",
    "    # Fit bounds to avoid spurious solutions\n",
    "    constraints = (\n",
    "        (-np.inf, 0, -np.inf, 0, 0, 0, 0, 0),\n",
    "        (0, np.inf, np.inf, np.inf, np.inf, np.inf, np.inf, np.inf),\n",
    "    )\n",
    "\n",
    "    maxval = np.max(detrend_flux)\n",
    "\n",
    "    # Some sensible initialisations for the fits\n",
    "    test_pars = [\n",
    "        spectrum.expected_vel,\n",
    "        maxval / 2,\n",
    "        4642,\n",
    "        maxval,\n",
    "        maxval,\n",
    "        maxval,\n",
    "        maxval,\n",
    "        maxval,\n",
    "    ]\n",
    "\n",
    "    res = least_squares(\n",
    "        bowen_residual_vector,\n",
    "        test_pars,\n",
    "        args=(detrend_wav, detrend_flux, detrend_fluxerr),\n",
    "        bounds=constraints,\n",
    "        jac=jac_bowen_residual_vector,\n",
    "        method=\"trf\",\n",
    "    )\n",
    "\n",
    "    # can recycle the optimal res.fun here since we're using normed resids\n",
    "    rchsq = np.sum(np.square(res.fun)) / (len(res.fun) - len(res.x))\n",
    "\n",
    "    cov = cov_from_jacobian(res.jac)\n",
    "    res_errs = np.sqrt(np.diag(cov) * rchsq)\n",
    "\n",
    "    # Let's save our fitted velocity and errors, and apply our barycentric corrections\n",
    "    # here\n",
    "    spectra_table.at[i, \"vel\"] = res.x[0] + spectrum.barycor\n",
    "    spectra_table.at[i, \"vel_err\"] = np.sqrt(\n",
    "        res_errs[0] ** 2 + 200**2\n",
    "    )  # ensure we get non-zero residuals, ~= systematic\n",
    "\n",
    "    # As a diagnostic, plot every 5th spectrum to check things are working as expected\n",
    "    if not (i % 5):\n",
    "        plt.plot(spectrum.wav, spectrum.flux - bg_poly(spectrum.wav))\n",
    "        plt.plot(plot_wav, bowen_model(plot_wav, res.x))\n",
    "        plt.xlim(bowen_roi[0] - 10, bowen_roi[1] + 10)\n",
    "        plt.axvline(bowen_roi[0], ls=\"--\", c=\"k\")\n",
    "        plt.axvline(bowen_roi[1], ls=\"--\", c=\"k\")\n",
    "        plt.xlabel(\"Wavelength $(\\AA)$\")\n",
    "        plt.ylabel(\"Flux\")\n",
    "        plt.ylim(-20, spectrum.flux[wavmask].max() + 50 - spectrum.flux[wavmask].min())\n",
    "        plt.axhline(0, c=\"k\", ls=\"--\")\n",
    "\n",
    "        for c in continuum_regions:\n",
    "            plt.fill_betweenx(\n",
    "                [0, 1200], [c[0], c[0]], [c[1], c[1]], color=\"r\", alpha=0.2\n",
    "            )\n",
    "\n",
    "        # plt.plot(plot_wav, bowen_model(plot_wav, test_pars))\n",
    "        plt.show()"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Fitting complete!\n",
    "---\n",
    "We now have the key data needed for the ephemeris - radial velocities of the donor in Sco X-1, and estimates of uncertainty."
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "spectra_table"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "plt.errorbar(spectra_table.jd.values, spectra_table.vel.values, yerr=spectra_table.vel_err.values, fmt='.k')\n",
    "plt.xlabel(\"Time (JD)\")\n",
    "plt.ylabel(\"Velocity (m/s)\")"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Determining the orbital parameters and ephemeris\n",
    "===="
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# local tweak to ensure we're in the right region of parameter space\n",
    "local_polish = least_squares(\n",
    "    circ_orbit_resids,\n",
    "    x0=[2455444.7284, 0.7873143, -113.7e3, 74.6e3],\n",
    "    args=(\n",
    "        spectra_table.jd.values,\n",
    "        spectra_table.vel.values,\n",
    "        spectra_table.vel_err.values\n",
    "    ),\n",
    "    jac=jac_circ_orbit_resids,\n",
    "    method=\"trf\",\n",
    "    x_scale=\"jac\",\n",
    ")\n",
    "\n",
    "# Phase on the best-fit ephemeris\n",
    "t0_init, P_init, gam_init, K_init = local_polish.x\n",
    "init_phase = ((spectra_table.jd.values - t0_init) / P_init) % 1\n",
    "\n",
    "times_plot = np.linspace(spectra_table.jd.min(), spectra_table.jd.max(), 10000)\n",
    "plt.errorbar(spectra_table.jd, spectra_table.vel, spectra_table.vel_err, fmt='.k')\n",
    "plt.plot(times_plot, circular_orbit(times_plot, local_polish.x))\n",
    "plt.show()\n",
    "\n",
    "plt.errorbar(init_phase, spectra_table.vel, yerr=spectra_table.vel_err, fmt='.k')\n",
    "plt.xlabel(\"Phase\")\n",
    "plt.ylabel(\"Velocity (m/s\")\n",
    "plot_phase = np.linspace(0, 1, 1000)\n",
    "folded_model = circular_orbit(plot_phase, [0, 1, gam_init, K_init])\n",
    "plt.plot(plot_phase, folded_model)\n",
    "plt.show()\n",
    "\n",
    "plt.hist(local_polish.fun, bins=\"auto\")\n",
    "plt.xlabel(\"Normalised residual\")\n",
    "plt.ylabel(\"N\")\n",
    "plt.show()\n",
    "print(local_polish.x)\n",
    "\n",
    "alldata_ephem = local_polish.x\n",
    "print(\n",
    "    np.sum(np.square(local_polish.fun)) / (len(local_polish.fun) - len(local_polish.x))\n",
    ")"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Parameter inference with MCMC\n",
    "-----\n"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# To avoid multimodality:\n",
    "tmin = t0_init - 3 * P_init\n",
    "tmax = t0_init + 3 * P_init\n",
    "\n",
    "# Definie the probabilistic model + priors\n",
    "def joint_circular_orbit_generative_model(x, y, yerr):\n",
    "    T0 = numpyro.sample(\"T0\", dist.Uniform(tmin, tmax))\n",
    "    P = numpyro.sample(\"P\", dist.Uniform(0.7, 0.8))\n",
    "    gamma = numpyro.sample(\"gamma\", dist.Uniform(-200e3, 0))\n",
    "    K = numpyro.sample(\"K\", dist.Uniform(0, 100e3))\n",
    "\n",
    "    efac = numpyro.sample(f\"efac\", dist.Uniform(0.5, 10))\n",
    "\n",
    "    pars = [T0, P, gamma, K]\n",
    "    model = circular_orbit(x, pars)\n",
    "\n",
    "    obs = numpyro.sample(\"obs\", dist.Normal(model, yerr*efac), obs=y)\n",
    "    return obs"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "rng_key = random.PRNGKey(0)\n",
    "rng_key, _rng_key = random.split(rng_key)\n",
    "\n",
    "# start in right region of parameter space\n",
    "paramnames = [\"T0\", \"P\", \"gamma\", \"K\"]\n",
    "init_location = {n: v for n, v in zip(paramnames, local_polish.x)}\n",
    "\n",
    "\n",
    "init_strategy = numpyro.infer.init_to_value(values=init_location)\n",
    "\n",
    "kernel = numpyro.infer.NUTS(\n",
    "    joint_circular_orbit_generative_model, init_strategy=init_strategy, dense_mass=True\n",
    ")\n",
    "\n",
    "mcmc = numpyro.infer.MCMC(\n",
    "    kernel,\n",
    "    num_warmup=3000,\n",
    "    num_samples=1000,\n",
    "    chain_method=\"parallel\",\n",
    "    num_chains=1,\n",
    "    jit_model_args=True,\n",
    ")\n",
    "\n",
    "mcmc.run(\n",
    "    rng_key=rng_key,\n",
    "    x=spectra_table.jd.values,\n",
    "    y=spectra_table.vel.values,\n",
    "    yerr=spectra_table.vel_err.values,\n",
    "    extra_fields=[\n",
    "        \"energy\",\n",
    "        \"potential_energy\",\n",
    "        \"diverging\",\n",
    "        \"accept_prob\",\n",
    "        \"num_steps\",\n",
    "    ],\n",
    ")\n",
    "\n",
    "mcmc.print_summary()"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "samples = mcmc.get_samples()\n",
    "\n",
    "# trace plots\n",
    "fig, ax = plt.subplots(len(samples), 1, dpi=120, sharex=True)\n",
    "plt.subplots_adjust(hspace=0)\n",
    "\n",
    "for i, (name, s) in enumerate(samples.items()):\n",
    "    ax[i].plot(s)\n",
    "    ax[i].set_ylabel(name.replace(\"_\", \"\\_\"))\n",
    "\n",
    "ax[-1].set_xlabel(\"Steps\")\n",
    "plt.show()\n",
    "\n",
    "chainwise_samples = mcmc.get_samples(group_by_chain=True)\n",
    "chainwise_diag_fields = mcmc.get_extra_fields(group_by_chain=True)\n",
    "\n",
    "chain_marginal_energy = chainwise_diag_fields[\"energy\"]\n",
    "e_bfmi = np.mean(np.square(np.diff(chain_marginal_energy, axis=1)), axis=1) / np.var(\n",
    "    chain_marginal_energy, axis=1, ddof=1\n",
    ")\n",
    "print(e_bfmi)\n",
    "dE = np.diff(chain_marginal_energy, axis=1) - np.diff(\n",
    "    chain_marginal_energy, axis=1\n",
    ").mean(axis=1, keepdims=True)\n",
    "E = chain_marginal_energy - np.mean(chain_marginal_energy, axis=1, keepdims=True)\n",
    "\n",
    "fig = plt.figure(dpi=120)\n",
    "fig.patch.set_facecolor(\"white\")\n",
    "\n",
    "plt.hist(dE.flatten(), bins=30, label=\"Transition energy\", range=(-20, 20))\n",
    "plt.hist(E.flatten(), bins=30, label=\"Marginal energy\", range=(-20, 20), alpha=0.5)\n",
    "plt.xlabel(\"$E$ - $<E>$\")\n",
    "plt.legend(edgecolor=None)\n",
    "plt.show()"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "samples_name_mapper = {\n",
    "    \"K\": \"$K$\",\n",
    "    \"P\": \"$P$\",\n",
    "    \"T0\": \"$T_0$\",\n",
    "    \"efac\": \"$\\varepsilon_\\mathrm{VLT}$\",\n",
    "    \"gamma\": \"$\\gamma$\",\n",
    "}\n",
    "\n",
    "test_plot_samples =    {\n",
    "        samples_name_mapper[key]: value[0][:1000]\n",
    "        for key, value in chainwise_samples.items()\n",
    "    }\n",
    "samples_rectified = {samples_name_mapper[key]: value for key, value in samples.items()}\n",
    "\n",
    "fig = plt.figure(dpi=300, figsize=(7.05556, 1.2 * 7.05556), clear=True)\n",
    "\n",
    "fig = corner.corner(\n",
    "    test_plot_samples,\n",
    "    use_math_text=True,\n",
    "    color=\"C0\",\n",
    "    label_kwargs=dict(fontsize=10),\n",
    "    max_n_ticks=3,\n",
    "    labelpad=0.1,\n",
    "    fig=fig,\n",
    ")\n",
    "plt.show()"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Minimising covariance between $T_0$ and $P$\n",
    "-----"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "ncycles = np.arange(-2000, 2000)[:, np.newaxis]\n",
    "\n",
    "T_n = samples['T0'] + ncycles*samples['P']\n",
    "\n",
    "cycle_offset_corr = [np.abs(np.corrcoef(T, samples['P'])[0][1]) for T in T_n]"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1, dpi=120)\n",
    "fig.patch.set_facecolor('white')\n",
    "ax.plot(ncycles, cycle_offset_corr, label='R')\n",
    "# ax.plot(ncycles, [np.abs(np.cov(T, samples['P'])[0][1]) for T in T_n], label=\"cov($T_0, P$)\")\n",
    "plt.yscale(\"log\")\n",
    "plt.xlabel(\"Cycles from least-squares $T_0$\")\n",
    "plt.ylabel(\"Magnitude\")\n",
    "plt.legend(frameon=False)"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "T0_new_samples = T_n[np.argmin(cycle_offset_corr)]\n",
    "Tasc_new_samples = T0_new_samples - 0.25*samples['P']"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "plt.hist(Tasc_new_samples, bins='auto', label=f\"RMS: {np.std(Tasc_new_samples)*86400:.2f}s\")\n",
    "plt.title(\"$T_{asc}$ from PEGS workshop\")\n",
    "plt.legend(frameon=False)\n",
    "plt.show()\n",
    "\n",
    "plt.hist(samples['P']*86400, bins='auto', label=f\"RMS: {np.std(samples['P'])*86400:.2f}s\")\n",
    "plt.legend(frameon=False)\n",
    "plt.show()"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [],
   "metadata": {
    "collapsed": false
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}