{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import emcee\n",
    "import corner\n",
    "from astropy import cosmology\n",
    "\n",
    "# Sample MCMC from emcee\n",
    "\n",
    "def lnprob(x, ivar):\n",
    "    return -0.5 * np.sum(ivar * x ** 2)\n",
    "\n",
    "ndim, nwalkers = 10, 100\n",
    "ivar = 1. / np.random.rand(ndim)\n",
    "p0 = [np.random.rand(ndim) for i in range(nwalkers)]\n",
    "\n",
    "sampler = emcee.EnsembleSampler(nwalkers, ndim, lnprob, args=[ivar])\n",
    "sampler.run_mcmc(p0, 1000)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Some theory:\n",
    "\n",
    "$\\mu(t)=\\frac{F_{\\nu,obs}(t)}{F_{\\nu,gs}}=\\frac{u(t)^2 + 2}{u(t)\\sqrt{u(t)^2 + 4}} $\n",
    "\n",
    "where $u(t)$ is the angular distance between the source and the lens in units of the Einstein angle, and $F_{\\nu,obs}(t)$ and $F_{\\nu, gs}$ are the observed flux density at time $t$ and the mean flux density in the ground state, respectively."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Some basic astronomy definitions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cosmo = cosmology.default_cosmology.get()\n",
    "G = 6.6738e-11\n",
    "c = 299792458.\n",
    "M_star = 1.9891e30\n",
    "Mpc = 3.0857e19\n",
    "\n",
    "\n",
    "def magnification(u):\n",
    "  # Magnification due to gravitational microlensing (low optical depth, Paczynski 1986, ApJ, 304, 1)\n",
    "  mag = (u * u + 2.) / (u * np.sqrt(u * u + 4.))\n",
    "  return mag\n",
    "\n",
    "\n",
    "def einstein_angle(m, z_l, z_s):\n",
    "  # Einstein angle (in arcsec) for source (z1) and lens (z2) with mass (m) in M_star\n",
    "  d_s = cosmo.angular_diameter_distance(z_s)\n",
    "  d_l = cosmo.angular_diameter_distance(z_l)\n",
    "  d_ls = cosmo.angular_diameter_distance_z1z2(z_l, z_s)\n",
    "  theta_e = np.sqrt(4. * G * M_star * m * d_ls / (c * c * d_l * d_s * Mpc)) * 3600.\n",
    "  return theta_e.value\n",
    "\n",
    "\n",
    "def angular_distance(theta, z_s, t):\n",
    "  # Angular distance between source and lens as a function of time (days), u(t), in terms of\n",
    "  # minimum impact parameter, transverse velocity (km/s), and redshift of lens\n",
    "  # in units of Einstein angle per second\n",
    "  m_l, v_t, t_0, z_l, u_0 = theta\n",
    "  d_l = cosmo.angular_diameter_distance(z_l).value * Mpc\n",
    "  theta_e = einstein_angle(m_l, z_l, z_s)\n",
    "  u = np.sqrt(u_0 * u_0 + np.power(v_t * (t - t_0) * 3.1104e11 / (d_l * theta_e), 2))\n",
    "  return u\n",
    "\n",
    "\n",
    "def fluxratio(t, theta, z_s):\n",
    "  u = angular_distance(theta, z_s, t)\n",
    "  flux = magnification(u)\n",
    "  return flux"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now the priors:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def lnlike(theta, data, z_s):\n",
    "  m_l, v_t, t_0, z_l, u_0 = theta\n",
    "  # Constrain results (really done by priors)\n",
    "  t, m, merr = data\n",
    "  model = fluxratio(t, theta, z_s) \n",
    "  inv_sigma2 = 1.0 / (merr ** 2.) # + model ** 2.)\n",
    "  return -0.5 * (np.sum((m - model) ** 2. * inv_sigma2 - np.log(inv_sigma2)))\n",
    "\n",
    "\n",
    "def mlprior(m_l):\n",
    "  # Log-prior on the lens mass (from the mass-weighted Chabrier (2003) IMF)\n",
    "  f1 = lambda m, a, m_c, sig: a * np.exp(-((np.log10(m) - np.log10(m_c)) ** 2.) / (2. * sig * sig))\n",
    "  f2 = lambda m, a, x: a * np.power(m, -x)\n",
    "  if m_l < 0.072:\n",
    "    type = 'bd'\n",
    "  elif m_l > 0.072 and m_l <= 1.:\n",
    "    type = 'lms'\n",
    "  elif m_l > 1. and m_l <= 9.:\n",
    "    type = 'ims'\n",
    "  else:\n",
    "    type = 'hms'\n",
    "  disk = {'bd': 0.04, 'lms': 0.41, 'ims': 0.35, 'hms': 0.2, 'rem': 0.0}\n",
    "  spheroid = {'bd': 0.0, 'lms': 0.48, 'ims': 0.34, 'hms': 0.18, 'rem': 0.0}\n",
    "  if m_l > 1.e-3 and m_l < 0.7:\n",
    "    p = f1(m_l, 0.158, 0.079, 0.69) * disk[type] + f1(m_l, 3.6e-4, 0.22, 0.33) * spheroid[type]\n",
    "  elif m_l >= 0.7 and m_l < 1.:\n",
    "    p = f1(m_l, 0.158, 0.079, 0.69) * disk[type] + f2(m_l, 7.1e-5, 1.3) * spheroid[type]\n",
    "  elif m_l >= 1. and m_l < 100.:\n",
    "    p = f2(m_l, 0.044, 1.3) * disk[type] + f2(m_l, 7.1e-5, 1.3) * spheroid[type]\n",
    "  else:\n",
    "    return -np.inf\n",
    "  return np.log(p) # Chabrier gives p(log m) and p(m) = p(log m) / (ln 10 * m) \n",
    "  \n",
    "\n",
    "def vtprior(v_t):\n",
    "  # Log-prior on the transverse velocity\n",
    "  if np.abs(v_t) > 1000:\n",
    "    return -np.inf\n",
    "  else:\n",
    "    return 0.\n",
    "\n",
    "\n",
    "def t0prior(t_0):\n",
    "  # Log-prior on the time delay\n",
    "  return 0.\n",
    "\n",
    "\n",
    "def zlprior(z_l, z_s):\n",
    "  # Log-prior on the lens redshift\n",
    "  if z_s == 0.: z_s = 1.309 # Fiducial value\n",
    "  if z_l > 0. and z_l < z_s:\n",
    "    return 0.\n",
    "  else:\n",
    "    return -np.inf\n",
    "\n",
    "\n",
    "def u0prior(u_0):\n",
    "  # Log-prior on the minimum impact parameter\n",
    "  if u_0 < 0.063 or u_0 > 1.:\n",
    "    return -np.inf\n",
    "  else:\n",
    "    return 0.\n",
    "\n",
    "\n",
    "def lnprior(theta, z_s):\n",
    "  m_l, v_t, t_0, z_l, u_0 = theta\n",
    "  lp = mlprior(m_l) + vtprior(v_t) + t0prior(t_0) + zlprior(z_l, z_s) + u0prior(u_0)\n",
    "  return lp\n",
    "\n",
    "\n",
    "def lnprob(theta, data, z_s):\n",
    "  lp = lnprior(theta, z_s)\n",
    "  if not np.isfinite(lp):\n",
    "    return -np.inf\n",
    "  return lp + lnlike(theta, data, z_s)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Get the data from: http://www.astro.caltech.edu/~mjg/fitdata.dat"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Run the fit:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def runmcmc(lnpost, result, data, zs, plotmcmc = False):\n",
    "  # MCMC: lnpost is posterior log-likelihood, result is initial guess for params\n",
    "  # data is data, zs is source redshift\n",
    "  ndim, nwalkers = 5, 100\n",
    "  pos = [result + 1.e-4 * np.random.random(ndim) for i in xrange(nwalkers)]\n",
    "  sampler = emcee.EnsembleSampler(nwalkers, ndim, lnpost, args = (data, zs))\n",
    "  pos, prob, state = sampler.run_mcmc(pos,500) # burn-in\n",
    "  sampler.reset()\n",
    "  sampler.run_mcmc(pos, 1000)\n",
    "  # Plot MCMC\n",
    "  samples = sampler.chain[:, 50:, :].reshape((-1, ndim))\n",
    "  params = sampler.flatchain\n",
    "  probs = sampler.flatlnprobability\n",
    "  if plotmcmc:\n",
    "    fig = corner.corner(samples, labels = [\"$m_l$\", \"$v_t$\", \"$t_0$\", \"$z_l$\", \"$u_0$\"], truths = result)\n",
    "    plt.show()\n",
    "    fig.savefig(\"triangle.png\")\n",
    "  return samples, params, probs\n"
   ]
  }
 ],
 "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.6.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}