NaN masking ALMA maps

mask_alma_maps.py

#!/usr/bin/env python

import numpy as np

from astropy.io import fits


files = ['HE0433_mom0.fits', 'HE0433_mom1.fits', 'HE0433_mom2.fits']
outnames = ['mom0_masked.fits', 'mom1_masked.fits', 'mom2_masked.fits']

thresholds = [0.02, 0.02, 0.02]


def maskimage(image_data, thresh):
    mask = (image_data < thresh)
    image_data[mask] = np.nan

    return image_data


def main():
    for image, thresh, outname in zip(files, thresholds, outnames):

        hdulist = fits.open(image)
        image_data = hdulist[0].data
        header = hdulist[0].header
        hdulist.close()

        masked_data = maskimage(image_data, thresh)
        newhdu = fits.PrimaryHDU(masked_data, header)

        newhdulist = fits.HDUList([newhdu])
        newhdulist.writeto(outname, clobber=True)


if __name__ == '__main__':
    main()