1# (c) Stefan Countryman 2019 

2 

3""" 

4Functions for working with point sources, applying point spread functions 

5(PSFs) thereto, and making those PSFs compatible with other HEALPix skymaps. 

6""" 

7 

8from abc import ABC, abstractproperty, abstractmethod 

9from collections import namedtuple 

10from math import pi 

11from llama.utils import memoize 

12from llama.files.healpix.utils import ( 

13 nest2uniq, 

14 dangle_rad, 

15 uniq2nside, 

16 uniq_intersection, 

17 uniq_rasterize, 

18 sky_area_deg, 

19) 

20 

21 

22def resolutions(): 

23 """Get an array mapping from HEALPix NSIDE values to corresponding 

24 resolutions in square degrees.""" 

25 import healpy as hp 

26 return tuple(hp.nside2resol(2**i)*180/pi for i in range(30)) 

27 

28 

29class Psf(ABC): 

30 """A point spread function that can be applied to a HEALPix skymap.""" 

31 

32 @abstractproperty 

33 def min_scale(self): 

34 """The rough scale [deg] of the smallest feature in this PSF. A 

35 discrete grid on which the PSF is being calculated should have pixel 

36 distances no greater than this.""" 

37 

38 @abstractproperty 

39 def max_scale(self): 

40 """The angular size of a circle in which nearly all of the probability 

41 will be contained [deg].""" 

42 

43 @abstractmethod 

44 def calculate(self): 

45 """Calculate a Point Spread Function (PSF) over a bunch of NUNIQ 

46 HEALPix pixels. 

47 

48 Returns 

49 ------- 

50 densities : array 

51 An array of densities for this skymap. See ``Psf.indices`` for the 

52 corresponding pixel indices. 

53 """ 

54 

55 @abstractproperty 

56 def ra(self): # pylint: disable=invalid-name 

57 """The right ascension of the point having a spread function applied to 

58 it [degrees].""" 

59 

60 @abstractproperty 

61 def dec(self): 

62 """The declination of the point having a spread function applied to 

63 it [degrees].""" 

64 

65 @property 

66 @memoize(method=True) 

67 def pix_area(self): 

68 """The area per-pixel in [deg**2] of the returned indices (note that 

69 they are at the same resolution, so this is a scalar).""" 

70 import healpy as hp 

71 return hp.nside2pixarea(self.min_nside) * (180*180/pi/pi) 

72 

73 @property 

74 @memoize(method=True) 

75 def min_nside(self): 

76 """Get the minimum HEALPix NSIDE required to provide adequate 

77 resolution for this PSF. Just find the NSIDE that will give pixel 

78 spaces smaller than ``self.min_scale``.""" 

79 import numpy as np 

80 return 2**(np.array(resolutions()) < self.min_scale).nonzero()[0][0] 

81 

82 @property 

83 @memoize(method=True) 

84 def max_nside(self): 

85 """Get the maximum HEALPix NSIDE scale, i.e. the scale at which 

86 ``1.2*self.max_scale`` (the maximum relevant radius of the PSF, times a 

87 conservative factor to account for the fact that pixels are not 

88 perfectly square) will be smaller than the average width of a pixel. At 

89 this scale, the pixel containing the point-source and its 8 neighbors 

90 will completely contain all of the relevant PSF pixels at 

91 ``self.min_nside`` resolution returned by ``self.indices``.""" 

92 import numpy as np 

93 return 2**(np.array(resolutions()) > 1.2*self.max_scale).nonzero()[0][-1] 

94 

95 @property 

96 @memoize(method=True) 

97 def max_scale_indices(self): 

98 """Get the 9 large pixel indices in NUNIQ ordering at the NSIDE scale 

99 defined by ``self.max_nside`` that will completely contain the PSF.""" 

100 import numpy as np 

101 import healpy as hp 

102 center = hp.ang2pix(self.max_nside, self.ra, self.dec, nest=True, 

103 lonlat=True) 

104 return nest2uniq( 

105 np.append( 

106 hp.get_all_neighbours(self.max_nside, center, nest=True), 

107 center 

108 ), 

109 self.max_nside 

110 ) 

111 

112 @property 

113 def indices(self): 

114 """Return the indices of the *important* pixels in this PSF in NUNIQ 

115 ordering; all other pixels will be assumed to be zero. Indices are 

116 sorted in ascending order and are guaranteed to be unique.""" 

117 return self._indices_and_dangles()[0] 

