while loops for processing FITS files(python)

Question

I am new at programming in python and am in the process of trying to create a setup of processing thousands of files with one piece of code in python. I created a practice folder to do this in. In it are two FITS files (FITS1.fits and FITS2.fits). I did the following to put them both in a .txt file:

ls > practice.txt

Here is what I did next:

$ python
import numpy
import pyfits
import matplotlib.pyplot as plt
from matplotlib import pylab
from pylab import *
import asciidata

a = asciidata.open('practice.txt')

print a[0][0] #To test to see if practice.txt really contains my FITS files FITS1.fits
i = 0
while i <=1 #Now I attempt a while loop to read data from columns in FITS files, plot the numbers desired, save and show the figures. I chose i <=1 because there are only two FITS files in the text(also because of zero-indexing). 

    b = pyfits.getdata(a[0][i]) # "i" will be the index used to use a different file when the while loop gets to the end

    time = b['TIME'] #'TIME' is a column in the FITS file
    brightness = b['SAP_FLUX']
    plt.plot(time, brightness)
    xlabel('Time(days)')
    ylabel('Brightness (e-/s)')
    title(a[0][i])

    pylab.savefig('a[0][i].png') #Here am I lost on how to get the while loop to name the saved figure something different every time. It takes the 'a[0][i].png' as a string and not as the index I am trying to make it be.

    pylab.show()

    i=i+1 # I placed this here, hoping that when the while loop gets to this point, it would start over again with a different "i" value

After pressing enter twice, I see the first figure as expected. Then I will close it and see the second. However, only the first figure is saved. Does anyone have any suggestions on how I can change my loop to do what I need it to?

Answer 1

In your code the i is being treated as the letter i, not the variable. If you wanted to keep this naming you could do something like:

FileName = 'a[0][%s].png' % i
pylab.savefig(FileName)

Answer 2

You should use glob to automatically get a the fits files as a list, from there using a for loop will let you iterate of the names of the files directly instead of using an index. When you call plt.savefig, you need to construct the file name you want to save it as. Here is the code cleaned up and put together:

from glob import glob
import pyfits
from matplotlib import pyplot as plt

files = glob('*.fits')

for file_name in files:
    data = pyfits.getdata(file_name)
    name = file_name[:-len('.fits')] # Remove .fits from the file name

    time       = data['TIME']
    brightness = data['SAP_FLUX']

    plt.plot(time, brightness)

    plt.xlabel('Time(days)')
    plt.ylabel('Brightness (e-/s)')
    plt.title(name)

    plt.savefig(name + '.png')
    plt.show()