Creating and manipulating FITS files

Question

I wrote a program that manipulated data from VLA observations and created FITS files to be used in creating spectral energy distributions of high redshift radio galaxies. I really tried to be as efficient as possible. I made lists of the variables and lists I had to use. I looped everything where I could, but it still came out horribly convoluted.

#!/usr/bin/python

import math
from astropy.io import ascii
import astropy.cosmology as cosmo
import pyfits as fits
import matplotlib.pyplot as plt
from matplotlib import lines
import numpy as np
import os

# removing any duplicates from the current directory
try:
    os.system("rm -f sed_and_alpha.png")
except:
    pass
try:
    os.system("rm -f sed_and_alpha.fits")
except:
    pass
#constants
CBAND = 4.86*10**9
XBAND = 8.46*10**9
CF = (10**-26)

#functions
def get_luminosity_distance(redshift):
    '''finds luminosity distance'''
    lum_dist = cosmo.luminosity_distance(redshift)
    tmp = str(lum_dist)
    tmp = tmp.replace("Mpc","")
    tmp = float(tmp)
    return tmp*3.09*10**22
    print lum_dist

def get_alpha(s1,s2,v1,v2):
    '''gets alpha from two flux densities'''
    top = math.log10(s1/s2)
    bottom = math.log10(v2/v1)
    return top/bottom

def get_luminosity(D_l,z,a,nu1,nu2,Sv2):
    '''gets luminosity from alpha'''
    Sv2 = (Sv2/1000)*CF
    lum = (4*math.pi*((D_l**2)))*(1/((1+z)**(1+a)))*((nu1/nu2)**a)*(Sv2)
    return lum


#lists
available_redshifts = [] # For making sure there are no blank spaces
luminosity_distribution = [] # which calculated luminosities are available
alpha_list = [] # available alphas
new_luminosity_distribution = [] # luminosity distr. after str[-2:]
default_data = {}
list22 = []
list23 = []
list24 = []
list25 = []
list26 = []
list27 = []
list28 = []
exponent_list = [
                    list22,
                    list23,
                    list24,
                    list25,
                    list26,
                    list27,
                    list28,
                ]
exponent_list_names = [
                        "list22",
                        "list23",
                        "list24",
                        "list25",
                        "list26",
                        "list27",
                        "list28",
                    ]
range_column_list = [10.0**22, 10.0**23, 10.0**24,10.0**25,10.0**26,10.0**27,10.0**28]
ylist = []
num_galaxies_list = []
filename = "/home/vhx/Documents/code/python/fits/radio_template.fits/apj384116t3_mrt.txt"
table_data = ascii.read(filename) # getting the table data
# Determining which redshifts are available
for i in xrange(246):
    blank = True
    tmp = (table_data["z"][i])
    if tmp != 'NO' and table_data["SX"][i] != 'NO' and table_data["SC"][i] != 'NO':
        available_redshifts.append(i)
    else:
        pass
# get the spectral indices 
for i in range(len(available_redshifts)-1):
    alpha_list.append(get_alpha(table_data["SX"][available_redshifts[i]],table_data["SC"][available_redshifts[i]],XBAND,CBAND))
    print table_data["z"][available_redshifts[i]], table_data["SX"][available_redshifts[i]], table_data["SC"][available_redshifts[i]]


# solving for luminosity
for i in range(len(available_redshifts)-1):
    x = get_luminosity_distance(table_data["z"][available_redshifts[i]])
    luminosity_distribution.append(get_luminosity(x,table_data["z"][available_redshifts[i]],alpha_list[i],CBAND,XBAND,table_data["SX"][available_redshifts[i]]))
# removing the units from the luminosity distribution so it can be converted
for z in range(len(luminosity_distribution)-1):
    tmp = str(luminosity_distribution[z])
    new_tmp = tmp.replace("Mpc","")
    print new_tmp
    try:
        new_luminosity_distribution.append(float(new_tmp))
    except ValueError:
        print 'Found blank space!'
        pass
new_luminosity_distribution.pop(0)
# checking stuff
alpha_list[:] = [x for x in alpha_list if math.isnan(x) != True]
print alpha_list
k = 0
for item in alpha_list:
    k += item
print k/(len(alpha_list)-1)
print max(alpha_list),min(alpha_list)
print '\n'*5
for item in new_luminosity_distribution:
    print '%e' % item
print '%e, %e' % (max(new_luminosity_distribution), min(new_luminosity_distribution))
sum = 0
for item in new_luminosity_distribution:
    sum += item
