Create full Stokes cube from individual Stokes images

make_stokes_cube.py

#! /usr/bin/env python

from astropy.io import fits
import numpy as np

from argparse import ArgumentParser

def get_args():
    ps = ArgumentParser("Create a full Stokes cube from individual Stokes images.")
    ps.add_argument("i_image", help="Stokes I image. Must be 2-dimensional (i.e. not existing Stokes or frequency axis.")
    ps.add_argument("-v", "--v_image", default=None, help="Stokes V image. If not specified data are nan.")
    ps.add_argument("-q", "--q_image", default=None, help="Stokes Q image. If not specified data are nan.")
    ps.add_argument("-u", "--u_image", default=None, help="Stokes U image. If not specified data are nan.")
    ps.add_argument("-o", "--outname", default=None, help="Output file name. If not specified, Stokes I image name is used (and overwritten).")
    args = ps.parse_args()
    return args

def make_cube(image_i, 
    image_v=None, 
    image_q=None, 
    image_u=None, 
    outname=None):

    image = fits.open(image_i)
    data = np.full((4, image[0].data.shape[0], image[0].data.shape[1]), np.nan)
    data[0, :, :] = image[0].data
    stokes = [image_q, image_u, image_v]
    for i, s in enumerate(stokes):
        if s is not None:
            data[i+1, :, :] = fits.getdata(s)
    
    header = image[0].header.copy()
    header["CTYPE3"] = "STOKES"
    header["CRVAL3"] = 1.0
    header["CDELT3"] = 1.0
    header["CRPIX3"] = 1.0

    if outname is None:
        outname = image_i

    fits.writeto(outname, data, header, overwrite=True)


def cli():
    args = get_args()
    make_cube(
        image_i=args.i_image,
        image_v=args.v_image,
        image_q=args.q_image,
        image_u=args.u_image,
        outname=args.outname,
    )

if __name__ == "__main__":
    cli()