match_io_DC2.py

import numpy as np
import glob
import fitsio
import GCRCatalogs
import GCR
from sklearn.neighbors import KDTree
import astropy.table
import healpy as hp
import os
import argparse

def spatial_closest_mag_1band(ra_data,dec_data,mag_data,
                              ra_true,dec_true,mag_true,true_id,
                              rmax=3,max_deltamag=1.):
    """
    Function to return the closest match in magnitude within a user-defined radius within certain
    magnitude difference.
    
    ***Caveats***: This method uses small angle approximation sin(theta)
    ~ theta for the declination axis. This should be fine to find the closest
    neighbor. This method does not use any weighting.
    
    Args:
    -----
    
    ra_data: Right ascension of the measured objects (degrees).
    dec_data: Declination of the measured objects (degrees).
    mag_data: Measured magnitude of the objects.
    ra_true: Right ascension of the true catalog (degrees).
    dec_true: Declination of the true catalog (degrees).
    mag_true: True magnitude of the true catalog.
    true_id: Array of IDs in the true catalog.
    rmax: Maximum distance in number of pixels to perform the query.
    max_deltamag: Maximum magnitude difference for the match to be good.
    
    Returns:
    --------
    
    dist: Distance to the closest neighbor in the true catalog. If inputs are
    in degrees, the returned distance is in arcseconds.
    true_id: ID in the true catalog for the closest match.
    matched: True if matched, False if not matched.
    """
    X = np.zeros((len(ra_true),2))
    X[:,0] = ra_true
    X[:,1] = dec_true
    tree = KDTree(X,metric='euclidean')
    Y = np.zeros((len(ra_data),2))
    Y[:,0] = ra_data
    Y[:,1] = dec_data
    ind,dist= tree.query_radius(Y,r=rmax*0.2/3600,return_distance=True)
    matched = np.zeros(len(ind),dtype=bool)
    ids = np.zeros(len(ind),dtype=true_id.dtype)
    dist_out = np.zeros(len(ind))
    for i, ilist in enumerate(ind):
        if len(ilist)>0:
            dmag = np.fabs(mag_true[ilist]-mag_data[i])
            good_ind = np.argmin(dmag)
            ids[i]=true_id[ilist[good_ind]]
            dist_out[i]=dist[i][good_ind]
            if np.min(dmag)<max_deltamag:
                matched[i]=True
            else:
                matched[i]=False
        else:
            ids[i]=-99
            matched[i]=False
            dist_out[i]=-99.
    return dist_out*3600., ids,matched
parser = argparse.ArgumentParser()
parser.add_argument('--init', dest='init', type=int, default=0, help='Tract to start')
parser.add_argument('--end', dest='end', type=int, default=168, help='Tract to end')
parser.add_argument('--data-catalog', dest='data_catalog', type=str, default='dc2_object_run2.1i_dr4',
                   help='(data) GCR catalog name to query')
args = parser.parse_args()
galaxies = GCRCatalogs.load_catalog(args.data_catalog)
cosmoDC2 = GCRCatalogs.load_catalog('cosmoDC2_v1.1.4_image')
star_table = fitsio.read('/global/projecta/projectdirs/lsst/groups/SSim/DC2/all_stars_DC2_run2.1i.fits.gz')

columns = ['ra','dec','mag_u_cModel','mag_g_cModel',
          'mag_r_cModel','mag_i_cModel','mag_z_cModel','mag_y_cModel', 'blendedness',
         'extendedness', 'objectId']
columns_true = ['galaxy_id','ra', 'dec', 'mag_u_lsst', 'mag_g_lsst', 'mag_r_lsst', 'mag_i_lsst',
                'mag_z_lsst', 'mag_y_lsst','shear_2_treecorr', 'redshift']
tracts = np.array(galaxies.available_tracts)
if 'dr1b' in args.data_catalog:
    tracts = tracts[tracts!=2897]