118 

119 @property 

120 def dangle(self): 

121 """Return the angular distances of the *important* pixels in this PSF 

122 from the center of the PSF in units of [degrees]; the indices of these 

123 pixels are given by the ``indices`` property in NUNIQ ordering. Omitted 

124 pixels are too far from the center of the PSF to be important.""" 

125 import numpy as np 

126 return np.degrees(self._indices_and_dangles()[1]) 

127 

128 @memoize(method=True) 

129 def _indices_and_dangles(self): 

130 """Return the NUNIQ indices and angular distances from the center of 

131 the PSF for the pixels in this function that contribute significantly 

132 to the PSF. Used to calculate both the ``indices`` and ``dangle`` 

133 properties (since they are calculated together anyway). Both indices 

134 and delta-angles are sorted before being returned. 

135 """ 

136 import numpy as np 

137 import healpy as hp 

138 # get central pixel and neighbors in NEST ordering; we will later use 

139 # this to find the relevant NUNIQ ordering. See: 

140 # https://healpix.sourceforge.io/pdf/intro.pdf 

141 # and: 

142 # dcc.ligo.org/public/0149/G1800186/003/healpix-multiresolution.pdf 

143 # pixels start in boundary (of the searched region) and are then sorted 

144 # into the interior or exterior of the max_scale region. the boundary 

145 # is then reevaluated as the neighbors of the previous iterations 

146 # boundary that have not already been evaluated (i.e. not in the 

147 # interior nor the exterior). 

148 ra_rad = self.ra * pi / 180. 

149 dec_rad = self.dec * pi / 180. 

150 max_scale_rad = self.max_scale * pi / 180. 

151 # store as numpy arrays 

152 boundary = np.array([hp.ang2pix(self.min_nside, self.ra, self.dec, 

153 nest=True, lonlat=True)]) 

154 index_in_interior = [True] # dummy value to start the while loop 

155 interior = np.array([], dtype=int) 

156 exterior = np.array([], dtype=int) 

157 dangles = np.array([], dtype=int) 

158 while np.any(index_in_interior): 

159 # angles for delta angle calculation; already in radians 

160 colat_rad, lon_rad = hp.pix2ang(self.min_nside, boundary, 

161 nest=True, lonlat=False) 

162 # dangle_rad takes ra, dec in units of radians; returns radians 

163 new_dangles = dangle_rad(ra_rad, dec_rad, lon_rad, pi/2 - 

164 colat_rad) 

165 index_in_interior = new_dangles <= max_scale_rad 

166 # add pixels within self.max_scale to interior 

167 interior = np.concatenate((interior, boundary[index_in_interior])) 

168 dangles = np.concatenate((dangles, new_dangles[index_in_interior])) 

169 # add pixels outside of self.max_scale to exterior 

170 exterior = np.concatenate(( 

171 exterior, 

172 boundary[np.logical_not(index_in_interior)] 

173 )) 

174 # find the new boundary (i.e. unsorted pixels) 

175 boundary = np.setdiff1d( 

176 np.unique( 

177 hp.get_all_neighbours( 

178 self.min_nside, 

179 boundary, 

180 nest=True 

181 ).flatten() 

182 ), 

183 np.concatenate((interior, exterior)), 

184 assume_unique=True 

185 ) 

186 # set pixels to NUNIQ ordering, u = p+4*NSIDE**2 

187 indices = nest2uniq(interior, self.min_nside) 

188 sort = indices.argsort() 

189 indices = indices[sort] 

190 assert np.all(np.unique(indices) == indices) # uniquness 

191 return indices, dangles[sort] 

192 

193 def skymap_product(self, indices, values, density=True): 

