post qlens parameter derivation-corecusp- from k0, gamma, n, a, s, get mvir, c_hat_200, r_core and rm2.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Gravitational Lensing Research\n",
    "# Program to take a qlens-generated sample chain and derive alternative parameters\n",
    "# d5: This version takes the corecusp format (k0, a, s) and derives mass, NEW concentration \n",
    "# and core radius, (mvir, c_hat_200, r_core)\n",
    "# where c_hat is defined as r_vir divided by the radius at which the log slope is -2,\n",
    "# and r_core is defined as the radius at which the log slope is -1.\n",
    "# Uses python 3.6\n",
    "# n fixed to 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "FlatLambdaCDM(H0=70 km / (Mpc s), Om0=0.3, Tcmb0=0 K, Neff=3.04, m_nu=None, Ob0=None)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'%.5g'"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import time\n",
    "import matplotlib.pyplot as plt\n",
    "import sys\n",
    "import numpy as np\n",
    "import os\n",
    "import pandas as pd\n",
    "import timeit\n",
    "from IPython.core.debugger import set_trace\n",
    "import datetime\n",
    "from shutil import copyfile\n",
    "from astropy import constants as const\n",
    "from astropy.cosmology import FlatLambdaCDM\n",
    "import astropy.units as u\n",
    "cosmo = FlatLambdaCDM(H0=70 * u.km / u.s /u.Mpc, Om0=0.3)\n",
    "print(cosmo)\n",
    "from scipy.integrate import quad, dblquad\n",
    "from scipy.optimize import minimize, minimize_scalar\n",
    "from numba import jit\n",
    "%precision %.5g"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import multiprocessing as mp\n",
    "nproc = 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/Users/kevin/CloudStation/KEVIN/UCI/Research/gravlensing/Python\n"
     ]
    }
   ],
   "source": [
    "# The current working directory should CONTAIN the \"chains_...\" subdirectory for the run of interest.\n",
    "print(os.getcwd())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/Users/kevin/CloudStation/KEVIN/UCI/Research/gravlensing/A2537\n"
     ]
    }
   ],
   "source": [
    "label = 'A2537.cc.v9'\n",
    "newflag = False# Use this to account for the \".new\" in names.\n",
    "if newflag:\n",
    "    extra = \".new\"\n",
    "else:\n",
    "    extra = ''\n",
    "rootpath = '/Users/kevin/CloudStation/KEVIN/UCI/Research/gravlensing/A2537' \n",
    "# rootpath = '/home/kea/KEVIN/UCI/Research/gravlensing/A383'\n",
    "os.chdir(rootpath)\n",
    "print(os.getcwd())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# verify that the key files exist\n",
    "assert os.path.exists('chains_' + label)\n",
    "for ext in ['', '.nparam', '.paramnames', '.latex_paramnames', '.ranges']:\n",
    "    assert os.path.exists('chains_' + label + '/' + label + extra + ext)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "path = rootpath + '/chains_' + label\n",
    "label = label + extra"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create the new directory, if it doesn't already exist\n",
    "app = '_d7'\n",
    "newlabel = label + app\n",
    "if not os.path.exists(os.getcwd() + '/chains_' + newlabel):\n",
    "        os.makedirs(os.getcwd() + '/chains_' + newlabel)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# copy key files\n",
    "for ext in ['', '.nparam', '.paramnames', '.latex_paramnames', '.ranges']:\n",
    "    copyfile(path + '/' + label + ext, 'chains_' + newlabel + '/' + newlabel + ext) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create 'newpath', for convenience\n",
    "newpath = rootpath + '/chains_' + newlabel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# we are deriving scale radius and core radius\n",
    "addl_pnames = ['m_vir', 'c_hat_200', 'r_core', 'rm2']\n",
    "addl_latex_pnames = ['m_{vir}', '\\hat{c}_{200}', 'r_{core}', 'r_{-2}']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "zlens = 0.189\n",
    "zsrc = 2.0  # try to consistently use this"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "163.77280884946157 solMass / kpc3\n",
      "3.1556470486647226 kpc / arcsec\n",
      "rho_200 =  1029288.3583223283 solMass / arcsec3\n"
     ]
    }
   ],
   "source": [
    "rho_crit = cosmo.critical_density(zlens).to(u.M_sun / u.kpc**3)\n",
    "print(rho_crit)\n",
    "kpc_per_arcsec = cosmo.kpc_proper_per_arcmin(zlens) * u.arcmin / (60. * u.arcsec)\n",
    "print (kpc_per_arcsec)\n",
    "rho_200 = rho_crit * 200 * kpc_per_arcsec**3  # units of M_sun arcsec^-3\n",
    "print('rho_200 = ', rho_200)\n",
    "# get rid of units now\n",
    "rho_200 = rho_200.value"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$\\Sigma _{\\text{crit}}=\\frac{c^2 D_S}{4 \\pi  G D_L D_{\\text{SL}}}$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3003559185.4506283 solMass / kpc2\n"
     ]
    }
   ],
   "source": [
    "sigma_crit = (const.c**2 * cosmo.angular_diameter_distance(zsrc) / (4. * np.pi * const.G *\\\n",
    "                           cosmo.angular_diameter_distance(zlens) * cosmo.angular_diameter_distance_z1z2(zlens, zsrc))\\\n",
    "             ).to(u.M_sun / u.kpc**2)\n",
    "print(sigma_crit)   \n",
    "# get rid of units now\n",
    "sigma_crit = sigma_crit.value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "# read in the chain and append the new parameter names\n",
    "samples = np.loadtxt(newpath + '/' + newlabel, comments=\"#\", delimiter=None, unpack=False)\n",
    "with open(newpath + '/' + newlabel + '.paramnames') as afile:\n",
    "    paramnames = afile.readlines()\n",
    "with open(newpath + '/' + newlabel + '.latex_paramnames') as afile:\n",
    "    latex_paramnames = afile.readlines()\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['k0\\n', 'gamma\\n', 'a\\n', 's\\n', 'q\\n', 'theta\\n', 'xc\\n', 'yc\\n', 'shear1\\n', 'shear2\\n', 'mtot1\\n', 'mtot2\\n', 'raw_chisq\\n']\n",
      "['k0\\t\\\\kappa_{0}\\n', 'gamma\\t\\\\gamma\\n', 'a\\ta\\n', 's\\ts\\n', 'q\\tq\\n', 'theta\\t\\\\theta\\n', 'xc\\tx_{c}\\n', 'yc\\ty_{c}\\n', 'shear1\\t\\\\gamma_{1}\\n', 'shear2\\t\\\\gamma_{2}\\n', 'mtot1\\tM_{tot,1}\\n', 'mtot2\\tM_{tot,2}\\n', 'raw_chisq\\t\\\\chi^2\\n']\n"
     ]
    }
   ],
   "source": [
    "# special fix for Abell_611_v28.cc: change parameter name for perturber 2\n",
    "# for i in range(len(paramnames)):\n",
    "#     if paramnames[i].find('mass2d(100000)') != -1:\n",
    "#         paramnames[i] = 'mass_p2\\n'\n",
    "# with open(newpath + '/' + newlabel + '.paramnames', 'w') as afile:\n",
    "#     afile.write(''.join(paramnames) + '')\n",
    "    \n",
    "print(paramnames)\n",
    "    \n",
    "for i in range(len(latex_paramnames)):\n",
    "    if latex_paramnames[i].find('mass2d(100000)') != -1:\n",
    "        latex_paramnames[i] = 'mass_p2\\tM_{p2}\\n'\n",
    "        \n",
    "with open(newpath + '/' + newlabel + '.latex_paramnames', 'w') as afile:\n",
    "    afile.write(''.join(latex_paramnames) + '')\n",
    "\n",
    "print(latex_paramnames)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "# append the new parameter names\n",
    "for nm, latex_nm in zip(addl_pnames, addl_latex_pnames):\n",
    "    paramnames.append(nm + '\\n')\n",
    "    latex_paramnames.append(nm + '\\t' + latex_nm + '\\n')\n",
    "    with open(newpath + '/' + newlabel + '.paramnames', 'a') as afile:\n",
    "        afile.write(nm + '\\n')\n",
    "    with open(newpath + '/' + newlabel + '.latex_paramnames', 'a') as afile:\n",
    "        afile.write(nm + '\\t' + latex_nm + '\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/Users/kevin/CloudStation/KEVIN/UCI/Research/gravlensing/A2537/chains_A2537.cc.v9_d7\n"
     ]
    }
   ],
   "source": [
    "# change the working directory to the new path\n",
    "os.chdir(newpath)\n",
    "print(os.getcwd())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.0  derived parameters, adding  4\n"
     ]
    }
   ],
   "source": [
    "# increase the number of derived parameters in the .nparam file\n",
    "pnum, dpnum = np.loadtxt(newlabel + '.nparam', comments=\"#\", delimiter=None, unpack=False)\n",
    "addl_np = len(addl_pnames)\n",
    "print(dpnum, ' derived parameters, adding ', addl_np)\n",
    "dpnum += addl_np\n",
    "np.savetxt(newlabel + '.nparam', [np.int(pnum), np.int(dpnum)], fmt='%i')\n",
    "\n",
    "# add new ranges to the .ranges file. Not sure if these should be calculated, just using 0 and 1e30.\n",
    "with open(newlabel + '.ranges', 'a') as afile:\n",
    "    for nm in addl_pnames:\n",
    "        afile.write('0  1e30\\n')\n",
    "\n",
    "# Strip off the \\n character.\n",
    "paramnames = [item[:-1] for item in paramnames] \n",
    "latex_paramnames = [item[:-1] for item in latex_paramnames]  \n",
    "samples = np.array(samples)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(120367, 15)\n",
      "['k0', 'gamma', 'a', 's', 'q', 'theta', 'xc', 'yc', 'shear1', 'shear2', 'mtot1', 'mtot2', 'raw_chisq', 'm_vir', 'c_hat_200', 'r_core', 'rm2']\n"
     ]
    }
   ],
   "source": [
    "ns, p = samples.shape # note that 'p' here is two more than the number of parameters.\n",
    "# as it counts the first column (weight) and the last (chisquare).\n",
    "print(samples.shape)\n",
    "print(paramnames)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "# extend the samples to have addl_np additional columns\n",
    "samples2 = np.append(samples, np.zeros((ns, addl_np)), axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "# move the chisquare to be the last column\n",
    "samples2[:,-1] = samples2[:, p-1]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can use the equation for virial mass and solve for the virial radius:\n",
    "$m_{200} = \\frac{4}{3} \\pi r_{200}^3 \\rho_{200}$, and $r_s = r_{200} / c$,\n",
    "Therefore\n",
    "$$ r_s = (\\frac{3}{4\\pi})^{1/3} (m_{200} / (c  \\rho_{200}))^{1/3}$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.62035"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(3 / (np.pi * 4.))**(1./3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "# calculate the new parameters\n",
    "[k0, gamma, r_s, r_c, q] = [col for col in samples2[:, [1,2, 3, 4, 5]].T] \n",
    "# for the cc model with n fixed, we need cols [1,2, 3, 4, 5], which are k0, gamma, a1, s and q"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import mpmath\n",
    "# def AppellF1(a, b, c, d , x, y):\n",
    "#     return mpmath.hyper2d({'m+n':[a], 'm':[b], 'n':[c]}, {'m+n':[d]}, x, y)\n",
    "# # test\n",
    "# AppellF1(1,2,3,4,.25,.5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "# def m200_b(r200, k0, n, gamma, r_s, r_c):\n",
    "#     return (1.68793e9 * r200**3 * r_c**-gamma * r_s**(-1 + gamma) * k0\\\n",
    "#             * AppellF1(3/2., gamma/2., (n - gamma)/2., 5/2., -(r200**2/r_c**2), -(r200**2/r_s**2)))\n",
    "# # test\n",
    "# m200_b(1470.22, 1.1, 3., 1.,  50.*4.333,15.*4.333)\n",
    "# I abandoned this method because of convergence issues."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set tolerance for relative error in integration to be reported\n",
    "tol = 1.e-2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 177 µs, sys: 6 µs, total: 183 µs\n",
      "Wall time: 187 µs\n",
      "mass = 3.05655e+14\n"
     ]
    }
   ],
   "source": [
    "# trying numerical integral instead of AppellF1, due to convergence issues\n",
    "# This version has no correction for axis ratio.\n",
    "def m200_b1(r200_as, k0, n, gamma, r_s_as, r_c_as, dummy):  # note distances must be in arc seconds\n",
    "    r200 = r200_as * kpc_per_arcsec.value\n",
    "    r_c = r_c_as * kpc_per_arcsec.value\n",
    "    r_s = r_s_as * kpc_per_arcsec.value\n",
    "    res = quad(lambda r: 4. * np.pi * r**2 *  (k0 * sigma_crit *  r_s**n)\\\n",
    "                /( 2. * np.pi * r_s * (r**2 + r_c**2)**(gamma/2.) * (r**2 + r_s**2)**((n - gamma)/2.)), 0, r200)\n",
    "    if res[0] !=0.:\n",
    "        if res[0]> .01 and np.abs(res[1] / res[0]) > tol:\n",
    "            print('Warning: relative error is above threshold.')\n",
    "            print('Result = {0:1.6e}, Error = {1:1.6e}, relative error = {2:2%}'.format(res[0], res[1], res[1]/res[0]))\n",
    "    return res[0]\n",
    "# test\n",
    "%time test = m200_b1(339, 1.1, 3., 1., 50., 15., 5e99)\n",
    "print('mass = {0:1.5e}'.format(test))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 18.4 ms, sys: 1.35 ms, total: 19.8 ms\n",
      "Wall time: 19.3 ms\n",
      "mass = 3.05655e+14\n"
     ]
    }
   ],
   "source": [
    "# now trying numerical double integral incorporating axis ratio, to account for elliptical nature\n",
    "def m200_b2(r200_as, k0, n, gamma, r_s_as, r_c_as, q):  # note distances must be in arc seconds\n",
    "    r200 = r200_as * kpc_per_arcsec.value\n",
    "    r_c = r_c_as * kpc_per_arcsec.value\n",
    "    r_s = r_s_as * kpc_per_arcsec.value\n",
    "    res = dblquad(lambda r, th: r**2 * np.sin(th) *  (k0 * sigma_crit *  r_s**(n-1)\\\n",
    "                /( ((r * np.sin(th))**2 + r_c**2 + (r * np.cos(th))**2/ q**2)**(gamma/2.)\\\n",
    "                  * ((r * np.sin(th))**2 + r_s**2 + (r * np.cos(th))**2/ q**2)**((n - gamma)/2.)))\\\n",
    "                , 0, np.pi, lambda r: 0., lambda r: r200)\n",
    "    if res[0] !=0.:\n",
    "        if res[0]> .01 and np.abs(res[1] / res[0]) > tol:\n",
    "            print('Warning: error is above threshold.')\n",
    "            print('Result = {0:1.6e}, Error = {1:1.6e}, relative error = {2:2%}'.format(res[0], res[1], res[1]/res[0]))\n",
    "    return res[0]\n",
    "# test\n",
    "%time test = m200_b2(339, 1.1, 3., 1., 50., 15., 1)\n",
    "print('mass = {0:1.5e}'.format(test))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 287 µs, sys: 2 µs, total: 289 µs\n",
      "Wall time: 292 µs\n",
      "mass = 3.05655e+14\n"
     ]
    }
   ],
   "source": [
    "# now trying the approximation method\n",
    "def m200_b3(r200_as, k0, n, gamma, r_s_as, r_c_as, q):  # note distances must be in arc seconds\n",
    "    r200 = r200_as * kpc_per_arcsec.value \n",
    "    r_c = r_c_as * kpc_per_arcsec.value * np.sqrt(q)\n",
    "    r_s = r_s_as * kpc_per_arcsec.value * np.sqrt(q)\n",
    "    # Manoj thought a factor of q**(-1./3) helped here, but I don't see that.\n",
    "    res = quad(lambda r:  4. * np.pi * r**2 *  (k0 * sigma_crit *  r_s**n)\\\n",
    "                /( 2. * np.pi * r_s * (r**2 + r_c**2)**(gamma/2.) * (r**2 + r_s**2)**((n - gamma)/2.)), 0, r200)\n",
    "    if res[0] !=0.:\n",
    "        if res[0]> .01 and np.abs(res[1] / res[0]) > tol:\n",
    "            print('Warning: relative error is above threshold.')\n",
    "            print('Result = {0:1.6e}, Error = {1:1.6e}, relative error = {2:2%}'.format(res[0], res[1], res[1]/res[0]))\n",
    "    return res[0]\n",
    "# test\n",
    "%time test = m200_b3(339, 1.1, 3., 1., 50., 15., 1)\n",
    "print('mass = {0:1.5e}'.format(test))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1.6797e+14"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "4/3. * np.pi * 339**3. * rho_200"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "virial radius = 265.44, virial mass = 8.06321e+13\n",
      "CPU times: user 13.1 ms, sys: 370 µs, total: 13.5 ms\n",
      "Wall time: 13.4 ms\n",
      "virial radius = 433.71, virial mass = 3.51734e+14\n",
      "CPU times: user 6.18 ms, sys: 193 µs, total: 6.38 ms\n",
      "Wall time: 6.2 ms\n"
     ]
    }
   ],
   "source": [
    "# find the root where the masses are equal\n",
    "def find_virial(k0, n, gamma, r_s, r_c, q):  # r_c and r_s should be in arc seconds\n",
    "    r = minimize_scalar(lambda r200: np.abs(m200_b1(r200, k0, n, gamma, r_s, r_c, q) - 4/3. * np.pi * r200**3.\\\n",
    "                        * rho_200), method = 'bounded', bounds=(0, 5000))  # r200 in units of arc seconds\n",
    "    # according to Manoj, don't have to do any correction now.\n",
    "    if r.success!=True:\n",
    "        print('Optimize failure.')\n",
    "        print(r)\n",
    "        print('k0: {0:1.5f}  gamma: {1:1.5f}  r_s:{2:1.5f}  r_c:{3:1.5f}  q:{4:1.5f}'.format(k0, gamma, r_s, r_c, q))\n",
    "    return r.x, 4/3. * np.pi * r.x**3. * rho_200\n",
    "# test\n",
    "%time print('virial radius = {0:1.2f}, virial mass = {1:1.5e}'.format(*find_virial(1.34, 3., 1.858, 17.7, .67, .8)))\n",
    "%time print('virial radius = {0:1.2f}, virial mass = {1:1.5e}'.format(*find_virial(0.167, 3., 2.353, 133.656, 19.21, .8)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "-2"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Now need a function to find c_hat. First, the Corecusp log slope.\n",
    "def cclogslope(r, gamma, r_s, r_c): # assumes n=3 and rho_0 = 1.\n",
    "    r2 = r * r\n",
    "    rs2 = r_s * r_s\n",
    "    rc2 = r_c * r_c\n",
    "    return - (r2 * (3 * r2 - rc2 * (gamma-3) + gamma * rs2)) / ((r2 + rc2) * (r2 + rs2))\n",
    "# test:\n",
    "cclogslope(15, 1, 15, 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "71.152"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# A function to find r_{beta}, the radius at which the log slope is 'beta'.\n",
    "def rbeta(gamma, r_s, r_c, beta):\n",
    "    rb = minimize_scalar(lambda r: np.abs(beta - cclogslope(r, gamma, r_s, r_c)), method='bounded', bounds=(0,5000))\n",
    "    if rb.success !=True:\n",
    "        print('Optimize Failure')\n",
    "        print(rb)\n",
    "        print(gamma, r_s, r_c, beta)\n",
    "    return rb.x\n",
    "# test\n",
    "rbeta(1.61, 104.6, 12.7, -2) # should be 71.1519 per Mathematica"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "# non-parallel version, works\n",
    "mvir = []\n",
    "c_hat = []\n",
    "r_core = []\n",
    "rm2 = []\n",
    "ns = len(k0)\n",
    "n = np.ones(ns) * np.array([3.])  # n=3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Calculation is complete.\n"
     ]
    }
   ],
   "source": [
    "for i, (k0p, gammap, npr, r_cp, r_sp, qp) in enumerate(zip(k0, gamma, n, r_c, r_s, q)):\n",
    "    res = find_virial(k0p, npr, gammap, r_sp, r_cp, qp)\n",
    "    mvir.append(res[1])\n",
    "    rm2p = rbeta(gammap, r_sp, r_cp, -2)\n",
    "    rm2.append(rm2p)\n",
    "    c_hat.append(res[0] / rm2p)\n",
    "    r_core.append(rbeta(gammap, r_sp, r_cp, -1))\n",
    "    sys.stdout.write('\\r{0:5.1%} complete.'.format(float(i) / ns))\n",
    "        \n",
    "sys.stdout.write(\"\\rCalculation is complete.\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# beep\n",
    "duration = 0.2  # second\n",
    "freq = 440  # Hz\n",
    "os.system('play --no-show-progress --null --channels 1 synth %s sine %f' % (duration, freq))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# # parallel version - need to fix to give r_core, rm2\n",
    "# def virial_and_rbeta(k0, n, gamma, r_s, r_c, q, beta):\n",
    "#     rv, mv = find_virial(k0, n, gamma, r_s, r_c, q)\n",
    "#     rb = rbeta(gamma, r_s, r_c, beta)\n",
    "#     return rv, mv, rb\n",
    "# pool = mp.Pool(processes=nproc)\n",
    "# ns = len(k0)\n",
    "# mvir = np.zeros(ns)\n",
    "# rvir = np.zeros(ns)\n",
    "# c_hat = np.zeros(ns)\n",
    "# results = []\n",
    "# for i, (k0p, gammap, r_sp, r_cp, qp) in enumerate(zip(k0, gamma, r_s, r_c, q)):\n",
    "#     results.append((i, pool.apply_async(virial_and_rbeta, args=(k0p, 3., gammap, r_sp, r_cp, qp, -2,))))\n",
    "# print(\"Started.\")\n",
    "# pool.close()   # now nothing else can be added to the pool\n",
    "# output = []\n",
    "# tstart = timeit.default_timer()\n",
    "# old_i = -1\n",
    "# for (i,k) in results:\n",
    "#     output.append(k.get())\n",
    "#     (result0, result1, result2) = output[-1]\n",
    "#     mvir[i] = result1\n",
    "#     rvir[i] = result0\n",
    "#     c_hat[i] = result0 / result2\n",
    "#     if old_i +1 != i:\n",
    "#         print('Out of order at i=', i)\n",
    "#     old_i = i\n",
    "#     tavg = (timeit.default_timer() - tstart) / (i+1.)\n",
    "#     proj_fin =  tavg * (ns -i -1)\n",
    "#     sys.stdout.write(\"\\r{0:5.2%} complete. Avg time per iteration: {1:6.2f} sec. Projected finish in {2:4.2f} min.\"\\\n",
    "#                      .format((float(i) / ns), tavg,  proj_fin/60))\n",
    "# sys.stdout.write(\"\\rCalculation is complete.                     \\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "samples2[:, p-1:p-1+addl_np] = np.array([mvir, c_hat, r_core, rm2]).T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "446358542791823.7 14.894334222673423\n"
     ]
    }
   ],
   "source": [
    "# confirm the correct value for the last sample\n",
    "ind = -1\n",
    "vir_test = find_virial(samples2[ind, 1], 3., *samples2[ind,2:6])[1]\n",
    "if vir_test != samples2[ind, -4]:\n",
    "    print(vir_test, samples2[ind, -3])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([ 7.28790000e-06,  2.34949000e+00,  6.58374000e-01,  3.60380000e+01,\n",
       "        1.11954000e-04,  6.21670000e-01,  2.79472000e+01,  4.30655000e-01,\n",
       "        2.33500000e+00,  3.89054000e-03, -3.13350000e-02,  6.52295000e+12,\n",
       "        2.12861000e+12,  1.45178000e+01,  4.46358543e+14,  1.12488890e+01,\n",
       "        1.48943342e+01,  4.17423052e+01, -5.73624000e+00])"
      ]
     },
     "execution_count": 41,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "samples2[-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(120367, 19)"
      ]
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "samples2.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [],
   "source": [
    "# write the revised samples back to the file\n",
    "np.savetxt(newlabel, samples2, fmt=\"%.6g\", delimiter='   ')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Best fit chi squared is  7.72813\n"
     ]
    }
   ],
   "source": [
    "for i, name in enumerate(paramnames):\n",
    "    if name == 'raw_chisq':\n",
    "        print('Best fit chi squared is ', np.amin(samples2[:,i+1]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'A2537.cc.v9_d7'"
      ]
     },
     "execution_count": 45,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "newlabel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [
    {
     "ename": "SystemExit",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "An exception has occurred, use %tb to see the full traceback.\n",
      "\u001b[0;31mSystemExit\u001b[0m\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/kevin/anaconda/envs/py36/lib/python3.6/site-packages/IPython/core/interactiveshell.py:2918: UserWarning: To exit: use 'exit', 'quit', or Ctrl-D.\n",
      "  warn(\"To exit: use 'exit', 'quit', or Ctrl-D.\", stacklevel=1)\n"
     ]
    }
   ],
   "source": [
    "sys.exit()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [],
   "source": [
    "# If there is a transform that you want to make, and a dat file for conversion, specify it here\n",
    "transform =True\n",
    "t_file = 'v8_transform.dat'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Success\n"
     ]
    }
   ],
   "source": [
    "# make the error file using mkdist\n",
    "if transform:\n",
    "    code = os.system('mkdist ' + newlabel + ' -E -T:../' + t_file + ' | tee ' + newlabel + '.err')\n",
    "else:\n",
    "    code = os.system('mkdist ' + newlabel + ' | tee ' + newlabel + '.err')    \n",
    "if code != 0:\n",
    "    print('Error! Code returned ', code)\n",
    "else:\n",
    "    print('Success')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Success\n"
     ]
    }
   ],
   "source": [
    "# make the various distributions and plots\n",
    "# having trouble with 'ranges much smaller than actual ranges'\n",
    "# so try adding the option '-B5e-4'\n",
    "if transform:\n",
    "    code = os.system('mkdist ' + newlabel + ' -n40 -N40  -B5e-4 -T:../' + t_file + ' |tee tmp.txt')\n",
    "else:\n",
    "    code = os.system('mkdist ' + newlabel + ' -n40 -N40  -B5e-4'+ ' |tee tmp.txt')\n",
    "if code != 0:\n",
    "    print('Error! Code returned ', code)\n",
    "else:\n",
    "    with open('tmp.txt', 'r') as file:\n",
    "        tmp = file.readlines()\n",
    "    if 'WARNING' in tmp or 'Error:' in tmp:\n",
    "        print('Completed with WARNINGS or Error')\n",
    "    elif len(tmp) == 1:\n",
    "        print('Warning: output was only one line:')\n",
    "        print(tmp)\n",
    "    else:\n",
    "        print('Success')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sys.exit()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Success\n"
     ]
    }
   ],
   "source": [
    "# to run the python triangle plot pdf maker\n",
    "code = os.system('python ' + newlabel + '_tri.py')\n",
    "if code != 0:\n",
    "    print('Error! Code returned ', code)\n",
    "else:\n",
    "    print('Success')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "Python [conda env:py36]",
   "language": "python",
   "name": "conda-env-py36-py"
  },
  "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}