sdssDR9findingChart

#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# This file is part of DownloadSDSSFindingCharts.
#
# This 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/>.

"""
Usage:
    sdssDR9findingChart [-h|--help]
    sdssDR9findingChart [options] fileWithAstroObjects

fileWithAstroObjects is a text file with an astronomical source name
per line. Non-valid names are ignored. The result is a finding chart
from SDSS DR9 as a JPEG downloaded in the directory from which the
sdssDR9findingChart script is called.

Use "sdssDR9findingChart -h" for the whole list of supported options.

See details of the DS9 Finding chart web service at:
  http://skyserver.sdss3.org/dr9/en/tools/chart/chart.asp
"""

from __future__ import division # so that 1/2 == 0.5, instead of 1/2==0

import multiprocessing
import os
import sys
import time
import urllib
from optparse import OptionParser

sesameQueryUrl = 'http://vizier.u-strasbg.fr/viz-bin/nph-sesame/-op/NSVA?%(name)s'
sdssQueryUrl   = 'http://skyservice.pha.jhu.edu/DR9/ImgCutout/getjpeg.aspx?'

def download_chart(args):
    """ Download the finding chart for an astronomical object (first element of
    the two-element tuple 'args') to a path (second element of the tuple)"""
    
    object_name, ra, dec, img_path, size, width, height, scale, rescale_velocity, sdss_opt = args
    
    # Perform the Sesame query
    sesameResponse = None
    
    if ra == None or dec == None:
        try:
            sesameQuery = sesameQueryUrl % {'name': urllib.quote(object_name.strip())}
            print "Processing Sesame query for %s" % object_name
            
            sesameResponse = urllib.urlopen(sesameQuery).readlines()
            ra, dec = filter(lambda x: x.find('%J') == 0, sesameResponse)[0].split(' ')[1:3]
        except Exception, e:
            # We raise the exception if we were not able to come up with a coordinate resolution
            if ra == None or dec == None:
                raise e
        # If we are here, we did not have ra and dec, and now we have it, so we have a valid sesameResponse
        if sesameResponse != None and rescale_velocity != None:
            try:
                velocityList = filter(lambda x: x.find('%V') == 0, sesameResponse)
                if len(velocityList) > 0:
                    velElements = velocityList[0].split(' ')
                    # We can use either velocity in km/s, or z 
                    if velElements[1] == 'v':
                        vel = float(velElements[2])
                    elif velElements[1] == 'z':
                        vel = 299792.458*float(velElements[2])
            except ValueError, e:
                vel = None
            if vel != None:
                scale = scale  * rescale_velocity / vel # Consider the scale 
    else: # We have been provided with ra and dec
        if size: # We have also been provided with a size
            # We want the R50 scale to occupy min(width/size, height/size pixels)
            scale = size*8/min(width/2, height/2)
    
    # SDSS DR9
    sdssParamsDict = {
        'ra':     ra,
        'dec':    dec,
        'width':  width,
        'height': height,
        'scale':  scale, # Native scale 0.396127
        'opt':    sdss_opt, # defaults to 'G'
        'query': 'SR(10,10)'
    }
    
    try:
        print "Querying SDSS for coordinates %s %s (%s)" % (ra, dec, object_name)
        url = sdssQueryUrl + urllib.urlencode(sdssParamsDict)
        
        def remove(path):
            """ Remove a partially downloaded file, ignore errors """
            try: os.unlink(path)
            except: pass
        
        try:
            urllib.urlretrieve(url, filename = img_path)
            print "Download of %s completed" % img_path
        except KeyboardInterrupt:
            remove(img_path)
            sys.exit(1)
        except:
            remove(img_path)
            print "Download of %s failed" % img_path
    
    except Exception, e:
        print "Error. Coordinates %s %s not found in SDSS" % (ra, dec)

