generate_priority_maps.py

import matplotlib.pyplot as plt
import numpy as np
import healpy as hp
from os import path
from sys import argv
from astropy import units as u
#from astropy_healpix import HEALPix
from astropy.coordinates import Galactic, TETE, SkyCoord
from astropy.io import fits
import config_utils
from pylab import cm


class PriorityMap:
    def __init__(self,npix=1,filter=None):
        self.vote_map = np.zeros(npix)
        self.filter = filter
        self.map_dict = {}

def build_priority_maps():

    # Fetch configuration:
    use_hp_maps = True
    if len(argv) < 2:
        config_file = input('Please enter the path to the configuration file: ')
    else:
        config_file = argv[1]
    config = config_utils.read_config(config_file)

    # Initialze the HEALpix map and storage array - this includes an extra
    # column for the combined_map and the maps summed over all filters
    #ahp = HEALPix(nside=config['NSIDE'], order='ring', frame=TETE())
    NPIX = hp.nside2npix(config['NSIDE'])

    # Build vote maps in each filter based on the stellar density of each HEALpix
    vote_maps = {}
    map_titles = ['combined_map']
    for code in config['map_codes']:
        map_title = config[code]['file_root_name']
        map_titles.append(map_title)

    for filter in config['filter_list']:
        map = PriorityMap(npix=NPIX,filter=filter)

        # Read in component maps for different science cases and add them
        # to the combined map for this filter
        for code in config['map_codes']:
            file_name = path.join(config['output_dir'],
                                    config[code]['file_root_name']+'_'+str(filter)+'.fits')

            # map.map_dict contains the original science maps with the
            # computed priority per-healpix, since this is calculated on
            # various different scientific grounds
            if use_hp_maps:
                pix_data = hp.read_map(file_name)
                map_title = config[code]['file_root_name']
            else:
                hdul = fits.open(file_name)
                pix_data = hdul[1].data['HEALpix_values']
                map_title = hdul[0].header['MAPTITLE']

            # vote_map contains the sum of all of the maps.
            # Note that no renormalization is performed at this stage because
            # this has been done for the maps generated by generate_sky_maps.
            map.vote_map += pix_data * config[code]['map_weight']
            map.map_dict[map_title] = pix_data

        # The final vote_map for this filter is normalized to 1.
        map.vote_map = map.vote_map/map.vote_map.max()

        vote_maps[filter] = map

    # Sum the priority per HEALpix of each science map over all filters
    summed_maps = {}
    for map_title in map_titles:
        sum_map = np.zeros(NPIX)
        for filter in config['filter_list']:
            if map_title == 'combined_map':
                sum_map += vote_maps[filter].vote_map
            else:
                sum_map += vote_maps[filter].map_dict[map_title]
        summed_maps[map_title] = sum_map

    # Output the priority map plots in both PNG and FITS formats, and the data
    # as a FITS binary table

    for filter in config['filter_list']:
        map = vote_maps[filter]

        fig = plt.figure(3,(10,10))
        plot_max = map.vote_map.max()
        if np.isnan(plot_max):
            plot_max = 1.0
        hp.mollview(map.vote_map, title="Regions of Interest "+str(filter)+"-band",
                    min=0.0, max=plot_max)
        hp.graticule()
        plt.tight_layout()
        plt.savefig(path.join(config['output_dir'],config['output_file_name']+'_plot_'+str(filter)+'.png'))
        plt.close(3)
        hp.write_map(path.join(config['output_dir'],config['output_file_name']+'_plot_'+str(filter)+'.fits'), map.vote_map, overwrite=True)

        # Header
        hdr = fits.Header()
        hdr['NSIDE'] = config['NSIDE']
        hdr['NPIX'] = hp.nside2npix(config['NSIDE'])
        hdr['MAPTITLE'] = 'Combined priority survey footprint'
        for i,map_title in enumerate(map_titles):
            hdr['SCIMAP'+str(i)] = map_title
        hdr['VERSION'] = config['VERSION']
        phdu = fits.PrimaryHDU(header=hdr)

        hdu_list = [phdu]

        # First table extension: combined vote map
        col_list = []
        col_list.append( fits.Column(name='combined_map', array=map.vote_map, format='E') )

        # Subsequent extensions contain the component maps for different science cases:
        for map_title, data in map.map_dict.items():
            col_list.append( fits.Column(name=map_title, array=data, format='E') )

        hdu = fits.BinTableHDU.from_columns(col_list)
        hdu_list.append(hdu)
        hdul = fits.HDUList(hdu_list)

        hdul.writeto(path.join(config['output_dir'],
                                config['output_file_name']+'_data_'+str(filter)+'.fits'),
                                overwrite=True)

    # Output summed maps
    hdr = fits.Header()
    hdr['NSIDE'] = config['NSIDE']
    hdr['NPIX'] = hp.nside2npix(config['NSIDE'])
    hdr['MAPTITLE'] = 'Total priority survey footprint'
    hdr['VERSION'] = '2.0.0'
    phdu = fits.PrimaryHDU(header=hdr)
    hdu_list = [phdu]

    col_list = []
    for map_title in map_titles:
        col_list.append( fits.Column(name=map_title, array=summed_maps[map_title], format='E') )

    hdu = fits.BinTableHDU.from_columns(col_list)
    hdu_list.append(hdu)
    hdul = fits.HDUList(hdu_list)
    hdul.writeto(path.join(config['output_dir'],
                            config['output_file_name']+'_data_sum.fits'),
                            overwrite=True)

if __name__ == '__main__':
    build_priority_maps()