process_wds_ids.py

#!/usr/bin/env python3
"""Extracts a clean list of Washington Double Star IDs from the output
of star_query.py, and then filters that using the main WDS data file
to generate a list of likely stellar companions.

This program is free software: you can redistribute it and/or modify it under
the terms of the GNU General Public License as published by the Free Software
Foundation, either version 3 of the License, or (at your option) any later
version.

This program is distributed in the hope that it will be useful, but WITHOUT
ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details.

You should have received a copy of the GNU General Public License along with
this program. If not, see <http://www.gnu.org/licenses/>.
"""

import argparse
import os
import os.path
import sys
from astropy.io import fits
from astropy.table import Table, Column, vstack
import DavesAstropyUtils as dapu
import warnings
from astropy.utils.exceptions import AstropyUserWarning, AstropyWarning
import WDS

__author__ = "Dave Strickland"
__copyright__ = "Copyright 2018, Dave Strickland"
__date__ = "2018/03/21"
__deprecated__ = False
__email__ =  "dave.strickland@gmail.com"
__license__ = "GPLv3"
__version__ = "0.2.0"

def command_line_opts():
    # TODO: environment variable for WDS location
    p_wds='data/WDS/B_wds.fits.gz'
    p_filter = 'negative'
    p_magdiff = 6.0
    p_wds_detail_table = 'wds_detail.html'
    p_fmt = ' Format is determined from file name.'
    p_css = 'darkTable.css'
    p_inp = 'INPUT_STAR_QUERY_OUTPUT.fits'

    parser = argparse.ArgumentParser()
    # required command line arguments
    parser.add_argument(dest='fitsfile', metavar=p_inp,
        default=None,
        help='Fits table that was generated by star_query.py')
    parser.add_argument(dest='output_table', 
        default=None, metavar='OUT_TABLE',
        help='Name for table of processed targets with'+
            ' their WDS identifiers.'+
            p_fmt)

    # optional
    parser.add_argument('--wds-detail', 
        dest='wds_detail', default=None, metavar=p_wds_detail_table,
        help='Name for optional output of WDS informational data on processed'+
            ' WDS components that passed filtering.'+
            p_fmt)
    parser.add_argument('--css',
        dest='cssfile', default=p_css, metavar='table_style.css',
        help='CSS table style for output HTML (default: {})'.format(p_css))
    parser.add_argument('-w', '--wdsfile',
        dest='wdsfile', default=p_wds,
        help='Location of WDS data table. (default: {})'.format(p_wds))
    parser.add_argument('--filter',
        dest='filter', default=p_filter,
        help='Type of binary star filtering to apply.'+
            ' Valid entries are "abc", "negative", or "positive" (default: {})'.format(p_filter)+
            ' "abc" selects A, B and C components.'+
            ' "positive" selects only physically likely companions.'+
            ' "negative" deselects unphysical components and components with large magnitude differences.')

    parser.add_argument('--magdiff',
        dest='magdiff', default=p_magdiff, type=float,
        help='Maximum magnitude difference allowed in negative filter (default: {})'.format(p_magdiff))

    parser.add_argument('-v', '--verbose',
        dest='verbose', action='store_true',
        help='Verbose output for each object processed. Useful for debugging purposes.')
        
    args = parser.parse_args()
    return args

def main():
    warnings.simplefilter('ignore', category=AstropyUserWarning, append=True)

    p_args = command_line_opts()

    # Don't have to worry about race conditions for this type of work
    if not os.path.isfile(p_args.fitsfile):
        print('Error: Input file {} not found. Current dir: {}'.format(p_args.fitsfile, os.getcwd()))
        sys.exit(1)

    # Read data from star_query...
    p_idata = dapu.read_table(p_args.fitsfile, p_args.verbose)

    # Create WDS class to handle WDS-related data collection
    p_wds = WDS.WDS(p_args.wdsfile, p_args.magdiff, p_args.verbose)
    
    # Process each input in turn maintaining link between user ID
    # and WDS id, if present. Store output tables for later combination.
    output_tables = [] # store returned table of WDS data
    processed_targets = []
    processed_wds_ids = []
    num_targets = len(p_idata)
    num_found = 0
    num_skipped = 0
    
    # Some useful lists of objects that had no WDS-like IDS in the input
    # and those that had all components filtered out.
    inputs_no_wds_list=[]
    all_wds_filtered_out_list=[]
    
    print('Processing {} targets from {}'.format(num_targets, p_args.fitsfile))
    for idx in range(num_targets):
        p_target = p_idata['Star'][idx]
        p_iwds = p_idata['WDS'][idx]
        p_owds = WDS.wds_id_from_simbad_wds(p_iwds)
        if p_owds is None:
            # inform user and skip to next target
            print('  Target #{} {} did not have WDS-like ID. Skipping.'.format(idx, p_target))
            num_skipped += 1
            inputs_no_wds_list.append(p_target)
            continue
        
        print('  Target #{} {} obtained WDS-like ID {} from {}'.format(idx, p_target, p_owds, p_iwds))

        # Get the like component data from the WDS
        p_table, p_ids = p_wds.get_likely_components(p_owds, p_args.filter)
        if p_table is None:
            print('    Target #{} {} has no likely WDS components after filtering'.format(idx, p_target))
            num_skipped += 1
            all_wds_filtered_out_list.append(p_target)
            continue

        # Otherwise we got some valid data.
        num_found += 1
        output_tables.append(p_table)
        for num in range(len(p_ids)):
            processed_targets.append(p_target)
            processed_wds_ids.append(p_ids[num])
        
    # summarize loop
    print('Found WDS IDs for {} input targets, skipped {}'.format(num_found, num_skipped))

    # create an output table from the data we have
    p_otable = make_output_table(processed_targets, processed_wds_ids)
    dapu.write_table(p_otable, p_args.output_table, p_args.cssfile)
    print('Wrote filtered WDS component for input targets to {}'.format(p_args.output_table))

    # combine data for components into final table
    detail_table = vstack(output_tables, join_type='outer')
    if p_args.wds_detail is None:
        if p_args.verbose:
            print('Select WDS detail information follows:')
            print(detail_table.info)
            print('Note this information could be written to disk using --wds_detail')
    else:
        dapu.write_table(detail_table, p_args.wds_detail, p_args.cssfile)
        print('Wrote WDS component detail info to {}'.format(p_args.wds_detail))

    print('Information on targets with no WDS or all WDS components filtered out.')
    print('  {} input targets with no WDS info: {}'.format(len(inputs_no_wds_list), inputs_no_wds_list))
    print('  {} input targets where {} filtering removed all components: {}'.format(len(all_wds_filtered_out_list), 
        p_args.filter,
        all_wds_filtered_out_list))
    return

def make_output_table(target_list, wds_id_list):
    """Construct an astropy Table from the target ID
    list and the filtered WDS component list
    """
    output_table = Table()
    target_column = Column(data=target_list,
        description='Original user-specified target identifier',
        name='Star',
        format='{}')
    wds_column = Column(data=wds_id_list,
        description='WDS identifiers for selected components',
        name='WDS',
        format='{}')
    output_table.add_columns(cols=[target_column, wds_column])
    return output_table

if __name__ == "__main__":
    main()