Documentation

from astropy.io import fits
import matplotlib.pyplot as plt
#https://github.com/Caltech-IPAC/Montage/tree/main/python/MontagePy
from MontagePy.main import mImgtbl, mMakeHdr,mAdd,mOverlaps,mDiffExec,mFitExec,mBgModel,mBgExec, mProject, mProjectQL
from io import BytesIO
import numpy as np 
import pandas as pd
import os
import re
import glob 
import requests
from astropy.io.votable import parse_single_table
from contextlib import contextmanager
import multiprocessing as mp

#A function to create a set of 10 arcmin images (pointings) to fill a given image size
def get_pointings(qra,qdec,size):
    dtor = np.pi/180.0
    overlap = 0.1
    #size of each tile - 10 arcmin (max allowed by SkyMapper SIA at present)
    diam = 10.0/60 

    #separation between tiles
    delta_ra = diam*(1-overlap)/np.cos(qdec*dtor)
    delta_dec = diam*(1-overlap)

    #number of tiles needed to create image of at least size across
    #this is not a perfect calculation, but should be fine for most applications
    ntiles = int(np.ceil(size/delta_dec))

    #lists to store tile centres
    x = []
    y = []

    sra = None
    sdec = None
    #if we have an even number of tiles
    if(ntiles % 2 == 0):
        #the field centre is at the centre of a grid of tiles (not in the centre of one of these tiles)
        #here we want the start ra,dec of the first tile to be [0][0] in the [tiles][tiles] matrix
        sra = qra - ntiles/2*delta_ra + delta_ra*0.5
        sdec = qdec - ntiles/2*delta_dec +delta_dec*0.5
    else:
        #odd number of tiles - easier as ra/dec already in centre of tile in centre of mosaic...
        nshift = np.floor(ntiles/2)
        sra = qra - delta_ra*nshift
        sdec = qdec - delta_dec*nshift

    #calculate the tile positions now that we have the first tile position
    for idx in range(0,ntiles):
        ra = sra + idx*delta_ra
        for jdx in range(0,ntiles):
            dec = sdec + jdx*delta_dec
            x.append(ra)
            y.append(dec)

    return (x,y)


#This is a helper function to ensure the multiprocessing pool starts and exits cleanly
#Kindly provided by James Tocknell (Data Central)
#Sourced from: https://github.com/aragilar/DiscSolver
@contextmanager
def nicer_mp_pool(*args, **kwargs):
    """
    Wrapper around multiprocessing.Pool - maybe look at forks or alternatives?
    """
    try:
        pool = mp.Pool(*args, **kwargs)
        yield pool
    finally:
        #terminate is usually called, which can break stuff if there's a problem
        pool.close()
        pool.join()

#A function to download files from the SIA service
#Science images are downloaded as-iss
#Mask images (bitmasks) are inverted to ensure they are Montage compatible
def dload(url,fname,mask=False):
    resp = requests.get(url)
    if(resp.status_code < 400):
        if(not mask):
            data = resp.content
            fptr = open(fname,'wb')
            fptr.write(data)
            fptr.close()
        else:
            #need to invert the masks for use with Montage
            #since the masks are floats, we convert to bool and back before inverting them
            hdu = fits.open(BytesIO(resp.content),mode='update',scale_back=True)
            dat = hdu[0].data
            dat = dat.astype(np.bool).astype(np.float64).astype(np.bool)
            dat = np.invert(dat).astype(np.float64)
            hdu[0].data = dat 
            hdu.writeto(fname)
            hdu.close()

