Xmatch: allow astropy tables input iwht more than 2 columns

hamogu · hamogu · commit db597e6f48f5 · 2016-01-14T20:03:09.000-05:00
Before this patch astropy tables as input could only have exactly
two columns that would be exactly in the order RA, DEC. That is a bad
design for two reasons:
- What is the order of the columns is different.
- After cross-matching, how do I know which object is which?
  Two RA/DEC columns are fine if I want to answer the question
  "What are the counterparts in catalog 2?" but not if I want
  to retrieve a joint catalog (e.g. I have my own optical data
  in cat1 and cross-match with 2MASS, then I want all mags in the
  returned catalog.
diff --git a/astroquery/xmatch/core.py b/astroquery/xmatch/core.py
@@ -33,10 +33,7 @@ def query(self, cat1, cat2, max_distance, colRA1=None,
             If the table is uploaded or accessed through a URL, it must be
             in VOTable or CSV format with the positions in J2000
             equatorial frame and as decimal degrees numbers.
-            Note: If the passed argument is an AstroPy table, the column names
-            are extracted from this object and the parameters `colRA1` and
-            `colDec1` are ignored!
-        cat2 : str or file
+         cat2 : str or file
             Identifier of the second table. Follows the same rules as *cat1*.
         max_distance : `~astropy.units.arcsec`
             Maximum distance in arcsec to look for counterparts.
@@ -84,57 +81,49 @@ def query_async(
             'RESPONSEFORMAT': 'csv',
         }
         kwargs = {}
-        if isinstance(cat1, six.string_types):
-            payload['cat1'] = cat1
-        elif isinstance(cat1, Table):
-            payload['colRA1'], payload['colDec1'] = cat1.colnames
+
+        self._prepare_sending_table(1, payload, kwargs, cat1, colRA1, colDec1)
+        self._prepare_sending_table(2, payload, kwargs, cat2, colRA2, colDec2)
+
+        response = commons.send_request(
+            self.URL, payload, self.TIMEOUT, **kwargs)
+        return response
+
+    def _prepare_sending_table(self, i, payload, kwargs, cat, colRA, colDec):
+        '''Check if table is a string, a `astropy.table.Table`, etc. and set
+        query parameters accordingly.
+        '''
+        catstr = 'cat{0}'.format(i)
+        if isinstance(cat, six.string_types):
+            payload[catstr] = cat
+        elif isinstance(cat, Table):
             # write the Table's content into a new, temporary CSV-file
             # so that it can be pointed to via the `files` option
             # file will be closed when garbage-collected
             fp = six.StringIO()
-            cat1.write(fp, format='ascii.csv')
+            cat.write(fp, format='ascii.csv')
             fp.seek(0)
-            kwargs['files'] = {'cat1': fp}
+            kwargs['files'] = {catstr: fp}
         else:
             # assume it's a file-like object, support duck-typing
-            kwargs['files'] = {'cat1': cat1}
-        if not self.is_table_available(cat1) and\
-                payload.get('colRA1') is None or\
-                payload.get('colDec1') is None:
+            kwargs['files'] = {catstr: cat}
+        if not self.is_table_available(cat):
+            if ((colRA is None) or (colDec is None)):
+                raise ValueError('Specify the name of the RA/Dec columns in' +
+                                 ' the input table.')
             # if `cat1` is not a VizieR table,
             # it is assumed it's either a URL or an uploaded table
-            payload['colRA1'] = colRA1
-            payload['colDec1'] = colDec1
-        if isinstance(cat2, six.string_types):
-            payload['cat2'] = cat2
-        elif isinstance(cat2, Table):
-            payload['colRA2'], payload['colDec2'] = cat2.colnames
-            # write the Table's content into a new, temporary CSV-file
-            # so that it can be pointed to via the `files` option
-            # file will be closed when garbage-collected
-            fp = six.StringIO()
-            cat1.write(fp, format='ascii.csv')
-            fp.seek(0)
-            kwargs['files'] = {'cat1': fp}
-        else:
-            # assume it's a file-like object, support duck-typing
-            kwargs['files'] = {'cat2': cat2}
-        if not self.is_table_available(cat2) and\
-                payload.get('colRA2') is None or\
-                payload.get('colDec2') is None:
-            # if `cat2` is not a VizieR table,
-            # it is assumed it's either a URL or an uploaded table
-            payload['colRA2'] = colRA2
-            payload['colDec2'] = colDec2
-        response = commons.send_request(
-            self.URL, payload, self.TIMEOUT, **kwargs)
-        return response
+            payload['colRA{0}'.format(i)] = colRA
+            payload['colDec{0}'.format(i)] = colDec
+
 
     def is_table_available(self, table_id):
         """Return True if the passed CDS table identifier is one of the
         available VizieR tables, otherwise False.
 
         """
+        if isinstance(table_id, six.string_types) and (table_id[:7] == 'vizier:'):
+            table_id = table_id[7:]
         return table_id in self.get_available_tables()
 
     def get_available_tables(self):
diff --git a/astroquery/xmatch/tests/data/posList.csv b/astroquery/xmatch/tests/data/posList.csv
@@ -1,7 +1,7 @@
-ra,dec
-267.22029,-20.35869
-274.83971,-25.42714
-275.92229,-30.36572
-283.26621,-8.70756
-306.01575,33.86756
-322.493,12.16703
+ra,dec,my_id
+267.22029,-20.35869,1
+274.83971,-25.42714,2
+275.92229,-30.36572,3
+283.26621,-8.70756,4
+306.01575,33.86756,5
+322.493,12.16703,6
diff --git a/astroquery/xmatch/tests/data/query_res.csv b/astroquery/xmatch/tests/data/query_res.csv
@@ -1,12 +1,12 @@
-angDist,ra,dec,2MASS,RAJ2000,DEJ2000,errHalfMaj,errHalfMin,errPosAng,Jmag,Hmag,Kmag,e_Jmag,e_Hmag,e_Kmag,Qfl,Rfl,X,MeasureJD
-1.352044,267.22029,-20.35869,17485281-2021323,267.220049,-20.358990,0.150,0.110,16,9.931,8.822,7.550,0.239,0.241,,EEU,226,2,2450950.8609
-1.578188,267.22029,-20.35869,17485288-2021328,267.220348,-20.359125,0.140,0.120,158,8.868,7.784,8.530,,,0.128,UUB,662,2,2450950.8609
-3.699368,267.22029,-20.35869,17485264-2021294,267.219344,-20.358171,0.130,0.120,33,9.562,8.525,9.445,,,0.119,UUB,662,2,2450950.8609
-3.822922,267.22029,-20.35869,17485299-2021279,267.220825,-20.357754,0.080,0.070,7,10.756,9.725,9.287,0.139,0.127,0.103,EBA,222,2,2450950.8609
-4.576677,267.22029,-20.35869,17485255-2021326,267.218994,-20.359064,0.140,0.130,87,10.431,9.348,7.926,0.159,0.316,,CEU,226,2,2450950.8609
-0.219609,274.83971,-25.42714,18192154-2525377,274.839773,-25.427162,0.060,0.060,90,9.368,8.431,7.919,0.024,0.044,0.036,AAA,211,0,2451407.5033
-1.633225,275.92229,-30.36572,18234133-3021582,275.922233,-30.366171,0.080,0.080,55,12.947,12.334,12.145,0.156,0.221,0.127,EEE,222,2,2451021.7212
-0.536998,283.26621,-8.70756,18530390-0842276,283.266284,-8.707690,0.060,0.060,45,12.182,11.534,11.380,0.057,0.071,0.063,AAA,222,0,2451301.7945
-1.178542,306.01575,33.86756,20240382+3352021,306.015944,+33.867275,0.060,0.060,90,13.575,12.684,12.321,0.025,0.027,0.026,AAA,222,0,2450948.9708
-0.853178,322.493,12.16703,21295836+1210007,322.493171,+12.166862,0.100,0.080,179,9.798,9.339,9.176,0.109,0.150,0.100,EEA,222,0,2451080.6935
-4.503950,322.493,12.16703,21295861+1210023,322.494242,+12.167332,0.100,0.080,1,10.057,9.720,9.483,0.077,0.136,0.088,EEE,222,0,2451080.6935
+angDist,ra,dec,my_id,2MASS,RAJ2000,DEJ2000,errHalfMaj,errHalfMin,errPosAng,Jmag,Hmag,Kmag,e_Jmag,e_Hmag,e_Kmag,Qfl,Rfl,X,MeasureJD
+1.352044,267.22029,-20.35869,1,17485281-2021323,267.220049,-20.358990,0.150,0.110,16,9.931,8.822,7.550,0.239,0.241,,EEU,226,2,2450950.8609
+1.578188,267.22029,-20.35869,1,17485288-2021328,267.220348,-20.359125,0.140,0.120,158,8.868,7.784,8.530,,,0.128,UUB,662,2,2450950.8609
+3.699368,267.22029,-20.35869,1,17485264-2021294,267.219344,-20.358171,0.130,0.120,33,9.562,8.525,9.445,,,0.119,UUB,662,2,2450950.8609
+3.822922,267.22029,-20.35869,1,17485299-2021279,267.220825,-20.357754,0.080,0.070,7,10.756,9.725,9.287,0.139,0.127,0.103,EBA,222,2,2450950.8609
+4.576677,267.22029,-20.35869,1,17485255-2021326,267.218994,-20.359064,0.140,0.130,87,10.431,9.348,7.926,0.159,0.316,,CEU,226,2,2450950.8609
+0.219609,274.83971,-25.42714,2,18192154-2525377,274.839773,-25.427162,0.060,0.060,90,9.368,8.431,7.919,0.024,0.044,0.036,AAA,211,0,2451407.5033
+1.633225,275.92229,-30.36572,3,18234133-3021582,275.922233,-30.366171,0.080,0.080,55,12.947,12.334,12.145,0.156,0.221,0.127,EEE,222,2,2451021.7212
+0.536998,283.26621,-8.70756,4,18530390-0842276,283.266284,-8.707690,0.060,0.060,45,12.182,11.534,11.380,0.057,0.071,0.063,AAA,222,0,2451301.7945
+1.178542,306.01575,33.86756,5,20240382+3352021,306.015944,+33.867275,0.060,0.060,90,13.575,12.684,12.321,0.025,0.027,0.026,AAA,222,0,2450948.9708
+0.853178,322.493,12.16703,6,21295836+1210007,322.493171,+12.166862,0.100,0.080,179,9.798,9.339,9.176,0.109,0.150,0.100,EEA,222,0,2451080.6935
+4.503950,322.493,12.16703,6,21295861+1210023,322.494242,+12.167332,0.100,0.080,1,10.057,9.720,9.483,0.077,0.136,0.088,EEE,222,0,2451080.6935
diff --git a/astroquery/xmatch/tests/test_xmatch.py b/astroquery/xmatch/tests/test_xmatch.py
@@ -69,8 +69,8 @@ def test_xmatch_is_avail_table(monkeypatch):
 
     assert xm.is_table_available('II/311/wise')
     assert xm.is_table_available('II/246/out')
-    assert not xm.is_table_available('vizier:II/311/wise')
-
+    assert xm.is_table_available('vizier:II/311/wise')
+    assert not xm.is_table_available('blablabla')
 
 def test_xmatch_query_local(monkeypatch):
     xm = XMatch()
@@ -87,7 +87,7 @@ def test_xmatch_query_local(monkeypatch):
     table = ascii.read(response.text, format='csv')
     assert isinstance(table, Table)
     assert table.colnames == [
-        'angDist', 'ra', 'dec', '2MASS', 'RAJ2000', 'DEJ2000',
+        'angDist', 'ra', 'dec', 'my_id', '2MASS', 'RAJ2000', 'DEJ2000',
         'errHalfMaj', 'errHalfMin', 'errPosAng', 'Jmag', 'Hmag', 'Kmag',
         'e_Jmag', 'e_Hmag', 'e_Kmag', 'Qfl', 'Rfl', 'X', 'MeasureJD']
     assert len(table) == 11
@@ -103,15 +103,15 @@ def test_xmatch_query_cat1_table_local(monkeypatch):
             request_mockreturn(request_type, url, data, **kwargs))
     with open(data_path('posList.csv')) as pos_list:
         input_table = Table.read(pos_list.readlines(),
-                                 names=['ra', 'dec'],
                                  format='ascii.csv',
                                  guess=False)
     response = xm.query_async(
-        cat1=input_table, cat2='vizier:II/246/out', max_distance=5 * arcsec)
+        cat1=input_table, cat2='vizier:II/246/out', max_distance=5 * arcsec,
+        colRA1='ra', colDec1='dec')
     table = ascii.read(response.text, format='csv')
     assert isinstance(table, Table)
     assert table.colnames == [
-        'angDist', 'ra', 'dec', '2MASS', 'RAJ2000', 'DEJ2000',
+        'angDist', 'ra', 'dec', 'my_id', '2MASS', 'RAJ2000', 'DEJ2000',
         'errHalfMaj', 'errHalfMin', 'errPosAng', 'Jmag', 'Hmag', 'Kmag',
         'e_Jmag', 'e_Hmag', 'e_Kmag', 'Qfl', 'Rfl', 'X', 'MeasureJD']
     assert len(table) == 11
