plot_source.py

import numpy as np
import argparse
import pandas as pd
import datetime as dt
import matplotlib.pyplot as plt
import matplotlib.colors as mc
import matplotlib.cm as cm
import os
import gdown
from astropy import units as u
from astropy.coordinates import SkyCoord
from astropy.io import fits
from astropy.table import Table
from astropy.wcs import WCS
from glob import glob
from subprocess import call

def match_catalogs(position, max_separation=3.5):
    """
    Find sources in point source catalogs that are within some input 
    separation.

    inputs
    ======
    position [SkyCoord] : center position 
    max_separation [float] : maximum separation from input position 
                             to include [arcminutes]

    returns
    =======
    matched_sources [dict] : 
    """
    
    catalogs = {
        'xmm' : 'catalogs/4XMM_abr_cat.fits',
        'chandra' : "catalogs/chandra.fits",
        'erosita' : "catalogs/erosita.fits",
        'rosat' : "catalogs/rosat.fits"
        }


    flag = False
    input_source = pd.DataFrame({"flux" : [np.inf],
                                 "separation" : [0],
                                 "ra" : position.ra.to('deg').value,
                                 "dec" : position.dec.to('deg').value})
    matched_sources = None 
    for catalog in catalogs.keys():
        if flag:
            continue
        cat = Table.read(catalogs[catalog],format='fits')

        if catalog == 'xmm':
            cat[f"PN_8_FLUX"].name = 'flux'
            cat = cat[~cat[f"flux"].mask]
            cat = cat[cat['flux'] > 0]
            obs_cnd = np.array(cat[f"PN_SUBMODE"] != 'UNDEFINED') & \
                      np.array(cat[f"PN_SUBMODE"] != 'PrimePartialRFS') & \
                      np.array(cat[f"PN_SUBMODE"] != 'PrimePartialW2') & \
                      np.array(cat[f"PN_SUBMODE"] != 'PrimePartialW3') & \
                      np.array(cat[f"PN_SUBMODE"] != 'PrimePartialW4') & \
                      np.array(cat[f"PN_SUBMODE"] != 'PrimePartialW5')
                      #np.array(cat[f"PN_SUBMODE"] != 'PrimeSmallWindow') & \
        else:
            obs_cnd = np.ones(len(cat)).astype(bool)
        matching_survey = catalog
        posns = SkyCoord(cat['ra'].value,cat['dec'].value,unit='deg')
        seps = position.separation(posns)
        sep_cnd = np.array(seps.to('arcmin').value < max_separation) & \
                  np.array(seps.to('arcsec').value > 5)

        if not np.any(sep_cnd&obs_cnd):
            continue

        fluxes = cat['flux'][sep_cnd&obs_cnd]
        sort_cnd = np.argsort(fluxes)[::-1]
        matched_sources = pd.DataFrame({
                            'flux' : cat['flux'][sep_cnd&obs_cnd][sort_cnd],
                            'separation' : seps[sep_cnd&obs_cnd][sort_cnd],
                            'ra' : posns.ra.to('deg').value[sep_cnd&obs_cnd][sort_cnd],
                            'dec' : posns.dec.to('deg').value[sep_cnd&obs_cnd][sort_cnd]})
        flag = True
    if matched_sources is not None:
        matched_sources = pd.concat([input_source,matched_sources],ignore_index=True)
    else:
        matched_sources = input_source


    return matched_sources, matching_survey

def download_image(position,max_separation=3.5,obsid=None):
    """
    Code to search XMM for observations around a given position. 

    Inputs
    ======
    position [SkyCoord] : position of candidate of interest
    max_separation [float] : radius of circular region in arcminutes
    obsid [str] : XMM observation ID to use (otherwise, finds the nearest one)

    Returns
    =======
    image_path [str] : points to downloaded image
    """

    xmm_cat = Table.read("catalogs/4XMM_abr_cat.fits")
    if obsid is None:
        obs_cnd = np.array(xmm_cat[f"PN_SUBMODE"] != 'UNDEFINED') & \
                  np.array(xmm_cat[f"PN_SUBMODE"] != 'PrimeSmallWindow') & \
                  np.array(xmm_cat[f"PN_SUBMODE"] != 'PrimePartialRFS') & \
                  np.array(xmm_cat[f"PN_SUBMODE"] != 'PrimePartialW2') & \
                  np.array(xmm_cat[f"PN_SUBMODE"] != 'PrimePartialW3') & \
                  np.array(xmm_cat[f"PN_SUBMODE"] != 'PrimePartialW4') & \
                  np.array(xmm_cat[f"PN_SUBMODE"] != 'PrimePartialW5')
        xmm_cat = xmm_cat[obs_cnd]
        s_cat = SkyCoord(xmm_cat['ra'],xmm_cat['dec'],unit='deg')
        seps = s_cat.separation(position)
        if seps.min() > max_separation*u.arcmin:
            return

        obsid = xmm_cat['obsid'][seps == seps.min()][0]
        

    if len(glob(f"images/{obsid}.fits")) == 0:
        print(f"Downloading image for ID {obsid}")
        cmd = f"curl -o {obsid}.tar " + \
              f"'https://nxsa.esac.esa.int/nxsa-sl/servlet/data-action-" + \
              f"aio?obsno={obsid}&level=PPS" + \
               "&name=3COLIM'"
        call(cmd,shell=True)
        call(f"tar xf {obsid}.tar",shell=True)
        if not os.path.isdir("images"):
            call("mkdir images",shell=True)
        call(f"mv {obsid}/pps/*.FTZ images/{obsid}.fits",shell=True)
        call(f"rm -rf {obsid} {obsid}.tar",shell=True)

    return f"images/{obsid}.fits"

