query.py | searchcode

/query.py

https://bitbucket.org/iraklis_k/samidb
Python | 543 lines | 482 code | 12 blank | 49 comment | 10 complexity | 87e518d054e545762842808d5f2b3a9e MD5 | raw file
Possible License(s): BSD-3-Clause

"""
query.py
--------
Python module within samiDB.

Description: 
  Query codes for the SAMI Galaxy Survey data archive. The main code is queryMultiple(), which uses PyTables to execute a series of individual queries and combine their outcomes. The rest supports this code, which has now superseded the old queryMaster() code on which it was based.

Development:
  08.10.2013, Iraklis Konstantopoulos (iraklis@aao.gov.au). Contributed codes acknowledged in the correspoding documentation strings.

Dependencies:
  tables, utils, os, h5py, sys, numpy
"""

import tables

# ~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~
def print_sami(s, idfile, queryText, outFile=True, verbose=True):
# ~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~
    """ Define a Row Iterator for screen output. 
    
    This bit that executes a PyTables query. All the coordination for this is done in queryMultiple() or queryMaster(), but the row iterator actually does the heavy lifting, one row (s) at a time. It can also output the ID list returned from the query code as an ascii file.

    Args:
        s         [str]: An individual row of the table being queried.
        
        idfile    [str]: Filename for output ID list (ascii).
        
        queryText [str]: Query being executed for screen output.

    Keyword args:
        outFile   [boo]: File output toggle.

        verbose   [boo]: Toggles diagnostic and declarative verbosity.
    """

    # Prepare some variables. 
    counter = 0
    idlist = []

    # Iterate over all supplied rows. 
    if outFile:
            f = open(idfile, 'w')
  
    for tables.row in s:
        name, z = tables.row['CATID'], tables.row['z_spec']
        if verbose:
            print("  Found SAMI galaxy %s at redshift z=%g" % (name, z))
        counter += 1
        idlist.append(name)
        if outFile: 
            f.write(str(name)+'\n')
    if verbose:
        print("\n  Found "+str(counter)+" galaxies satisfying query:\n  "+
              queryText)

    if outFile: 
        f.close()
    
    return(idlist)


# ~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~
def makeTable(table, tabIndex, quickLookDir):
# ~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~
    """ Generate an html table for website output.

    Outputs an html table containing the rows satisfying a query (or list of queries). The quickLookDir arg will be phased out as soon as the single-object viewer can be generated on-the-fly, although it will remain in use by the flat backup filesystem.

    Args:
        table        [str]: The name of the table being parsed (within h5 block).
        
        tabIndex     [str]: Index array for the rows of a table to be display.
        
        quickLookDir [str]: URL where quicklook PDFs can be foind on server.
    """

    print('TEST')
    print(tabIndex)
    print('TEST')

    # Write html preamble. 
    htmlTab = '<html><body><table>'

    # Write Row 0, column headings. 
    htmlTab = htmlTab+\
              "<tr>"+\
              "<td>Quicklook</td>"+\
              "".join(["<td>"+str(s)+"</td>" for s in table.colnames])+\
              "</tr>"

    # Populate the table. 
    for tables.row in table[tabIndex]:
        
        hlink = "<td><a href='"+quickLookDir+str(tables.row[0])+".pdf'>" +\
                "View</a></td>"
        try:
            htmlTab = htmlTab+\
                      "<tr>" + hlink+\
                      "".join(["<td>"+str(s)+"</td>" for s in tables.row])+\
                      "</tr>"
        except:
            print('Nah, mate [2]')

    # Wrap up html, return table.
    htmlTab = htmlTab +"</table></body></html>"
    return(htmlTab)


    """ This used to write an html file. Keeping old code here for now. 
    # Open a file to write the table, write html preamble. 
    f = open(tableOut, 'w')

    f.write('<html><body><table>')
    f.write("<tr>" + 
            "<td>Quicklook</td>" + 
            "".join(["<td>"+str(s)+"</td>" for s in table.colnames]) + 
            "</tr>")

    # Where do the quick-look plots live? 
    baseURL ='file:///Users/iraklis/Data/SAMI/datasheets/GAMA/'

    # Do the deed. 
    for tables.row in table[tabIndex]:

        hlink = "<td><a href='"+baseURL+str(tables.row[0])+".pdf'>" +\
                "View</a></td>"
        f.write("<tr>" + hlink + 
                "".join(["<td>"+str(s)+"</td>" for s in tables.row]) + 
                "</tr>")

    # Wrap up html. 
    f.write("</table></body></html>")
    f.close()
    """ 


# ~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~
def queryMaster(h5file, queryIn, version='', idfile='sami_query.lis', 
                verbose=False, returnID=False, tabulate=True, 
                quickLookDir='./', overwrite=True):
