redshift-plot.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Make a plot with both redshift and universe age axes using astropy.cosmology\n",
    "\n",
    "## Authors\n",
    "Neil Crighton, Stephanie T. Douglas\n",
    "\n",
    "## Learning Goals\n",
    "* Plot relationships using `matplotlib`\n",
    "* Add a second axis to a `matplotlib` plot\n",
    "* Relate distance, redshift, and age for two different types of cosmology using `astropy.cosmology`\n",
    "\n",
    "## Keywords\n",
    "units, physics, cosmology, matplotlib\n",
    "\n",
    "## Summary\n",
    "\n",
    "Each redshift corresponds to an age of the universe, so if you're plotting some quantity against redshift, it's often useful show the universe age too.  The relationship between the two changes depending the type of cosmology you assume, which is where `astropy.cosmology` comes in. In this tutorial we'll show how to use the tools in `astropy.cosmology` to make a plot like this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set up matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import Image\n",
    "Image(filename=\"ang_dist.png\", width=500)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We start with a cosmology object. We will make a flat cosmology (which means that the curvature density $\\Omega_k=0$)  with a hubble parameter of $70$ km/s/Mpc and matter density $\\Omega_M=0.3$ at redshift 0. The `FlatLambdaCDM` cosmology then automatically infers that the dark energy density $\\Omega_\\Lambda$ must $=0.7$, because $\\Omega_M + \\Omega_\\Lambda + \\Omega_k = 1$. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from astropy.cosmology import FlatLambdaCDM\n",
    "import astropy.units as u\n",
    "\n",
    "# In this case we just need to define the matter density \n",
    "# and hubble parameter at z=0.\n",
    "\n",
    "# Note the default units for the hubble parameter H0 are km/s/Mpc. \n",
    "# We will pass in a `Quantity` object with the units specified. \n",
    "\n",
    "cosmo = FlatLambdaCDM(H0=70*u.km/u.s/u.Mpc, Om0=0.3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that we could instead use one of the built-in cosmologies, like `WMAP9` or `Planck13`, in which case we would just redefine the `cosmo` variable.\n",
    "\n",
    "Now we need an example quantity to plot versus redshift. Let's use the angular diameter distance, which is the physical transverse distance (the size of a galaxy, say) corresponding to a fixed angular separation on the sky. To calculate the angular diameter distance for a range of redshifts:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "zvals = np.arange(0, 6, 0.1)\n",
    "dist = cosmo.angular_diameter_distance(zvals)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that we passed an array of redshifts to `cosmo.angular_diameter_distance` and it produced a corresponding array of distance values, one for each redshift. Let's plot them:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(6,4))\n",
    "ax = fig.add_subplot(111)\n",
    "ax.plot(zvals, dist)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To check the units of the angular diameter distance, take a look at the unit attribute:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dist.unit"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's put some age labels on the top axis. We're going to pick a series of round age values where we want to place axis ticks. You may need to tweak these depending on your redshift range to get nice, evenly spaced ticks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ages = np.array([13, 10, 8, 6, 5, 4, 3, 2, 1.5, 1.2, 1])*u.Gyr"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To link the redshift and age axes, we have to find the redshift corresponding to each age. The function `z_at_value` does this for us."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from astropy.cosmology import z_at_value\n",
    "ageticks = [z_at_value(cosmo.age, age) for age in ages]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we make the second axes, and set the tick positions using these values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(6,4))\n",
    "ax = fig.add_subplot(111)\n",
    "ax.plot(zvals, dist)\n",
    "ax2 = ax.twiny()\n",
    "ax2.set_xticks(ageticks)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We have ticks on the top axis at the correct ages, but they're labelled with the redshift, not the age. We can fix this by setting the tick labels by hand. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(6,4))\n",
    "ax = fig.add_subplot(111)\n",
    "ax.plot(zvals, dist)\n",
    "ax2 = ax.twiny()\n",
    "ax2.set_xticks(ageticks)\n",
    "ax2.set_xticklabels(['{:g}'.format(age) for age in ages.value])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We need to make sure the top and bottom axes have the same redshift limits. They may not line up properly in the above plot, for example, depending on your setup (the age of the universe should be ~13 Gyr at z=0). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(6,4))\n",
    "ax = fig.add_subplot(111)\n",
    "ax.plot(zvals, dist)\n",
    "ax2 = ax.twiny()\n",
    "ax2.set_xticks(ageticks)\n",
    "ax2.set_xticklabels(['{:g}'.format(age) for age in ages.value])\n",
    "zmin, zmax = 0.0, 5.9\n",
    "ax.set_xlim(zmin, zmax)\n",
    "ax2.set_xlim(zmin, zmax)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We're almost done. We just need to label all the axes, and add some minor ticks. Let's also tweak the y axis limits to avoid putting labels right near the top of the plot."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(6,4))\n",
    "ax = fig.add_subplot(111)\n",
    "ax.plot(zvals, dist)\n",
    "ax2 = ax.twiny()\n",
    "ax2.set_xticks(ageticks)\n",
    "ax2.set_xticklabels(['{:g}'.format(age) for age in ages.value])\n",
    "zmin, zmax = 0, 5.9\n",
    "ax.set_xlim(zmin, zmax)\n",
    "ax2.set_xlim(zmin, zmax)\n",
    "ax2.set_xlabel('Time since Big Bang (Gyr)')\n",
    "ax.set_xlabel('Redshift')\n",
    "ax.set_ylabel('Angular diameter distance (Mpc)')\n",
    "ax.set_ylim(0, 1890)\n",
    "ax.minorticks_on()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now for comparison, let's add the angular diameter distance for a different cosmology, from the Planck 2013 results. And then finally, we save the figure to a png file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from astropy.cosmology import Planck13\n",
    "dist2 = Planck13.angular_diameter_distance(zvals)\n",
    "\n",
    "fig = plt.figure(figsize=(6,4))\n",
    "ax = fig.add_subplot(111)\n",
    "ax.plot(zvals, dist2, label='Planck 2013')\n",
    "ax.plot(zvals, dist, label=\n",
    "        '$h=0.7,\\ \\Omega_M=0.3,\\ \\Omega_\\Lambda=0.7$')\n",
    "ax.legend(frameon=0, loc='lower right')\n",
    "ax2 = ax.twiny()\n",
    "ax2.set_xticks(ageticks)\n",
    "ax2.set_xticklabels(['{:g}'.format(age) for age in ages.value])\n",
    "zmin, zmax = 0.0, 5.9\n",
    "ax.set_xlim(zmin, zmax)\n",
    "ax2.set_xlim(zmin, zmax)\n",
    "ax2.set_xlabel('Time since Big Bang (Gyr)')\n",
    "ax.set_xlabel('Redshift')\n",
    "ax.set_ylabel('Angular diameter distance (Mpc)')\n",
    "ax.minorticks_on()\n",
    "ax.set_ylim(0, 1890)\n",
    "fig.savefig('ang_dist.png', dpi=200, bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`bbox_inches='tight'` automatically trims any whitespace from around the plot edges. \n",
    "\n",
    "And we're done!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Well, almost done.  Notice that we calculated the times on the upper axis using the original cosmology, not the new cosmology based on the Planck 2013 results. So strictly speaking, this axis applies only to the original cosmology, although the difference between the two is small. As an exercise, you can try plot two different upper axes, slightly offset from each other, to show the times corresponding to each cosmology. Take a look at the first answer to [this question on Stack Overflow](http://stackoverflow.com/questions/7733693/matplotlib-overlay-plots-with-different-scales) for some hints on how to go about this."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "astropy-tutorials": {
   "author": "Neil Crighton <neilcrighton@gmail.com>",
   "date": "April 2015",
   "description": "Demonstrates how to make a plot with both redshift and universe age axes, using astropy.cosmology and matplotlib.",
   "link_name": "Making a plot with both redshift and universe age axes",
   "name": "Making a plot with both redshift and universe age axes",
   "published": true
  },
  "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.6.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}