ctools

#!/usr/bin/env python

# author: Nathan Kelley-Hoskins

# date  : Oct, 2015

#

# Run this script with no arguments to see the usage, 

# or scroll to the bottom of this script.

#

import sys, gammalib, math, numpy

from scipy import ndimage

from astropy.io import fits

def denoise_tv_chambolle_2d(im, weight=50, eps=2.e-4, n_iter_max=200):

    """Perform total-variation denoising on 2D images.

    Parameters

    ----------

    im : ndarray

        Input data to be denoised.

    weight : float, optional

        Denoising weight. The greater `weight`, the more denoising (at

        the expense of fidelity to `input`)

    eps : float, optional

        Relative difference of the value of the cost function that determines

        the stop criterion. The algorithm stops when:

            (E_(n-1) - E_n) < eps * E_0

    n_iter_max : int, optional

        Maximal number of iterations used for the optimization.

    Returns

    -------

    out : ndarray

        Denoised array of floats.

    Notes

    -----

    ***Function taken from the python module scikit-image***

    ***https://github.com/scikit-image/scikit-image/blob/master/skimage/restoration/_denoise.py***

    The principle of total variation denoising is explained in

    http://en.wikipedia.org/wiki/Total_variation_denoising.

    This code is an implementation of the algorithm of Rudin, Fatemi and Osher

    that was proposed by Chambolle in [1]_.

    References

    ----------

    .. [1] A. Chambolle, An algorithm for total variation minimization and

           applications, Journal of Mathematical Imaging and Vision,

           Springer, 2004, 20, 89-97.

    """

    px = numpy.zeros_like(im)

    py = numpy.zeros_like(im)

    gx = numpy.zeros_like(im)

    gy = numpy.zeros_like(im)

    d = numpy.zeros_like(im)

    i = 0

    while i < n_iter_max:

        d = -px - py

        d[1:]    += px[:-1]

        d[:, 1:] += py[:, :-1]

        out = im + d

        E   = (d ** 2).sum()

        gx[:-1]    = numpy.diff(out, axis=0)  # [], gradient/variation in x direction

        gy[:, :-1] = numpy.diff(out, axis=1)  # [], gradient/variation in y direction 

        norm  = numpy.sqrt(gx ** 2 + gy ** 2) # total variation here

        E    += weight * norm.sum()           # adjust by weight between sameness and variation

        norm *= 0.5 / weight

        norm += 1

        px   -= 0.25 * gx

        px   /= norm

        py   -= 0.25 * gy

        py   /= norm

        E    /= float(im.size)

        if i == 0:

            E_init     = E

            E_previous = E

        else:

            if numpy.abs(E_previous - E) < eps * E_init:

                break

            else:

                E_previous = E

        i += 1

    return out

def background2fits( data, fname ) :

  """Function just for taking final background data, and writing it to a fits file.

  Does no binning or smearing.

  **Args:**

    data  : dictionary containing info about background, see write_background_to_fits for more info.

    fname : output fits filename to write background to

  """

  # figure out basics of the fits table

  ndetx   = len( data['DETX_LO'] )

  ndety   = len( data['DETY_LO'] )

  nenez   = len( data['BGD'    ] )

  formx   = '%dE' % ndetx

  formz   = '%dE' % nenez

  formxyz = '%dE' % ( ndetx * ndety * nenez )

  bgddim  = '(%d,%d,%d)' % ( ndetx, ndety, nenez ) 

  # setup table columns

  #print

  cols = []

  for key in [ 'DETX_LO', 'DETX_HI', 'DETY_LO', 'DETY_HI' ] :

    #print 'adding col %s with %d elements: ' % ( key, len(data[key]) ), [ '%.2f'%d for d in data[key] ]

    cols.append( fits.Column( name=key , format=formx  , unit='deg'       , array=[ data[key] ]             ) )

  for key in [ 'ENERG_LO', 'ENERG_HI' ] :

    #print 'adding col %s' % key

    encol = [ en.TeV() for en in data[key] ] # convert from list of GEnergy's to list of TeV floats

    cols.append( fits.Column( name=key , format=formz  , unit='TeV'       , array=[ encol     ]             ) )

  for key in [ 'BGD' ] :

    #print 'adding col %s' % key

    cols.append( fits.Column( name=key , format=formxyz, unit='1/s/MeV/sr', array=[ data[key] ], dim=bgddim ) )

  # setup fits table

  tablehdu      = fits.BinTableHDU.from_columns( cols )

  # name the table and make sure the dimensionality of the BGD column is set correctly

  tablehdu.name = 'BACKGROUND'

  tablehdu.header.set( 'TDIM7', bgddim )

  # write table to file

  print 'writing fits table to %s' % fname

  tablehdu.writeto( fname, clobber=True)