print sum/len(new_luminosity_distribution)
new_luminosity_distribution_array = np.array(new_luminosity_distribution)
alpha_list_array = np.array(alpha_list)
numbers = np.random.normal(size = 1000)
width = 1
zed = 0
#making new_luminosity_dist. and alpha_list equal in length
print len(new_luminosity_distribution), len(alpha_list)
alpha_list.pop()
alpha_list.pop()

# checking the range of exponent values
for i in range(len(new_luminosity_distribution)-1):
    tmp = str(new_luminosity_distribution[i])
    tmp = tmp[-2:]
    if int(tmp) == 22:
        list22.append(i)
    elif int(tmp) == 23:
        list23.append(i)
    elif int(tmp) == 24:
        list24.append(i)
    elif int(tmp) == 25:
        list25.append(i)
    elif int(tmp) == 26:
        list26.append(i)
    elif int(tmp) == 27:
        list27.append(i)
    elif int(tmp) == 28:
        list28.append(i)
for lists in exponent_list:
    print lists

# grouping average alphas
for i in range(len(exponent_list)-1):
    sum = 0
    for item in exponent_list[i]:
        sum += alpha_list[item]
    try:
        sum /= len(exponent_list[i])
        ylist.append(sum)
    except ZeroDivisionError:
        sum = 0
    average_dictionary["average of "+str(exponent_list_names[i])] = sum
print average_dictionary
# making the bar graph
for zed in range(len(exponent_list)-1):
    current_list = exponent_list[zed]
    item = 0
    try:
        x = str(new_luminosity_distribution[current_list[item]])
        x = int(x[-2:])
    except:
        x = 0
    height = len(current_list)
    plt.bar(x,height,width,color = 'gray')

for item in exponent_list:
    num_galaxies_list.append(len(item))

#annotating the alphas
plt.annotate("alpha = \n"+str(average_dictionary['average of list22']),xy = (22,5),xytext = (22.3,10),rotation = 90)
plt.annotate("alpha = \n"+str(average_dictionary['average of list24']),xy = (24,5),xytext = (24.3,15),rotation = 90)
plt.annotate("alpha = \n"+str(average_dictionary['average of list25']),xy = (25,5),xytext = (25.3,20),rotation = 90)
plt.annotate("alpha = \n"+str(average_dictionary['average of list26']),xy = (22,5),xytext = (26.3,20),rotation = 90)
plt.annotate("alpha = \n"+str(average_dictionary['average of list22']),xy = (22,5),xytext = (27.3,20),rotation = 90)
plt.savefig("sed_and_alpha")

# Putting everything into a FITS file
hdu = fits.PrimaryHDU()
hdulist = fits.HDUList([hdu])
range_column = fits.Column(name = 'Range', format = 'FLOAT',
                                        array = range_column_list)
num_galaxies = fits.Column(name = 'Number of Galaxies', format = 'FLOAT',array = num_galaxies_list)
mean_alpha_column = fits.Column(name = 'mean_alpha', format = 'FLOAT', array = ylist)
mean_alpha = fits.ColDefs([range_column,mean_alpha_column,num_galaxies])
alpha_change = fits.new_table(mean_alpha)
alpha_change.header.update('EXTNAME','ALPHACHANGE','mapping the change of spectral index')
primary_header = fits.Header()
primary_header.append('LUMINOS')
for item in range_column_list:
    primary_header.append(str(item))
primary_hdu = fits.PrimaryHDU(header = primary_header)
thdulist = fits.HDUList([primary_hdu,alpha_change])
thdulist.writeto("sed_and_alpha.fits")
print '<====DONE====>'

Questions:

• Is there anything that can be done to make it more efficient?
• Can list comprehension be used to make anything shorter?
• Is it good code? Have I improved from my past posts?

Note: nothing can be done to shorten the plt.annotate block--the values are arbitrary (but not random).

Answer 1

There is a lot going on here, I don't really know where to start from… Maybe point out some mistakes you were already told about in your previous questions:

• exception handling;
• white space;
• list-comprehensions.

Exception handling

Never use bare except clauses. You must know what kind of exception to expect before thinking about using a try. If you don't, then there is no need to try to recover from an "error" since you won't know how to handle said "error".

Even if you know which kind of exception to expect, if you don't know how to handle it, there is no point in trying to do so. Better let the exception bubble up and let some other part of the code handle it if it knows what to do…

These are obviously generic comments but the point is that you should use a try ... except clause only if you both now what to expect and how to handle it.

White space

Giveyourcodesomefreshair. Noonelikestoreaddenseblocksofcode. Use spaces to make reading easier and to delimit logical sections on a single (complex) expression.

Python has PEP8 as its official coding style. Give it a read, it has some parts about whitespaces.

List-comprehensions

