/
xmatch.py
198 lines (161 loc) · 7.03 KB
/
xmatch.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
# Licensed under a 3-clause BSD style license - see LICENSE.rst
from __future__ import absolute_import, division, print_function, unicode_literals
import logging
from collections import OrderedDict
import numpy as np
from astropy.extern import six
from astropy.coordinates import Angle
from astropy.table import Table, Column
from astropy.table import vstack as table_vstack
from astropy.table import hstack as table_hstack
from astropy.coordinates import SkyCoord
from .utils import skycoord_from_table
__all__ = [
'catalog_xmatch_circle',
'catalog_xmatch_combine',
'table_xmatch_circle_criterion',
'table_xmatch',
]
log = logging.getLogger(__name__)
def catalog_xmatch_circle(catalog, other_catalog,
radius='Association_Radius',
other_radius=Angle(0, 'deg')):
"""Find associations within a circle around each source.
This is convenience function built on `~astropy.coordinates.SkyCoord.search_around_sky`,
extending it in two ways:
1. Each source can have a different association radius.
2. Handle source catalogs (`~astropy.table.Table`) instead of `~astropy.coordinates.SkyCoord`.
Sources are associated if the sum of their radii is smaller than their separation on the sky.
Parameters
----------
catalog : `~astropy.table.Table`
Main source catalog
other_catalog : `~astropy.table.Table`
Other source catalog of potential associations
radius, other_radius : `~astropy.coordinates.Angle` or `str`
Main source catalog association radius.
For `str` this must be a column name (in degrees if without units)
Returns
-------
associations : `~astropy.table.Table`
The list of associations.
"""
if isinstance(radius, six.text_type):
radius = Angle(catalog[radius])
if isinstance(other_radius, six.text_type):
other_radius = Angle(other_catalog[other_radius])
skycoord = skycoord_from_table(catalog)
other_skycoord = skycoord_from_table(other_catalog)
association_catalog_name = other_catalog.meta.get('name', 'N/A')
# Compute associations as list of dict and store in `Table` at the end
associations = []
for source_index in range(len(catalog)):
# TODO: check if this is slower or faster than calling `SkyCoord.search_around_sky` here!?
separation = skycoord[source_index].separation(other_skycoord)
max_separation = radius[source_index] + other_radius
other_indices = np.nonzero(separation < max_separation)[0]
for other_index in other_indices:
association = OrderedDict(
Source_Index=source_index,
Source_Name=catalog['Source_Name'][source_index],
Association_Index=other_index,
Association_Name=other_catalog['Source_Name'][other_index],
Association_Catalog=association_catalog_name,
# There's an issue with scalar `Quantity` objects to init the `Table`
# https://github.com/astropy/astropy/issues/3378
# For now I'll just store the values without unit
Separation=separation[other_index].degree,
)
associations.append(association)
# Need to define columns if there's not a single association
if len(associations) == 0:
log.debug('No associations found.')
table = Table()
table.add_column(Column([], name='Source_Index', dtype=int))
table.add_column(Column([], name='Source_Name', dtype=str))
table.add_column(Column([], name='Association_Index', dtype=int))
table.add_column(Column([], name='Association_Name', dtype=str))
table.add_column(Column([], name='Association_Catalog', dtype=str))
table.add_column(Column([], name='Separation', dtype=float))
else:
log.debug('Found {} associations.'.format(len(associations)))
table = Table(associations, names=associations[0].keys())
return table
def table_xmatch_circle_criterion(max_separation):
"""An example cross-match criterion for `table_xmatch` that reproduces `catalog_xmatch_circle`.
TODO: finish implementing this and test it.
Parameters
----------
max_separation : `~astropy.coordinates.Angle`
Maximum separation
Returns
-------
xmatch : function
Cross-match function to be passed to `table_xmatch`.
"""
def xmatch(row1, row2):
skycoord1 = SkyCoord(row1['RAJ2000'], row1['DEJ2000'], unit='deg')
skycoord2 = SkyCoord(row2['RAJ2000'], row2['DEJ2000'], unit='deg')
separation = skycoord1.separation(skycoord2)
if separation < max_separation:
return True
else:
return False
return xmatch
def table_xmatch(table1, table2, xmatch_criterion, return_indices=True):
"""Cross-match rows from two tables with a cross-match criterion callback.
Note: This is a very flexible and simple way to find matching
rows from two tables, but it can be very slow, e.g. if you
create `~astropy.coordinates.SkyCoord` objects or index into them
in the callback cross-match criterion function:
https://github.com/astropy/astropy/issues/3323#issuecomment-71657245
Parameters
----------
table1, table2 : `~astropy.table.Table`
Input tables
xmatch_criterion : callable
Callable that takes two `~astropy.table.Row` objects as input
and returns True / False when they match / don't match.
return_indices : bool
If `True` this function returns a Table with match indices
``idx1`` and ``idx2``, if False it stacks the matches in a table using
`~astropy.table.hstack`.
Returns
-------
matches : `~astropy.table.Table`
Match table (one match per row)
"""
matches = Table(names=['idx1', 'idx2'], dtype=[int, int])
for row1 in table1:
for row2 in table2:
if xmatch_criterion(row1, row2):
matches.add_row([row1.index, row2.index])
if return_indices:
return matches
else:
raise NotImplementedError
# TODO: need to sub-set table1 and table1 using the matches
table = table_hstack([matches, table1, table2])
return table
def catalog_xmatch_combine(associations):
"""Combine (vertical stack) association tables.
Parameters
----------
associations : dict or (str, `~astropy.table.Table`)
Associations
Returns
-------
combined_associations : `~astropy.table.Table`
Combined associations table.
"""
# Add a column to each table with the catalog name
for name, table in associations.items():
log.debug('{:10s} has {:5d} rows'.format(name, len(table)))
if len(table) != 0:
table['Association_Catalog'] = name
table = table_vstack(list(associations.values()))
# Sort table columns the way we like it
names = ['Source_Name', 'Association_Catalog', 'Association_Name', 'Separation']
table = table[names]
log.debug('Combined number of associations: {}'.format(len(table)))
return table