.py

#! /usr/bin/env python
from __future__ import print_function, division

import argparse
import os.path

from astropy.io import fits
from numba import njit, float64, prange
import numpy as np


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('fitsfile', nargs='+')
    parser.add_argument('--weightsuffix')
    parser.add_argument('--out', default='myadd.fits')
    args = parser.parse_args()

    header = None
    imagesum = None
    weightsum = None

    for i, fitsfile in enumerate(args.fitsfile):
        print("Adding file %d/%d..." % (i+1, len(args.fitsfile)))

        try:
            image = fits.open(fitsfile)[0]
            if args.weightsuffix:
                path, ext = os.path.splitext(fitsfile)
                weights = fits.open(path + args.weightsuffix + ext)[0].data
            else:
                weights = image.data.copy()
                weights[:] = 1

            weights[~np.isfinite(weights)] = 0
            weights[~np.isfinite(image.data)] = 0
            image.data[~np.isfinite(image.data)] = 0
            image.data *= weights

            if imagesum is not None:
                imagesum += image.data
            else:
                imagesum = image.data

            if weightsum is not None:
                weightsum += weights
            else:
                weightsum = weights

            if header is None:
                header = image.header

        except IOError:
            print("Could not open %s (or its associated rms file)" % fitsfile)

    # Normalise
    imagesum /= weightsum

    hdu = fits.PrimaryHDU(imagesum, header=header)
    hdu.writeto(args.out, overwrite=True)

    hdu = fits.PrimaryHDU(weightsum, header=header)
    path, ext = os.path.splitext(args.out)
    hdu.writeto(path + '-weights' + ext, overwrite=True)



if __name__ == '__main__':
    main()
    
    
++++++++++

#! /usr/bin/env python
from __future__ import print_function, division

import argparse

from astropy.io import fits
import numpy as np

from astrobits.coordinates import fits_to_radec
from astrobits.mwabeam import MWABeam


parser = argparse.ArgumentParser()
parser.add_argument('--metafits', required=True)
parser.add_argument('--template', required=True)
parser.add_argument('--output', required=True)
args = parser.parse_args()

template = fits.open(args.template)[0]
ras, decs = fits_to_radec(template)
xx, yy = MWABeam(args.metafits).power(ras.flatten(), decs.flatten(), freq=template.header['CRVAL3'])
xx, yy = np.reshape(xx, template.data.shape), np.reshape(yy, template.data.shape)
template.data = xx + yy
template.writeto(args.output, overwrite=True)
print("Done")