def option_parsing_setup():
    """Use optparse to get the arguments"""
    # Constants for the defaults
    sdss_native_scale = 0.396127
    pixel_size = 512
    
    # Creating the options
    op = OptionParser(usage="sdssDR9findingChart [options] file") 
    op.add_option(
        "-x", "--x-size", dest="xsize", metavar="PIXEL", type="int",
        default = pixel_size,
        help="The finding chart will be PIXEL pixels wide. Default value, %s." % pixel_size
    )
    op.add_option(
        "-y", "--y-size", dest="ysize", metavar="PIXEL", type="int",
        default = pixel_size,
        help="The finding chart will be PIXEL pixels high. Default value, %s." % pixel_size
    )
    op.add_option(
        "-s", "--scale", dest="scale", metavar="SCALE", type="float",
        default = sdss_native_scale,
        help=("The finding chart pixel resolution will be SCALE arc-seconds per pixel. "+
        "Defaults to %s arc-seconds per pixel (native SDSS pixel resolution).") % sdss_native_scale
    )
    op.add_option(
        "-z", "--zoom", dest="zoom", metavar="ZOOM-RATIO", type="float",
        help=("The finding chart pixel resolution will be zoomed by ZOOM-RATIO from specified "+
        "scale ratio (native scale ratio if -s is not used). Values above 1 "+
        "imply zooming in, values below 1 zooming out. Used alone, is equivalent to the -s "+
        "option with a SCALE value of %s/ZOOM-RATIO.") % sdss_native_scale
    )
    op.add_option(
        "-r", "--rescale", dest="rescale_velocity", metavar="VELOCITY", type="float",
        help="When this option is set, the size of the image is chosen so that "+
        "most of the object can fit the image. Works more reliably for objects "+
        "with recessional velocities >1500 km/s. If you want to know the actual "+
        "scale, use it with the -L option. The autosizing mechanism will try to "+
        "assume a scaling proportional to the ration of the given VELOCITY and "+
        "the object velocity as reported by Sesame. It cannot work with the "+
        "SDSS_COORDINATES option. Use the --size-column option instead."
    )
    op.add_option(
        "-c", "--size-column", dest="size_column", action="store_true",
        default=False,
        help="This option can only be used with the SDSS_COORDINATES option. "+
        "If present, the file must contain a column named size, which shows "+
        "the size of the object in arcseconds. "+
        "The script will use the size column to compute the scale that will "+
        "keep the galaxy centered and zoomed in."
    )
    op.add_option(
        "-G", "--grid", dest="sdss_opt", action="append_const", 
        const="G",
        help="Corresponds with the G option in the SDSS web chart tool; "+
        "when used, the image sports a grid with  and crosshairs."
    )
    op.add_option(
        "-L", "--label", dest="sdss_opt", action="append_const", 
        const="L",
        help="Corresponds with the L option in the SDSS web chart tool; "+
        "when used, the image has a label overlay indicating the chart parameters."
    )
    op.add_option(
        "-I", "--invert", dest="sdss_opt", action="append_const", 
        const="I",
        help="Corresponds with the I option in the SDSS web chart tool; "+
        "when set, the image color map is inverted."
    )
    op.add_option(
        "-P", "--photometric-objects", dest="sdss_opt", action="append_const", 
        const="P",
        help="Corresponds with the P option in the SDSS web chart tool; "+
        "shows light-blue circles around SDSS photometric objects."
    )
    op.add_option(
        "-S", "--spectroscopic-objects", dest="sdss_opt", action="append_const", 
        const="S",
        help="Corresponds with the S option in the SDSS web chart tool; "+
        "shows red squares around SDSS spectroscopic objects."
    )
    op.add_option(
        "-O", "--outlines", dest="sdss_opt", action="append_const", 
        const="O",
        help="Corresponds with the O option in the SDSS web chart tool; "+
        "shows green outlines around SDSS photometric objects, showing "+
        "where the area for which the photometry is calculated."
    )
    op.add_option(
        "-B", "--bounding-boxes", dest="sdss_opt", action="append_const", 
        const="B",
        help="Corresponds with the B option in the SDSS web chart tool; "+
        "shows pink squares that bound SDSS photometric objects."
    )
    op.add_option(
        "-F", "--fields", dest="sdss_opt", action="append_const", 
        const="F",
        help="Corresponds with the F option in the SDSS web chart tool; "+
        "shows gray outlines demarcating the SDSS fields."
    )
    op.add_option(
        "-M", "--masks", dest="sdss_opt", action="append_const", 
        const="M",
        help="Corresponds with the M option in the SDSS web chart tool; "+
        "shows pink footprints around masked objects, such as bright objects "+
        "and data artifacts."
    )
    op.add_option(
        "-Q", "--plates", dest="sdss_opt", action="append_const", 
        const="Q",
        help="Corresponds with the Q option in the SDSS web chart tool; "+
        "shows circular plates (in lavender) used for spectra collection."
    )
    op.add_option(
        "-m", "--mode", dest="mode", metavar="MODE",
        type="choice",
        choices=["OBJECT_LIST", "SDSS_COORDINATES"],
        default="OBJECT_LIST",
        help="MODE can be either OBJECT_LIST or SDSS_COORDINATES. If set to "+
             "OBJECT_LIST, each line is expected to have an object name, "+
             "and Sesame will be used to obtain RA and DEC. This is the default. "+
             "If set to SDSS_COORDINATES, the file is supposed to be either "+
             "comma-, blank-, or tab-separated, with at least columns for RA y Dec. "+
             "If a column with name size is present, the tool will try to fit "+
             "the size into the frame."
    )
    
    return op

