Coverage for /usr/local/lib/python3.9/site-packages/lccalib/match.py: 9%

163 statements  

« 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 

6 

7deg2rad = deg2rad = np.pi / 180 

8 

9#pylint: disable=dangerous-default-value, consider-using-enumerate, too-many-locals 

10 

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 ) 

17 

18 

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 ) 

26 

27 

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 

42 

43 

44class NearestNeighAssoc: 

45 """Solve the fixed-radius nearest neighbor search on a 2D cartesian 

46 lattice 

47 """ 

48 

49 def __init__(self, first=[], extension=[], radius=1): 

50 self.belongs = {} 

51 self.clusters = [] 

52 self.radius = radius 

53 

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] 

74 

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 ) 

86 

87 i = np.digitize(x, self.x_bins) 

88 j = np.digitize(y, self.y_bins) 

89 index = np.zeros(len(i)) 

90 

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 

120 

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 

152 

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 

160 

161 

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 

176 

177 

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) 

183 

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 

204 

205 

206def radius_to_nside(radius): 

207 """Find the smallest nside for wich a circle of radius radius is 

208 contained within an healpy pixel. 

209 

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 

215 

216 

217class NearestNeighAssocHealpy: 

218 """Solve the fixed-radius nearest neighbor search on a 2D Healpy 

219 lattice. Good for full-sky match. 

220 

221 """ 

222 

223 def __init__(self, first=[], radius=np.pi / 180.0 / 3600.0): 

224 

225 self.belongs = {} 

226 self.clusters = [] 

227 self.radius = radius 

228 self.nside = radius_to_nside(radius) 

229 

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] 

236 

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") 

243 

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 

269 

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 

297 

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