{
 "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": null,
   "metadata": {},
   "outputs": [],
   "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": null,
   "metadata": {},
   "outputs": [],
   "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": null,
   "metadata": {},
   "outputs": [],
   "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": null,
   "metadata": {},
   "outputs": [],
   "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.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}