# Define __main__ environment
if __name__=='__main__':
    op = option_parsing_setup() # Returns a suitably setup option parser
    (options, args) = op.parse_args()
    
    if options.zoom != None:
        options.scale = options.scale/options.zoom
    
    # If no options, 
    sdss_opt = ""
    if options.sdss_opt:
        sdss_opt = "".join(options.sdss_opt)
    
    objects_file = args[0] # We get the first file; we could loop on this, but not for the time being.
    
    # Use a pool of workers and download the images in parallel.
    # As there is more latency than processing, we don't use CPU counts; instead, we use a number that does not cause Sesame to throttle our client 
    pool = multiprocessing.Pool(8)
    map_async_args = None
    
    with open(objects_file,'r') as f:
        # Get only non-empty lines, stripped of trailing blanks and \n
        if options.mode == "OBJECT_LIST":
            objects = filter(lambda x: len(x)>0, map(lambda x: x.strip(), f.readlines()))
            map_async_args = (
                (name, None, None, name.strip() + '.jpg', None,
                 options.xsize, options.ysize, options.scale, options.rescale_velocity,
                 sdss_opt)
                   for name in objects
            )
        elif options.mode == "SDSS_COORDINATES":
            # Use astropy to read coordinates
            try:
                from astropy.table import Table, Column
                import numpy as np
            except ImportError, e:
                print "You need astropy and numpy for parsing SDSS_COORDINATES files"
                raise e
            
            objectArray    = Table.read(f, format='ascii')
            columnNames    = objectArray.colnames
            raColNames     = filter(lambda x: x.find("ra") == 0, columnNames)
            decColNames    = filter(lambda x: x.find("dec") == 0, columnNames)
            sizeColNames   = filter(lambda x: x.find("size") == 0, columnNames)
            objectColNames = filter( # we try to use columns starting with obj, name, source or target
                                lambda x: x.find("obj") == 0 
                                       or x.find("name") == 0
                                       or x.find("source") == 0
                                       or x.find("target") == 0,
                                columnNames
            )
            
            raCol     = None
            decCol    = None
            objectCol = None
            sizeCol   = None
            
            if not raColNames or not decColNames:
                raise ValueError("There sould be a ra and dec columns")
            else:
                raCol = raColNames[0]
                decCol = decColNames[0]
            
            if not objectColNames:
                numRows = len(np.array(objectArray))
                objectArray.add_column(
                              Column( name='objid', data=range(1,numRows+1) ),
                              index=0
                )
                objectColNames = ['objid']
            
            # One way or the other, we have an object column
            objectCol = objectColNames[0]
            
            if options.size_column and not sizeColNames:
                raise ValueError(
                    "You must provide a file with a size column "+
                    "if you use the --size-column option"
                )
            if sizeColNames and options.size_column:
                sizeCol = sizeColNames[0]
            else: # We add a fake column both if there is no size column, and if we don't want to use the provided size column
                numRows = len(np.array(objectArray))
                objectArray.add_column(
                            Column( name = 'fake_size', data=[None for i in range(numRows)])
                )
                sizeCol = 'fake_size'
            
            map_async_args = (
                (row[objectCol], row[raCol], row[decCol], "%s.jpg" % row[objectCol], row[sizeCol],
                options.xsize, options.ysize, options.scale, options.rescale_velocity,
                sdss_opt)
                for row in objectArray
            )
        else:
            raise ValueError("--mode should be OBJECT_LIST or SDSS_COORDINATES")
        
        # With the pool and map_async_args we get our results
        result = pool.map_async(download_chart, map_async_args)
        while not result.ready():
            time.sleep(1)
        result.get() # reraise exceptions of the remote call, if any