#returns the filters available for a given position
def get_filters(ra,dec,size):
    #url for accessing skymapper data from WITHIN Australia
    url = 'https://api.skymapper.nci.org.au/aus/siap/dr3/query?POS=%s,%s&SIZE=%s&RESPONSEFORMAT=VOTABLE&FORMAT=image/fits&INTERSECT=CENTER&VERB=3' % (ra,dec,size)
    #url for accessing skymapper data from OUTSIDE Australia
    #url = 'https://api.skymapper.nci.org.au/public/siap/dr2/query?POS=%s,%s&SIZE=%s&RESPONSEFORMAT=VOTABLE&FORMAT=image/fits&INTERSECT=CENTER&VERB=3' % (ra,dec,size)
    resp = requests.get(url,timeout=10)
    if(resp.status_code < 400):
        table = parse_single_table(BytesIO(resp.content)).to_table(use_names_over_ids=True)
        names = [name for name in table.colnames if len(table[name].shape) <= 1]
        df = table[names].to_pandas()
        #restrict our focus to the 'main' survey images
        df = df[df['image_type'] == 'main'].reset_index(drop=True)
        #optionally filter on the seeing in the image (here <= 3.0 arcsec)
        #df = df[(df['image_type'] == 'main') & (df['mean_fwhm'] <= 3.0)].reset_index(drop=True)
        return df['band'].unique()
    else:
        return []

def get_images(odir,im,ra,dec,size,band):
    #url for accessing skymapper data from WITHIN Australia
    url = 'https://api.skymapper.nci.org.au/aus/siap/dr3/query?POS=%s,%s&SIZE=%s&BAND=%s&RESPONSEFORMAT=VOTABLE&FORMAT=image/fits&INTERSECT=CENTER&VERB=3' % (ra,dec,size,band)
    #url for accessing skymapper data from OUTSIDE Australia
    #url = 'https://api.skymapper.nci.org.au/public/siap/dr2/query?POS=%s,%s&SIZE=%s&BAND=%s&RESPONSEFORMAT=VOTABLE&FORMAT=image/fits&INTERSECT=CENTER&VERB=3' % (ra,dec,size,band)
    resp = requests.get(url,timeout=10)
    if(resp.status_code < 400):
        table = parse_single_table(BytesIO(resp.content)).to_table(use_names_over_ids=True)
        names = [name for name in table.colnames if len(table[name].shape) <= 1]
        df = table[names].to_pandas()
        #restrict our focus to the 'main' survey images
        df = df[df['image_type'] == 'main'].reset_index(drop=True)
        print (df)
        #optionally filter on the seeing in the image (here <= 3.0 arcsec)
        #df = df[(df['image_type'] == 'main') & (df['mean_fwhm'] <= 3.0)].reset_index(drop=True)
        for idx, row in df.iterrows():
            ofname='%s/%d-%d.fits' % (odir,im,idx)
            print ('Downloading ',ofname)
            dload(row['get_fits'],ofname)

            ofname='%s/masks/%d-%d.fits' % (odir,im,idx)
            print ('Downloading ',ofname)
            dload(row['get_mask'],ofname,mask=True)
        return 1
    return 0

#A wrapper function for the projection, in case we want to call different methods
#This is mostly to help pass different parameters that cannot easily be passed using zip()
def WrapperProject(inp,output,template,weight):
    #most thorough projection method
    mProject(inp,output,template,weight_file=weight)
    #alternative quicker way to do projection - have to set noAreas to false 
    #to ensure area files are generated (they are not generated by default).
    #mProjectQL(inp,output,template,weight_file=weight,noAreas=False)

