write code to generate moment maps

[aakepley]

Rough code to do this below. Need to make sure that smoothing accounts for edges properly.

[aakepley]

# launch with ipython --matplotlib='qt' to get quicklook to work.
# first index is channel (opposite of what I expect).

# creating moment images using spectral-cube.

def mass_produce_moment_images(dataDir, **kwargs):

    '''
    Create moment images for all cubes in a specific directory in the current working directory
    '''

    import glob 
    import os.path
  
    cubelist = glob.glob(os.path.join(dataDir,"*rebase*smooth*hanning*.fits"))

    for cube in cubelist:
        create_moment_images(cube, **kwargs)

def create_moment_images(imagename,smoothFactor=3.0,specSmoothFactor=3.0,snr=0.9,plotIt=True):

    '''
    Create integrated intensity (mom0) and peak intensity (mom8) maps for DEGAS data

    rmsFrac: fraction of channels to use for RMS estimate. The RMS estimate comes from the beginning rmsFrac*nchan channels and end rmsFrac*chan

    smoothFactor: multiple of the original beam size to smooth masking image by spatially.

    specSmoothFactor: how many channels to boxcar smooth together

    rmsCut: S/N to cut smoothed mask at.

    The images will be written in the current working directory.

    '''

    import os.path
    import glob
    import math
    import numpy as np
    import matplotlib.pyplot as plt
    import re
    
    
    from spectral_cube import SpectralCube
    import radio_beam
    from astropy import units as u
    from astropy.convolution import Box1DKernel
    from astropy.stats import median_absolute_deviation as mad
    from astropy.io import fits

    MADtoRMS = 1.4826

    # strip out the directly name so that the new images are in the cwd.
    baseimagename = os.path.basename(imagename)
    dataDir = os.path.split(imagename)[0]

    (galaxy,line) = baseimagename.split('_')[0:2]

    print "Creating moments maps for " + galaxy + ' ' + line

    # read in cube
    cube = SpectralCube.read(imagename)

    # calculate the noise properties of the cube and then create a noise cube.
    nchan = len(cube.spectral_axis)
    cubeData = cube.unmasked_data[:,:,:].value

    noise2dImage = MADtoRMS * mad(cubeData,axis=0,ignore_nan=True)
    noiseCubeData = np.repeat(noise2dImage[np.newaxis,:,:],axis=0,repeats=nchan)

    noiseCube = SpectralCube(noiseCubeData,cube.wcs.copy(),beam=radio_beam.Beam(cube.beam.major,cube.beam.minor,cube.beam.pa),read_beam=False,mask=cube.mask)

    # smooth cube spatially by some multiple of the beam
    smoothBeam = radio_beam.Beam(major=cube.beam.major * smoothFactor, minor=cube.beam.minor * smoothFactor, pa=cube.beam.pa)
    
    ### WHAT I HAVE DISCOVERED IS THAT THE CONVOLVE_TO FUNCTION DOESN'T DEAL WITH EDGES AND NANS VERY WELL, LEADING TO CRAZY NOISE MAPS.
    # I've ripped off the code inside spectral-cube to do it better. It's a fairly simple addition.
    smoothCube = betterSmooth(cube,smoothBeam)
    smoothNoiseCube = betterSmooth(noiseCube,smoothBeam)
    
    # smooth cube on spectral axis. 
    kernel = Box1DKernel(specSmoothFactor) # This might not be the best kernel to use ever, but I understand what it does.
    specSmoothCube = smoothCube.spectral_smooth(kernel)
    specSmoothCubeData = specSmoothCube.unmasked_data[:,:,:].value

    specSmoothNoiseCube = smoothNoiseCube.spectral_smooth(kernel)
    specSmoothNoiseCubeData = specSmoothNoiseCube.unmasked_data[:,:,:].value

    snrSmoothImage = specSmoothCubeData/specSmoothNoiseCubeData
    smoothMask = (snrSmoothImage > snr)  ## not quite sure the units are working out right because I have to choose a pretty low snr to get the relevant regions to be masked.

    #  Create the masked cube
    maskedCube = cube.with_mask(smoothMask,inherit_mask=True)

    # create the mom0 image.
    mom0 = maskedCube.moment(order=0)
    if plotIt: mom0.quicklook()
    mom0.write(baseimagename.replace('.fits','.mom0.fits'),format='fits',overwrite=True)

    # create the masked peak intensity image
    masked_mom8 = maskedCube.max(axis=0,how='cube')
    if plotIt: masked_mom8.quicklook()
    masked_mom8.write(baseimagename.replace('.fits','.mom8.masked.fits'),format='fits',overwrite=True)

    # create the peak intensity image
    mom8 = cube.max(axis=0,how='cube')
    if plotIt: mom8.quicklook()
    mom8.write(baseimagename.replace('.fits','.mom8.fits'),format='fits',overwrite=True)


def betterSmooth(oldCube,smoothBeam):

    '''
    The SpectralCube.convolve_to() function currently doesn't do very well with Nans. I'm writing a quick improved version.
    '''

    from astropy import wcs
    import astropy.units as u
    from astropy.convolution import convolve
    from spectral_cube import SpectralCube

    pixscale = wcs.utils.proj_plane_pixel_area(oldCube.wcs.celestial)**0.5 * u.deg

    convolution_kernel = \
        smoothBeam.deconvolve(oldCube.beam).as_kernel(pixscale)

    # For some reason, this only works per channel for me. I don't know what they are doing in spectral_cube 
    # to do all the channels at once.
    
    newData = oldCube.unmasked_data[:,:,:].value.copy()
    nchan = len(oldCube.spectral_axis)
    for i in range(nchan):
        newData[i,:,:] = convolve(oldCube.unmasked_data[i,:,:].value, convolution_kernel,
                           normalize_kernel=True,boundary='extend')

    # make the output a cube.
    newCube = SpectralCube(newData,oldCube.wcs.copy(),beam=smoothBeam,read_beam=False,mask=oldCube.mask)

    return newCube
'''