Create some fake all-sky polarization data in HEALpix format

dummydata.py

import numpy as np
import healpy as hp
from astropy.io import fits
import astropy.constants as const


def compscreen(lsq, phi_0, psi_0):
    '''
    Get complex polarization for a Faraday screen.
        Inputs:
                lsq - lambda^2 vector
                phi_0 - Center of Burn slab
                psi_0 - Initial polarisation angle
        Output:
                Complex polarisation (Q+iU)
    '''
    POL = np.exp(2j * (np.radians(psi_0) + phi_0 * lsq))
    return POL

def savecube(data, f, Nside, stokes):
    '''
    Save data to FITS file
    '''
    hdu0 = fits.PrimaryHDU(data)
    new_hdul = fits.HDUList([hdu0])

    hdr1 = new_hdul[0].header

    hdr1.set('BUNIT', 'pseudo K', 'Arbitrary brightness units')
    hdr1.set('STOKES', stokes)
    hdr1.set('CTYPE1', 'FREQ', 'Frequency')
    hdr1.set('CUNIT1', 'Hz')
    hdr1.set('CRPIX1', f[0])
    hdr1.set('CDELT1', f[1] - f[0])

    hdr1.set('CTYPE2', 'pix', 'HEALpix pixel')
    hdr1.set('CRPIX2', 0)
    hdr1.set('CDELT2', 1)
    hdr1.set('NSIDE', Nside, 'HEALpix NSIDE')

    outfile = 'dummycube.' + stokes + '.' + str(Nside) + '.fits'
    new_hdul.writeto(outfile, overwrite = True)
    print 'Saved to ./' + outfile

def savemap(data, plane, Nside, stokes):
    outfile = 'dummymap.' + str(freq[plane]/1e6) + '.' + stokes + '.' + str(Nside) + '.fits'

    if stokes == 'q':
        column = 'Q_POLARISATION'
    if stokes == 'u':
        column = 'U_POLARISATION'
    hp.write_map(outfile, data, column_names = [column], coord = 'G')
    print 'Saved to ./' + outfile

if __name__ == "__main__":
    # Choose HEALpix NSIDE
    nside = 32 # must be power of 2
    N_pix = hp.nside2npix(nside) # No. of pixels (12*NSIDE^2)

    # Choose frequency range (e.g. GMIMS-LBS)
    f_0 = 300e6
    f_1 = 480e6
    d_f = 0.5e6

    freq = np.arange(f_0, f_1 + d_f, d_f) # frequency vector in Hz

    C = const.c.to_value()
    lsq = (C / freq)**2 # \lambda^2 vector

    N_freq = len(freq)

    # Generate range of Faraday depths and initial angles
    depths = np.arange(-99.,99.+1.,1.)
    angles = np.arange(-np.pi/2, +np.pi/2+np.pi/100, np.pi/100)

    # Init cubes
    qcube = np.zeros((N_pix, N_freq)) * np.nan
    ucube = np.zeros((N_pix, N_freq)) * np.nan

    # Generate random data
    for i in range(N_pix):
        phi_0 = np.random.choice(depths)
        psi_0 = np.random.choice(angles)
        POL = compscreen(lsq, phi_0, psi_0)
        qcube[i] = np.real(POL[::-1]) # reverse array to match frequency order
        ucube[i] = np.imag(POL[::-1]) # reverse array to match frequency order

    # Save 2D array to a FITS file -- This is useful but not a standard
    savecube(qcube, freq, nside, 'q')
    savecube(ucube, freq, nside, 'u')

    # Save single frequency map to HEALpix FITS file --  This is standard but can
    # only handle one frequency slice at a time
    plane = 0 # Choose first frequency slice,
    # you could loop over all planes, but this would generate many files
    savemap(qcube[:,plane], plane, nside, 'q')
    savemap(ucube[:,plane], plane, nside, 'u')