Power spectrum plotting in CASA

powerspecplot_in_casa.py

from turbustat.statistics import psds
import pylab as pl
import numpy as np

def plotpsd(image, pixscale=None, mask=None):
    if isinstance(image, str):
        ia.open(image)
        data = ia.getchunk().squeeze()
        ia.close()
        
        pixscale = (np.abs(imhead(image)['incr'][:2]*3600*180/np.pi)).mean()
    else:
        data = image
        assert pixscale is not None
        
    if mask is not None:
        assert mask.shape == data.shape
        ft = np.fft.fftshift(np.fft.fft2(data * mask)**2)
    else:
        ft = np.fft.fftshift(np.fft.fft2(data)**2)
        
    pfreq, pspec = psds.pspec(ft)
    
    # filter out bad data
    pfreq = pfreq[np.isfinite(pspec)] / pixscale
    pspec = pspec[np.isfinite(pspec)]
    
    pl.loglog(pfreq_skymodel, np.abs(pspec_skymodel), 'k')
    pl.loglog(pfreq, np.abs(pspec))
    
    pl.xlabel("Frequency (1/arcsec)")
    pl.ylabel("Power")