Hertzsprung-Russell.ipynb

{
 "metadata": {
  "name": "Hertzsprung-Russell"
 },
 "nbformat": 3,
 "nbformat_minor": 0,
 "worksheets": [
  {
   "cells": [
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "%pylab inline"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "with open('data/hygfull.csv') as f:\n",
      "    print f.read(440) + '...'"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "import csv\n",
      "with open('data/hygfull.csv') as f:\n",
      "    stars = list(csv.DictReader(f))"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "stars[123]"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "tau = 2.0 * pi"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "from spectral_classification import build_color_chart\n",
      "color_chart = build_color_chart('starcolors.txt')"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "def paren(spectral_class):\n",
      "    letter = spectral_class[:2]\n",
      "    roman = spectral_class[2:] or 'V'\n",
      "    return '{}({})'.format(letter, roman)\n",
      "\n",
      "print paren('G8III')\n",
      "print paren('B9')"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "stars2 = [\n",
      "    star for star in stars\n",
      "    if star['ColorIndex'] and\n",
      "       paren(star['Spectrum']) in color_chart\n",
      "    ]\n",
      "print len(stars2), 'of', len(stars), 'have good colors'\n",
      "\n",
      "star_magnitudes = array(\n",
      "    [ float(star['AbsMag']) for star in stars2 ])\n",
      "\n",
      "star_colors = array(\n",
      "    [ float(star['ColorIndex']) for star in stars2 ])\n",
      "\n",
      "star_spectra = array(\n",
      "    [ star['Spectrum'] for star in stars2 ])"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "print star_magnitudes[:5]\n",
      "print star_colors[:5]\n",
      "print star_spectra[:5]"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "hr_axes = [-0.5, 2, 15, -5]"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "scatter(star_colors, star_magnitudes, s=1,\n",
      "             c=(0, 0, 0), linewidths=0)\n",
      "axis(hr_axes)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "star_rgbs = [\n",
      "    color_chart[paren(spectral_class)]\n",
      "    for spectral_class in star_spectra\n",
      "    ]"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "scatter(star_colors, star_magnitudes, s=1,\n",
      "        c=star_rgbs, linewidths=0.0)\n",
      "axis(hr_axes)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "hr_range = [[-5, 15], [-0.5, 2]]"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "heatmap, yedges, xedges = histogram2d(\n",
      "    star_magnitudes, star_colors, bins=50,\n",
      "    range=hr_range)\n",
      "\n",
      "extent = [xedges[0], xedges[-1], yedges[-1], yedges[0]]\n",
      "imshow(heatmap, aspect=0.08, extent=extent)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "heatmap[10][10]"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "from matplotlib.cm import get_cmap"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "imshow(heatmap, cmap=get_cmap('gray'),\n",
      "       aspect=0.08, extent=extent)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "What if we want some color?"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "colorgrid = zeros((50, 50, 3))\n",
      "\n",
      "ymin, ymax, xmin, xmax = -5, 15, -0.5, 2\n",
      "y_range = ymax - ymin\n",
      "x_range = xmax - xmin\n",
      "\n",
      "for i in range(len(stars2)):\n",
      "    xbin = 50.0 * (star_colors[i] - xmin) / x_range\n",
      "    ybin = 50.0 * (star_magnitudes[i] - ymin) / y_range\n",
      "\n",
      "    if not 0.0 <= xbin < 50.0:\n",
      "        continue\n",
      "    if not 0.0 <= ybin < 50.0:\n",
      "        continue\n",
      "\n",
      "    xbin = int(xbin)\n",
      "    ybin = int(ybin)\n",
      "    color = color_chart[paren(star_spectra[i])]\n",
      "    colorgrid[ybin][xbin] += color\n",
      "\n",
      "print heatmap[10][10]\n",
      "print colorgrid[25][20:25]\n",
      "colorgrid /= colorgrid.max()\n",
      "print colorgrid[25][20:25]"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "imshow(colorgrid, aspect=0.08, extent=extent)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "x = linspace(0.0, 1.2, 100)\n",
      "plot(x, x ** 2)\n",
      "plot(x, x)\n",
      "plot(x, x ** 0.5)\n",
      "axes().set_aspect(1.0)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "plot = plt.figure().add_subplot(111)\n",
      "plot.imshow(colorgrid ** 0.5, aspect=0.08,\n",
      "            extent=extent)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [],
     "language": "python",
     "metadata": {},
     "outputs": []
    }
   ],
   "metadata": {}
  }
 ]
}