# ~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~
    """ Read a SAMI master table and perform a query 

    This code queries the SAMI target table. It is now obsolete, as the more advanced queryMultiple() code can query the target table in as little time. It has not, however, been deleted, as an application for it may arise when we deploy the SAMI server interface. See above for an explanation of the quickLookDir argument. The rest should be self-explanatory, but drop the maintainer a line if you don't think so. The output of all query codes is restricted to lists of SAMI identifiers (long integers). This way the lists can be read in by the export code and be modular. This also helps with keeping transactions between the Python code and the server as simple as possible.

    Args:
        h5file       [str]: Name of SAMI archive file.
        
        queryIn      [str]: Query text or name of file containing the query string.

    Keyword args:
        version      [str]: Data version to query. Blank string for latest version.
        
        idfile       [str]: Name of output ID list ascii file.
        
        returnID     [boo]: Toggle ascii file output.
        
        tabulate     [boo]: Toggle html table output (always true for server).
        
        quickLookDir [str]: URL where quicklook PDFs can be foind on server.
        
        overwrite    [boo]: Clobber flag for 'idfile'.
        
        verbose      [boo]: Toggles diagnostic and declarative verbosity.
    """

    import utils
    import os
    import h5py as h5

    # Interpret the 'query' argument (look for a filename). 
    if os.path.isfile(queryIn):
        fq = open(queryIn, 'r')
        queryText = fq.readlines()[0]
    else:
        queryText = queryIn

    # Check that the nominated h5 file exists. 
    if not os.path.isfile(h5file):
        raise SystemExit("Cannot find nominated HDF5 file ('"+h5file+"').")
        
    # Get latest data version, if not supplied
    hdf0 = h5.File(h5file, 'r')
    version = utils.getVersion(h5file, hdf0, version)
    hdf0.close()

    # Open-read h5file. 
    hdf = tables.openFile(h5file, 'r')

    # Optionally open an ascii file to write IDs returned by the query. 
    if returnID:
        # Check if the file exists, check overwrite flag:
        if os.path.isfile(idfile):
            if not overwrite:
                hdf.close()
                raise SystemExit("The nominated output file ('"+idfile+"') "+
                                 "already exists. Please raise the 'overwrite'"+
                                 " flag or enter a different filename. ")

    # Identify the SAMI master table, assumed to live in the Table directory
    g_table = hdf.getNode('/SAMI/'+version+'/Table/')
    master = g_table.SAMI_MASTER

    # Run the row iterator.
    try: 
        idlist = print_sami(master.where(queryText), idfile, 
                            queryText, outFile=returnID, verbose=verbose)
    except:
        hdf.close()
        raise SystemExit("Oops! Your query was not understood. Please "+
                         "check the spelling of the chosen variable.")

    # Generate a table, if requested, given a row index and the table.
    if tabulate:
        tabIndex = [row.nrow for row in master.where(queryText)]
        madeTab = makeTable(master, tabIndex, quickLookDir)
        hdf.close()
        return(madeTab)

    # Otherwise close h5 file and return list of query results. 
    if not tabulate:
        hdf.close()
        return(idlist)



# ~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~
def queryMultiple(h5file, queryIn, writeFile=False, outFile='multipleQuery.lis',
                  overwrite=False, tabulate=True, quickLookDir='./', 
                  verbose=True, version=''):
