I have many reasons to like Python: rapid prototyping, powerful libraries, great interactive tools like the IPython shell and notebook, and the ability to write beautiful code. I will admit, the last one definitely is in the eye of the beholder. Brandon Rhodes gave a good talk at PyCon Canada a couple of year ago in which he explained what makes Python beautiful to him. One of the points he made, and one that I wholeheartedly agree with, is that you can make Python code look just like the math you are coding up. We know how to write math; we know how to typeset equations so they look beautiful and are easy to read, and thus easier to understand. We can use that knowledge to write more readable Python code.
But I’m not a mathematician, I’m an astrophysicist. Can I take this further and make my code look like physics, not like math?
“Wait”, you might say, “math is the language of physics. If you can just write down equations in Python, what is missing?” The answer to this is “units”. We drill our students to understand that reporting a quantity without units is meaningless. We educate them about the power of dimensional analysis and its importance to make sure that their results make sense. Yet all of this is strangely absent from the codes we write.
I have probably wasted many hours because I forgot to document whether some function in my code expected lengths in Mpc or pc, critical density in $\mathrm{g\,cm}^{-3}$ or $M_\odot\,\mathrm{Mpc}^{-3}$, or simply messed up the conversion of the gravitational constant to astronomical units. This is where astropy‘s Units and Quantities changed the way I write code.
import numpy as np
from astropy import cosmology
from astropy import units as u
cosmology.set_current(cosmology.WMAP9)
rho_crit_0 = cosmology.critical_density(0)
print "Critical density at redshift 0:", rho_crit_0
Look, the output has units! Astropy added them automatically. What if I need the critical density in $M_\odot\,\mathrm{Mpc}^{-3}$? Simple, conversion between compatible units is built is:
print "Critical density at redshift 0:", rho_crit_0.to(u.solMass / u.Mpc**3)
Also one of the most common sources of errors in astronomical codes, forgetting to convert degree to radian before calling trigonometrical functions, is now automatically taken care of:
phi = 30 * u.deg
print "sin(%.1f %s) = %.4f" % (phi.value, phi.unit, np.sin(phi))
phi = phi.to(u.rad)
print "sin(%.5f %s) = %.4f" % (phi.value, phi.unit, np.sin(phi))
phi = phi.to(u.arcsec)
print "sin(%.0f %s) = %.4f" % (phi.value, phi.unit, np.sin(phi))
I am not going to repeat the entire Units and Quantities documentation here. The astropy website does a great job documenting their modules. My point here is that this has a profound impact on the way we can write code that does physical computations. We do not anymore have to care about the units variables have when we pass them to functions. Now there is a system that keeps track of it for us. No longer do we have to specify in code comments which units a function expects, no longer are we deprived of automated and reliable means to check that the return values of functions have units that make sense. Take as an illustration this simple function computing the gravitational force a body of mass $m$ experiences at a distance $r$ from the Sun:
from astropy.constants import G
from astropy.units import imperial
def grav_force(m, r):
"""
Compute the gravitational force on a body.
Parameters:
-----------
m: astropy.Quantity, mass of the body
r: astropy.Quantity, distance from the Sun
Returns:
--------
f: astropy.Quantity, gravitational force according to Newton's law.
"""
m_sun = 1 * u.solMass
f = G * m_sun * m / r**2
return f.to(u.N)
print "Gravitational force on a body of mass 1 lb at 1 Earth orbit from the Sun:", grav_force(1 * imperial.lb, 1 * u.AU)
print "Units of the gravitational constants in astropy:", G.unit
The wonderful thing is, the Python code above is completely agnostic with respect to the units passed to it. We pass the mass of the body in imperial pounds (this insanity is just for illustration!), multiply it by one Solar mass and the gravitational constant in SI units, and divide that by the square of an astronomical unit. Not only can we completely remove any error prone code to convert units to a common base system from our own code, we also have built-in checks that our units are consistent:
try:
grav_force(20 * u.m, 500 * u.kg)
except u.UnitsError, e:
print e
Writing code that uses the astropy Quantity objects forces you to use them consistently throughout your code. Addition and subtraction of Quantities and numbers will raise an exception. Multiplication of quantities with numbers that should have units but do not will lead to inconsistencies later on. Some will undoubtedly see the requirement to use units throughout the code as a burden. I think it is a tremendous aid in writing physics code. This requirement is exactly what we ask our students to do when they hand in their homework or their lab reports. We put that “burden” on our students for good reasons and exactly the same reasons apply to us when we write code: Knowing the scale of our numerical values and making sure the units coming out of the combination of physical quantities make sense. Astropy does this in a wonderful way, which enables us to have these advantages and at the same time relieve of us the tedious and error prone task of unit conversion!
I will not hide that not everything is perfect. There are a few numpy and scipy functions and modules that do not play nicely with Quantities (yet). However, I think that this is such a tremendous addition to the way we write physics code in Python that it would be a shame to have this “hidden” in a library used by astronomers only. I hope this part of astropy will make it to a more general science library.