And a few things about iteration

A few times in your code, you have constructs that looks like

a_list = []
for i in range(len(some_other_list)-1):
    value = call_function(some_other_list[i])
    a_list.append(value)

There are 3 things wrong with that:

• range(x) generate integer from 0 (inclusive) to x (exclusive). I hardly see the need of the -1 at the end since it removes the last item each time you do it.

• Python's for loops are meant to iterate over the elements of a sequence (well, any iterable, if we want to be pedantic, but sequences are iterables). You don't need indices to get the elements one by one. You don't even want indices as accessing elements using their index will be slower than having the for loop providing them to you:

  a_list = []
  for element in some_other_list:
      a_list.append(call_function(element))
  

  If you think you need the index to iterate over several sequence together… think again, and learn about zip. If you truly need indexes and elements, use enumerate.

• appending to a list is slow, better builld the list directly using a list-comprehension:

  a_list = [call_function(e) for e in some_other_list]
  

Write functions

As it stand, your code is almost only a big giant pile of instructions without any associated meaning. Organizing things into functions could lead to your top-level code being something along the lines of

remove_old_files()
alpha, luminosity = process_redshifts("/home/vhx/Documents/code/python/fits/radio_template.fits/apj384116t3_mrt.txt")
exponent_list = separate(luminosity, alpha)
plot(exponent_list, "sed_and_alpha.png")
save(exponent_list, "sed_and_alpha.fits")

This is a rough sketch, I’m not sure this is exactly what is being computed here but having this kind of code at top-level gives much more information about what is going on.

Moreover, top-level code is better put under an if __name__ == '__main__': statement. It avoid running the code in the event you import your file in an interactive session (or any other testing environment).

Using functions also have the advantage of giving more meaning to parameters than having them floatting around the code as global variables.

Summarysing what I said so far, we could write process_redshifts the following way:

from itertools import izip

def process_redshifts(input_file_name):
    data = ascii.read(input_file_name)
    # Canot generate two lists at once with a list-comprehension
    alpha = []
    luminosity = []

    for z, sx, sc in izip(data['z'], data['SX'], data['SC']):
        if 'NO' not in (z, sx, sc):
            alpha_value = get_alpha(sx, sc, XBAND, CBAND)
            alpha.append(alpha_value)
            x = get_luminosity_distance(z)
            luminosity.append(get_luminosity(z, x, CBAND, XBAND, sx))

    return alpha, luminosity

I’m using izip instead of zip here to process triplets one at a time and not build a whole list of triplets up-front.

os operations

You want to remove a file from your computer? You should use os.remove then, which is built exactly for that purpose.

def remove_old_files():
    for filename in ('sed_and_alpha.png', 'sed_and_alpha.fits'):
        try:
            os.remove(filename)
        except FileNotFoundError:
            pass

Use astropy builtin capabilities

I’m not versed in astropy at all, but I had a glance at their code on their github. It turns out that the objects you are converting to floats using str and removing 'Mpc' are astropy.units.Quantity objects.

These objects define their __str__ method by formating self.value in front of their unit. They also provide a __float__ method. So you have two builtin ways of getting the float value you are looking for:

1. the_value_i_want = astropy_object.value
2. the_value_i_want = float(astropy_object)

Stop using str for such conversions.

Prints

There are a lot of them, and they all look like debug messages. At this point you should have tested your code long enough to be confident in what it is doing and you thus are able to remove them. Their sheer number makes understanding the code much harder for reviewers.

You even happened to have a print line after a return in one of your functions.

There is also this special one print '<====DONE====>' at the end… It's completely useless as the end of the script will be signified to the user by their prompt appearing again.

Answer 2

The first, really simple thing I would do is tidy up the way you're using scientific notation. Rather than writing 4.86*10**9, you can write 4.86e9, and so on. It's easier to read and less error-prone.

Second, some of the things you're trying to do are duplicated in library functions. For example, your get_luminosity_distance function can be replaced with cosmo.luminosity_distance(redshift).si.value. (at least, in the current version of astropy. I don't know which version you're using).

I would store your exponent_lists in a dictionary; you're trying to do that already with the exponent_list_names list. To check the range of exponent values, you could then do something like:

for idx, luminosity in enumerate(new_luminosity_distribution):
    exponent_dict[int(np.log10(luminosity))].append(idx)

which is much more compact than

# checking the range of exponent values
for i in range(len(new_luminosity_distribution)-1):
    tmp = str(new_luminosity_distribution[i])
    tmp = tmp[-2:]
    if int(tmp) == 22:
        list22.append(i)
    elif int(tmp) == 23:
        list23.append(i)
    elif int(tmp) == 24:
        list24.append(i)
    elif int(tmp) == 25:
        list25.append(i)
    elif int(tmp) == 26:
        list26.append(i)
    elif int(tmp) == 27:
        list27.append(i)
    elif int(tmp) == 28:
        list28.append(i)

