agwb.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# We need\n",
    "\n",
    "- `HubbleInvMpcNow = cosmo.Hubble(0) # Hubble constant in Mpc-1`\n",
    "- ```\n",
    "  z_points = cosmo.get_background()['z']\n",
    "  chi_points = cosmo.get_background()['comov. dist.']\n",
    "  ```\n",
    "  in order to get a function chi(z)\n",
    "- ```\n",
    "  matterPk = cosmo.get_pk_array(k_vec, z_vec, \n",
    "                                  len(k_vec), len(z_vec), \n",
    "                                  nonlinear=isnonlinear)\n",
    "  ```\n",
    "- ```\n",
    "  cosmo = Class()\n",
    "  cosmo.set(params_cosmo)\n",
    "  cosmo.compute()\n",
    "  ```\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax\n",
    "# enable 64-bit precision\n",
    "from jax.config import config\n",
    "config.update(\"jax_enable_x64\", True)\n",
    "import jax.numpy as jnp\n",
    "import jax_cosmo as jc\n",
    "from jax_cosmo.power import linear_matter_power, nonlinear_matter_power\n",
    "\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cosmo = jc.Planck15()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cosmo.__dict__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cosmo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "constants = {\n",
    "    'JoulestoErg'      : 1.0E7,\n",
    "    'metersTocm'       : 1.0E2,\n",
    "    'metersToMpc'      : 3.2408E-23,\n",
    "    'gravityConstantG' : 6.673E-11, # G in m^3/ s^2 / kg\n",
    "    'speedOfLightC'    : 299792458.0, # c in m / s\n",
    "    'cMpcInvSec'       : 9.72E-15 # Speed of light in Mpc/sec\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def rho_crit(cosmo):\n",
    "    \"\"\"Calculates the critical energy density at present time\n",
    "    density in erg/cm^3 for a given cosmology\n",
    "    rho_crit = 3*H0**2 * c^2 /(8 pi G)\n",
    "\n",
    "    Args:\n",
    "        cosmo (Class): An instance of the computed cosmology\n",
    "\n",
    "    Returns:\n",
    "        float: Critical density at present time in erg/cm^3\n",
    "    \"\"\"\n",
    "    HubbleInvMpcNow = cosmo.h * 100 * constants['metersToMpc'] / constants['cMpcInvSec']\n",
    "    conversions = (constants['metersToMpc']**2 / constants['metersTocm']**3) * constants['JoulestoErg']\n",
    "    rhoCritical = conversions * 3.0 * HubbleInvMpcNow**2 * constants['speedOfLightC']**4 / (8.0 * jnp.pi * constants['gravityConstantG']) \n",
    "    return rhoCritical"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rho_crit(cosmo)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def chi_of_z(cosmo, z):\n",
    "    \"\"\"Calculates the comoving distance in Mpc for a given redshift\n",
    "\n",
    "    Args:\n",
    "        cosmo (Class): An instance of the computed cosmology\n",
    "\n",
    "    Returns:\n",
    "        float: Comoving distance in Mpc\n",
    "    \"\"\"\n",
    "    chi = jc.background.radial_comoving_distance(cosmo, 1.0 / (1.0 + z))\n",
    "    return chi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "z = jnp.linspace(0.0, 10.0, 1000)\n",
    "chi = chi_of_z(cosmo, z)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(z, chi)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def z_from_chi(cosmo, chi_vec):\n",
    "    \"\"\"Calculates the comoving distance in Mpc for a given redshift\n",
    "\n",
    "    Args:\n",
    "        cosmo (Class): An instance of the computed cosmology\n",
    "\n",
    "    Returns:\n",
    "        float: Comoving distance in Mpc\n",
    "    \"\"\"\n",
    "    from jax_cosmo.scipy.interpolate import interp\n",
    "    # create an array of z values to interpolate\n",
    "    z = jnp.linspace(0, 10, 1000)\n",
    "    chi = chi_of_z(cosmo, z)\n",
    "    z_vec = interp(chi_vec, chi, z)\n",
    "    return z_vec"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "chi = jnp.linspace(0.0, 6700.0, 1000)\n",
    "z = z_from_chi(cosmo, chi)\n",
    "plt.plot(chi, z)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# #Let's have a look at the linear power\n",
    "# k = jnp.logspace(-3,-0.5, 50).reshape(-1,1)\n",
    "# a = jnp.linspace(0.1, 1, 10).T\n",
    "# print(k.shape)\n",
    "# print(a.shape)\n",
    "\n",
    "# pk = linear_matter_power(cosmo, k/cosmo.h, a)\n",
    "# pk_nonlin = nonlinear_matter_power(cosmo, k/cosmo.h, a)\n",
    "\n",
    "# print(pk.shape)\n",
    "# print(pk_nonlin.shape)\n",
    "# plt.figure(figsize=(7,5))\n",
    "# for i in range(a.shape[0]):\n",
    "#     plt.loglog(k,jnp.sqrt(pk[:,i]/cosmo.h**3), label=f'a={a[i]:.2f}')\n",
    "# plt.legend()\n",
    "# plt.xlabel('k [Mpc]')\n",
    "# plt.ylabel('delta_k');\n",
    "# plt.title('Non-linear power spectrum')\n",
    "# plt.show()\n",
    "\n",
    "# plt.figure(figsize=(7,5))\n",
    "# for i in range(a.shape[0]):\n",
    "#     plt.loglog(k,jnp.sqrt(pk_nonlin[:,i]/cosmo.h**3), label=f'a={a[i]:.2f}')\n",
    "# plt.legend()\n",
    "# plt.xlabel('k [Mpc]')\n",
    "# plt.ylabel('delta_k');\n",
    "# plt.title('Non-linear power spectrum')\n",
    "# plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from jax_gw.signal.agwb import compute_cl, parser_with_arguments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "parser = parser_with_arguments()\n",
    "# LIGO BAND:\n",
    "# A_max = 6E-38\n",
    "# z_peak = 0.6\n",
    "# z_sigma = 0.7\n",
    "# LISA BAND:\n",
    "# A_max = \n",
    "# z_peak = \n",
    "# z_sigma = \n",
    "args_data = parser.parse_args(\"./jax_gw/data/stochastic_GW/ --preBessel --overwriteKernel\".split())\n",
    "cls = compute_cl(jnp.array([6E-38, 0.6, 0.7]), args_data, f_value=63.1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cls[2], cls[3], cls[4]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ell_arr = jnp.arange(0, len(cls[1][0]))\n",
    "plt.loglog(ell_arr[1:], ell_arr[1:]*(ell_arr[1:]+1)*cls[1][0][1:]/(2*jnp.pi), label='LIGO band')\n",
    "plt.xlabel(r'$\\ell$')\n",
    "plt.ylabel(r'$\\ell(\\ell+1)C_\\ell/(2\\pi)$')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from jax import random\n",
    "import jax.numpy as jnp\n",
    "import numpyro\n",
    "# import numpyro handlers\n",
    "from numpyro import handlers\n",
    "import numpyro.distributions as dist\n",
    "from numpyro.infer import MCMC, NUTS, Predictive, HMC"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_cl_data(A_max, z_peak, z_sigma):\n",
    "    samples = jnp.array([A_max, z_peak, z_sigma])\n",
    "    args_data = parser.parse_args(f\"./jax_gw/data/stochastic_GW/ --preBessel --overwriteKernel\".split())\n",
    "    cls = compute_cl(samples, args_data, f_value=63.1)\n",
    "    # add noise\n",
    "    cls_01 = cls[1][0] + cls[1][1]\n",
    "    return cls_01\n",
    "\n",
    "cl_data = generate_cl_data(6E-38, 0.5, 0.6)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_loglkl_from_cls(data_cl, theory_cl, l_vec):\n",
    "    # Combining theory_cl and data_cl to calculate the likelihood\n",
    "    # Using equation (3) of arXiv 1811.11584\n",
    "    # Note that most of the expression below can be precomputed if needed\n",
    "    # Note that data_cl already have noise inside\n",
    "\n",
    "    chi2_l = (2.0 * l_vec + 1.0) * \\\n",
    "                ( (data_cl / theory_cl) + jnp.log(theory_cl) ) \\\n",
    "            - (2.0 * l_vec - 1.0) * jnp.log(data_cl) \n",
    "    # Exclude l = 0 \n",
    "    if l_vec[0]==0:\n",
    "        chi2_l = chi2_l[1:]\n",
    "    chi2 = jnp.sum(chi2_l)\n",
    "    loglklhood = - 0.5 * chi2\n",
    "    return loglklhood\n",
    "\n",
    "\n",
    "def likelihood_fn(A_max=None, z_peak=None, z_sigma=None):\n",
    "    \"\"\"Likelihood function for the astrpphysical GW stochastic background (AGWB)\n",
    "\n",
    "    The likelihood is a Wishart distribution with a covariance given by the AGWB power spectrum\n",
    "\n",
    "    Args:\n",
    "        A_max (float): Maximum amplitude of the AGWB\n",
    "        z_peak (float): Redshift of the peak of the AGWB\n",
    "        z_sigma (float): Width of the AGWB\n",
    "    \"\"\"\n",
    "    # Sample the parameters\n",
    "    with handlers.seed(rng_seed=0):\n",
    "        A_max_sample = numpyro.sample(\"A_max\", dist.Uniform(1E-40, 1E-35))\n",
    "        z_peak_sample = numpyro.sample(\"z_peak\", dist.Uniform(0.0, 2.0))\n",
    "        z_sigma_sample = numpyro.sample(\"z_sigma\", dist.Uniform(0.0, 3.0))\n",
    "    str_formatted = f\"./jax_gw/data/stochastic_GW/ --preBessel --overwriteKernel\"\n",
    "    str_formatted_splitted = str_formatted.split()\n",
    "    args_data = parser.parse_args(str_formatted_splitted)\n",
    "    cls = compute_cl(\n",
    "        jnp.array([A_max_sample, z_peak_sample, z_sigma_sample]),\n",
    "        args_data, \n",
    "        f_value=63.1)\n",
    "    ell_arr = jnp.arange(0, len(cls[1][0]))\n",
    "\n",
    "    # Compute the log likelihood\n",
    "    loglklhood = compute_loglkl_from_cls(cl_data, cls[1][0], ell_arr)\n",
    "    numpyro.factor(\"loglklhood\", loglklhood)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rng_key = random.PRNGKey(0)\n",
    "rng_key, rng_key_ = random.split(rng_key)\n",
    "mcmc = MCMC(NUTS(likelihood_fn), num_warmup=1000, num_samples=1000, num_chains=1)\n",
    "mcmc.run(rng_key=rng_key_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_sparse(l_max, l_min=0, intervals=[30,40,240,1000], sample_distances=[1,10,20,40]):\n",
    "    \"\"\"If intervals is not sorted it will be sorted.\n",
    "    TODO: The default values probably need refinement.\n",
    "\n",
    "    Args:\n",
    "        l_vec (ndarray): The array of all ell at which the likelihood is evaluated\n",
    "        intervals (list, optional): The upper limit of each interval.\n",
    "                                    Each interval has different sampling rate.\n",
    "                                    Defaults to [30,40,240,1000].\n",
    "        sample_distances (list, optional): The distance between samples in each interval. \n",
    "                                           Defaults to [1,10,20,40].\n",
    "\n",
    "    Raises:\n",
    "        ValueError: If the specifications of the interval are inconsistent.\n",
    "\n",
    "    Returns:\n",
    "        ndarray: array of sparse ell.\n",
    "    \"\"\"\n",
    "    if (not isinstance(intervals, list))  or (not isinstance(sample_distances, list)):\n",
    "        raise ValueError(\"intervals and samples_distances should both be lists.\")\n",
    "    if len(intervals) != len(sample_distances):\n",
    "        raise ValueError(\"Arrays `intervals` and `sample_distances` should have same length\")\n",
    "    if any(x < 1 for x in sample_distances):\n",
    "        raise ValueError(\"All sample distances have to be greater than 1.\")\n",
    "    \n",
    "    # if np.max(np.array(intervals)).astype(int) < l_max.astype(int):\n",
    "    #     raise ValueError(\"Not sure how to deal with the interval from {} to {}\".format(np.max(intervals),l_max))\n",
    "    \n",
    "    intervals = sorted(intervals)    \n",
    "    \n",
    "    ell_value = l_min\n",
    "    ell_list = []\n",
    "    interv_idx = 0\n",
    "    while interv_idx < len(intervals):\n",
    "        upper_bound = intervals[interv_idx]\n",
    "        while ell_value < upper_bound and ell_value <= l_max:\n",
    "            ell_list.append(ell_value)\n",
    "            ell_value += sample_distances[interv_idx]\n",
    "        interv_idx += 1\n",
    "    if l_max not in ell_list:\n",
    "        ell_list.append(l_max)\n",
    "\n",
    "    l_sparse_vec = jnp.array(ell_list)\n",
    "    \n",
    "    return l_sparse_vec"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "make_sparse(1000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "jitted_make_sparse = jax.jit(make_sparse, static_argnums=(0,))\n",
    "\n",
    "jitted_make_sparse(1000)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Decoupling precompute from likelihood"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# samples contains the values of the parameters of the kernel\n",
    "# args contains the arguments of the code\n",
    "# transfer the samples to args\n",
    "from jax_gw.signal.agwb import *\n",
    "parser = parser_with_arguments()\n",
    "args = parser.parse_args(f\"./src/jax_gw/data/stochastic_GW/ --preBessel --overwriteKernel\".split())\n",
    "samples = np.array([1E-38, 0.7, 0.6])\n",
    "f_value = None\n",
    "f_ref = 63.1\n",
    "f_min = 20\n",
    "f_max = 500\n",
    "verbose = True\n",
    "\n",
    "\n",
    "args.A_max = samples[0]\n",
    "args.mean_z = samples[1]\n",
    "args.sigma_z = samples[2]\n",
    "nonlinear = 'Halofit'\n",
    "# Write here cosmological parameters used to calculate the data\n",
    "params_cosmo = {\n",
    "        'output': 'mPk',\n",
    "        'z_pk': '0., 3.0, 7.0, 10.0',\n",
    "        'P_k_max_1/Mpc': '70',\n",
    "        'non linear': nonlinear\n",
    "#         'gauge' : 'Newtonian' #TODO: commented this as it should be the same. Check!\n",
    "    }\n",
    "\n",
    "if not args.output_path: # Assign default name\n",
    "    if args.overwriteKernel:\n",
    "        args.output_path = os.path.join(args.input_dir,'data_cl_f_l_GAUSS.fits')\n",
    "    else:\n",
    "        args.output_path = os.path.join(args.input_dir,'data_cl_f_l_TABLES.fits')\n",
    "    \n",
    "if not args.bessel_path: # Assign default name\n",
    "    args.bessel_path = os.path.join(args.input_dir,'sph_bessel_k_z_l_TEST_gwtools.fits')\n",
    "\n",
    "# Assign absolute path for all files\n",
    "input_dir  = os.path.abspath(args.input_dir)\n",
    "path = {\n",
    "    'f'      : os.path.join(input_dir, args.f_fname),\n",
    "    'z'      : os.path.join(input_dir, args.z_fname),\n",
    "    'A'      : os.path.join(input_dir, args.A_fname),\n",
    "    'output' : os.path.abspath(args.output_path),\n",
    "    'bessel' : os.path.abspath(args.bessel_path)\n",
    "}\n",
    "\n",
    "### NOTE: temporary comment out. put back in when interpolating from files\n",
    "\n",
    "# Check existence of input files and warning before overwriting output\n",
    "# assert os.path.isfile(path['f']), 'File {} not found!'.format(path['f'])\n",
    "# assert os.path.isfile(path['z']), 'File {} not found!'.format(path['z'])\n",
    "# assert os.path.isfile(path['A']), 'File {} not found!'.format(path['A'])\n",
    "# if os.path.isfile(path['output']) and args.storeCl:\n",
    "#     print('WARNING! I am going to overwrite a pre-existing data file!')\n",
    "    \n",
    "\n",
    "\n",
    "# # Import files\n",
    "# # initially 71 redshifts and 141 frequencies. Kernel A in erg/cm^3.\n",
    "# f_vec  = np.genfromtxt(path['f'], delimiter='\\t')\n",
    "# z_vec = np.genfromtxt(path['z'], delimiter='\\t')\n",
    "# A_z_f  = np.genfromtxt(path['A'])\n",
    "# # Check that the imported arrays have the right dimensions and consistent with input parameters\n",
    "# assert z_vec.shape+f_vec.shape==A_z_f.shape, 'The dimensions of the imported arrays are wrong!'\n",
    "# assert args.z_min>=z_vec.min() and args.z_max<=z_vec.max(), 'Check redshift boundaries!'\n",
    "# # Interpolate the kernel and check that the arguments are z and f in this order\n",
    "# A_z_f_interp = interp2d(z_vec, f_vec, A_z_f.T, kind='cubic')\n",
    "# assert z_vec.min()==A_z_f_interp.x_min and z_vec.max()==A_z_f_interp.x_max\n",
    "# assert f_vec.min()==A_z_f_interp.y_min and f_vec.max()==A_z_f_interp.y_max\n",
    "l_vec = np.arange(args.l_max+1)\n",
    "if args.full_ell:\n",
    "    l_compute = np.arange(args.l_max+1)\n",
    "else:\n",
    "    l_compute = jax.jit(make_sparse, static_argnums=(0,))(args.l_max)\n",
    "\n",
    "### DENSE\n",
    "\n",
    "x_vec = get_x_full(ell_max=args.l_max, x_min=args.x_min, \n",
    "                    after=args.num_after_max, points_pp=args.points_pp)\n",
    "\n",
    "# Create vectors\n",
    "k_num = int(args.k_density * (np.log10(args.k_max) - np.log10(args.k_min)))\n",
    "k_vec = create_array(args.k_min, args.k_max, k_num, \"log\")\n",
    "\n",
    "\n",
    "### SPARSE\n",
    "k_sparse = create_array(args.k_min, args.k_max+1, args.k_sparse_num, 'log')\n",
    "z_sparse = create_array(args.z_min, args.z_max, args.z_sparse_num, 'log')    \n",
    "\n",
    "intermediate_grids = get_intermediate_grids(k_vec, x_vec, k_sparse)\n",
    "\n",
    "#     print(\"NOTE: choosing very narrow frequency interval\")\n",
    "if f_value is None:\n",
    "    f_vec = create_array(f_min, f_max, args.f_num, args.f_spacing)\n",
    "else:\n",
    "    f_vec = [f_value,]\n",
    "\n",
    "# Calculate the matter power spectrum. This is frequency independent.\n",
    "# This is the only place were we need Class\n",
    "# cosmo = Class()\n",
    "# cosmo.set(params_cosmo)\n",
    "# cosmo.compute()\n",
    "\n",
    "cosmo = jc.Planck15()\n",
    "\n",
    "\n",
    "b_eff, deltaM_eff, assorted_grids = \\\n",
    "                get_cosmo_eff(cosmo, z_sparse, intermediate_grids, \n",
    "                                args, nonlinear)\n",
    "\n",
    "# used in the evaluation of noise\n",
    "chi_mid = chi_from_z(cosmo, assorted_grids[\"z_mid\"])\n",
    "    \n",
    "    \n",
    "# Precompute Spherical Bessel Function\n",
    "try: \n",
    "    if verbose:\n",
    "        print(\"Checking for pre-computed Spherical Bessel\")\n",
    "    assert os.path.isfile(path['bessel'])\n",
    "    _l = read_data_from_fits(path['bessel'], 'l')\n",
    "    _x = read_data_from_fits(path['bessel'], 'x')\n",
    "    assert _x.size == x_vec.size\n",
    "    assert _x.min() == x_vec.min()\n",
    "    assert _x.max() == x_vec.max()\n",
    "    # TODO: re-implement this check in jax\n",
    "    # assert set(_l) >= set(l_compute)\n",
    "except AssertionError:\n",
    "    print(\"Could not find consistent precomputed Bessel\")\n",
    "    if args.preBessel:\n",
    "        print('WARNING! I am going to overwrite the precomputed bessel file!')\n",
    "        print(\"Writing Bessel. This might take a while time and it might require a lot of memory\")\n",
    "\n",
    "        after_func = lambda l: max(l, args.min_after_nu)\n",
    "        write_sph_bessel(path['bessel'], l_compute, x_vec=x_vec,\n",
    "                            before=args.num_before_nu, after=after_func)\n",
    "        print('Finished writing Bessel')\n",
    "    else:\n",
    "        path['bessel'] = None\n",
    "        print('Not going to use precomputed Bessel Function')\n",
    "        print(\"Use --preBessel True to store and use spherical Bessel Functions for these k and z vectors\")\n",
    "else:\n",
    "    if verbose:\n",
    "        print(\"Found Precomputed Bessels\")\n",
    "if verbose:\n",
    "    print(\"Recovering Bessel\")\n",
    "bessel_x_l = get_bessel_x_l(path['bessel'], l_compute)[None,...]\n",
    "\n",
    "not_chi_mask_nonzero = (~assorted_grids[\"chi_mask\"]).nonzero()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f_vec.min(), f_vec.max(), f_vec.size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def cl_narrowband_from_grids(params, f, args,):\n",
    "    f_len = 1\n",
    "    clustering = np.zeros((f_len, len(l_vec)))\n",
    "    noise = np.zeros(f_len)\n",
    "    data = np.zeros((f_len, len(l_vec)))\n",
    "    nf = 0\n",
    "    print(f\"\\r{nf} {f:.4f}\\tHz\", end=\" \")\n",
    "    sys.stdout.flush()\n",
    "        \n",
    "    if args.overwriteKernel:\n",
    "        A_eff, A_sparse = compute_kernel_on_grid(\n",
    "                                            params,\n",
    "                                            freq=f, \n",
    "                                            assorted_grids=assorted_grids, \n",
    "                                            args=args, \n",
    "                                            not_chi_mask_nonzero=not_chi_mask_nonzero,\n",
    "                                            A_kernel_interp2d=None)\n",
    "        A_eff = A_eff * f\n",
    "        A_sparse = A_sparse * f\n",
    "    else:\n",
    "        A_eff, A_sparse = compute_kernel_on_grid(freq=f, \n",
    "                                            assorted_grids=assorted_grids, \n",
    "                                            args=args, \n",
    "                                            A_kernel_interp2d=A_z_f_interp)\n",
    "        \n",
    "    noise = noise.at[nf].set(compute_spatial_shot_noise(cosmo,\n",
    "                                            A_z=A_sparse, \n",
    "                                            chi_vec=chi_mid, \n",
    "                                            n_G=args.n_G))\n",
    "    clustering_l = compute_clustering_cl(cosmo, A_eff, b_eff, deltaM_eff, \n",
    "                                            bessel_x_l, intermediate_grids[\"chi_grid\"], \n",
    "                                            k_vec)\n",
    "    clustering = clustering.at[nf].set(\n",
    "        interpolate_cl(\n",
    "            clustering_l, \n",
    "            l_compute, \n",
    "            l_vec, \n",
    "    ))\n",
    "        \n",
    "    # data[nf,:] = clustering[nf,:] + noise[nf]\n",
    "    data = data.at[nf].set(clustering[nf,:] + noise[nf])\n",
    "    return data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def cl_broadband_from_grids(params, f_vec, f_ref, args,):\n",
    "    f_len = 1\n",
    "    clustering = np.zeros((f_len, len(l_vec)))\n",
    "    noise = np.zeros(f_len)\n",
    "    data = np.zeros((f_len, len(l_vec)))\n",
    "    # print(f\"\\r{0} {f_vec[0]:.4f}-{f_vec[-1]:.4f}\\tHz\", end=\" \")\n",
    "    sys.stdout.flush()\n",
    "    nf=0\n",
    "    if args.overwriteKernel:\n",
    "        A_eff, A_sparse = compute_kernel_on_grid(\n",
    "                                            params,\n",
    "                                            freq=f_ref, \n",
    "                                            assorted_grids=assorted_grids, \n",
    "                                            args=args, \n",
    "                                            not_chi_mask_nonzero=not_chi_mask_nonzero,\n",
    "                                            A_kernel_interp2d=None)\n",
    "        f_slope = 2/3\n",
    "        broadband = 1 / (f_slope + 2) / f_ref**f_slope * \\\n",
    "            (f_vec[-1]**(f_slope+2) - f_vec[0]**(f_slope+2))\n",
    "        A_eff = A_eff * broadband\n",
    "        A_sparse = A_sparse * broadband\n",
    "    else:\n",
    "        A_eff, A_sparse = compute_kernel_on_grid(freq=f, \n",
    "                                            assorted_grids=assorted_grids, \n",
    "                                            args=args, \n",
    "                                            A_kernel_interp2d=A_z_f_interp)\n",
    "        \n",
    "    noise = noise.at[0].set(compute_spatial_shot_noise(cosmo,\n",
    "                                            A_z=A_sparse, \n",
    "                                            chi_vec=chi_mid, \n",
    "                                            n_G=args.n_G))\n",
    "    clustering_l = compute_clustering_cl(cosmo, A_eff, b_eff, deltaM_eff, \n",
    "                                            bessel_x_l, intermediate_grids[\"chi_grid\"], \n",
    "                                            k_vec)\n",
    "    clustering = clustering.at[nf].set(\n",
    "        interpolate_cl(\n",
    "            clustering_l, \n",
    "            l_compute, \n",
    "            l_vec, \n",
    "    ))\n",
    "        \n",
    "    # data[nf,:] = clustering[nf,:] + noise[nf]\n",
    "    data = data.at[nf].set(clustering[nf,:] + noise[nf])\n",
    "    return data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from jax import random\n",
    "import jax.numpy as jnp\n",
    "import numpyro\n",
    "# import numpyro handlers\n",
    "from numpyro import handlers\n",
    "import numpyro.distributions as dist\n",
    "from numpyro.infer import MCMC, NUTS, Predictive, HMC"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_cl_data(A_max, z_peak, z_sigma, f_vec, f_ref):\n",
    "    samples = jnp.array([A_max, z_peak, z_sigma])\n",
    "    str_formatted = f\"./src/jax_gw/data/stochastic_GW/ --preBessel --overwriteKernel\"\n",
    "    str_formatted_splitted = str_formatted.split()\n",
    "    args_data = parser.parse_args(str_formatted_splitted)\n",
    "    A_max_sample, z_peak_sample, z_sigma_sample = samples\n",
    "    cls = cl_broadband_from_grids(\n",
    "        jnp.array([\n",
    "            A_max_sample*1E-37,\n",
    "            z_peak_sample, \n",
    "            z_sigma_sample\n",
    "        ]),\n",
    "        f_vec, \n",
    "        f_ref,\n",
    "        args=args_data,\n",
    "    )\n",
    "    # add noise\n",
    "    cls_01 = cls[1][0] + cls[1][1]\n",
    "    return cls_01\n",
    "\n",
    "f_min, f_max = 20, 500\n",
    "f_ref = 63.1\n",
    "f_vec = np.linspace(f_min, f_max, 10)\n",
    "cl_data = generate_cl_data(0.6, 0.5, 0.6, f_vec, f_ref)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(cl_data.min(), cl_data.max())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_loglkl_from_cls(data_cl, theory_cl, l_vec):\n",
    "    # Combining theory_cl and data_cl to calculate the likelihood\n",
    "    # Using equation (3) of arXiv 1811.11584\n",
    "    # Note that most of the expression below can be precomputed if needed\n",
    "    # Note that data_cl already have noise inside\n",
    "\n",
    "    chi2_l = (2.0 * l_vec + 1.0) * \\\n",
    "                ( (data_cl / theory_cl) + jnp.log(theory_cl) ) \\\n",
    "            - (2.0 * l_vec - 1.0) * jnp.log(data_cl) \\\n",
    "                - 2.0 * jnp.log(data_cl)\n",
    "\n",
    "    \n",
    "    # Exclude l = 0 from the sum\n",
    "    chi2_l = chi2_l[1:]\n",
    "    chi2 = jnp.sum(chi2_l)\n",
    "    loglklhood = - 0.5 * chi2 + 1225505.9\n",
    "    return loglklhood\n",
    "\n",
    "def likelihood_fn(A_max=None, z_peak=None, z_sigma=None):\n",
    "    \"\"\"Likelihood function for the astrpphysical GW stochastic background (AGWB)\n",
    "\n",
    "    The likelihood is a Wishart distribution with a covariance given by the AGWB power spectrum\n",
    "\n",
    "    Args:\n",
    "        A_max (float): Maximum amplitude of the AGWB\n",
    "        z_peak (float): Redshift of the peak of the AGWB\n",
    "        z_sigma (float): Width of the AGWB\n",
    "    \"\"\"\n",
    "    # Sample the parameters\n",
    "    # with handlers.seed(rng_seed=0):\n",
    "    A_max_sample = numpyro.sample(\"A_max\", dist.Uniform(0.4, 0.8))\n",
    "    z_peak_sample = numpyro.sample(\"z_peak\", dist.Uniform(0.2, 0.8))\n",
    "    z_sigma_sample = numpyro.sample(\"z_sigma\", dist.Uniform(0.2, 0.8))\n",
    "        \n",
    "\n",
    "    str_formatted = f\"./src/jax_gw/data/stochastic_GW/ --preBessel --overwriteKernel\"\n",
    "    str_formatted_splitted = str_formatted.split()\n",
    "    args_data = parser.parse_args(str_formatted_splitted)\n",
    "    f_min, f_max = 20, 500\n",
    "    f_ref = 63.1\n",
    "    f_vec = np.linspace(f_min, f_max, 10)\n",
    "    cls = cl_broadband_from_grids(\n",
    "        jnp.array([\n",
    "            A_max_sample*1E-37,\n",
    "            z_peak_sample, \n",
    "            z_sigma_sample\n",
    "        ]),\n",
    "        f_vec, \n",
    "        f_ref,\n",
    "        args=args_data,\n",
    "    )\n",
    "    ell_arr = jnp.arange(0, len(cls[0]))\n",
    "\n",
    "    # Compute the log likelihood\n",
    "    loglklhood = compute_loglkl_from_cls(cl_data, cls[0], ell_arr)\n",
    "    print(loglklhood)\n",
    "    numpyro.factor(\"loglklhood\", loglklhood)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from numpyro.infer import init_to_feasible, init_to_median, init_to_sample\n",
    "rng_key = random.PRNGKey(0)\n",
    "rng_key, rng_key_ = random.split(rng_key)\n",
    "kernel = NUTS(\n",
    "    likelihood_fn,\n",
    "    init_strategy=init_to_median(),\n",
    "    )\n",
    "mcmc = MCMC(kernel, num_warmup=100, num_samples=100, num_chains=1)\n",
    "mcmc.run(rng_key=rng_key_)\n",
    "mcmc.print_summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from numpyro.infer import init_to_feasible, init_to_median, init_to_sample\n",
    "rng_key = random.PRNGKey(0)\n",
    "rng_key, rng_key_ = random.split(rng_key)\n",
    "kernel = NUTS(\n",
    "    likelihood_fn,\n",
    "    init_strategy=init_to_median(),\n",
    "    )\n",
    "mcmc = MCMC(kernel, num_warmup=100, num_samples=100, num_chains=1)\n",
    "mcmc.run(rng_key=rng_key_)\n",
    "mcmc.print_summary()\n",
    "posterior = mcmc.get_samples()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Instrumental Noise"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Runs\n",
    "\n",
    "- One run with f=0.001"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.9.16 ('jax-gwb')",
   "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.9.16"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "9015ef353a66a28b1cf10b4768999d3c6c5ace8f5466ddae6f166734f75ffc03"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}