# ~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~
    """ Query multiple tables and combine results 
    
    Very much like queryMaster(), in fact it is an extension of that code. So please have a look at the queryMaster() documentation string for some basic information on how SAMI queries are executed. This code runs a series of individual queries and then performs an intersection of the resultant ID lists. No effort has been made to immitate the speedup advantage offered by a proper table 'join', as this seems to be on the development list for PyTables proper. The code is not slow, however, and should satisfy even the most demanding queries.

    Args:
        h5file       [str]: Name of SAMI archive file.
        
        queryIn      [str]: Query text or name of file containing the query string.
        
    Keyword args:
        writeFile    [boo]: Toggle ascii file output.
        
        outFile      [str]: Name of output ID list ascii file.
        
        version      [str]: Data version to query. Blank string for latest version.
        
        tabulate     [boo]: Toggle html table output (always true for server).
        
        quickLookDir [str]: URL where quicklook PDFs can be foind on server.
        
        overwrite    [boo]: Clobber flag for 'outFile'.
        
        verbose      [boo]: Toggles diagnostic and declarative verbosity.
    """

    import utils
    import os
    import h5py as h5
    
    # Interpret the 'query' argument (look for a filename). 
    if os.path.isfile(queryIn):
        queryInFile = True
    else:
        queryInFile = False

    # Check that the nominated h5 file exists. 
    if not os.path.isfile(h5file):
        raise SystemExit("Cannot find nominated HDF5 file ('"+h5file+"').")

    # Is the query being read from a file? 
    if queryInFile:
        # Open and read the query file line-per-line (even=table,odd=query). 
        counter = 0
        tabs, queries = [], []
        with open(queryIn) as f:
            for i, l in enumerate(f):
                # Read the file in odd even lines
                if i%2==0:
                    tabs.append(l.strip())
                else:
                    queries.append(l.strip())
                    counter += 1
        f.close()
    
    # Or is it being supplied as a multi-line string? 
    else:
        # Split odd-even lines. 
        tabs = queryIn.split('\n')[::2]
        queries = queryIn.split('\n')[1::2]
        
        # Get rid of any blank lines...
        while '' in tabs: 
            tabs.remove('')
        while '' in queries: 
            queries.remove('')
        
        # ...and any whitespace. 
        tabs = [s.strip() for s in tabs]
        query = [s.strip() for s in queries]

        # Keep a counter of tables+queries for loops to come.
        counter = len(tabs) + len(queries)
    
    # Check if the length of the two lists is equal. 
    if len(tabs) != len(queries):
        raise SystemExit("Table-query mismatch. Please input an equal "+
                         "number of 'table' and 'query' lines. ")
    else: 
        print("Read %g queries" % (counter/2))
        if verbose: 
            for i in range(counter/2):
                print("In table '"+tabs[i]+"' query: "+queries[i])
            print("")

    # Get latest data version, if not supplied
    hdf0 = h5.File(h5file, 'r')
    version = utils.getVersion(h5file, hdf0, version)
    hdf0.close()

    # Now run all queries: 

    # Open h5 file with PyTabs: 
    hdf = tables.openFile(h5file, 'r')

    # Read all tables, append to list
    h5tabs = []
    for i in range(counter/2):
        try:
            h5tabs.append(hdf.getNode('/SAMI/'+version+'/Table/', tabs[i]))
        except:
            hdf.close()
            raise SystemExit("Oops! Your query was not understood. Please "+
                             "check the spelling of table '"+tabs[i]+"'.")

    # OK, have the tables defined as variables, now need to query them. 

    # Run each query: 
    all_lists = []
    for i in range(counter/2):
        try:
            idlist = print_sami(h5tabs[i].where(queries[i]), outFile, 
                                queries[i], outFile=False, verbose=False)
            all_lists.append(idlist)
        except:
            hdf.close()
            raise SystemExit("Oops! Your query was not understood. Please "+
                             "check the spelling of query '"+queries[i]+"'.")
        if verbose: 
            print("Query "+str(i+1)+": Found "+str(len(idlist))+
                  " galaxies satisfying "+queries[i])

    # Finally, intersect all idlists within all_lists container return, exit.
    final_idlist = set(all_lists[0]).intersection(*all_lists)
    if verbose: 
            print("\n ------- Found "+str(len(final_idlist))+
                  " galaxies satisfying all queries simultaneously.")
    if writeFile:
        # Check if the file exists. 
        if (os.path.isfile(outFile)) and (not overwrite):
                print("\nFile already exists, please choose other filename or "+
                      "raise 'overwrite' flag.")

        if (not os.path.isfile(outFile)) or overwrite:
            f = open(outFile, 'w')
            [f.write(s) for s in str(list(final_idlist))]
            f.close()
    
    if tabulate:
        # Identify the SAMI master table, assumed to live in the Table directory
        g_table = hdf.getNode('/SAMI/'+version+'/Table/')
        master = g_table.SAMI_MASTER
        
        # Iterate over Master, locate CATID in final_idlist. 
        tabIndex = []
        for row in master:
            if row['CATID'] in final_idlist:
                tabIndex.append(row.nrow)

        madeTab = makeTable(master, tabIndex, quickLookDir)
        hdf.close()
        return(madeTab)
    else:
        hdf.close()
        return(final_idlist)
        

# ~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~
def test_contents(h5file):
# ~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~
    """ Check if h5file contains cubes and RSS datasets 
    
    This development code is largely obsolete now, but can morph into a quality assessment tool down the track.

    Args:
        h5file [str]: HDF5 archive to be scanned.
    """

    ### UNDER CONSTRUCTION ###

    import h5py as h5
    import sys
    
    f = h5.File(h5file)
    if 'SAMI' in f.keys():

        if ('Targets' in f['SAMI'].keys()):
            print('File is SAMI-formatted')
            
            targ_gr = f['SAMI/Targets']
            targets = targ_gr.keys()

        else: sys.exit('File is not SAMI-formatted')

        for i in range(len(targets)):

            # Check contents of each group:
            all_cube = ['Blue_cube_data', 'Blue_cube_variance', 
                        'Blue_cube_weight', 'Red_cube_data', 
                        'Red_cube_variance', 'Red_cube_weight']
            
            all_rss = ['Blue_RSS_data', 'Blue_RSS_variance', 
                       'Red_RSS_data', 'Red_RSS_variance'] 
                        
            if all_cube[:] in f['SAMI/Targets/'+targets[i]+'/v01'].keys():
                print(targets[i], 'is complete')

    else: sys.exit('File is not SAMI-formatted')


