{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Cosmology\n",
    "\n",
    "## RUBIX cosmology module\n",
    "\n",
    "For RUBIX we implemented a cosmology modlue that is similar to the astropy cosmology module and is jax compatible. Therefore, we follow https://github.com/ArgonneCPAC/dsps/blob/main/dsps/cosmology/flat_wcdm.py\n",
    "\n",
    "We assume Planck15 cosmology: The present day matter density Om0 is set to 0.3089. The present day dark energy equation of state w0 is set to -1. The dark energy equation of state parameter wa is set to 0.0. The Hubble constant h is set to 0.6774."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "BaseCosmology(Om0=f32[], w0=f32[], wa=f32[], h=f32[])\n",
      "FlatLambdaCDM(name=\"Planck15\", H0=67.74 km / (Mpc s), Om0=0.3075, Tcmb0=2.7255 K, Neff=3.046, m_nu=[0.   0.   0.06] eV, Ob0=0.0486)\n"
     ]
    }
   ],
   "source": [
    "# NBVAL_IGNORE_OUTPUT\n",
    "from rubix.cosmology import PLANCK15 as rubix_cosmo\n",
    "\n",
    "# Compare to astropy\n",
    "from astropy.cosmology import Planck15 as astropy_cosmo\n",
    "\n",
    "print(rubix_cosmo)\n",
    "print(astropy_cosmo)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Comparison to astropy.cosmology\n",
    "\n",
    "We can now compare our RUBIX cosmology module with the astropy cosmology module. We show the comparison for the angular diameter distance, the comoving distance, the lookback to z and the age."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Angular Diameter Distance\n",
      "rubix cosmo:  702.5322\n",
      "astropy cosmo:  702.3747610737071 Mpc\n",
      "Comoving Distance\n",
      "rubix cosmo:  843.0387\n",
      "astropy cosmo:  842.8497132884485 Mpc\n",
      "lookback to z\n",
      "rubix cosmo:  2.5104787\n",
      "astropy cosmo:  2.509878627257186 Gyr\n",
      "Age\n",
      "rubix cosmo:  [11.310789]\n",
      "astropy cosmo:  11.287737269639198 Gyr\n"
     ]
    }
   ],
   "source": [
    "# NBVAL_IGNORE_OUTPUT\n",
    "z = 0.2\n",
    "print(\"Angular Diameter Distance\")\n",
    "print(\"rubix cosmo: \",rubix_cosmo.angular_diameter_distance_to_z(z))\n",
    "print(\"astropy cosmo: \",astropy_cosmo.angular_diameter_distance(z))\n",
    "\n",
    "\n",
    "print(\"Comoving Distance\")\n",
    "print(\"rubix cosmo: \",rubix_cosmo.comoving_distance_to_z(z))\n",
    "print(\"astropy cosmo: \",astropy_cosmo.comoving_distance(z))\n",
    "\n",
    "print(\"lookback to z\")\n",
    "print(\"rubix cosmo: \",rubix_cosmo.lookback_to_z(z))\n",
    "print(\"astropy cosmo: \",astropy_cosmo.lookback_time(z))\n",
    "\n",
    "print(\"Age\")\n",
    "print(\"rubix cosmo: \",rubix_cosmo.age_at_z(z))\n",
    "print(\"astropy cosmo: \",astropy_cosmo.age(z))\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "9.5\n"
     ]
    }
   ],
   "source": [
    "from rubix.cosmology.utils import trapz\n",
    "import jax.numpy as jnp\n",
    "\n",
    "x = jnp.array([0, 1, 2, 3])\n",
    "y = jnp.array([0, 1, 4, 9])\n",
    "print(trapz(x, y))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.0\n",
      "1946.5874\n"
     ]
    }
   ],
   "source": [
    "from rubix.cosmology import PLANCK15 as rubix_cosmo\n",
    "import jax.numpy as jnp\n",
    "\n",
    "#from rubix.cosmology.base import scale_factor_to_redshift\n",
    "scale_factor = jnp.array(0.5)\n",
    "result = rubix_cosmo.scale_factor_to_redshift(jnp.array(0.5))\n",
    "print(result)  # Output: 1.0\n",
    "\n",
    "result2 = rubix_cosmo.comoving_distance_to_z(0.5)\n",
    "print(result2)  # Output: 0.0"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "rubix",
   "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.12.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}