194 """Take the pixelwise product of ``self.calculate`` with another 

195 HEALPix NUNIQ-ordered skymap. Will cleverly downselect and resize the 

196 other skymap's pixels to efficiently use memory and CPU resources. 

197 

198 Parameters 

199 ---------- 

200 indices : array 

201 Indices of other skymap's HEALPix pixels in NUNIQ ordering. 

202 values : array 

203 Corresponding values of the other skymap's pixels. 

204 density : bool, optional 

205 Whether the other skymap's pixel values are some sort of density. 

206 If True (the default), the pixel values will not be modified during 

207 rescaling (since they don't depend on pixel area). Otherwise, 

208 *THE PIXEL VALUES WILL BE CONVERTED TO DENSITIES*. 

209 

210 Returns 

211 ------- 

212 product_indices : array 

213 The HEALPix NUNIQ indices of the pixels of the product of the two 

214 skymaps. 

215 product_values : array 

216 The pixel values of the product of the two skymaps corresponding to 

217 the ``product_indices``. 

218 

219 Examples 

220 -------- 

221 The product between a PSF and a lower-resolution density value of 

222 1 should be the same as the original PSF: 

223 

224 >>> import numpy as np 

225 >>> identity_inds = np.arange(4, 16) # NSIDE = 1 

226 >>> identity_values = np.ones_like(identity_inds) 

227 >>> psf = PsfGaussian(30, 10, 1) # use a Gaussian PSF as an example 

228 >>> inds, values = psf.skymap_product(identity_inds, identity_values, 

229 ... density=True) 

230 >>> np.all(psf.indices == inds) 

231 True 

232 >>> np.all(psf.calculate() == values) 

233 True 

234 """ 

235 # downselect the input skymap to the pixels that overlap with the PSF 

236 # search region; we only care about which skymap pixels are in the PSF, 

237 # so get the unique list and ignore the fact that a single skymap pixel 

238 # may correspond to multiple PSF pixels. 

239 import numpy as np 

240 search_inds = np.unique(uniq_intersection(indices, 

241 self.max_scale_indices)[0]) 

242 indices = indices[search_inds] 

243 values = values[search_inds] 

244 # get the pixel values and indices for the PSF 

245 psf_indices = self.indices 

246 psf_values = self.calculate() # ALREADY SORTED BY INDEX 

247 # if the PSF skymap is at lower resolution than the other skymap, 

248 # rasterize it to that higher resolution 

249 nsides = uniq2nside(indices) 

250 psf_nside = self.min_nside 

251 if nsides.max() > psf_nside: 

252 psf_nside = nsides.max() 

253 # uniq_rasterize also SORTS ON INDEX 

254 psf_indices, psf_values = uniq_rasterize(psf_indices, psf_values, 

255 density=True, 

256 nside=psf_nside) 

257 # rasterize the other skymap's pixels so that they are at the same 

258 # resolution as the PSF skymap; ALREADY SORTED ON INDEX AND CHECKED FOR 

259 # HEALPIX OVERLAP (AND HENCE UNIQUENESS) by uniq_rasterize 

260 indices, values = uniq_rasterize(indices, values, density=density, 

261 nside=psf_nside) 

262 # if the other skymap is not a density, convert it to one now 

263 if not density: 

264 import healpy as hp 

265 values /= sky_area_deg() / hp.nside2npix(psf_nside) 

266 # find the indices in ``indices`` that correspond to indices in 

267 # ``psf_indices`` and slice those values from ``values``, then take the 

268 # product between ``psf_values`` and that ordered slice of ``values``. 

269 intersection = np.isin(indices, psf_indices, assume_unique=True) 

270 # The PSF pixels should be a strict subset of the skymap pixels 

271 assert np.all(indices[intersection] == psf_indices) 

272 psf_values *= values[intersection] # in-place 

273 return psf_indices, psf_values 

274 

275 

276class PsfGaussian(namedtuple("GaussTuple", ("ra", "dec", "sigma")), Psf): 

277 """ 

278 A Point Spread Function (PSF) for a Gaussian distribution. 

279 

280 Parametars 

281 ---------- 

282 ra : float 

283 Right Ascension of the center of the distribution [degrees] 

284 dec : float 

285 Declination of the center of the distribution [degrees] 

286 sigma : float 

287 Standard deviation of the distribution [degrees] 

288 """ 

289 

290 @property 

291 def min_scale(self): 

292 return self.sigma / 10 

293 

294 @property 

295 def max_scale(self): 

296 return self.sigma * 5 

297 

298 def calculate(self): 

299 """Return a 2-D Gaussian PDF function to be used as a PSF (point spread 

300 function) on the sphere. Note that this Gaussian PDF *does not* take 

301 account of curvature and is accurate only for small angular spreads. 

302 Angles must be in degrees. 

303 """ 

304 import numpy as np 

305 # TODO Finish, check for correctness 

306 two_sigma2 = 2*self.sigma**2 

307 densities = np.exp(-self.dangle**2/two_sigma2) / (np.pi*two_sigma2) 

308 # if degrees: 

309 # return Quantity(densities, "1/deg^2").to("1/steradian").value 

310 return densities