#This is needed in order to run the multiprocessing pool
if __name__ == '__main__':

    #10 arcmin tiles
    qsize = 10.0/60
    #0.5 degree mosaic size
    qmosaicsize = 0.5

    obj = 'Centaurus A'
    #recommended to use multiprocessing as projection of all images takes some time
    USE_MULTIPROCCESSING = 1
    #USE_MULTIPROCCESSING = 0

    #Skymapper filters:
    #u v g r i z 
    #get_filters queries the SIA service to find out what filters are available
    #If you want to specify just one filter, this can be done as:
    #for mnum in range(88,110):
    #Use the CDS SIMBAD name resolver service to get RA and DEC of the object.
    url = "https://simbad.u-strasbg.fr/simbad/sim-nameresolver?Ident=%s&output=json" % obj
    simbad = requests.get(url).json()
    qra = simbad[0]['ra']
    qdec = simbad[0]['dec']

    (ra_cen,dec_cen) = get_pointings(qra,qdec,qmosaicsize)
    
    #if you want to create multiple mosaics    
    #filters = get_filters(qra,qdec,qsize)
    #to create just a single band mosaic, in this case i-band
    filters = ['i']
    for qband in filters:
        qdir=re.sub(" ","",obj) + qband
        print (qdir)
        imtable = qdir + '/images.tbl'
        header = qdir + '/images.hdr'
        projdir = qdir + '/proj'
        pimtable = projdir + '/images.tbl'
        mosaic = qdir + '/' + qband + '.fits'
        difftable = projdir + '/diff.tbl'
        fitstable = projdir + '/fits.tbl'
        corrtable = projdir + '/corr.tbl'
        diffdir = projdir + '/diff'
        corrdir = projdir + '/corr'
        maskdir = qdir + '/masks'
        if(not os.path.exists(qdir)):
            os.mkdir(qdir)
        if(not os.path.exists(projdir)):
            os.mkdir(projdir)
        if(not os.path.exists(diffdir)):
            os.mkdir(diffdir)
        if(not os.path.exists(corrdir)):
            os.mkdir(corrdir)
        if(not os.path.exists(maskdir)):
            os.mkdir(maskdir)

        #download all the data...
        for idx in range(0,len(ra_cen)):
            rtn = get_images(qdir,idx,ra_cen[idx],dec_cen[idx],qsize,qband)

        #We use Montage to do the mosaicking: see
        #http://montage.ipac.caltech.edu/docs/index.html
        #and
        #https://github.com/Caltech-IPAC/Montage/tree/main/python/MontagePy
        #and
        #https://github.com/Caltech-IPAC/MontageNotebooks

        #create a table of the images to be worked on
        rtn = mImgtbl(qdir, imtable,showCorners=True)
        print (rtn)
        #create a header that encompasses all the images
        rtn = mMakeHdr(imtable,header,northAligned=1)
        print (rtn)
        images = glob.glob(qdir+'/*.fits')

        if(USE_MULTIPROCCESSING):
            #create args for mProject as lists - needed for calling with starmap 
            #as part of multiprocessing pool
            img_src = []
            img_proj = []
            img_header = []
            img_weight = []
            for im in images:
                im = re.sub(".*\/","",im)
                img_src.append(qdir+'/'+im)
                img_proj.append(projdir+'/'+im)
                img_header.append(header)
                img_weight.append(maskdir+'/'+im)
            #number of processes to run
            nworkers = min(len(images),mp.cpu_count()-1)
            print ("Projecting images using %d workers" % (nworkers))
            with nicer_mp_pool(nworkers) as pool:
                #project the images using the bitmask as the weight_file
                for result in pool.starmap(WrapperProject,zip(img_src,img_proj,img_header,img_weight)):
                    #we are only interested in running the function, so we do nothing here
                    pass
            #we are done with the pool...        

        else:
            for im in images:
                print ("Projecting ",im)
                im = re.sub(".*\/","",im)
                #project the images using the bitmask as the weight_file
                rtn = mProject(qdir+'/'+im,projdir+'/'+im,header,weight_file=maskdir+'/'+im)

        #create a table of all the projected images    
        rtn = mImgtbl(projdir, pimtable)
        print (rtn)
        #steps necessary for proper background handling...
        rtn = mOverlaps(pimtable,difftable)
        print (rtn)
        rtn = mDiffExec(projdir,difftable,header,diffdir) 
        print (rtn)
        rtn = mFitExec(difftable,fitstable,diffdir) 
        print (rtn)
        rtn = mBgModel(pimtable,fitstable,corrtable)
        print (rtn)
        rtn = mBgExec(projdir,pimtable,corrtable,corrdir)
        print (rtn)
        print ("Creating final mosaic: ",mosaic)
        #finally create the final mosaic
        rtn = mAdd(corrdir,pimtable,header,mosaic)
        print (rtn)
        #if the above steps for background handling (mOverlaps..mBgExec) are skipped, 
        #then mAdd can be run on projdir alone
        #rtn = mAdd(projdir,pimtable,header,mosaic)
        #print (rtn)