def write_background_to_fits( obs, fname, binwidth=0.4, nbufferbin=5, nenergybins=5, energyprange=[0.02,0.99], zeroval=1e-5, smoothmode='gauss', smooth=0.4 ) :

  """Create and write a background from a GObservations object to a fits file.

  **Arguments:**

    obs   : GObservations object

    fname : str, filename of output fits file

  **Options:**

    background binning options:

      binwidth     : float, width in degrees of one side of a bin

      nbufferbin   : int, number of extra (empty) buffer bins to add to the 

                     edge of the camera before smoothing.  If 0, will only

                     create bins out to event with largest camera offset.

      nenergybins  : int, number of equal-statistics energy bins to create

      energyprange : list of 2 floats, [min,max]. Only use the events whose 

                     energy is between the min and max percentiles. [0.0,1.0] 

                     for all events.

      zeroval      : float, after binning events, replace all 0-event bins 

                     with this value

      smoothmode   : str, type of smoothing to apply to background

                     'gauss' : smooth each bin with a 'smooth'-degree-wide gaussian

                     'totalvariation' : apply total-variation denoising, with

                      weight 'smooth'.  See denoise_tv_chambolle_2d() for

                      more information.

                     'none' for no smoothing

      smooth       : float, if smoothmode is 'gauss', then smooth determines

                     the gaussian width [deg] of the smoothing.  if smoothmode is

                     'totalvariation', smooth is the weight given to the function

                     denoise_tv_chambolle_2d(..., weight=smooth)

  **Returns:**

    dictionary with info about background produced, keys are:

      ['enbins'    ] : list of pairs of energies, one pair for each energy bin

      ['detx'      ] : list of detx pairs of angles for camera bin edges

      ['dety'      ] : list of detx pairs of angles for camera bin edges

      ['bgd'       ] : background counts 3D array, [ienergy][idetx][idety]

      ['bgd_smooth'] : same as 'bgd', except smoothed

      ['bgd_norm'  ] : same as 'bgd_smooth', except normalized for exposure

"""

  print

  print 'Using %s smoothing!' % smoothmode

  print

  # init some vars

  quart_min = min( energyprange )

  quart_max = max( energyprange )

  # figure out min/max energy range, max offset

  nev = 0

  e_min, e_max = 0.0, 0.0

  maxoffset = 0.0

  for run in obs :

    for ev in run.events() :

      en = ev.energy()

      dx = ev.dir().detx() * gammalib.rad2deg

      dy = ev.dir().dety() * gammalib.rad2deg

      offset = math.sqrt( dx*dx + dy*dy )

      if offset > maxoffset :

        maxoffset = offset

      if nev == 0 :

        e_min = en

        e_max = en

      if en > e_max : e_max = en

      if en < e_min : e_min = en

      nev += 1

  binr    = maxoffset + ( nbufferbin * binwidth )

  # push e_max up by a tiny amount,

  # to include the highest event in the binning

  e_max = e_max + gammalib.GEnergy(1,'GeV')

  print 'total: %d events from %.3f to %.1f TeV' % ( nev, e_min.TeV(), e_max.TeV() )

  # setup fine-energy binning

  nfinebins      = 1000

  uselogbinning  = True

  enfinespec     = gammalib.GEbounds( nfinebins, e_min, e_max, uselogbinning )

  enfinebins     = numpy.zeros( nfinebins )

  enfinebinwidth = (e_max.log10TeV() - e_min.log10TeV() ) / nfinebins

  # loop over events, adding to fine energy bins

  for run in obs :

    for ev in run.events() :

      enfinebins[ enfinespec.index( ev.energy() ) ] += 1

  # only use events that are within 

  # quart_min and quart_max of the distribution

  # search for energy that contains the given fraction of events at

  # that energy or lower

  def frac2en( fraclist ) :

    """search binned histogram enfinespec (defined a few lines up) for the 

   energies that have the given fraction of events at or below that energy.

    **Example:**

      frac2en( [0.4 ] ) # find what energy has 40% of events below that energy

      frac2en( [0.88] ) # find what energy has 88% of events below that energy

    **Args:**

      fraclist : float or list of floats, each float from 0.0 to 1.0

    **Returns:**

      list of gammalib.GEnergy objects, each element corrosponding 

        to each fraction in 'fraclist'

  """

    # force input to be a list

    if type(fraclist) is not type( [] ) :

      fraclist = [ fraclist ]

    # init variables

    enlist = [ 0 for i in range(len(fraclist)) ]

    last_nfrac = 0.0

    # search histogram bins for given fraction

    for i in range(enfinespec.size()) :

      nfrac = numpy.sum(enfinebins[:i]) / float(nev)

      # loop over input fractions

      for i_frac, frac in enumerate( fraclist ) :

        # when found, 

        if last_nfrac < frac and nfrac > frac :

          enlist[i_frac] = enfinespec.elogmean(i)

      last_nfrac = nfrac

    return enlist

  # calc median/quartile energies

  quart_min_en, quart_median_en, quart_max_en = frac2en( [ quart_min, 0.5, quart_max ] )

  # calculate quartile boundaries for each energy bin

  quartbins = []

  quart_width = ( quart_max - quart_min ) / nenergybins

  for i in range(nenergybins) :

    quartbins.append( [quart_min + (quart_width*i), quart_min + (quart_width*(i+1))] )

  newquartedges = sorted(list(set([ j for i in quartbins for j in i ])))[1:-1]

  # calculate energies of the new edges

  newquartedges_en = frac2en( newquartedges ) 

  # figure out list of detx/y bins

  detbins = []

  j = 0 

  while True :

    inner = j     * binwidth

    outer = (j+1) * binwidth

    if inner > binr :

      break

    detbins.append(    [    i for i in [inner,outer] ] )

    detbins.insert( 0, [ -1*i for i in [outer,inner] ] )

    j += 1

  # setup background dictionary

  print 'energy bins, ~%d events (%d%% of all events) per bin' % ( nev*(quartbins[0][1]-quartbins[0][0]), 100*(quartbins[0][1]-quartbins[0][0] ) )

  background = {}

  background['DETX_LO' ] = background['DETY_LO'] = [ det[0] for det in detbins ]

  background['DETX_HI' ] = background['DETY_HI'] = [ det[1] for det in detbins ]

  background['ENERG_LO'] = []

  background['ENERG_HI'] = []

  background['BGD'     ] = []

  # setup energy bins

  enbins = []

  for fracrange in quartbins :

    energybin = {}

    emin, emax = frac2en( fracrange )

    enbins.append( [emin, emax] )

    print '%5.3f - %5.3f TeV' % ( emin.TeV(), emax.TeV() )

    background['ENERG_LO'].append(emin)

    background['ENERG_HI'].append(emax)

    logmean = pow(10, ( emin.log10TeV() + emax.log10TeV() ) / 2.0 )

  # add backgrounds full of zeros

  ndet = len(detbins)

  nen  = len( enbins)

  bgd = numpy.zeros( (len(enbins),ndet,ndet), dtype='int32')

  # fill background 

  mindet = min( background['DETX_LO'] )

  def det2ind(val) :

    return int( ( val - mindet ) / binwidth )

  def en2ind(en) :

    for i, erange in enumerate( enbins ) :

      if en > erange[0] and en < erange[1] :

        return i

  livetime = 0

  for run in obs :

    livetime += run.livetime()

    for ev in run.events() :

      idetx = det2ind( ev.dir().detx() * gammalib.rad2deg )

      idety = det2ind( ev.dir().dety() * gammalib.rad2deg )

      ien   = en2ind(  ev.energy()                        )

      if type(ien) is type(3) :

        bgd[ien][idety][idetx] += 1

  # convert from array of ints (from numpy.zeros) to array of floats

  bgd = bgd.astype(numpy.float32)

  # set all bins with zeroes to something really small

  bgd[ bgd == 0 ] = zeroval

  # smooth background values via a gaussian kernel with a width of 'smooth' degrees

  bgd_smooth = []

  for i in range(nen) :

    if smoothmode == 'gauss' :

      # gaussian smoothing shamelessly stolen from stackoverflow

      # http://stackoverflow.com/questions/32155488/smoothing-a-2d-array-along-only-one-axis

      bgd_smooth.append( ndimage.gaussian_filter( bgd[i], smooth/binwidth ) ) # convert gaussw to units of '# bins'

    elif smoothmode == 'totalvariation' :

      presmoothsum = numpy.sum( bgd[i] )

      bgd_smooth.append( denoise_tv_chambolle_2d( bgd[i], weight=smooth   ) )

      postsmoothsum = numpy.sum( bgd_smooth[-1] )

      # normalize by the integral of all events

      bgd_smooth[-1] *= presmoothsum

      bgd_smooth[-1] /= postsmoothsum

    elif smoothmode == 'none' :

      bgd_smooth.append( bgd[i] )

    else :

      raise ValueError( "smoothmode is currently %s, must be one of ['gauss', 'totalvariation', 'none']" % repr(smoothmode) )

  # find min/max background values

  zmin = numpy.min( [bgd,bgd_smooth] )

  zmax = numpy.max( [bgd,bgd_smooth] )

  # apply normalization factor to account for exposure

  bin_sr = binwidth * binwidth * gammalib.deg2rad * gammalib.deg2rad

  bgd_norm = []

  for i in range(nen) :

    emin, emax = enbins[i]

    en_norm    = emax.MeV() - emin.MeV()

    normbin    = bin_sr * livetime * en_norm

    bgd_norm.append( bgd_smooth[i] / normbin )

  norm_min = numpy.min( bgd_norm )

  norm_max = numpy.max( bgd_norm )

  # plots of exposure-normalized bin values

  mevnorm = [ en[1].MeV() - en[0].MeV() for en in enbins ]

  # write background to a fits file

  background['BGD'] = bgd_norm

  # create temporary fits filename and write our fits table to it

  #fname = os.path.join( tempfile._get_default_tempdir(), next(tempfile._get_candidate_names())+'.fits' )

  background2fits( background, fname )  

  ret = {}

  ret['maxoffset' ] = maxoffset

  ret['enbins'    ] = enbins

  ret['detx'      ] = detbins

  ret['dety'      ] = detbins

  ret['bgd'       ] = bgd

  ret['bgd_smooth'] = bgd_smooth

  ret['bgd_norm'  ] = bgd_norm

  return ret

if __name__ == '__main__' :

  if len(sys.argv) == 1 :

    print 'Script for converting a GObservations xml file into a GCTABackground3D fits file.'

    print 'Usage is:'

    print ' $ %s xmlfile.xml outputbackground.fits' % __file__

    sys.exit(0)

  if len(sys.argv) != 3 :

    print 'Incorrect number of arguments.'

    print

    print 'Usage is:'

    print ' $ %s xmlfile.xml outputbackground.fits' % __file__

    sys.exit(1)

  xmlfile  = sys.argv[1]

  fitsfile = sys.argv[2]

  obs = gammalib.GObservations()

  obs.load( xmlfile )

  write_background_to_fits( obs, fitsfile )