diff --git a/astroquery/xmatch/tests/test_xmatch_remote.py b/astroquery/xmatch/tests/test_xmatch_remote.py
@@ -32,7 +32,8 @@ def test_xmatch_avail_tables(xmatch):
 def test_xmatch_is_avail_table(xmatch):
     assert xmatch.is_table_available('II/311/wise')
     assert xmatch.is_table_available('II/246/out')
-    assert not xmatch.is_table_available('vizier:II/311/wise')
+    assert xmatch.is_table_available('vizier:II/311/wise')
+    assert not xmatch.is_table_available('blablabla')
 
 
 @remote_data
@@ -43,7 +44,7 @@ def test_xmatch_query(xmatch):
             colRA1='ra', colDec1='dec')
     assert isinstance(table, Table)
     assert table.colnames == [
-        'angDist', 'ra', 'dec', '2MASS', 'RAJ2000', 'DEJ2000',
+        'angDist', 'ra', 'dec', 'my_id', '2MASS', 'RAJ2000', 'DEJ2000',
         'errHalfMaj', 'errHalfMin', 'errPosAng', 'Jmag', 'Hmag', 'Kmag',
         'e_Jmag', 'e_Hmag', 'e_Kmag', 'Qfl', 'Rfl', 'X', 'MeasureJD']
     assert len(table) == 11
@@ -52,13 +53,13 @@ def test_xmatch_query(xmatch):
 @remote_data
 def test_xmatch_query_astropy_table(xmatch):
     datapath = os.path.join(DATA_DIR, 'posList.csv')
-    input_table = Table.read(datapath, names=['ra', 'dec'],
-                             format='ascii.csv')
+    input_table = Table.read(datapath, format='ascii.csv')
     table = xmatch.query(
-        cat1=input_table, cat2='vizier:II/246/out', max_distance=5 * arcsec)
+        cat1=input_table, cat2='vizier:II/246/out', max_distance=5 * arcsec,
+        colRA1='ra', colDec1='dec')
     assert isinstance(table, Table)
     assert table.colnames == [
-        'angDist', 'ra', 'dec', '2MASS', 'RAJ2000', 'DEJ2000',
+        'angDist', 'ra', 'dec', 'my_id', '2MASS', 'RAJ2000', 'DEJ2000',
         'errHalfMaj', 'errHalfMin', 'errPosAng', 'Jmag', 'Hmag', 'Kmag',
         'e_Jmag', 'e_Hmag', 'e_Kmag', 'Qfl', 'Rfl', 'X', 'MeasureJD']
     assert len(table) == 11
