Coverage for src/lccalib/match.py: 0%
163 statements
« prev ^ index » next coverage.py v7.6.9, created at 2024-12-23 12:30 +0000
« prev ^ index » next coverage.py v7.6.9, created at 2024-12-23 12:30 +0000
1"""
2Tools to match object from their sky (ra,dec) coordinates.
3"""
4import numpy as np
5import healpy
7deg2rad = deg2rad = np.pi / 180
9#pylint: disable=dangerous-default-value, consider-using-enumerate, too-many-locals
11def euclidian(x1, y1, x2, y2):
12 """ return euclidian distance.
13 """
14 return np.sqrt(
15 (np.array(x1) - np.array(x2)) ** 2 + (np.array(y1) - np.array(y2)) ** 2
16 )
19def haversine(ra1, dec1, ra2, dec2):
20 """ return haversine distance.
21 """
22 dlambda = np.array(ra1 - ra2)
23 return np.arccos(
24 np.sin(dec1) * np.sin(dec2) + np.cos(dec1) * np.cos(dec2) * np.cos(dlambda)
25 )
28def gnomonic_projection(ra, dec, center=None):
29 """ ra,dec to x,y.
30 """
31 ra_rad, dec_rad = ra * deg2rad, dec * deg2rad
32 if center is None:
33 center = ra_rad.mean(), dec_rad.mean()
34 else:
35 center = center[0] * deg2rad, center[1] * deg2rad
36 s, c = np.sin(dec_rad), np.cos(dec_rad)
37 c2 = np.cos(center[0] - ra_rad)
38 d = np.sin(center[1]) * s + np.cos(center[1]) * c * c2
39 x = (np.cos(center[1]) * np.sin(center[0] - ra_rad)) / d
40 y = (np.sin(center[1]) * c - np.cos(center[1]) * s * c2) / d
41 return x, y
44class NearestNeighAssoc:
45 """Solve the fixed-radius nearest neighbor search on a 2D cartesian
46 lattice
47 """
49 def __init__(self, first=[], extension=[], radius=1):
50 self.belongs = {}
51 self.clusters = []
52 self.radius = radius
54 if extension:
55 xmin, xmax, ymin, ymax = extension
56 self.x_bins = np.arange(xmin - 0.01 * radius, xmax + 0.01 * radius, radius)
57 self.y_bins = np.arange(ymin - 0.01 * radius, ymax + 0.01 * radius, radius)
58 elif first:
59 firstx, firsty = first
60 xmin, xmax, ymin, ymax = (
61 firstx.min(),
62 firstx.max(),
63 firsty.min(),
64 firsty.max(),
65 )
66 self.x_bins = np.arange(xmin - 0.01 * radius, xmax + 0.01 * radius, radius)
67 self.y_bins = np.arange(ymin - 0.01 * radius, ymax + 0.01 * radius, radius)
68 self.clusters = list(zip(firstx, firsty))
69 i = np.digitize(firstx, self.x_bins)
70 j = np.digitize(firsty, self.y_bins)
71 for k in range(len(i)): #pylint: disable=consider-using-enumerate
72 ik, jk = i[k], j[k]
73 self.belongs[(ik, jk)] = self.belongs.get((ik, jk), []) + [k]
75 def append(self, x, y, metric=haversine):
76 """ populate clusters.
77 """
78 if not hasattr(self, "x_bins"):
79 xmin, xmax, ymin, ymax = x.min(), x.max(), y.min(), y.max()
80 self.x_bins = np.arange(
81 xmin - 0.01 * self.radius, xmax + 0.01 * self.radius, self.radius
82 )
83 self.y_bins = np.arange(
84 ymin - 0.01 * self.radius, ymax + 0.01 * self.radius, self.radius
85 )
87 i = np.digitize(x, self.x_bins)
88 j = np.digitize(y, self.y_bins)
89 index = np.zeros(len(i))
91 for k in range(len(i)): #pylint: disable=consider-using-enumerate
92 ik, jk = i[k], j[k]
93 # gather the list of clusters in the neighborhood
94 candidates = sum(
95 [
96 self.belongs.get((i_n, j_n), [])
97 for i_n in (ik - 1, ik, ik + 1)
98 for j_n in (jk - 1, jk, jk + 1)
99 ],
100 [],
101 )
102 if candidates:
103 distance = metric(
104 x[k],
105 y[k],
106 [self.clusters[l][0] for l in candidates],
107 [self.clusters[l][1] for l in candidates],
108 )
109 l = distance.argmin()
110 if candidates and distance[l] < self.radius:
111 m = candidates[l]
112 index[k] = m
113 self.clusters[m][2] += 1
114 else:
115 clu_i = len(self.clusters)
116 index[k] = clu_i
117 self.clusters.append([x[k], y[k], 1])
118 self.belongs[(ik, jk)] = self.belongs.get((ik, jk), []) + [clu_i]
119 return index
121 def match(self, x, y, metric=haversine):
122 """ choose min distance candidate from cluters.
123 """
124 i = np.digitize(x, self.x_bins)
125 j = np.digitize(y, self.y_bins)
126 index = np.zeros(len(i), dtype="int")
127 for k in range(len(i)):
128 ik, jk = i[k], j[k]
129 # gather the list of clusters in the neighborhood
130 candidates = sum(
131 [
132 self.belongs.get((i_n, j_n), [])
133 for i_n in (ik - 1, ik, ik + 1)
134 for j_n in (jk - 1, jk, jk + 1)
135 ],
136 [],
137 )
138 if candidates:
139 distance = metric(
140 x[k],
141 y[k],
142 [self.clusters[l][0] for l in candidates],
143 [self.clusters[l][1] for l in candidates],
144 )
145 l = distance.argmin()
146 if candidates and distance[l] < self.radius:
147 m = candidates[l]
148 index[k] = m
149 else:
150 index[k] = -1
151 return index
153 def get_cat(self):
154 """ convert clusters into catalog.
155 """
156 clusters = np.rec.fromrecords(self.clusters, names=["ra", "dec", "n"])
157 clusters["ra"] /= deg2rad
158 clusters["dec"] /= deg2rad
159 return clusters
162def assoc(catalog, project=True, xy=False, radius=3e-4 * deg2rad):
163 """ Return an index and a catalog of associated objects.
164 """
165 if project:
166 x, y = gnomonic_projection(np.array(catalog["ra"]), np.array(catalog["dec"]))
167 # x, y = gnomonic_projection(catalog['ra'], catalog['dec'])
168 elif xy:
169 x, y = np.array(catalog["x"]), np.array(catalog["y"])
170 else:
171 x, y = catalog["ra"] * deg2rad, catalog["dec"] * deg2rad
172 _assoc = NearestNeighAssoc(radius=radius)
173 index = _assoc.append(x, y, metric=euclidian if (project or xy) else haversine)
174 clusters = _assoc.get_cat()
175 return index.astype(int), clusters
178def match(refcat, cat, project=True, xy=False, arcsecrad=1):
179 """ Return an index that matches 2 catalogs.
180 """
181 def _arr(x):
182 return np.array(x)
184 if project:
185 xref, yref = gnomonic_projection(
186 _arr(refcat["ra"]),
187 _arr(refcat["dec"]),
188 center=[refcat["ra"].mean(), refcat["dec"].mean()],
189 )
190 x, y = gnomonic_projection(
191 _arr(cat["ra"]),
192 _arr(cat["dec"]),
193 center=[refcat["ra"].mean(), refcat["dec"].mean()],
194 )
195 elif xy:
196 xref, yref = _arr(refcat["x"]), _arr(refcat["y"])
197 x, y = _arr(cat["x"]), _arr(cat["y"])
198 else:
199 xref, yref = _arr(refcat["ra"]) * deg2rad, _arr(refcat["dec"]) * deg2rad
200 x, y = _arr(cat["ra"]) * deg2rad, _arr(cat["dec"]) * deg2rad
201 _assoc = NearestNeighAssoc(first=[xref, yref], radius=arcsecrad / 3600.0 * deg2rad)
202 index = _assoc.match(x, y, metric=euclidian if (project or xy) else haversine)
203 return index
206def radius_to_nside(radius):
207 """Find the smallest nside for wich a circle of radius radius is
208 contained within an healpy pixel.
210 radius: in radian
211 """
212 npix = 4 * np.pi / (2 * radius) ** 2
213 order = int(np.ceil(np.log2(np.sqrt(npix / 12))))
214 return 2**order
217class NearestNeighAssocHealpy:
218 """Solve the fixed-radius nearest neighbor search on a 2D Healpy
219 lattice. Good for full-sky match.
221 """
223 def __init__(self, first=[], radius=np.pi / 180.0 / 3600.0):
225 self.belongs = {}
226 self.clusters = []
227 self.radius = radius
228 self.nside = radius_to_nside(radius)
230 if first:
231 firstra, firstdec = first[0] * deg2rad, first[1] * deg2rad
232 self.clusters = list(zip(firstra, firstdec))
233 i = healpy.ang2pix(self.nside, np.pi / 2 - firstdec, firstra)
234 for k in range(len(i)):
235 self.belongs[i[k]] = self.belongs.get(i[k], []) + [k]
237 def append(self, ra, dec, metric=haversine):
238 """ Populate cluters.
239 """
240 ra_rad, dec_rad = np.array(ra * deg2rad), np.array(dec * deg2rad)
241 i = healpy.ang2pix(self.nside, np.pi / 2 - dec_rad, ra_rad)
242 index = np.zeros(len(i), dtype="int")
244 for k in range(len(i)):
245 ik = i[k]
246 # gather the list of clusters in the neighborhood
247 neighpix = list(healpy.get_all_neighbours(self.nside, ik)) + [ik]
248 candidates = sum(
249 [self.belongs.get(i_n, []) for i_n in neighpix if neighpix != -1], []
250 )
251 if candidates:
252 distance = metric(
253 ra_rad[k],
254 dec_rad[k],
255 [self.clusters[l][0] for l in candidates],
256 [self.clusters[l][1] for l in candidates],
257 )
258 l = distance.argmin()
259 if candidates and distance[l] < self.radius:
260 m = candidates[l]
261 index[k] = m
262 self.clusters[m][2] += 1
263 else:
264 clu_i = len(self.clusters)
265 index[k] = clu_i
266 self.clusters.append([ra_rad[k], dec_rad[k], 1])
267 self.belongs[ik] = self.belongs.get(ik, []) + [clu_i]
268 return index
270 def match(self, ra, dec, metric=haversine):
271 """ choose min distance candidate from cluters.
272 """
273 ra_rad, dec_rad = np.array(ra * deg2rad), np.array(dec * deg2rad)
274 i = healpy.ang2pix(self.nside, np.pi / 2 - dec_rad, ra_rad)
275 index = np.zeros(len(i), dtype="int")
276 for k in range(len(i)):
277 ik = i[k]
278 # gather the list of clusters in the neighborhood
279 neighpix = list(healpy.get_all_neighbours(self.nside, ik)) + [ik]
280 candidates = sum(
281 [self.belongs.get(i_n, []) for i_n in neighpix if neighpix != -1], []
282 )
283 if candidates:
284 distance = metric(
285 ra_rad[k],
286 dec_rad[k],
287 [self.clusters[l][0] for l in candidates],
288 [self.clusters[l][1] for l in candidates],
289 )
290 l = distance.argmin()
291 if candidates and distance[l] < self.radius:
292 m = candidates[l]
293 index[k] = m
294 else:
295 index[k] = -1
296 return index
298 def get_cat(self):
299 """ convert clusters into catalog.
300 """
301 clusters = np.rec.fromrecords(self.clusters, names=["ra", "dec", "n"])
302 clusters["ra"] /= deg2rad
303 clusters["dec"] /= deg2rad
304 return clusters