Avoiding Loops: a homework example

In previous post, Tim has shown us that using array is much efficient than loops. Here I'm going to give a real example about how to use numpy array to solve one of our homework problem.

The example is CA4, problem3. We have a data set of the masses of known transit planets and want to find the power index of the mass distribution function.


##########################################################

from numpy import *
import asciitable
import pylab as py

# read the transit data
all_transits = asciitable.read('transits.csv',delimiter=",")
transit_mass = all_transits['MSINI']   

# find the data with 1 < m <20
ind = ((transit_mass >= 1)& (transit_mass <=20)).nonzero()[0]
mass_data = transit_mass[ind]

n =  mass_data.size  

# the total number of planets in the given mass range

# build an array for a possible alpha range
alpha_grid = arange(1.1,3.5,0.001)

# here we use numpy array rather than loop
# (1) we build 2D meshgrid arrays for alpha and mass
# NumPy provides a convenient function for generating meshgrid arrays,
# called, appropriately, meshgrid.

(alpha,mass) = meshgrid(alpha_grid,mass_data)
#print alpha
#print mass

#(2) Calculate the probability at each alpha and m
# Attention: array math operates on an element by element basis,
# i. .e, unlike matrix math.
mass_max =  mass_data.max()
mass_min = mass_data.min()
const = 1./(1-alpha)*(mass_max**(1.-alpha)-mass_min**(1.-alpha))
log_pdf = log(1./const*mass**(-1.*alpha))

# (3) we sum the log_pdf in the mass dimension
log_pdf_alpha = log_pdf.sum(axis = 0)
plot(alpha_grid,exp(log_pdf_alpha-log_pdf_alpha.max()))

##########################################################

The example code and transit planet data can be downloaded from here:

I find this introduction of numpy calculation very useful. It is totally worth your time.

http://www.physics.nmt.edu/~raymond/software/python_notes/paper003.html