def parse_header(hdu,output_filename):
    """
    Prints info from an HDU header to an output .log file. 

    Inputs
    ======
    hdu [astropy.io.fits.image.PrimaryHDU] : image HDU
    output_filename [str] : base name for output .log file
    """

    keys = ["INSTRUME", "EXPOSURE", "CONTENT", "DATE"]

    with open(output_filename+'.log','w') as fil:
        fil.write(f"Header info read from {hdu.fileinfo()['file'].name}\n")

        for key in keys:
            fil.write(f"{key} : {hdu.header[key]}\n")


def plot_matches(matched_sources,catalog,max_separation=3.5,image=None,fn=None):
    """
    Plots both the source of interest and the matching sources from a given
    catalog. Will also plot the X-ray image (if available).

    Inputs
    ======
    matched_sources [DataFrame] : output of match_catalogs, contains both the
                                  source of interest and the matching catalog
                                  sources
    catalog [str] : string identifier for catalog ['xmm','chandra','erosita','rosat']
    max_separation [float] : separation from source of interest in arcminutes
    image [str] : points to .fits image
    fn [str] : output filename
    """
    
    if len(matched_sources) <= 1:
        return 
    fig,ax = plt.subplots(figsize=(10,10))
    plt.xlabel('RA [deg]')
    plt.ylabel('Dec [deg]')
    ax.set_facecolor('black')
    if image is not None:
        hdu = fits.open(image)[0]
        parse_header(hdu,image.split('/')[-1].split('.')[0]+'_info')
        w = WCS(hdu.header)
        im = np.sum(hdu.data,axis=0) 

        # plot the image using pixels 
        ax.imshow(im,cmap=cm.Greys_r,norm=mc.LogNorm())

        # need to calculate pixel values for position offsets
        center = SkyCoord(matched_sources['ra'][0],matched_sources['dec'][0],unit='deg')
        center_pix = np.array(center.to_pixel(w)).astype(float) 
        ax.set_xlim([center_pix[0] - 0.1*len(im[0]),center_pix[0] +  0.1*len(im[0])])
        ax.set_ylim([center_pix[1] -  0.1*len(im[0]), center_pix[1] +  0.1*len(im[0])])

        # tick units are now correct, but labels are still in pixels (and unformatted)
        yticks = plt.yticks()[0]
        xticks = plt.xticks()[0][:len(yticks)]
        yticks = yticks[:len(xticks)]
        xticks_txt, yticks_txt = np.array(w.pixel_to_world_values(np.array([xticks,
                                                                            yticks,
                                                                            xticks]).T)).T[:2]
        plt.xticks(xticks,
                   [f"{np.floor(xt):2.0f}h"+ \
                           f"{(xt-np.floor(xt))*60:2.1f}m" for xt in xticks_txt*u.deg.to('hourangle')],
                   rotation=45)
        yt_formatted = []
        for yt in yticks_txt:
            if yt > 0:
                yt_formatted.append(f"{np.floor(yt*u.deg).value:2.0f}"+r"$\degree$"+ \
                          f"{(yt-np.floor(yt))*u.deg.to('arcmin'):2.2f}'")
            else:
                yt_formatted.append(f"-{np.floor(abs(yt)*u.deg).value:2.0f}"+r"$\degree$"+ \
                          f"{60-(yt-np.floor(yt))*u.deg.to('arcmin'):2.2f}'")
        plt.yticks(yticks,yt_formatted)


        # need to calculate the pixel positions of the sources as well
        pix_ra, pix_dec = np.array(w.world_to_pixel_values(np.array([matched_sources['ra'],
                                                                     matched_sources['dec'],
                                                                     matched_sources['dec']]).T)).T[:2]
        ax.scatter(pix_ra[0],pix_dec[0],marker='*',facecolor='yellow',
                   edgecolor='black',s=300,label='Source',zorder=2)
        ax.scatter(pix_ra[1:],pix_dec[1:],marker='+',color='red',label=catalog,
                    #s=(matched_sources['flux'][1:]/matched_sources['flux'].min())*50)
                   s=500*matched_sources['flux'][1:]/matched_sources['flux'][1:].max())

        circ = w.world_to_pixel_values(
                [[center.ra.to('deg').value,
                   (center.dec+max_separation*u.arcmin).to('deg').value,
                   1]])[0]

        radius = np.sqrt((center_pix[0]-circ[0])**2 + (center_pix[1]-circ[1])**2)
        circle = plt.Circle((pix_ra[0],pix_dec[0]),
                            radius, 
                            color='lime',lw=3,fill=False)
        ax.add_patch(circle)

    else:
        if len(matched_sources['ra'] < 2):
            plt.close()
            return
        ax.scatter(matched_sources['ra'][0],matched_sources['dec'][0],
                    facecolor='yellow',marker='*',s=300,label='Source',edgecolor='black',
                    zorder=2)
        ax.scatter(matched_sources['ra'][1:],matched_sources['dec'][1:],
                   s=500*matched_sources['flux'][1:]/matched_sources['flux'][1:].max(),
                    marker='+',color='red',label=catalog)
        ax.set_xlim([(matched_sources['ra'][0]*u.deg - 1.2*max_separation*u.arcmin).value,
                  (matched_sources['ra'][0]*u.deg + 1.2*max_separation*u.arcmin).value])
        ax.set_ylim([(matched_sources['dec'][0]*u.deg - 1.2*max_separation*u.arcmin).value,
                  (matched_sources['dec'][0]*u.deg + 1.2*max_separation*u.arcmin).value])
        circle = plt.Circle((matched_sources['ra'][0],
                             matched_sources['dec'][0]),
                            max_separation*u.arcmin.to('deg'),
                            color='lime',lw=3,fill=False)
        ax.add_patch(circle)
        ax.tick_params(axis='x',rotation=45)
    ax.set_aspect('equal')
    plt.tight_layout()
    plt.legend(loc=(-0.1,-0.1))
    if fn is not None:
        plt.savefig(f"{fn}.png")
    plt.show()
    plt.close()


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
            description="Find sources within a given angular separation " + \
                        "from various instrument catalogs")
    parser.add_argument("-ra",help="Source RA [deg]",type=float)
    parser.add_argument("-dec",help="Source DEC [deg]",type=float)
    parser.add_argument("-rpp_name",help="Name of RPP source",type=str)
    parser.add_argument("-sep",help="Maximum separation [arcmin]",
                        type=float,default=3.5)
    parser.add_argument("-np",help="Don't plot",default=False,type=bool,
                        action=argparse.BooleanOptionalAction)
    parser.add_argument("-f",help="Output filename (.png), if None then plot" +\
                        "isn't saved",
                        type=str)
    parser.add_argument("-o",help="XMM Observation ID to use for image," +\
                        "otherwise use the closest to the input position",
                        type=str)

    args = parser.parse_args()
    if not os.path.isdir("catalogs"):
        gdown.download_folder("https://drive.google.com/drive/" + \
                              "folders/19iyl3ieItlPSRNnukeGczjhu7vqQd15k?usp=sharing")
    rpp = Table.read("catalogs/rpp_catalog.txt",format='ascii')

    if args.ra is not None and args.dec is not None:
        position = SkyCoord(args.ra,args.dec,unit='deg')
        s_rpp = SkyCoord(rpp['ra'],rpp['dec'],unit='deg')
        seps = position.separation(s_rpp)
        sep = seps.min()
        src = rpp[np.where(seps == sep)[0][0]]['name']
        if sep > 1*u.deg:
            print(f"closest RPP: {src} ({sep.to('deg'):2.2f})")
        elif sep < 1*u.arcmin:
            print(f"closest RPP: {src} ({sep.to('arcsec'):2.2f})")
        else:
            print(f"closest RPP: {src} ({sep.to('arcmin'):2.2f})")
    elif args.rpp_name is not None:
        if args.rpp_name in rpp['name']:
            src = rpp[rpp['name'] == args.rpp_name]
            position = SkyCoord(src['ra'],src['dec'],unit='deg')
        else:
            raise ValueError(f"{args.rpp_name} not in catalog")
    else:
        raise ValueError("Must provide either -ra and -dec or -rpp_name")

    matched_sources, matching_survey = match_catalogs(position,
                                                      max_separation=args.sep)
    today = dt.datetime.today()
    today =f"{str(today.month).zfill(2)}{str(today.day).zfill(2)}{today.year}"
    filename = f"{matching_survey}_{today}"
    if args.f is not None:
        filename = args.f+'_'+matching_survey
    if len(matched_sources) <= 1:
        print("No sources found")
    else:
        print(f"writing sources to {filename}.csv")
        matched_sources.to_csv(filename+'.csv',index=False)
        if not args.np:
            image = download_image(position,max_separation=args.sep,obsid=args.o)
            if args.o is None:
                if image is not None:
                    obsid = "_"+image.split('/')[-1].split('.')[0]
                else:
                    obsid = ''
            else:
                obsid = "_"+args.o
            plot_matches(matched_sources,matching_survey,max_separation=args.sep,image=image,fn=filename+obsid)