if args.end > len(tracts-1):
    args.end = len(tracts-1)
for tract in tracts[args.init:args.end]:
    if os.path.isfile(f'/global/projecta/projectdirs/lsst/groups/SSim/DC2/matched_ids_{args.data_catalog}_{tract}.fits.gz'):
        continue
    else:
        print('Generating match for tract:', tract)
        data_meas = galaxies.get_quantities(columns, native_filters= [f'tract == {tract}']) #tract 2897 missing/corrupt
        print('Got', len(data_meas['ra']), 'objects in the object catalog')
        max_ra = np.nanmax(data_meas['ra'])
        min_ra = np.nanmin(data_meas['ra'])
        max_dec = np.nanmax(data_meas['dec'])
        min_dec = np.nanmin(data_meas['dec'])
        vertices = hp.ang2vec(np.array([min_ra, max_ra, max_ra, min_ra]), np.array([min_dec, min_dec, max_dec, max_dec]), lonlat=True)
        print('vertices: ', min_ra, max_ra, min_dec, max_dec)
        ipix = hp.query_polygon(32, vertices, inclusive=True)
        print('healpixels to query', ipix)
        native_filter = f'(healpix_pixel == {ipix[0]})'
        for ipx in ipix:
            native_filter=native_filter+f' | (healpix_pixel == {ipx})'
        data_true = cosmoDC2.get_quantities(columns_true,
            filters=[f'ra >= {min_ra}',f'ra <={max_ra}', f'dec >= {min_dec}', f'dec <= {max_dec}',
                'mag_u_lsst < 30', 'mag_g_lsst < 30', 'mag_r_lsst < 30', 'mag_i_lsst < 30',
                'mag_z_lsst < 30', 'mag_y_lsst < 30'], native_filters=native_filter)
        print('Got', len(data_true['ra']), 'objects in cosmoDC2')
        obj_ids = np.concatenate([data_true['galaxy_id'], star_table['simobjid']])
        obj_ra = np.concatenate([data_true['ra'], star_table['ra']])
        obj_dec = np.concatenate([data_true['dec'], star_table['decl']])
        obj_mag = np.concatenate([data_true['mag_r_lsst'], star_table['rmag']]) # We match in r-band

        dist, ids, matched = spatial_closest_mag_1band(data_meas['ra'],
                                               data_meas['dec'],
                                               data_meas['mag_r_cModel'],
                                               obj_ra,
                                               obj_dec,
                                               obj_mag,
                                               np.arange(len(obj_mag)),
                                               max_deltamag=1, rmax=5)

        star = np.zeros(len(obj_mag), dtype=np.bool)
        star[len(data_true['ra'])-1:]=True
        umag = np.concatenate([data_true['mag_u_lsst'],star_table['umag']])
        gmag = np.concatenate([data_true['mag_g_lsst'],star_table['gmag']])
        imag = np.concatenate([data_true['mag_i_lsst'],star_table['imag']])
        zmag = np.concatenate([data_true['mag_z_lsst'],star_table['zmag']])
        ymag = np.concatenate([data_true['mag_y_lsst'],star_table['ymag']])
        redshift = np.concatenate([data_true['redshift'], np.zeros(len(star_table))])

        tab = astropy.table.Table([obj_ids[ids], data_meas['objectId'], matched, star[ids], obj_ra[ids], obj_dec[ids],
                           umag[ids], gmag[ids],obj_mag[ids], imag[ids], zmag[ids], ymag[ids], redshift[ids]],
                          names=('truthId','objectId','is_matched','is_star', 'ra', 'dec',
                                 'mag_u_lsst', 'mag_g_lsst','mag_r_lsst', 'mag_i_lsst', 'mag_z_lsst', 'mag_y_lsst', 'redshift_true'))
        tab.write(f'/global/projecta/projectdirs/lsst/groups/SSim/DC2/matched_ids_{args.data_catalog}_{tract}.fits.gz', overwrite=True)