[Matrix-SIG] FITS image file interface for NumPy

[David Buscher]

Please find enclosed a Python module to read and write FITS image files
to/from NumPy arrays. FITS is a format widely used in the astronomical
community for storage and interchange of data, especially multidimensional
image data. The code is so short, it does not seem worth the trouble of
putting it on an ftp archive. Others have probably written the same thing
themselves, but I couldn't find one anywhere, and I wanted to save others
needing to reinvent this wheel again. 

The module defines only a Read() and a Write() function. The comments in
the module should be sufficient for people to work out how to use it, but
please let me know if it is too obscure. I was going to add proper
docstrings and a test function, but waiting for the spare time to do that
has already taken 3 months. 

enjoy,
David

p.s. This code has been regularly used on various types of data, but
the usual caveats apply... 

----------------------------------------+-------------------------------------
 David Buscher                          |  Phone  +44 191 374 7462
 Dept of Physics, University of Durham, |  Fax    +44 191 374 3709
 South Road, Durham DH1 3LE, UK         |  Email  david.buscher@durham.ac.uk
----------------------------------------+-------------------------------------

#!/usr/bin/env python
# $Id: FITS.py,v 1.1 1998/08/24 14:49:57 dfb Exp dfb $
#
# Functions to read and write FITS image files
# For more information on the FITS format, see http://www.cv.nrao.edu/fits. 
#
# This module reads and writes only the image formats, not the 
# more complex binary and ascii tables. 
#
# It will handle 8, 16 and 32 bit integer and 32 bit float formats, but
# it does not handle the recently-proposed 64-bit format. This would be 
# trivial to do, I have just never come across any data in this format 
# that could act as a test input. It also does not handle the obscure 
# 'groups' format.


import string
from Numeric import *

error = 'FITS error'
#
# Read a FITS image file
#
# Returns a header, image tuple
# The header is returned in two forms, a raw list of header records, and
# a parsed dictionary for keyword access.
#
# By default, the image is returned as floats, using the BSCALE and BZERO
# keywords where available. By setting the asFloat keyword to zero it is 
# returned unscaled and in the (presumably more compact) numeric format
# that it was stored in
# 
def Read(filename, asFloat = 1) :
    file = open(filename, "r")
    header = {}
    rawHeader = []
    buffer = file.read(2880)
    if buffer[:6] != 'SIMPLE' :
	raise error, 'Not a simple fits file'
    while(1) :
	for char in range(0,2880,80) :
	    line = buffer[char:char+80]
	    rawHeader.append(line)
	    key = string.strip(line[:8])
	    if key :
		val = line[9:]
		val = string.strip(val)
		if val :
		    if val[0] == "'" :
			pos = string.index(val,"'",1)
			val = val[1:pos]
		    else :
			pos = string.find(val, '/')
			if pos != -1 :
			    val = val[:pos]
		header[key] = val
	if header.has_key('END') : break
	buffer = file.read(2880)
    naxis = string.atoi(header['NAXIS'])
    shape = []
    for i in range(1,naxis+1) :
	shape.append(string.atoi(header['NAXIS%d' % i]))
    shape.reverse()
    numPix = 1
    for i in shape :
	numPix = numPix * i
    bitpix = string.atoi(header['BITPIX'])
    if bitpix == 8 :
	type = UnsignedInt8
    elif bitpix == 16 :
	type = Int16
    elif bitpix == 32 :
	type = Int32
    elif bitpix == -32 :
	type = Float32
	bitpix = 32
    numByte = numPix * bitpix/8
    data = file.read(numByte)
    data = fromstring(data, type)
    data.shape = shape
    if LittleEndian : data = data.byteswapped()
    if asFloat :
	bscale = string.atof(header.get('BSCALE', '1.0'))
	bzero = string.atof(header.get('BZERO', '0.0'))
	data = data*bscale + bzero
    return( { 'raw' : rawHeader, 'parsed' : header},  data  )

#
# Save a numeric array to a FITS image file
#
# The FITS image dimensions are taken from the array dimensions.
# The data is saved in the numeric format of the array, or as 32-bit
# floats if a non-FITS data type is passed.
#
# A list of extra header records can be passed in extraHeader. The
# SIMPLE, BITPIX, NAXIS* and END records in this list are ignored.
# Header records are padded to 80 characters where necessary.
#
def Write(data, filename, extraHeader = None) :
    type = data.typecode()
    if type == UnsignedInt8 : bitpix = 8
    elif type == Int16 : bitpix = 16
    elif type == Int32 : bitpix = 32
    else :
	data = data.astype(Float32)
	bitpix = -32
    shape = list(data.shape)
    shape.reverse()
    naxis = len(shape)
    header = [ 'SIMPLE  = T', 'BITPIX  = %d' % bitpix, 'NAXIS   = %d' % naxis ]
    for i in range(naxis) :
	keyword = 'NAXIS%d' % (i + 1)
	keyword = string.ljust(keyword, 8)
	header.append('%s= %d' % (keyword, shape[i]))
    if extraHeader != None :
	for rec in extraHeader :
	    try :
		key = string.split(rec)[0]
	    except IndexError :
		pass
	    else :
		if key != 'SIMPLE' and \
		   key != 'BITPIX' and \
		   key[:5] != 'NAXIS' and \
		   key != 'END' :
		    header.append(rec)
    header.append('END')
    header = map(lambda x: string.ljust(x,80)[:80], header)
    header = string.join(header,'')
    numBlock = (len(header) + 2880 - 1) / 2880
    header = string.ljust(string.join(header,''), numBlock*2880)
    file = open(filename, 'w')
    file.write(header)
    if LittleEndian : data = data.byteswapped()
    data = data.tostring()
    file.write(data)
    numBlock = (len(data) + 2880 - 1) / 2880
    padding = ' ' * (numBlock*2880 - len(data))
    file.write(padding)