Broadband_tomography_with_castor.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Broadband Tomography with CASTOR\n",
    "\n",
    "This script implements a forecast for a broadband tomographic measurement of \n",
    "the UV background with the Cosmological Advanced Survey Telescope for \n",
    "Optical-UV Research (CASTOR). \n",
    "\n",
    "Broadband tomography (arxiv:) applies clustering redshift estimation (arxiv:) to derive\n",
    "a bias weighted distribution of photon emission that is, in principle, insensitive to \n",
    "foreground emission from galactic or solar system components. \n",
    "\n",
    "This script is organized as follows. It begins by importing the CASTOR filter curves, \n",
    "an emissivity model is then defined, followed by a forward model for the observed frame \n",
    "quantity is implemented. With this forward model, given an error model, we then define \n",
    "a data likelihood and conclude with an MCMC routine to generate estimates of the \n",
    "recovery of the posterior distributions for the input parameters. \n",
    "\n",
    "This script can be imported and delivers functions for modeling a broadband tomographic survey. \n",
    "Modifying the filter passbands allows one to model other surveys. \n",
    "\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import statements\n",
    "\n",
    "import numpy as np\n",
    "from scipy import interpolate\n",
    "from scipy.stats import norm\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# Emcee sampling, multiprocessing and analysis\n",
    "import time\n",
    "import corner\n",
    "from multiprocess import Pool\n",
    "\n",
    "# define background cosmology using the colossus package\n",
    "# Set to Planck 2018 best fit\n",
    "\n",
    "from colossus.cosmology import cosmology\n",
    "\n",
    "cosmo = cosmology.setCosmology('planck18')\n",
    "\n",
    "# define physical constants \n",
    "\n",
    "speed_of_light = 3*10**18 #in angstroms/s\n",
    "nu1500 = speed_of_light/1500\n",
    "nu1216 = speed_of_light/1216\n",
    "nu912 = speed_of_light/912\n",
    "\n",
    "# define observed frame wavelength and frequency range\n",
    "\n",
    "obs_waves = np.linspace(300, 6000, 6000-299) #note, this range more than the castor filter set. \n",
    "\n",
    "nu_obs = speed_of_light/obs_waves\n",
    "inverse_nu_obs = 1.0/nu_obs\n",
    "\n",
    "# import  sampler\n",
    "\n",
    "import emcee\n",
    "\n",
    "#set a global path variable to passband files, e.g. for running on a cluster\n",
    "\n",
    "#path = \"/rhome/bscot002/bigdata/passbands/\"\n",
    "path = \"passbands/\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Read in filter curves and interpolate in wavelength and frequency. \n",
    "\n",
    "class ThroughputCurve:\n",
    "    \"\"\"\n",
    "    A class representing a throughput curve.\n",
    "\n",
    "    ...\n",
    "\n",
    "    Attributes\n",
    "    ----------\n",
    "    None\n",
    "\n",
    "    Methods\n",
    "    -------\n",
    "    None \n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, dataframe):\n",
    "        \"\"\"\n",
    "        Parameters\n",
    "        ----------\n",
    "        waves : pandas dataframe\n",
    "            The wavelength in microns (galex) or angstroms (castor) which are related by a factor of 1e5\n",
    "        throughput : pandas dataframe\n",
    "            The unnormalized filter throughput, e.g. transmission fraction\n",
    "\n",
    "        \"\"\"\n",
    "        self.waves = dataframe[0]\n",
    "        self.throughput = dataframe[1]\n",
    "\n",
    "        \n",
    "# define names of filter curve files\n",
    "        \n",
    "castor_names = ['passband_castor_uv.csv', \n",
    "                'passband_castor_uv_dark.csv', \n",
    "                'passband_castor_u.csv',\n",
    "                'passband_castor_u_wide.csv', \n",
    "                'passband_castor_g.csv']\n",
    "\n",
    "galex_names = ['galex1500.res.txt', \n",
    "               'galex2500.res.txt']\n",
    "\n",
    "throughput_curve_names = castor_names + galex_names\n",
    "\n",
    "# and read filter curves to throughput curve object\n",
    "\n",
    "filter_waves = ThroughputCurve(pd.read_csv(path+castor_names[1], \n",
    "                                                      delim_whitespace=True, skiprows=0, header=None)).waves\n",
    "\n",
    "throughputs = np.zeros((5, 1001))\n",
    "for i,f in enumerate(castor_names): \n",
    "       throughputs[i,:] = ThroughputCurve(pd.read_csv(path+f, \n",
    "                                                      delim_whitespace=True, skiprows=0, header=None)).throughput\n",
    "\n",
    "castor_uv_throughput = throughputs[0,:]\n",
    "castor_uv_waves = filter_waves\n",
    "\n",
    "castor_uv = interpolate.interp1d(np.array(castor_uv_waves)*10000, np.array(castor_uv_throughput), bounds_error = False, fill_value = 0)\n",
    "castor_uv_dark = interpolate.interp1d(np.array(filter_waves)*10000, np.array(throughputs[1,:]), bounds_error = False, fill_value = 0)\n",
    "castor_u = interpolate.interp1d(np.array(filter_waves)*10000, np.array(throughputs[2,:]), bounds_error = False, fill_value = 0)\n",
    "castor_u_wide = interpolate.interp1d(np.array(filter_waves)*10000, np.array(throughputs[3,:]), bounds_error = False, fill_value = 0)\n",
    "castor_g = interpolate.interp1d(np.array(filter_waves)*10000, np.array(throughputs[4,:]), bounds_error = False, fill_value = 0)\n",
    "\n",
    "\n",
    "throughputs_galex_fuv = ThroughputCurve(pd.read_csv(path+galex_names[0], \n",
    "                                                      delim_whitespace=True, skiprows=0, header=None)).throughput\n",
    "        \n",
    "filter_waves_fuv = ThroughputCurve(pd.read_csv(path+galex_names[0], \n",
    "                                                      delim_whitespace=True, skiprows=0, header=None)).waves\n",
    "\n",
    "galex_fuv = interpolate.interp1d(np.array(filter_waves_fuv), np.array(throughputs_galex_fuv), bounds_error = False, fill_value = 0)\n",
    "\n",
    "throughputs_galex_nuv = ThroughputCurve(pd.read_csv(path+galex_names[1], \n",
    "                                                      delim_whitespace=True, skiprows=0, header=None)).throughput\n",
    "        \n",
    "filter_waves_nuv = ThroughputCurve(pd.read_csv(path+galex_names[1], \n",
    "                                                      delim_whitespace=True, skiprows=0, header=None)).waves\n",
    "\n",
    "galex_nuv = interpolate.interp1d(np.array(filter_waves_nuv), np.array(throughputs_galex_nuv), bounds_error = False, fill_value = 0)\n",
    "\n",
    "filters = np.array([castor_uv, castor_uv_dark, castor_u, castor_u_wide, castor_g, galex_fuv, galex_nuv])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# this cell initializes the parameter vector for later\n",
    "# use with the MCMC routine\n",
    "\n",
    "log_e1500z0_b1500z0 = 25.13  # note this degeneracy needs to be resolved\n",
    "gammae1500 = 2.06\n",
    "alpha1500z0 = -0.08\n",
    "Calpha1500 = 1.85\n",
    "alpha1100z0 = -3.71\n",
    "Calpha1100 = 0.50 \n",
    "EWLyAz03 = -6.17\n",
    "EWLyAz1 = 88.02\n",
    "gammabnu = -0.86\n",
    "gammabz = 0.79\n",
    "\n",
    "#normalizations\n",
    "\n",
    "b1500z0 = 0.32 #from \"breaking the degeneracy section\" --> follow up on this! \n",
    "e1500z0 = (10**(log_e1500z0_b1500z0))/b1500z0\n",
    "\n",
    "#assumed values\n",
    "\n",
    "alpha900 = -1.5\n",
    "logfLyCz1 = -0.53 \n",
    "logfLyCz2 = -0.84 \n",
    "\n",
    "fLyCz1 = np.exp(logfLyCz1)\n",
    "fLyCz2 = np.exp(logfLyCz2)\n",
    "\n",
    "theta =  gammabnu, gammabz, log_e1500z0_b1500z0, gammae1500, alpha1500z0, Calpha1500, alpha1100z0, Calpha1100, EWLyAz1, logfLyCz1, logfLyCz2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def CLyC(logfLyCz2 = -0.84, logfLyCz1 = -0.53):\n",
    "    \"\"\"Returns CLyC parameter used to model the evolution of the LyC escape fraction. \n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    logfLyCz2 : log10 of the LyC escape fraction at z=2, fiducial = -0.84\n",
    "    logfLyCz1 : log10 of the LyC escape fraction at z=1, fiducial = -0.53\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    float\n",
    "        the value of the CLyC parameter used to evaluate equation 7\n",
    "    \"\"\"\n",
    "    return (logfLyCz2 - logfLyCz1)/np.log10((1+2)/(1+1))\n",
    "\n",
    "# def logfLyC(z, CLyC, fLyCz2, fLyCz1):\n",
    "#     return CLyC(fLyCz2, fLyCz1) * np.log10((1+z)/(1+1)) + np.log10(fLyCz2)\n",
    "\n",
    "def logfLyC(z, logfLyCz2 = -0.84, logfLyCz1 = -0.53):\n",
    "    \"\"\" Returns the log of the Lyman-Continuum escape fraction evolutions with redshift \n",
    "    as a function of observed values at redshifts 1 and 2.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    z : The emission redshift \n",
    "    logfLyCz2 : log10 of the LyC escape fraction at z=2, fiducial = -0.84\n",
    "    logfLyCz1 : log10 of the LyC escape fraction at z=1, fiducial = -0.53\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    float\n",
    "        The value of the f_LyC parameter defined in equation 7\n",
    "    \"\"\"\n",
    "\n",
    "    return CLyC(logfLyCz2, logfLyCz1) * np.log10((1+z)/(1+1)) + logfLyCz2\n",
    "\n",
    "#Lyman alpha parameters \n",
    "\n",
    "def CLyA(EWLyAz1 = 88.02, EWLyAz03 = -6.17):\n",
    "    \n",
    "    \"\"\" Returns the log of the Lyman-Continuum escape fraction evolutions with redshift \n",
    "    as a function of observed values at redshifts 1 and 2.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    EWLyAz1 : The equivalent width of Ly-alpha at z = 1, fiducial = 88.02\n",
    "    EWLyA03 : The equivalent width of Ly-alpha at z = 0.3, fiducial = -6.17\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    float\n",
    "        the value of the CLyA parameter defined in order to evaluate equation 10\n",
    "    \"\"\"\n",
    "        \n",
    "    return (EWLyAz1 - EWLyAz03)/np.log10((1+1)/(1+0.3))\n",
    "\n",
    "def EWLyA(z, EWLyAz1 = 88.02, EWLyAz03 = -6.17):\n",
    "    \"\"\" Returns the equivalent width of Lyman-alpha as a function of redshift\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    z : The emission redshift\n",
    "    EWLyAz1 : The equivalent width of Ly-alpha at z = 1, fiducial = 88.02\n",
    "    EWLyA03 : The equivalent width of Ly-alpha at z = 0.3, fiducial = -6.17\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    float\n",
    "        the value of the Lyman alpha equivalent width as a function of redshift, see equation 10\n",
    "    \"\"\"\n",
    "        \n",
    "    return CLyA(EWLyAz1, EWLyAz03) * np.log10((1+z)/(1+0.3)) + EWLyAz03\n",
    "\n",
    "def e1500_evol(z, gammae1500 = 2.06):\n",
    "    \"\"\" Returns evolution of the 1500 Angstrom emissivity as a function of redshift. \n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    z : The emission redshift\n",
    "    gammae1500 : The (power law) evolution of the emission with redshift, fiducial = 2.06\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    float\n",
    "        The evolution of the emission at 1500 Angstrom as a function of redshift as defined in equation 12. \n",
    "    \"\"\"\n",
    "    \n",
    "    return (1+z)**(gammae1500)\n",
    "\n",
    "def alpha1500(z, Calpha1500=1.85, alpha1500z0=-0.08):\n",
    "    \"\"\" Returns evolution of the 1500 Angstrom slope as a function of redshift. \n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    z : The emission redshift\n",
    "    alpha1500z0 : The (power law) evolution of the emission with redshift, fiducial = -0.08\n",
    "    Calpha1500 : Normalization of the evolution with redshift, fiducial = 1.85\n",
    "    \n",
    "    Returns\n",
    "    -------\n",
    "    float\n",
    "        The evolution of the slope at 1500 Angstrom as a function of redshift as defined in equation 8. \n",
    "    \"\"\"\n",
    "    \n",
    "    return alpha1500z0 + Calpha1500*np.log10(1+z)\n",
    "\n",
    "def alpha1100(z, Calpha1100 = 0.50, alpha1100z0 = -3.71):\n",
    "    \"\"\" Returns evolution of the 1100 Angstrom slope as a function of redshift. \n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    z : The emission redshift\n",
    "    alpha1100z0 : The (power law) evolution of the emission with redshift, fiducial = -3.71\n",
    "    Calpha1100 : Normalization of the evolution with redshift, fiducial = 0.50\n",
    "    \n",
    "    Returns\n",
    "    -------\n",
    "    float\n",
    "        The evolution of the slope at 1100 Angstrom as a function of redshift as defined in equation 8. \n",
    "    \"\"\"\n",
    "    \n",
    "    #switched log10 --> log\n",
    "    return alpha1100z0 + Calpha1100*np.log10(1+z)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# begin by interpolating the filter curves to this wavelength range\n",
    "\n",
    "#nuv\n",
    "\n",
    "throughputs_galex_nuv = ThroughputCurve(pd.read_csv(path+galex_names[1], \n",
    "                                                      delim_whitespace=True, skiprows=0, header=None)).throughput\n",
    "        \n",
    "filter_waves_nuv = ThroughputCurve(pd.read_csv(path+galex_names[1], \n",
    "                                                      delim_whitespace=True, skiprows=0, header=None)).waves \n",
    "\n",
    "#filter waves is in angstroms, speed of light in angstroms/second, so this is in Hz \n",
    "\n",
    "filter_freqs_nuv = speed_of_light/filter_waves_nuv \n",
    "\n",
    "#interpolate\n",
    "\n",
    "nuv_freq = interpolate.interp1d(filter_freqs_nuv, throughputs_galex_nuv, bounds_error = False, fill_value = 0)\n",
    "\n",
    "#and for fuv\n",
    "\n",
    "throughputs_galex_fuv = ThroughputCurve(pd.read_csv(path+galex_names[0], \n",
    "                                                      delim_whitespace=True, skiprows=0, header=None)).throughput\n",
    "        \n",
    "filter_waves_fuv = ThroughputCurve(pd.read_csv(path+galex_names[0], \n",
    "                                                      delim_whitespace=True, skiprows=0, header=None)).waves\n",
    "\n",
    "filter_freqs_fuv = speed_of_light/filter_waves_fuv \n",
    "\n",
    "fuv_freq = interpolate.interp1d(filter_freqs_fuv, throughputs_galex_fuv, bounds_error = False, fill_value = 0)\n",
    "\n",
    "filter_freqs = speed_of_light/(filter_waves*10000) # to convert from angstroms\n",
    "\n",
    "# interpolating functions for filter curves\n",
    "\n",
    "castor_uv = interpolate.interp1d(np.array(filter_freqs), np.array(castor_uv_throughput), bounds_error = False, fill_value = 0)\n",
    "castor_u = interpolate.interp1d(np.array(filter_freqs), np.array(throughputs[2,:]), bounds_error = False, fill_value = 0)\n",
    "castor_g = interpolate.interp1d(np.array(filter_freqs), np.array(throughputs[4,:]), bounds_error = False, fill_value = 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def b(theta, nu, z):\n",
    "    \"\"\"\n",
    "    Returns evolution of the intensity map bias, b_im as a function of redshift and frequency \n",
    "    as described in equation 17. \n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    z : The emission redshift\n",
    "    nu: frequency in Hz\n",
    "    theta: A vector of ordered parameters expected by the MCMC routine, \n",
    "    theta[0] should be the power law evolution in frequency\n",
    "    theta[1] should be the power law evolution in redshift \n",
    "    \n",
    "    Returns\n",
    "    -------\n",
    "    float\n",
    "        The (unnormalized) evolution of the map bias in frequency and redshift \n",
    "    \"\"\"\n",
    "    \n",
    "    gammabnu = theta[0]\n",
    "    gammabz = theta[1]\n",
    "    \n",
    "    return (nu/nu1500)**(gammabnu) * (1+z)**(gammabz)\n",
    "\n",
    "def delta(nu):\n",
    "    \"\"\"\n",
    "    For a vector of frequencies, returns a value 1 at the frequency corresponding \n",
    "    to 1216 angstroms and 0 elsewhere.\n",
    "    \n",
    "    Note\n",
    "    ----------\n",
    "    no longer used. \n",
    "    \"\"\"\n",
    "    condlist = []\n",
    "    funclist = [1, 0]\n",
    "    return np.select(condlist, funclist)\n",
    "\n",
    "def prefactor(z): \n",
    "    \"\"\"\n",
    "    Returns the multiplicative pre-factor in the observed frame intensity forward model. \n",
    "    This is essentially the luminosity distance.\n",
    "    Calls and requires the colossus cosmology to be set globally. Colossus is a convenience package of cosmology \n",
    "    functions.  \n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    z : The emission redshift\n",
    "    \n",
    "    Returns\n",
    "    -------\n",
    "    float\n",
    "        The (unnormalized) evolution of the map bias in frequency and redshift \n",
    "        \n",
    "    Requires\n",
    "    --------\n",
    "    Calls and requires the colossus cosmology to be set globally. Colossus is a convenience package of cosmology \n",
    "    functions.  \n",
    "    \n",
    "    \"\"\"\n",
    "        \n",
    "    MPCtoCM = 3.086e24\n",
    "    speed_of_light_km_s = 3*10**5\n",
    "    return speed_of_light_km_s *(10**23)/(4*np.pi*(cosmology.Cosmology.Hz(cosmo, z)*(1+z))*(MPCtoCM**2))\n",
    "\n",
    "def optical_depth(nu, z):\n",
    "    \"\"\"\n",
    "    An analytic approximation to effective optical depth for Poisson distributed absorbers \n",
    "    to redshift z from Madau (2000) \n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    z : The emission redshift\n",
    "    \n",
    "    Returns\n",
    "    -------\n",
    "    float\n",
    "        The effective optical depth.  \n",
    "    \n",
    "    \"\"\"\n",
    "    \n",
    "    \n",
    "    A = 4*10**7\n",
    "    sigma_L = 6.3*10**(-18)\n",
    "    prefactor = (4.0/3.0)*np.sqrt(np.pi * sigma_L)\n",
    "    nu_L = 3.29 * 10**15 #Hz, Rydberg frequency\n",
    "    nu_term = (nu/nu_L)**(-1.5)\n",
    "    \n",
    "    #note redshift term assumes z0 = 0\n",
    "    \n",
    "    redshift_term = (1**(1.5))*((1+z)**(1.5)-(1)**(1.5))\n",
    "    \n",
    "    return np.exp(-prefactor * A * nu_term * redshift_term)\n",
    "\n",
    "\n",
    "def emissivity_model_optical_depth(theta, nu, z):\n",
    "    \"\"\"\n",
    "    Returns a vector of emissivities*bias*optical depth for a set of input frequencies at a given redshift z. This implements\n",
    "    the model described in Section 2 and multiplies by the optical depth as described in the forward model given in\n",
    "    equation 16.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    theta : a vector of input parameters describing the emissivity model\n",
    "    the parameters of theta, in order, are;  \n",
    "    the evolution of the bias in redshift and frequency,\n",
    "    the log10(normalization) of the 1500 angstrom emissivity times the bias,\n",
    "    the power law evolution parameter of the 1500 angstrom emissivity,\n",
    "    the power law evolution and amplitude of the slope as measured at redshift 0,\n",
    "    the same quantities at 1100 angstroms,\n",
    "    the lyman alpha equivalent width at z=1,\n",
    "    and the log10(LyC) escape fractions at z=1,2. \n",
    "    nu: a vector of frequencies to evaluate the emissivity over. The frequencies should correspond to \n",
    "    the CASTOR observed wavelength range\n",
    "    z : The emission redshift\n",
    "    \n",
    "    Returns\n",
    "    -------\n",
    "    np.array \n",
    "        The bias weighted emissivity as a function of frequency at redshift z times the optical depth to redshift z. \n",
    "    \n",
    "    \"\"\"\n",
    "    \n",
    "    \n",
    "    gammabnu, gammabz, log_e1500z0_b1500z0, gammae1500, alpha1500z0, Calpha1500, alpha1100z0, Calpha1100, EWLyAz1, logfLyCz1, logfLyCz2 = theta\n",
    "    \n",
    "    # note, we use np.select to vectorize our piecewise function in frequency. \n",
    "    \n",
    "    condlist = [\n",
    "                speed_of_light/nu > 1216.5,\n",
    "                (1215.5 > (speed_of_light/nu)) * (912 <= speed_of_light/nu),\n",
    "                (np.trunc(speed_of_light/nu)) >= 1215.5 * ((np.trunc(speed_of_light/nu)) <= 1216.5), \n",
    "                speed_of_light/nu < 912\n",
    "                ]\n",
    "    funclist = [\n",
    "                10**(log_e1500z0_b1500z0) * e1500_evol(z, gammae1500) * (nu/nu1500)**(alpha1500(z, Calpha1500, alpha1500z0)), \n",
    "        \n",
    "                10**(log_e1500z0_b1500z0) * (e1500_evol(z, gammae1500) *  (nu1216/nu1500)**(alpha1500(z, Calpha1500, alpha1500z0))*\n",
    "                ((nu/nu1216)**alpha1100(z, Calpha1100, alpha1100z0)))*optical_depth(nu, z), \n",
    "        \n",
    "                10**(log_e1500z0_b1500z0) * (e1500_evol(z, gammae1500) *  (nu1216/nu1500)**(alpha1500(z, Calpha1500, alpha1500z0))*\n",
    "                ((nu/nu1216)**alpha1100(z, Calpha1100, alpha1100z0) + EWLyA(z, EWLyAz1, EWLyAz03)))*optical_depth(nu, z), \n",
    "                #(nu**2)/speed_of_light),\n",
    "        \n",
    "                10**(logfLyC(z, logfLyCz2, logfLyCz1))*\n",
    "                10**(log_e1500z0_b1500z0) * e1500_evol(z, gammae1500) * (nu1216/nu1500)**alpha1500(z, Calpha1500, alpha1500z0) * \n",
    "                (nu912/nu1216)**alpha1100(z, Calpha1100, alpha1100z0) * (nu/nu912)**alpha900\n",
    "                ] \n",
    "    return np.select(condlist, funclist)\n",
    "\n",
    "def biased_weighted_emissivity(theta, z):\n",
    "    \"\"\"\n",
    "    Returns the observed frame quantity dJ/dz b_im as described in equation 16. Note that the name of this function is a\n",
    "    misnomer and retained for consistency. This is a bias weighted filter specific intensity that is seen in the observed\n",
    "    frame, while the emissivity is the rest frame underlying quantity whose model we seek to infer. \n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    theta : a vector of input parameters describing the emissivity model, see documentation for the emissivity model\n",
    "    for details on this vector. \n",
    "    z : The emission redshift\n",
    "    \n",
    "    Returns\n",
    "    -------\n",
    "    np.array \n",
    "        A 3xN vector of the observed frame quantity dJ/dz b_im for each castor filter. \n",
    "    \n",
    "    \"\"\"\n",
    "    \n",
    "    # compute each term in this integral, nu_obs is defined globally\n",
    "    inverse_nu_obs = 1.0/nu_obs \n",
    "    \n",
    "    #compute the bandpass\n",
    "    \n",
    "    R_castor_u = np.abs(np.trapz(castor_u(nu_obs)*inverse_nu_obs, nu_obs))\n",
    "    R_castor_uv = np.abs(np.trapz(castor_uv(nu_obs)*inverse_nu_obs, nu_obs))\n",
    "    R_castor_g = np.abs(np.trapz(castor_g(nu_obs)*inverse_nu_obs, nu_obs))\n",
    "    \n",
    "    # in this case, the galex nuv curve R_nuv (nu_obs)\n",
    "    \n",
    "    R_nu_obs_u = castor_u(nu_obs)\n",
    "    R_nu_obs_uv = castor_uv(nu_obs)\n",
    "    R_nu_obs_g = castor_g(nu_obs)\n",
    "\n",
    "    # the bias factor, evaluated in the rest frame nu = nu_obs (1+z)\n",
    "    \n",
    "    bias_nu = b(theta, nu_obs * (1+z), z)\n",
    "    \n",
    "    # and the emissivity model itself \n",
    "    \n",
    "    e_nu_z = emissivity_model_optical_depth(theta, nu_obs * (1+z), z)\n",
    "    \n",
    "    # and perform the integral\n",
    "    #syntax for np.trapz is (f(x), x)\n",
    "    \n",
    "    dJdz_bj_u = prefactor(z)/R_castor_u * np.abs(np.trapz(R_nu_obs_u * bias_nu * e_nu_z *inverse_nu_obs, nu_obs))\n",
    "    dJdz_bj_uv = prefactor(z)/R_castor_uv * np.abs(np.trapz(R_nu_obs_uv * bias_nu * e_nu_z *inverse_nu_obs, nu_obs))\n",
    "    dJdz_bj_g = prefactor(z)/R_castor_g * np.abs(np.trapz(R_nu_obs_g * bias_nu * e_nu_z *inverse_nu_obs, nu_obs))\n",
    "    \n",
    "\n",
    "    return np.array([dJdz_bj_uv, dJdz_bj_u, dJdz_bj_g])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def frac_shot_error(depth = 1):\n",
    "    \"\"\"\n",
    "    Returns the fractional shot noise error for the surveys we consider. \n",
    "    The shot noise term is computed separately from a collection of survey catalogs and loaded as a .npy\n",
    "    This implements the denominator of equation 22, where the full variance is determined in quadrature.  \n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    depth : a factor to multiply the suvey by in order to determine an optimal survey. By default this is 1 to use\n",
    "    the real survey catalog. \n",
    "    \n",
    "    Returns\n",
    "    -------\n",
    "    np.array \n",
    "        A 80x1 vector of shot noise terms for each redshift bin we consider in this work.  \n",
    "    \n",
    "    \"\"\"\n",
    "    \n",
    "    survey_shot = np.load('frac_errors_survey.npy')\n",
    "    return (1/depth)*survey_shot\n",
    "\n",
    "def photo_z_error(A_n, zs):\n",
    "    \"\"\"\n",
    "    Returns the fractional error due to photometric noise as it evolves with redshift as defined in equation 26. \n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    A_n : The overall noise amplitude. Typical values are 0.01 for CASTOR, 0.03 for GALEX. \n",
    "    zs : redshifts to evaluate the noise at it. \n",
    "    \n",
    "    Returns\n",
    "    -------\n",
    "    np.array \n",
    "        A 80x1 vector of fractional photometric zero-point errors for the redshift binning we consider in this work.  \n",
    "    \n",
    "    \"\"\"\n",
    "    \n",
    "    return A_n*(1+zs)**2\n",
    "\n",
    "def bias_error(sigma_z, alpha, zs):\n",
    "    \"\"\"\n",
    "    Returns the fractional error due to unmodeled evolution of the reference catalog bias. \n",
    "    This implements equation 25. \n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    sigma_z : gaussian width of the tracer catalog redshift distribution, we assume z = 0.5\n",
    "    alpha : power law describing the unmodeled evolution of the bias. \n",
    "    zs : redshifts to evaluate the noise at it. \n",
    "    \n",
    "    Returns\n",
    "    -------\n",
    "    np.array \n",
    "        A 80x1 vector of fractional bias errors for the redshift binning we consider in this work.  \n",
    "    \n",
    "    \"\"\"\n",
    "    \n",
    "    \n",
    "    error = []\n",
    "\n",
    "    for z0 in zs:\n",
    "        weights = norm.pdf(zs, z0, sigma_z)\n",
    "        error.append((np.average(zs, weights = (zs**alpha)*weights) - np.average(zs, weights = weights)))\n",
    "    return np.array(error)\n",
    "\n",
    "def tot_frac_error(zs, depth, model):\n",
    "    \"\"\"\n",
    "    Computes the total fractional error on the observed frame quantity dJ/dz b_im by adding in quadruture.  \n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    zs : redshifts to evaluate the noise at it. \n",
    "    depth : survey depth factor to scale the shot noise by\n",
    "    model : integer selecting either model 1, the optimal model which is photometric error dominated\n",
    "            or model 2, the full model incoroporating all effects above. \n",
    "    \n",
    "    Returns\n",
    "    -------\n",
    "    np.array \n",
    "        A 80x1 vector of fractional errors for the redshift binning we consider in this work.  \n",
    "    \n",
    "    \"\"\"\n",
    "    \n",
    "    if model == 1:\n",
    "        A = np.sqrt(frac_shot_error(5)**2 + 0.01**2)\n",
    "        \n",
    "    elif model == 2:\n",
    "        A = np.sqrt(frac_shot_error(1)**2 + photo_z_error(0.01, zs)**2 + bias_error(0.5, 0.01, zs)**2)\n",
    "        \n",
    "    return A\n",
    "\n",
    "\n",
    "\n",
    "def mk_data(theta, zmin = 0.01, zmax = 4, sample_size = 80, depth = 1, error_model = 1):  \n",
    "    \n",
    "    \"\"\"\n",
    "    This function makes the final data and error vectors.   \n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    theta : a vector of input parameters describing the emissivity model, see documentation for the emissivity model\n",
    "    for details on this vector.  \n",
    "    zmin : minimum redshift to compute model over, fiducial = 0.01, should be nonzero and positive\n",
    "    zmax : maximum redshift to compute model over, fiducial = 4, should be nonzero and positive\n",
    "    sample_size : number of redshift bins to use. fiducial is 80, corresponding to z_step = 0.05\n",
    "    error_model : parameter specifying which error model to use. 1 or 2, 1 corresponds to an optimal model, \n",
    "                  2 to the full model. \n",
    "    \n",
    "    Returns\n",
    "    -------\n",
    "    np.array \n",
    "        sim_data is the simulated data under our SED model evaluated at the bin spacing\n",
    "    np.array\n",
    "        error is the uncertainty on the sim data, computed as the fractional error * sim_data \n",
    "        and evaluated at the bin spacing\n",
    "    \n",
    "    \"\"\"\n",
    "    \n",
    "    z_int_sampling = np.linspace(zmin, zmax, sample_size)\n",
    "    sim_data = np.zeros((len(z_int_sampling), 3))\n",
    "    error = np.zeros((len(z_int_sampling), 3))\n",
    "    \n",
    "    for i, z in enumerate(z_int_sampling):\n",
    "        sim_data[i,:] = biased_weighted_emissivity(theta, z)\n",
    "        error[i,:] = sim_data[i,:]*tot_frac_error(z_int_sampling, depth, error_model)[i] #(z, n_galaxies, n_fields)\n",
    "        #dJdzb_error = np.random.multivariate_normal(mean=np.array([0, 0, 0]), cov=np.diag(error[i,:]))\n",
    "        #sim_data[i,:] = dJdzb + dJdzb_error\n",
    "\n",
    "    return sim_data, error\n",
    "\n",
    "def model(theta, z):\n",
    "    \"\"\"\n",
    "    Convenience wrapper for the biased_weighted_emissivity. See documentation for that function. \n",
    "    \"\"\"\n",
    "    return biased_weighted_emissivity(theta, z)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def lnprior(theta):\n",
    "    \"\"\"\n",
    "    This function returns the log of the prior on each parameter as summarized in Table 1.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    theta : a vector of input parameters describing the emissivity model, see documentation for the emissivity model\n",
    "    for details on this vector.  \n",
    "    \n",
    "    Returns\n",
    "    -------\n",
    "    np.array\n",
    "        log(prior)\n",
    "    \n",
    "    \"\"\"\n",
    "    \n",
    "    gammabnu, gammabz,log_e1500z0_b1500z0, gammae1500, alpha1500z0, Calpha1500, alpha1100z0, Calpha1100, EWLyAz1, logfLyCz2, logfLyCz2 = theta\n",
    "\n",
    "    \n",
    "    condlist = [-0.08 - (2*0.84)<=alpha1500z0<=-0.08+(2*1.28), 1.85-(2*1.28)<=Calpha1500<=1.85+(2*1.22),\n",
    "               -3.71 - (2*0.98)<=alpha1100z0<=-3.71 + (2*1.34), 0.50 - (2*1.44)<=Calpha1100<=0.50+(2*1.46),\n",
    "                88.02 - (2*48.87) <= EWLyAz1 <= 88.02 + (2*51.44), \n",
    "                logfLyCz1 <= 0.0, logfLyCz2 <= 0.0, \n",
    "               -0.86 - (2*1.29) <= gammabnu<= -0.86 + (2*0.83)]\n",
    "    \n",
    "    if not(all(condlist)):\n",
    "        return -np.inf\n",
    "    \n",
    "    mu1 = 2\n",
    "    sigma1 = 0.3\n",
    "    mu2 = 0.79\n",
    "    sigma2 = 0.33\n",
    "    mu3 = 25.13\n",
    "    sigma3 = 0.02\n",
    "    \n",
    "    return (np.log(1.0/(np.sqrt(2*np.pi)*sigma1))-0.5*(gammae1500-mu1)**2/sigma1**2)+(np.log(1.0/(np.sqrt(2*np.pi)*sigma2))-0.5*(gammabz-mu2)**2/sigma2**2)+(np.log(1.0/(np.sqrt(2*np.pi)*sigma3))-0.5*(log_e1500z0_b1500z0-mu3)**2/sigma3**2)\n",
    "\n",
    "\n",
    "def lnlike_parallel(theta, z, datum, datum_error):\n",
    "    \"\"\"\n",
    "    This function returns the log likelihood of the data given the model and is parallelized for use with the \n",
    "    emcee parallel sampler. \n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    theta : a vector of input parameters describing the emissivity model, see documentation for the emissivity model\n",
    "    for details on this vector.\n",
    "    z : redshift to compute the likelihood at\n",
    "    datum : individual simulated data point corresponding to the redshift z\n",
    "    datum_error : individual simulated data error corresponding to the data at redshift z\n",
    "    \n",
    "    Returns\n",
    "    -------\n",
    "    float\n",
    "        log(like)\n",
    "    \"\"\"\n",
    "    \n",
    "    diff = model(theta, z) - datum\n",
    "    icov = np.diag(np.array([1/datum_error[0]**2, 1/datum_error[1]**2, 1/datum_error[2]**2]))\n",
    "    \n",
    "    return -0.5* (np.dot(np.transpose(diff), np.dot(icov, diff)))\n",
    "\n",
    "\n",
    "def lnprob_parallel(theta, zs, data):\n",
    "    \"\"\"\n",
    "    This function returns the log posterior probability of the data given the model \n",
    "    and is parallelized for use with the emcee implementation of a parallel sampler. \n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    theta : a vector of input parameters describing the emissivity model, see documentation for the emissivity model\n",
    "    for details on this vector.\n",
    "    zs : redshift to evaluate at\n",
    "    data : data and errors computed from the make_data function. \n",
    "    \n",
    "    Returns\n",
    "    -------\n",
    "    float\n",
    "        log(posterior)\n",
    "    \n",
    "    \"\"\"\n",
    "    \n",
    "    lp = lnprior(theta)\n",
    "    \n",
    "    lks = np.zeros(len(zs))\n",
    "    \n",
    "    for i, (z, datum, datum_error) in enumerate(zip(zs, data[0], data[1])):\n",
    "        lks[i] = lnlike_parallel(theta, z, datum, datum_error)\n",
    "    lksum = np.sum(lks)\n",
    "    if not np.isfinite(lp):\n",
    "        return -np.inf\n",
    "    if not np.isfinite(lksum):\n",
    "        return -np.inf\n",
    "    return lp + lksum\n",
    "\n",
    "\n",
    "\n",
    "def main():\n",
    "    \"\"\"\n",
    "    main function of the script, sets and initializes various parameters, prints files to log and calls emcee \n",
    "    sampler as well as saving example corner plot locally. \n",
    "    \"\"\"\n",
    "    \n",
    "    \n",
    "    print('Starting emcee sampling')\n",
    "    nwalkers = 100\n",
    "    ndim = len(theta)\n",
    "    n_samps = 80\n",
    "    chain_len = 10**5\n",
    "    zmax = 4\n",
    "    zmin = 0.01\n",
    "    zs = np.linspace(zmin,zmax, n_samps) #define redshift sampling\n",
    "    print('nwalkers, ndim, n_samps, chain_len, min(zs), max(zs)')\n",
    "    print([nwalkers, ndim, n_samps, chain_len, min(zs), max(zs)])\n",
    "    \n",
    "    print('Initializing Walkers')\n",
    "\n",
    "    p0 = theta + 1e-3 * np.random.randn(nwalkers, len(theta))\n",
    "    \n",
    "    print('Making data')\n",
    "\n",
    "    data = mk_data(theta, zmin = zmin, zmax = zmax, sample_size = n_samps, error_model = 1)\n",
    "    \n",
    "    print(np.min(data[0]), np.max(data[0]), np.mean(data[0]), np.min(data[1]), np.max(data[1]), np.mean(data[1]))\n",
    "    \n",
    "    from multiprocess import cpu_count\n",
    "\n",
    "    ncpu = cpu_count()\n",
    "    print(\"{0} CPUs\".format(ncpu))\n",
    "    \n",
    "    filename = \"/rhome/bscot002/bigdata/error_model_1_1223_optimal.h5\"\n",
    "    backend = emcee.backends.HDFBackend(filename)\n",
    "    #backend.reset(nwalkers, ndim)\n",
    "\n",
    "    with Pool() as pool: \n",
    "        sampler = emcee.EnsembleSampler(nwalkers, ndim, lnprob_parallel, args = [zs, data], pool=pool, backend = backend)\n",
    "        #sampler = emcee.EnsembleSampler(nwalkers, ndim, lnprob_parallel, args = [zs, data], pool=pool)\n",
    "        start = time.time()\n",
    "        pos, prob, state = sampler.run_mcmc(p0, chain_len, progress = True)\n",
    "        end = time.time()\n",
    "        multi_time = end - start\n",
    "        print(\"Multiprocessing took {0:.1f} seconds\".format(multi_time))\n",
    "        \n",
    "    tau = sampler.get_autocorr_time()\n",
    "    print('sampling complete')\n",
    "    print('tau is:')\n",
    "    print(tau)\n",
    "    \n",
    "    flat_samples = sampler.get_chain(discard=100, thin=15, flat=True)\n",
    "    print(flat_samples.shape)\n",
    "    \n",
    "    fig = corner.corner(flat_samples);\n",
    "    \n",
    "    fig.savefig(path+'corner_test.png')\n",
    "        \n",
    "    pass\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    main()\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.8.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}