That actually brings up another issue: when you're looping over the elements of a list, rather than doing

for i in range(len(thelist)):
    ... etc. ...

it's much more pythonic to do:

for item in thelist:
    ... etc. ...

if you just need the items, or use enumerate (as I did above) if you need the index, as well.

Answer 3

I'm not a python developer but I will do my best.

First at all list22, list23, list24 is strange you should do something like this

exponent_list[0]  #for list22
exponent_list[1]  #for list23
exponent_list[2]  #for list24

etc... It is more abstract and useful In this way you save to do 3 or 4 arrays like you have right now

this code,

range_column_list = [10.0**22, 10.0**23, 10.0**24,10.0**25,10.0**26,10.0**27,10.0**28]

can be done like this

range_column_list = []
for i in range(22,28)
    range_column_list.append(10.0 ** i)

or better and 10.0 is a constant

for i in xrange(246):

in that line I think that you want go for all lines in file but 246 for me is a magic number and if tomorrow will change you need change the program

if is a requirement or always come like that put a constant however if not here you are a link for count the lines of a file https://stackoverflow.com/questions/845058/how-to-get-line-count-cheaply-in-python

first run a loop for all item in the file if is a redshift you add it in an array

after that loop over that array and do other thing why not together? I think that is the same code

# Determining which redshifts are available
for i in xrange(246):
    blank = True
    tmp = (table_data["z"][i])
    if tmp != 'NO' and table_data["SX"][i] != 'NO' and table_data["SC"][i] != 'NO':
        available_redshifts.append(i)
        #now you know that is a redshift so... I do more things
        # get the spectral indices 
        alpha_list.append(get_alpha(table_data["SX"][available_redshifts[i]],table_data["SC"][available_redshifts[i]],XBAND,CBAND))
        print table_data["z"][available_redshifts[i]], table_data["SX"][available_redshifts[i]], table_data["SC"][available_redshifts[i]]
        # now you iterate again for each redshift but I already have one, so... I don't need a loop
        # solving for luminosity
        x = get_luminosity_distance(table_data["z"][available_redshifts[i]])
        luminosity_distribution.append(get_luminosity(x,table_data["z"][available_redshifts[i]],alpha_list[i],CBAND,XBAND,table_data["SX"][available_redshifts[i]]))
        # I was done the last loop here too but I realise that all spectral indexes will not together so I do in a different loop like you    

    else:
        pass

You really need or verify the data before, I see a lot of assumptions, I understand that the data always come in the same format but you never know. check all array indexes please.

this code is a little confusing

new_luminosity_distribution.pop(0)

why do you remove the first item if you just append data?

there are a lot of this

print k/(len(alpha_list)-1)

I don't know what is it, put text for clarify the values Now this

alpha_list.pop()
alpha_list.pop()

you remove two items, it is magic for me

this

for item in new_luminosity_distribution:
    print '%e' % item
print '%e, %e' % (max(new_luminosity_distribution), min(new_luminosity_distribution))
sum = 0
for item in new_luminosity_distribution:
    sum += item

why repeat 2 times if you can do one time?

sum = 0
for item in new_luminosity_distribution:
    print '%e' % item
    sum += item
print '%e, %e' % (max(new_luminosity_distribution), min(new_luminosity_distribution))

Here the same and assuming that do the exponent_list instead list22 you have the next

# checking the range of exponent values
for i in range(len(new_luminosity_distribution)-1):
    tmp = str(new_luminosity_distribution[i])
    tmp = tmp[-2:]
    if int(tmp) == 22:
        list22.append(i)
    elif int(tmp) == 23:
        list23.append(i)
    elif int(tmp) == 24:
        list24.append(i)
    elif int(tmp) == 25:
        list25.append(i)
    elif int(tmp) == 26:
        list26.append(i)
    elif int(tmp) == 27:
        list27.append(i)
    elif int(tmp) == 28:
        list28.append(i)

like this

# checking the range of exponent values
for i in range(len(new_luminosity_distribution)-1):
    tmp = str(new_luminosity_distribution[i])
    tmp = tmp[-2:]
    exponent_list[tmp-22].append(i)
    # tmp-22 will give you the index
    # 0 for 22, 1 for 23, 2 for 24 etc like list22, list23, list24 etc

try to reduce the number of loops and do optimizations If you have two loops and are the same loop but the content is different look for if there are a way the do at the same time