ts_map.py

#!/usr/bin/env python
"""Compute a TS map using binned Poisson likelihood.
Usage: ts_map.py counts.fits psf.fits background.fits ts.fits

For testing use the output files of make_test_data.py.

TODO: Use a constrained minimizer, forcing the amplitudes >= 0!!!
"""
import sys
import logging
import numpy as np
from scipy.ndimage import convolve
from scipy.optimize import fmin, fmin_l_bfgs_b
from astropy.io import fits
from stats.likelihood import poisson_nll, TS_from_NLL

def compute_ts(x, y, counts, psf, background, 
               disp=False, xtol=1e-3, ftol=1e-3, bounded_optimization=False):
    """Compute TS for one position x, y"""
    logging.debug('Computing TS for position %d, %d' % (x, y))

    # Shift psf to position x, y
    #signal = scipy.ndimage.shift(signal, (x,y))

    logging.debug('Computing an expected counts image for the signal')
    signal = np.zeros_like(counts)
    signal[y, x] = 1 # Note: numpy arrays have y, x index order!
    signal = convolve(signal, psf, mode='constant')
    #fits.writeto('signal_shifted_%04d_%04d.fits' % (x, y), 
    #               signal, clobber=True)
    
    logging.debug('Fitting H0')
    def L0_func(p):
        """B = Background amplitude"""
        B = p[0]
        return poisson_nll(counts, B * background).sum()

    if bounded_optimization:
        # Doesn't work yet! Don't know why.
        B0, NLL0, results_dict =  fmin_l_bfgs_b(L0_func, [1], bounds = [(0, 1e3)]) 
        logging.debug(results_dict)
    else:
        results =  fmin(L0_func, [1], 
                        xtol=xtol, ftol=ftol,
                        disp=disp, full_output=True)
        B0, NLL0 = results[0][0], results[1]
        logging.debug(results)

    logging.debug('Fitting H1')
    def L1_func(p):
        """B = Background amplitude,
        S = Signal amplitude"""
        B, S = p[0], p[1]
        return poisson_nll(counts, B * background + S * signal).sum()
    results =  fmin(L1_func, [1, 0], 
                    xtol=xtol, ftol=ftol,
                    disp=disp, full_output=True)
    B1, S1, NLL1 = results[0][0], results[0][1], results[1]
    logging.debug(results)

    logging.debug('Computing TS')
    TS = TS_from_NLL(NLL0, NLL1)
    
    logging.info('%4d %4d %g %g %g %g %g %g' 
                 % (x, y, TS, NLL0, NLL1, B0, B1, S1))
    
    return TS, NLL0, NLL1, B0, B1, S1

def make_ts_map(counts, psf, background):
    TS = np.empty_like(counts)
    NLL0 = np.empty_like(counts)
    NLL1 = np.empty_like(counts)
    B0 = np.empty_like(counts)
    B1 = np.empty_like(counts)
    S1 = np.empty_like(counts)
    
    # Process one pixel at a time
    nx, ny = counts.shape
    for x in range(nx):
        for y in range(ny):
            (TS[x, y], NLL0[x, y], NLL1[x, y], 
            B0[x, y], B1[x, y], S1[x,y]) = compute_ts(x, y, counts, psf, background)
            
    return TS, NLL0, NLL1, B0, B1, S1


if __name__ == '__main__':
    logging.basicConfig(level=logging.INFO)
    
    logging.debug('Parsing command line options')
    if len(sys.argv) != 5:
        print __doc__
        exit()
    counts_file     = sys.argv[1]
    psf_file     = sys.argv[2]
    background_file = sys.argv[3]
    ts_file         = sys.argv[4]
    print 'counts_file:    ', counts_file
    print 'psf_file:    ', psf_file
    print 'background_file:', background_file
    print 'ts_file:        ', ts_file
    
    logging.debug('Reading data')
    header     = fits.getheader(counts_file)
    counts     = fits.getdata(counts_file)
    psf        = fits.getdata(psf_file)
    background = fits.getdata(background_file)
    
    logging.debug('Computating TS map')
    # @todo Time this part!
    #ts = compute_ts(30, 60, counts, psf, background)
    #print 'ts:', ts
    
    maps = make_ts_map(counts, psf, background)
    ts = maps[0]
    
    logging.debug('Writing ts_significance.fits')
    fits.writeto('ts.fits', ts, header, clobber=True)
    fits.writeto('ts_significance.fits', 
                   np.sqrt(np.abs(ts)), header, clobber=True)
    
    # Now make Li & Ma significance map for comparison!