# ~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~
def querycone(h5file, RAc, DECc, radius, version='', idfile='sami_query.lis', 
              outFile=True, verbose=True, returnID=True, overwrite=True):
# ~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~
    """ Performs simple cone search, returns ID list.

    Written by Luca Cortese, based on the queryMaster() function. This should be self explanatory to astronomers, please contact the maintainer if you would like to see some clarifications inserted here.

    Args:
        h5file    [str]: Name of SAMI archive file.
        RAc       [flt]: Right ascension at centre of cone in decimal degrees.
        DECc      [flt]: Declination at centre of cone in decimal degrees.
        radius    [flt]: Radius of cone in decimal degrees.

    Keyword args:
        version   [str]: Data version to query. Blank string for latest version.
        idfile    [str]: Name of output ID list ascii file.
        outFile   [boo]: Toggle ascii file output.
        returnID  [boo]: Toggle ascii file output.
        overwrite [boo]: Clobber flag for 'outfile'.
        verbose   [boo]: Toggles diagnostic and declarative verbosity.
    
    .. warning:: Discovered a redundancy between 'outFile' and 'returnID'. Also between the 'name' variable and 'idFile' argument. [14.04.2014]
    """
    import utils
    import os
    import h5py as h5
 
    # Get latest data version, if not supplied
    hdf0 = h5.File(h5file, 'r')
    version = utils.getVersion(h5file, hdf0, version)
    hdf0.close()

    # Open-read h5file. 
    hdf = tables.openFile(h5file, 'r')

    # Optionally open an ascii file to write IDs returned by the query. 
    if returnID:
        # Check if the file exists, check overwrite flag:
        if os.path.isfile(idfile):
            if not overwrite:
                hdf.close()
                raise SystemExit("The nominated output file ('"+idfile+"') "+
                                 "already exists. Please raise the 'overwrite'"+
                                 " flag or enter a different filename. ")

    # Identify the SAMI master table, assumed to live in the Table directory
    g_table = hdf.getNode('/SAMI/'+version+'/Table/')
    master = g_table.SAMI_MASTER

    # Prepare some variables. 
    counter = 0
    idlist = []

    if outFile:
            f = open(idfile, 'w')

   
    for tables.row in master:
    			ra, dec = tables.row['RA'], tables.row['Dec']
			dist=sph_dist(ra,dec,RAc,DECc)
			 
			if (dist<radius):
			        name, z = tables.row['CATID'], 
                                tables.row['z_spec']
        			
				if verbose:
            				print("  Found SAMI galaxy %s at redshift z=%g at a distance of %f degrees" % (name, z, dist))
        			
				counter += 1
        			idlist.append(name)
        			
				if outFile: 
            				f.write(str(name)+'\n')
    
    print("\n  Found "+str(counter)+" galaxies satisfying the cone search: RA=%f, DEC=%f, radius=%f \n" % (RAc, DECc, radius))

    if outFile: 
        f.close()
    

    # Close h5 file 
    hdf.close()


# ~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~
def sph_dist(ra1, dec1,ra2, dec2):
# ~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~-~
    """ Spherical distance calculator. 
        
    Written by Luca Cortese. Computes the spherical distance between 2 pairs of coordinates using the Haversine formula. This should be self-explanatory to astronomers, please contact the maintainer if you would like to see some clarifications inserted here.
    Args:
        ra1  [flt]: Right ascension of first coordinate pair in decimal degrees.
        ra2  [flt]: Right ascension of second coordinate pair in decimal degrees.
        dec2 [flt]: Declination of first coordinate pair in decimal degrees.
        dec1 [flt]: Declination of second coordinate pair in decimal degrees.
    """
    import numpy as np

    ra1_rad = np.radians(ra1)
    dec1_rad = np.radians(dec1)
    ra2_rad = np.radians(ra2)
    dec2_rad = np.radians(dec2)

    d = np.sin((dec1_rad-dec2_rad)/2)**2;
    d += np.sin((ra1_rad-ra2_rad)/2)**2 * np.cos(dec1_rad)*np.cos(dec2_rad)

    return np.degrees(2*np.arcsin(np.sqrt(d)))