generate_dv_nonlimber.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "185d081b",
   "metadata": {},
   "outputs": [],
   "source": [
    "year = 'Y1'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "baeb8a56",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import scipy\n",
    "import sacc\n",
    "import pyccl as ccl\n",
    "from scipy.interpolate import interp1d\n",
    "from scipy.special import erf"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "830c1795",
   "metadata": {},
   "source": [
    "# Fiducial cosmology\n",
    "* Planck 2018 TT,TE,EE+lowE+lensing\n",
    "* https://arxiv.org/abs/1807.06209 Table 2 second-to-last column\n",
    "* for our fiducial cosmology we have 3 degenerate massive neutrinos with equal mass, and a total mass of 0.1 eV (to include the effect of massive neutrinos without having to do calculations for multiple masses)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b672d502",
   "metadata": {},
   "outputs": [],
   "source": [
    "fiducial_cosmology = {'name': 'desy3',\n",
    "                      'h': 0.6736,\n",
    "                      'Omega_m': 0.3153,\n",
    "                      'Omega_b': 0.0493,\n",
    "                      'A_s': 2.092e-9,\n",
    "                      'n_s': 0.9656,\n",
    "                      'w0': -1,\n",
    "                      'wa': 0,\n",
    "                      #'tau': 0.0536,\n",
    "                      'sigma8': 0.8091,\n",
    "                      'S8': 0.828\n",
    "                     }"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2df1177",
   "metadata": {},
   "source": [
    "# $\\ell$ binning choices\n",
    "\n",
    "* decide whether to make bins smaller\n",
    "* look at hsc, des, kids choices?\n",
    "* I don't think there is any point in going to 15 000 so even if we want the same width we should be able to just make fewer bins"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3390434",
   "metadata": {},
   "outputs": [],
   "source": [
    "auto_galaxy_clustering_only = False\n",
    "\n",
    "nbin_ell = 20\n",
    "lmin_bins = 20\n",
    "lmax_bins = 2000\n",
    "\n",
    "twopoint_data = '3x2pt'\n",
    "    \n",
    "lbinning_str = str(nbin_ell)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5230c556",
   "metadata": {},
   "source": [
    "# Galaxy clustering systematics and analysis choices\n",
    "\n",
    "* nonlinear bias with Eulerian perturbation theory, redshift evolution (see bias challenge paper https://arxiv.org/abs/2307.03226, figure 7 of https://arxiv.org/pdf/2111.00501.pdf)\n",
    "* k_max 0.1-0.2 h/Mpc\n",
    "* redshift: shift and stretch?\n",
    "* magnification\n",
    "* redshift space distortions\n",
    "\n",
    "## To Do:\n",
    "* include nonlinear bias\n",
    "* values for magnification bias\n",
    "* include photo-z errors? or no errors in datavector but still marginalize in analysis?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4e72300",
   "metadata": {},
   "outputs": [],
   "source": [
    "if year=='Y1' or year==1 or year=='y1':\n",
    "    # galaxy clustering fiducial systematic parameters\n",
    "    linear_bias_lens = np.array([1.23852, 1.36452, 1.50445, 1.6498, 1.79915]) #add nonlinear bias\n",
    "    kmax = 0.2*fiducial_cosmology['h']\n",
    "    magnification_bias_lens = np.array([0.1,0.1,0.1,0.1,0.1]) #figure out fiducial values\n",
    "\n",
    "    # galaxy clustering redshift bins\n",
    "    nbin_z_lens = 5\n",
    "    z0_lens = 0.26\n",
    "    alpha_lens = 0.94\n",
    "    sig_z_lens = 0.03\n",
    "    z_min_lens = 0.2\n",
    "    z_max_lens = 1.2\n",
    "    \n",
    "    # survey-informed parameters --TODO: update these\n",
    "    Ngal_lens = 18 # Normalisation, galaxies/arcmin^2\n",
    "    sky_area = 11808.28 #square degrees, this is still srd values for now"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4030c06c",
   "metadata": {},
   "source": [
    "# Weak Lensing systematics and analysis choices\n",
    "\n",
    "* TATT for intrinsic alignments\n",
    "* Mean redshift shift and stretch\n",
    "* Multiplicative Shear Bias\n",
    "* using DESC SRD lmax cutoff, ell binning and z binning for now"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2d43260",
   "metadata": {},
   "outputs": [],
   "source": [
    "if year=='Y1' or year==1 or year=='y1':\n",
    "    # weak lensing shear (source sample)\n",
    "    lmax_src = 2000\n",
    "    \n",
    "    z0_src = 0.13\n",
    "    alpha_src = 0.78\n",
    "    sig_z_src = 0.05  #(0.05(1+z))\n",
    "    sig_e_src = 0.26\n",
    "    nbin_z_src = 5\n",
    "    \n",
    "    Ngal_src = 10"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee72ae49",
   "metadata": {},
   "source": [
    "# CCL cosmology object\n",
    "* Done: HMCode 2020 with baryons"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28c2f539",
   "metadata": {},
   "outputs": [],
   "source": [
    "cosmo = ccl.Cosmology(\n",
    "    Omega_c=fiducial_cosmology['Omega_m']-fiducial_cosmology['Omega_b'],\n",
    "    Omega_b=fiducial_cosmology['Omega_b'],\n",
    "    w0=-1.0,\n",
    "    wa=0.0,\n",
    "    A_s = fiducial_cosmology['A_s'],\n",
    "    n_s=fiducial_cosmology['n_s'],\n",
    "    h=fiducial_cosmology['h'],\n",
    "    Omega_k=0.0,\n",
    "    m_nu = 0.1,\n",
    "    mass_split = 'equal',\n",
    "    matter_power_spectrum = \"camb\",\n",
    "    extra_parameters = {\"camb\": {\"halofit_version\": \"mead2020_feedback\",\n",
    "                                     \"HMCode_logT_AGN\": 7.8}}\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "729ba705",
   "metadata": {},
   "outputs": [],
   "source": [
    "if year=='Y1' or year==1 or year=='y1':\n",
    "\n",
    "    #All other Systematic Parameters\n",
    "\n",
    "    # Bare CCL setup\n",
    "    a_1 = 1.0\n",
    "    eta_1 = -1.0\n",
    "    z_pivot = 0.62\n",
    "\n",
    "    #Shear Calibration\n",
    "    s_cal = np.array([0.0, 0.0, 0.0, 0.0, 0.0])\n",
    "\n",
    "    k_max = 0.1/ fiducial_cosmology['h']\n",
    "    \n",
    "    #THIS IS JUST APPROXIMATE\n",
    "    #need the n(z) distributions first to calculate true <z>\n",
    "    z_mean_lens = np.array([0.3, 0.5, 0.7, 0.9, 1.1])\n",
    "    a_mean_lens = 1/(1+z_mean_lens)\n",
    "    l_max_lens =  kmax * cosmo.comoving_radial_distance(a_mean_lens) + 0.5\n",
    "    print(l_max_lens)\n",
    "    \n",
    "    #1-Loop Galaxy Bias\n",
    "    b_1 = 1.05/cosmo.growth_factor(a_mean_lens)\n",
    "    print(b_1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "95ba88f8",
   "metadata": {},
   "source": [
    "# $\\ell$ binning\n",
    "\n",
    "* currently matches DESC SRD\n",
    "* consider changing binning (smaller lmax, narrower bins?)\n",
    "* should we use a more realistic ell binning function based on mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed41aa5e",
   "metadata": {},
   "outputs": [],
   "source": [
    "ell_edges = np.geomspace(lmin_bins, lmax_bins, nbin_ell+1)\n",
    "ell_centers = [lmin_bins**(1-(i+0.5)/nbin_ell)*lmax_bins**((i+0.5)/nbin_ell) for i in range(nbin_ell)]\n",
    "print(ell_centers)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8bf874e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# tophat window functions in ell (use more realistic window fns?)\n",
    "ell_unbinned = np.arange(0, lmax_bins+1)\n",
    "n_ell_unbinned = len(ell_unbinned)\n",
    "ell_windows = np.zeros([len(ell_unbinned), nbin_ell])\n",
    "for i in range(nbin_ell):\n",
    "    ell_windows[:,i] = (ell_unbinned>=ell_edges[i]) & (ell_unbinned<=ell_edges[i+1])\n",
    "    ell_windows[:,i]/=np.sum(ell_windows[:,i])\n",
    "    plt.plot(ell_unbinned, ell_windows[:,i], label=r'$\\ell$ window '+str(i+1))\n",
    "plt.xlim(10, lmax_bins)\n",
    "plt.xscale('log')\n",
    "plt.xlabel('$\\ell$')\n",
    "plt.title('window functions in $\\ell$')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6a7c0f91",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a SACC bandpower window object\n",
    "sacc_ell_windows = sacc.BandpowerWindow(ell_unbinned, ell_windows)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2621bfbe",
   "metadata": {},
   "source": [
    "# Galaxy clustering"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2aa312d2",
   "metadata": {},
   "source": [
    "## z binning\n",
    "\n",
    "* same as DESC SRD (5 equally spaced bins from z=0.2-1.2 for LSST y1)\n",
    "* no cross spectra between bins (look into whether crossing neighbouring bins adds signal)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "399abdf8",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_redshift_distribution(z, z0, alpha): # shape, not normalized\n",
    "    return (z / z0)**2. * np.exp(-(z / z0)**alpha)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f90746cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "def avg_z_int(z, splined):\n",
    "    return z * splined(z)\n",
    "\n",
    "def nz_int(z, sig_e, z_max, z_min, splined):\n",
    "    sig_e *= (1+z)\n",
    "    upper_lim = (z-z_max)/(sig_e*np.sqrt(2))\n",
    "    lower_lim = (z-z_min)/(sig_e*np.sqrt(2))\n",
    "    return  0.5*(erf(lower_lim)-erf(upper_lim))*splined(z)\n",
    "\n",
    "def get_nz_bins_from_zmid(z_arr, ngal, nz, sig_z, z_lower_bounds, z_upper_bounds, zmids=None, title=''):\n",
    "    nz_bins = []\n",
    "    z_avgs = []\n",
    "    nz_avgs = []\n",
    "    if zmids is None:\n",
    "        zmids = (z_lower_bounds + z_upper_bounds)/2\n",
    "    fig, ax = plt.subplots(1,2, figsize=(10,5))\n",
    "    ax[0].plot(z_arr, nz, label='total nz')\n",
    "    splined = interp1d(z_arr, nz)\n",
    "    for i, zmin in enumerate(z_lower_bounds):\n",
    "        zmax = z_upper_bounds[i]\n",
    "        zmid = zmids[i]        \n",
    "        nz_bin_i = np.zeros_like(z_arr)\n",
    "        nz_bin_i[(z_arr>zmin) & (z_arr<zmax)] = nz[(z_arr>zmin) & (z_arr<zmax)]\n",
    "        ax[0].plot(z_arr, nz_bin_i, label='z bin %d'%(i+1))\n",
    "        ax[0].set_xlabel('$z$')\n",
    "        ax[0].set_ylabel('$p(z)$')\n",
    "        ax[0].set_title('before gaussian binning')\n",
    "        ax[0].set_xlim(0,2)\n",
    "\n",
    "        ax[0].vlines(zmid, 0, np.amax(nz), color='k', linestyle=':')\n",
    "        norm, e = scipy.integrate.quad(nz_int, z_arr[0], z_arr[-1], args=(sig_z, zmax, zmin, splined))\n",
    "        nz_bin_i = nz_int(z_arr, sig_z, zmax, zmin, splined)/norm\n",
    "        nz_bins.append(nz_bin_i)\n",
    "        splined_2 = interp1d(z_arr, nz_bin_i)\n",
    "        z_avgs.append((scipy.integrate.quad(avg_z_int, zmin, zmax, args=(splined_2))[0])/(scipy.integrate.quad(splined_2, zmin, zmax)[0]))\n",
    "        nz_avgs.append(norm*ngal)\n",
    "        ax[1].plot(z_arr, nz_bin_i, label='z bin %d'%(i+1))\n",
    "\n",
    "        ax[1].set_xlabel('$z$')\n",
    "        ax[1].set_title('after gaussian binning')\n",
    "        ax[1].set_xlim(0,2)\n",
    "        ax[1].vlines(z_avgs[-1], 0, np.amax(nz_bin_i), color='k', linestyle=':')\n",
    "\n",
    "        \n",
    "    plt.legend()\n",
    "    plt.suptitle(title, fontsize=15)\n",
    "    plt.show()\n",
    "    return z_avgs, nz_avgs, nz_bins"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d11feb85",
   "metadata": {},
   "outputs": [],
   "source": [
    "z_arr = np.linspace(0., 5., 5001)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c8bbf44f",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# z binning galaxy clustering (lens sample): evenly spaced redshift bins\n",
    "bin_edges_lens = np.linspace(z_min_lens, z_max_lens, nbin_z_lens+1)\n",
    "z_bounds_low_lens = bin_edges_lens[:-1]\n",
    "z_bounds_high_lens = bin_edges_lens[1:]\n",
    "\n",
    "pz_lens = get_redshift_distribution(z_arr, z0_lens, alpha_lens)\n",
    "norm, e = scipy.integrate.quad(get_redshift_distribution, 0, 10, args=(z0_lens, alpha_lens))\n",
    "pz_lens = pz_lens/norm  \n",
    "z_avgs_lens, nz_avgs_lens, pz_bins_lens = get_nz_bins_from_zmid(z_arr, Ngal_lens, pz_lens, sig_z_lens, z_bounds_low_lens, z_bounds_high_lens, title='redshift distribution for lens galaxy bins')\n",
    "\n",
    "chis = ccl.comoving_radial_distance(cosmo, 1/(1+np.array(z_avgs_lens))) # Mpc\n",
    "lmax_lens_zbins = kmax*chis-0.5\n",
    "print(z_avgs_lens)\n",
    "print(nz_avgs_lens)\n",
    "print(z_bounds_low_lens)\n",
    "print(lmax_lens_zbins)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "921ba9ab",
   "metadata": {},
   "source": [
    "## tracers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63255d2c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# tracers for galaxy clustering\n",
    "s = sacc.Sacc()\n",
    "tracers_lens=[]\n",
    "tracers_lens_nobias = []\n",
    "\n",
    "for i in range(nbin_z_lens):\n",
    "    # galaxy clustering CCL tracer for computing mock datavector\n",
    "    tracers_lens.append(ccl.NumberCountsTracer(cosmo, \n",
    "                                          has_rsd=False, \n",
    "                                          dndz=(z_arr, pz_bins_lens[i]), \n",
    "                                          bias=(z_arr,b_1[i]*np.ones(len(z_arr))), \n",
    "                                          mag_bias=(z_arr, magnification_bias_lens[i]*np.ones(len(z_arr)))\n",
    "                                        ))  \n",
    "    # galaxy clustering SACC tracer for saving to file later\n",
    "    s.add_tracer('NZ', \"lens%d\"%i,  \n",
    "                 quantity='galaxy_density', \n",
    "                 z=z_arr,  \n",
    "                 nz=pz_bins_lens[i]) "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "008ce739",
   "metadata": {},
   "source": [
    "# Weak lensing"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "adbcba65",
   "metadata": {},
   "source": [
    "## source z bins\n",
    "equal number binning\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de2dfec7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_fraction_in_redshift_bin(zmin, zmax, z0, alpha):\n",
    "    i, e1 = scipy.integrate.quad(get_redshift_distribution, zmin, zmax, args=(z0, alpha))\n",
    "    i_all_z, e2 = scipy.integrate.quad(get_redshift_distribution, 0, 10, args=(z0, alpha))\n",
    "    fraction_in_bin = i/i_all_z # what fraction of the total galaxies are between redshifts zmin and zmax\n",
    "    return fraction_in_bin\n",
    "\n",
    "#need to get equal for GAUSSIAN bins\n",
    "def get_equal_number_binning_bounds(nbin_z, ngal, splined, z_arr, sig_e, z0, alpha):\n",
    "    z_bounds = [0]\n",
    "    i_all_z, e2 = scipy.integrate.quad(get_redshift_distribution, 0, 10, args=(z0, alpha))\n",
    "\n",
    "    for i in range(nbin_z-1):\n",
    "        # use cumulative redshift distribution to solve for zmax \n",
    "        # such that 1/nbin_z of the galaxies are in each bin\n",
    "        sol = scipy.optimize.root_scalar(lambda zmax, zmin, bin_number: \n",
    "            get_fraction_in_redshift_bin(zmin, zmax, z0, alpha) - (bin_number+1)/nbin_z,\n",
    "            args = (0,i), method='bisect', bracket=[0., 5.])\n",
    "        assert(sol.converged)\n",
    "        z_bounds.append(sol.root)\n",
    "\n",
    "    z_bounds.append(5) #technically infinity since there is no upper bound but this is a pain numerically, and redshift distribution falls before then\n",
    "\n",
    "    return z_bounds\n",
    "\n",
    "def nz_integrated(z, sig_e, zmax, zmin, splined):\n",
    "    return scipy.integrate.quad(nz_int, z[0], z[-1], args=(sig_e, zmax, zmin, splined))[0]\n",
    "\n",
    "def get_nz_bins_eq(z_arr, ngal, nz, sig_z, z_lower_bounds, z_upper_bounds, title=''):\n",
    "    nz_bins = []\n",
    "    z_avgs = []\n",
    "    nz_avgs = []\n",
    "    fig, ax = plt.subplots(1,2, figsize=(10,5))\n",
    "    splined = interp1d(z_arr, nz)\n",
    "\n",
    "    for i, zmin in enumerate(z_lower_bounds):\n",
    "        zmax = z_upper_bounds[i]\n",
    "        \n",
    "        nz_bin_i = np.zeros_like(z_arr)\n",
    "        nz_bin_i[(z_arr>zmin) & (z_arr<zmax)] = nz[(z_arr>zmin) & (z_arr<zmax)]\n",
    "        dz = z_arr[1]-z_arr[0]\n",
    "        ax[0].plot(z_arr, nz_bin_i, label='z bin %d'%(i+1))\n",
    "        ax[0].set_xlabel('$z$')\n",
    "        ax[0].set_ylabel('$p(z)$')\n",
    "        ax[0].set_title('before gaussian binning')\n",
    "        ax[0].set_xlim(0,2)\n",
    "        \n",
    "        norm, e = scipy.integrate.quad(nz_int, z_arr[0], z_arr[-1], args=(sig_z, zmax, zmin, splined))\n",
    "        nz_bin_i = nz_int(z_arr, sig_z, zmax, zmin, splined)/norm\n",
    "        nz_bins.append(nz_bin_i)\n",
    "\n",
    "        splined_2 = interp1d(z_arr, nz_bin_i)\n",
    "        z_avgs.append((scipy.integrate.quad(avg_z_int, zmin, zmax, args=(splined_2))[0])/(scipy.integrate.quad(splined_2, zmin, zmax)[0]))\n",
    "        nz_avgs.append(norm*ngal)\n",
    "        ax[1].plot(z_arr, nz_bin_i, label='z bin %d'%(i+1))\n",
    "        ax[1].set_xlabel('$z$')\n",
    "        ax[1].set_title('after gaussian binning')\n",
    "        ax[1].set_xlim(0,2)\n",
    "        ax[1].vlines(z_avgs[-1], 0, np.amax(nz_bin_i), color='k', linestyle=':')\n",
    "\n",
    "    plt.legend()\n",
    "    plt.suptitle(title, fontsize=15)\n",
    "    plt.show()\n",
    "\n",
    "    return z_avgs, nz_avgs, nz_bins"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4c50df7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# z binning weak lensing source sample: source galaxies use equal number binning\n",
    "pz_src = get_redshift_distribution(z_arr, z0_src, alpha_src)\n",
    "norm, e = scipy.integrate.quad(get_redshift_distribution, 0, 10, args=(z0_src, alpha_src))\n",
    "pz_src = pz_src/norm \n",
    "splined_src = interp1d(z_arr, pz_src)\n",
    "\n",
    "#get_equal_gaussian_bins(nbin_z_src, Ngal_src, splined_src, z_arr, sig_z_src, z0_src, alpha_src)\n",
    "bin_edges_src = get_equal_number_binning_bounds(nbin_z_src, Ngal_src, splined_src, z_arr, sig_z_src, z0_src, alpha_src)\n",
    "z_bounds_low_src = np.array(bin_edges_src[:-1])\n",
    "z_bounds_high_src = np.array(bin_edges_src[1:])\n",
    "\n",
    "\n",
    "z_avgs_src, nz_avgs_src, pz_bins_src = get_nz_bins_eq(z_arr, Ngal_src, pz_src, sig_z_src,\n",
    "                                     z_bounds_low_src, z_bounds_high_src,\n",
    "                                     title='redshift distribution for source galaxy bins')\n",
    "print(z_avgs_src)\n",
    "print(bin_edges_src)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "819a8e9b",
   "metadata": {},
   "source": [
    "## weak lensing ccl and sacc tracers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76d12dc0",
   "metadata": {},
   "outputs": [],
   "source": [
    "tracers_src=[]\n",
    "tracers_lensing_kernel = []\n",
    "for i in range(nbin_z_src):\n",
    "    # weak lensing CCL tracer for computing mock datavector\n",
    "    tracers_src.append(ccl.WeakLensingTracer(cosmo,\n",
    "                                             has_shear = True,\n",
    "                                             dndz=(z_arr, pz_bins_src[i]),\n",
    "                                             ia_bias = (z_arr, a_1*np.ones_like(z_arr)*((1+z_arr)/(1+z_pivot))**eta_1),\n",
    "                                             use_A_ia=True\n",
    "                                             ))\n",
    "    # weak lensing SACC tracer for saving to file later\n",
    "    s.add_tracer(\"NZ\", \"src%d\"%i, \n",
    "                 quantity='galaxy_shear',\n",
    "                 z=z_arr,\n",
    "                 nz=pz_bins_src[i])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ce13ebf",
   "metadata": {},
   "outputs": [],
   "source": [
    "ells = []\n",
    "for ell in ell_centers:\n",
    "    ells.append(int(ell))\n",
    "\n",
    "cl_ss_theory = []\n",
    "cl_gg_theory = []\n",
    "cl_gs_theory = []\n",
    "cl_gs_theory_limber = []\n",
    "cl_gg_theory_limber = []\n",
    "for i in range(nbin_z_src):\n",
    "    temp = []\n",
    "    for j in range(nbin_z_src):\n",
    "        temp.append(ccl.angular_cl(cosmo,tracers_src[i], tracers_src[j], ells))\n",
    "    cl_ss_theory.append(temp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "679e8ea8",
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(nbin_z_lens):\n",
    "    temp = []\n",
    "    for j in range(nbin_z_src):\n",
    "        temp.append(ccl.angular_cl(cosmo, tracers_lens[i], tracers_src[j], ells, l_limber=np.max(ells), non_limber_integration_method=\"FKEM\"))\n",
    "    cl_gs_theory.append(temp)\n",
    "\n",
    "for i in range(nbin_z_lens):\n",
    "    temp = []\n",
    "    for j in range(nbin_z_src):\n",
    "        temp.append(ccl.angular_cl(cosmo, tracers_lens[i], tracers_src[j], ells))\n",
    "    cl_gs_theory_limber.append(temp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e11d71dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(nbin_z_lens):\n",
    "    temp = []\n",
    "    for j in range(nbin_z_lens):\n",
    "        temp.append(ccl.angular_cl(cosmo, tracers_lens[i], tracers_lens[j], ells, l_limber=np.max(ells), non_limber_integration_method=\"FKEM\"))\n",
    "    cl_gg_theory.append(temp)\n",
    "\n",
    "\n",
    "for i in range(nbin_z_lens):\n",
    "    temp = []\n",
    "    for j in range(nbin_z_lens):\n",
    "        temp.append(ccl.angular_cl(cosmo, tracers_lens[i], tracers_lens[j], ells))\n",
    "    cl_gg_theory_limber.append(temp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2325ead8",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(len(cl_ss_theory), len(cl_ss_theory[0]), len(cl_ss_theory[0][0]))\n",
    "print(len(cl_gg_theory), len(cl_gg_theory[0]), len(cl_gg_theory[0][0]))\n",
    "print(len(cl_gg_theory_limber), len(cl_gg_theory_limber[0]), len(cl_gg_theory_limber[0][0]))\n",
    "print(len(cl_gs_theory_limber), len(cl_gs_theory_limber[0]), len(cl_gs_theory_limber[0][0]))\n",
    "print(len(cl_gs_theory), len(cl_gs_theory[0]), len(cl_gs_theory[0][0]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "faa06077",
   "metadata": {},
   "source": [
    "# Angular power spectra\n",
    "bin combos from DESC SRD\n",
    "* 5 density-density\n",
    "* 7 density-shear spectra\n",
    "* 15 shear-shear spectra (src-src)\n",
    "\n",
    "\n",
    "NOTE: if using SRD cosmolike covmat the order is shear-shear, shear-density, density-density"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7605b5d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# for plots\n",
    "colors = ['#e41a1c', '#ff7f00', '#4daf4a', '#377eb8', '#984ea3', '#a65628', '#f781bf', '#999999', '#dede00'] \n",
    "linestyles = ['-', '--', '-.', ':', ':']\n",
    "markerstyles = ['o', '^', 'v', '+', 'x']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "feb7c784",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(10,10))\n",
    "ell_centers = np.array(ell_centers)\n",
    "for i in range(nbin_z_lens):\n",
    "    \n",
    "    #plt.vlines(l_max_lens[i], -1.e-4, 1.e-4, color = 'k', linestyle=linestyles[i])\n",
    "    if not auto_galaxy_clustering_only:\n",
    "        for j in range(i, nbin_z_lens):\n",
    "            cl_binned = cl_gg_theory[i][j]\n",
    "            s.add_ell_cl('galaxy_density_cl',  \n",
    "             \"lens%d\" % (i),  \n",
    "             \"lens%d\" % (j), \n",
    "             ell_centers,  \n",
    "             cl_binned,  \n",
    "             window=sacc_ell_windows\n",
    "            )\n",
    "\n",
    "            plt.plot(ells, cl_gg_theory[i][j], color=colors[i], linestyle='-', marker = markerstyles[j], label='lens bins %d-%d'%(i+1,j+1))\n",
    "            plt.plot(ells, cl_gg_theory_limber[i][j], color=colors[i], linestyle=':', marker = markerstyles[j], label='limber lens bins %d-%d'%(i+1,j+1))\n",
    "            #plt.plot(ells, (cl_gg_theory[i][j]-cl_gg_theory_limber[i][j])/cl_gg_theory[i][j], color=colors[i], linestyle=linestyles[j], label='lens bins %d-%d'%(i+1,j+1))\n",
    "\n",
    "    else:\n",
    "        cl_binned = cl_gg_theory[i][i]\n",
    "        s.add_ell_cl('galaxy_density_cl',  \n",
    "         \"lens%d\" % (i),  \n",
    "         \"lens%d\" % (i), \n",
    "         ell_centers,  \n",
    "         cl_binned,  \n",
    "         window=sacc_ell_windows\n",
    "        )\n",
    "\n",
    "        plt.plot(ells, cl_gg_theory[i][i], color=colors[i], linestyle=linestyles[0], label='lens bins %d-%d'%(i+1,i+1))\n",
    "        plt.plot(ells, cl_gg_theory_limber[i][i], color=colors[i], linestyle=linestyles[1], label='limber lens bins %d-%d'%(i+1,i+1))\n",
    "        #plt.plot(ells, (cl_gg_theory[i][i]-cl_gg_theory_limber[i][i])/cl_gg_theory[i][i], color=colors[i], linestyle=linestyles[i], label='lens bins %d-%d'%(i+1,i+1))\n",
    "       \n",
    "if not auto_galaxy_clustering_only:\n",
    "    plt.yscale('symlog', linthresh=1.e-10)\n",
    "    plt.ylim(-1.e-6, 1.e-4)\n",
    "else:\n",
    "    plt.yscale('log')\n",
    "    plt.ylim(1.e-8, 1.e-4)\n",
    "plt.xscale('log')\n",
    "plt.legend(loc='lower center', ncols=3)\n",
    "plt.title('lens-lens', fontsize=20)\n",
    "plt.ylabel(r'$C_{g,g}(l)$', fontsize=20)\n",
    "plt.xlabel('ell', fontsize=20)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc428312",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(10,10))\n",
    "tracer_combos = [[0, 2], [0, 3], [0, 4], [1, 3], [1, 4], [2, 4], [3, 4]]\n",
    "for inds in tracer_combos:\n",
    "    i = inds[0]\n",
    "    j = inds[1]\n",
    "    plt.loglog(ell_centers, cl_gs_theory[i][j], color = colors[i], marker=markerstyles[j], label='lens bin %d and source bin %d'%(i+1,j+1))\n",
    "    plt.loglog(ell_centers, cl_gs_theory_limber[i][j], linestyle=':', color = colors[i], marker=markerstyles[j], label='limber lens bin %d and source bin %d'%(i+1,j+1))\n",
    "    #plt.vlines(l_max_lens[i], 10**-10, 10**-6, color = 'k', linestyle=linestyles[j],)\n",
    "    cl_binned = cl_gs_theory[i][j]\n",
    "    s.add_ell_cl('galaxy_shearDensity_cl_e',  \n",
    "                 \"lens%d\" % (i),  \n",
    "                 \"src%d\" % (j), \n",
    "                 ell_centers,  \n",
    "                 cl_binned,  \n",
    "                 window=sacc_ell_windows\n",
    "                )  \n",
    "plt.yscale('symlog', linthresh=1.e-7)\n",
    "plt.ylim(-1.e-5, 1.e-5)\n",
    "plt.xscale('log')\n",
    "plt.legend(loc='lower center', ncols=3)\n",
    "plt.title('lens-source', fontsize=20)\n",
    "plt.ylabel(r'$C_{g,s}(l)$', fontsize=20)\n",
    "plt.xlabel('ell', fontsize=20)\n",
    "plt.show()  \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "84dce0e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(10,10))\n",
    "\n",
    "for i in range(nbin_z_src):\n",
    "    for j in range(i, nbin_z_src):\n",
    "        plt.loglog(ell_centers, cl_ss_theory[i][j], color=colors[i], marker=markerstyles[j], label='z bins %d and %d'%(i+1,j+1))\n",
    "        cl_binned = cl_ss_theory[i][j]\n",
    "        s.add_ell_cl(\"galaxy_shear_cl_ee\", \n",
    "                     \"src%d\" % (i), \n",
    "                     \"src%d\" % (j), \n",
    "                     ell_centers,  \n",
    "                     cl_binned,  \n",
    "                     window=sacc_ell_windows\n",
    "                     )\n",
    "\n",
    "plt.yscale('log')\n",
    "plt.ylim(1.e-11, 1.e-7)\n",
    "plt.xscale('log')\n",
    "plt.legend(loc='lower center', ncols=3)\n",
    "plt.title('source-source', fontsize=20)\n",
    "plt.ylabel(r'$C_{s,s}(l)$', fontsize=20)\n",
    "plt.xlabel('ell', fontsize=20)\n",
    "plt.show() "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "928c6a03",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}