Plot the exoplanets.org database with Python/Matplotlib

exoplanetplotter.py

'''
Downloads, unpacks and "pickles" (saves a python binary) of the exoplanets.org
planet database, saves the data in nested dictionaries.

For exoplanets.org's parameter field keywords, see Table 1 of 
Wright et al. 2011: http://arxiv.org/pdf/1012.5676v3.pdf

Brett Morris
'''

import numpy as np
from glob import glob
import time
import os
import cPickle
from urllib import urlopen
from matplotlib import pyplot as plt

# Save the exoplanet.org database to the directory that this file is saved in
exodbPath = os.path.dirname(os.path.abspath(__file__))

def downloadAndPickle():
    pklDatabaseName = os.path.join(exodbPath,'exoplanetDB.pkl')	 ## Name of exoplanet database C-pickle
    pklDatabasePaths = glob(pklDatabaseName)   ## list of files with the name pklDatabaseName in cwd
    csvDatabaseName = os.path.join(exodbPath,'exoplanets.csv')  ## Path to the text file saved from exoplanets.org
    csvDatabasePaths = glob(csvDatabaseName)

    '''If there's a previously archived database pickle in this current working 
        directory then use it, if not, grab the data from exoplanets.org in one big CSV file and make one.
        If the old archive is >14 days old, grab a fresh version of the database from exoplanets.org.
        '''
    if csvDatabasePaths == []:
        print 'No local copy of exoplanets.org database. Downloading one...'
        rawCSV = urlopen('http://www.exoplanets.org/csv-files/exoplanets.csv').read()
        saveCSV = open(csvDatabaseName,'w')
        saveCSV.write(rawCSV)
        saveCSV.close()
    else: 
        '''If the local copy of the exoplanets.org database is >14 days old, download a new one'''
        secondsSinceLastModification = time.time() - os.path.getmtime(csvDatabaseName) ## in seconds
        daysSinceLastModification = secondsSinceLastModification/(60*60*24*30)
        if daysSinceLastModification > 7:
            print 'Your local copy of the exoplanets.org database is >14 days old. Downloading a fresh one...'
            rawCSV = urlopen('http://www.exoplanets.org/csv-files/exoplanets.csv').read()
            saveCSV = open(csvDatabaseName,'w')
            saveCSV.write(rawCSV)
            saveCSV.close()
        else: print "Your local copy of the exoplanets.org database is <14 days old. That'll do."

    if len(pklDatabasePaths) == 0:
        print 'Parsing '+os.path.split(csvDatabaseName)[1]+', the CSV database from exoplanets.org...'
        rawTable = open(csvDatabaseName).read().splitlines()
        labels = rawTable[0].split(',')
        exoplanetDB = {}
        planetNameColumn = np.arange(len(labels))[np.array(labels,dtype=str)=='NAME'][0]
        for row in range(1,len(rawTable)): 
            splitRow = rawTable[row].split(',')
            exoplanetDB[splitRow[planetNameColumn]] = {}	## Create dictionary for this row's planet
            for col in range(0,len(splitRow)):
                exoplanetDB[splitRow[planetNameColumn]][labels[col]] = splitRow[col]
        
        output = open(pklDatabaseName,'wb')
        cPickle.dump(exoplanetDB,output)
        output.close()
    else: 
        print 'Using previously parsed database from exoplanets.org...'
        inputFile = open(pklDatabaseName,'rb')
        exoplanetDB = cPickle.load(inputFile)
        inputFile.close()
    
    return exoplanetDB

def getmass(planet):
    '''Get mass in Jupiter masses, return zero if there is none'''
    if exoplanetDB[planet]['MASS'] == '':
        return 0.0
    else:
        return float(exoplanetDB[planet]['MASS'])

def getradius(planet):
    '''Get radius in Jupiter radii, rturn zero if there is none'''
    if exoplanetDB[planet]['R'] == '':
        return 0.0
    else:
        return float(exoplanetDB[planet]['R'])

# Load the database
exoplanetDB = downloadAndPickle()

# Get the names of all planets in the database
allplanets = exoplanetDB.keys()

# Get a list of the non-KOI planets, because the KOIs have false masses listed
nonKOIplanets = [planet for planet in allplanets if planet.startswith('KOI') == False]

# Get the masses (in Jupiter masses) and radii (in Jupiter radii) 
# of every non-KOI planet where the value are non-zero
masses = [getmass(planet) for planet in nonKOIplanets if getmass(planet) != 0 \
          and getradius(planet) != 0]
radii = [getradius(planet) for planet in nonKOIplanets if getmass(planet) != 0 \
          and getradius(planet) != 0]

plt.loglog(masses, radii, 'k.')
plt.xlabel('$M/M_J$')
plt.ylabel('$R/R_J$')
plt.title(str(len(masses))+' Exoplanets from exoplanets.org')
plt.show()