Coverage for /opt/anaconda/lib/python3.7/site-packages/llama/files/healpix/utils.py : 49%
Hot-keys on this page
r m x p toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1# (c) Stefan Countryman 2019
3"""
4Utility functions used across healpix_skymap classes.
5"""
7from math import pi
8import logging
9import gzip
10from tempfile import NamedTemporaryFile
11import binascii
13LOGGER = logging.getLogger(__name__)
14GZIP_BUFFSIZE = 10**5
17def sky_area():
18 """Get the area of the entire sky as an ``Astropy.units.Quantity``."""
19 from astropy.units import Quantity
20 return Quantity(4*pi, "rad2")
23def sky_area_deg():
24 """Get the area of the entire sky in square degrees."""
25 return sky_area().to("deg**2").value
28def dangle_rad(ra, dec, mapra, mapdec): # pylint: disable=invalid-name
29 """Get an array of angular distances in radians between the point specified
30 in ``ra``, ``dec`` and the points specified in ``mapra``, ``mapdec``. All
31 angles, including the return value, are specified in radians.
33 sin(DEC1)*sin(DEC2) + cos(DEC1)*cos(DEC2)*cos(RA2-RA1)
35 So we get the angle by taking the arccos of this function.
37 Parameters
38 ----------
39 ra : float
40 Right ascension of the single point [radians]
41 dec : float
42 Declination of the single point [radians]
43 mapra : array
44 Right ascensions whose angular distances to the single point will be
45 calculated; must be same length as ``mapdec``. [radians]
46 mapdec : array
47 Declinations whose angular distances to the single point will be
48 calculated; must be same length as ``mapra``. [radians]
50 Returns
51 -------
52 dangles : array
53 Angular distances between the single point provided and the arrays of
54 points [radians].
55 """
56 import numpy as np
57 delta_ra = mapra - ra
58 dot_prod = (np.sin(mapdec)*np.sin(dec) +
59 np.cos(mapdec)*np.cos(dec)*np.cos(delta_ra))
60 return np.arccos(dot_prod)
63def nest2uniq(indices, nside):
64 """Return the NUNIQ pixel indices for nested ``indices`` with
65 NSIDE=``nside``.
67 Parameters
68 ----------
69 indices : array
70 Indices of HEALPix pixels in NEST ordering.
71 nside : int
72 A valid NSIDE value not greater than any of the NSIDE values of the
73 ``indices``.
75 Returns
76 -------
77 uniq : array
78 An array of HEALPix pixels in NEST ordering corresponding to the input
79 ``indices`` and ``nside``.
81 Raises
82 ------
83 ValueError
84 If ``nside`` is an invalid NSIDE value, i.e. not a power of 2.
86 Examples
87 --------
88 >>> import numpy as np
89 >>> nest2uniq(np.array([284, 286, 287, 10, 8, 2, 279, 278, 285]), 8)
90 array([540, 542, 543, 266, 264, 258, 535, 534, 541])
91 """
92 check_valid_nside(nside)
93 return indices + 4*nside**2
96def check_valid_nside(nside):
97 """Checks whether ``nside`` is a valid HEALPix NSIDE value. Raises a
98 ValueError if it is not.
100 Parameters
101 ----------
102 nside : int or array
103 An integer or array of integers representing HEALPix NSIDE values.
105 Raises
106 ------
107 ValueError
108 If any of the values are not valid HEALPix NSIDE values.
109 """
110 import numpy as np
111 if np.any(2**np.floor(np.log2(nside)) != nside):
112 raise ValueError("Not a valid NSIDE value: {}".format(nside))
115def check_valid_nuniq(indices):
116 """Checks that ``indices`` are valid NUNIQ indices.
118 Raises
119 ------
120 ValueError
121 If ``indices`` are not valid NUNIQ indices, i.e. they are not integers
122 greater than 3.
123 """
124 import numpy as np
125 if not np.all(np.mod(indices, 1) == 0):
126 raise ValueError("NUNIQ indices must be integers.")
127 if not indices.min() > 3:
128 raise ValueError("NUNIQ indices start at 4; found smaller values.")
131def uniq2nside(indices):
132 """Get the NSIDE value of the given NUNIQ-ordered indices.
134 Raises
135 ------
136 ValueError
137 If ``indices`` are not valid NUNIQ indices, i.e. they are not integers
138 greater than 3.
139 """
140 import numpy as np
141 check_valid_nuniq(indices)
142 return 2**np.array(np.floor(np.log2(indices/4.)/2.), dtype=int)
145def uniq2nest_and_nside(indices):
146 """Parameters
147 ----------
148 indices : array
149 A scalar or numpy array of NUNIQ indices
151 Returns
152 -------
153 nest_indices : array
154 The indices in nest ordering at their respective resolutions
155 nside : array
156 The resolution expressed as NSIDE of the indices given in
157 ``nest_indices``
159 Examples
160 --------
161 >>> import numpy as np
162 >>> nuniq = np.array([540, 542, 543, 266, 264, 258, 535, 534, 541])
163 >>> uniq2nest_and_nside(nuniq)
164 (array([284, 286, 287, 10, 8, 2, 279, 278, 285]), array([8, 8, 8, 8, 8, 8, 8, 8, 8]))
166 Confirm that this is the inverse of nest2uniq:
167 >>> all(nest2uniq(*uniq2nest_and_nside(nuniq)) == nuniq)
168 True
169 """
170 # uniq2nside implicitly checks whether the indices are valid NUNIQ indices
171 nside = uniq2nside(indices)
172 return indices - 4*nside**2, nside
175def uniq2uniq_lowres(indices, nside):
176 """Find the indices at resolutions ``nside`` (each of which MUST be lower
177 resolution than that of the corresponding pixel in ``indices``) which
178 contain the higher resolution pixels described by the corresponding
179 ``indices``.
181 Parameters
182 ----------
183 indices : array
184 Indices of HEALPix pixels in NUNIQ ordering.
185 nside : array
186 A valid NSIDE value not greater than any of the NSIDE values of the
187 ``indices``. Each of these corresponds to an index in ``indices``.
189 Returns
190 -------
191 lowres_indices : array
192 Indices of HEALPix pixels in NUNIQ ordering with resolution ``nside``
193 which contain the corresponding pixels in ``indices``.
195 Raises
196 ------
197 ValueError
198 If the NSIDE resolution of any index in ``indices`` is smaller than
199 ``nside``, or if ``nside`` is an invalid NSIDE value, i.e. not a power
200 of 2.
202 Examples
203 --------
204 Take a few NUNIQ pixels at NSIDE=16 and an NUNIQ pixel at NSIDE=1024 and
205 find which NUNIQ pixels they correspond to at NSIDE=8. The first four
206 should all correspond to the same pixel at NSIDE=8.
208 >>> import numpy as np
209 >>> indices = np.array([1024, 1025, 1026, 11295603])
210 >>> uniq2uniq_lowres(indices, 8)
211 array([256, 256, 256, 689])
212 """
213 import numpy as np
214 check_valid_nside(nside)
215 # we will multiply by the following factor and take the floor to find the
216 # corresponding lowres pixels. the values of nside_factor must be strictly
217 # less-than or equal to 1. uniq2nside implicitly checks whether the indices
218 # are valid NUNIQ indices.
219 nside_factor = (uniq2nside(indices) / nside)**-2
220 if np.max(nside_factor) > 1.0:
221 raise ValueError(("nside must be lower than the NSIDE resolution of "
222 "corresponding input indices."))
223 # use the fact that, at NSIDE=N, pixel p corresponds to NSIDE=2N pixels
224 # 4p, 4p+1, 4p+2, 4p+3
225 return np.array(np.floor(indices*nside_factor), dtype=int)
228def uniq_intersection(indices1, indices2):
229 """Downselect the pixel indices given in ``indices1`` to the set that
230 overlaps with pixels in ``indices2`` and return pairs of indices into
231 both of the input index lists showing which pixel in ``indices2`` each
232 downselected pixel from ``indices1`` overlaps with. Use this to get rid of
233 pixels outside of a given region, or alternatively use it as part of a
234 calculation involving two multi-resolution skymaps whose pixel sizes are
235 non-identical.
236 Should have stable O(len(indices1)*len(indices2)) performance, so make
237 sure that your index lists are not both huge.
239 Parameters
240 ----------
241 indices1 : array
242 Indices of HEALPix pixels in NUNIQ ordering.
243 indices2 : array
244 Indices of HEALPix pixels in NUNIQ ordering.
246 Returns
247 -------
248 indices_into_indices1 : array
249 Indices *into* ``indices1`` that overlap with ``indices2``.
250 indices_into_indices2 : array
251 Corresponding indices *into* ``indices2`` that overlap with
252 ``indices1``.
254 Examples
255 --------
256 Some pixels with NSIDE of 16, 32, and 64, respectively:
257 >>> import numpy as np
258 >>> indices1 = np.array([1024, 4100, 44096])
260 Pixels at NSIDE = 32 that overlap with only the first and last pixels of
261 ``indices1``:
262 >>> indices2 = np.array([4096, 4097, 11024])
264 We should see correspondence between index 0 of ``indices1`` and indices 0,
265 1 of ``indices1``; and correspondence between index 2 of ``indices1`` and
266 index 2 of ``indices2``:
267 >>> uniq_intersection(indices1, indices2)
268 (array([0, 0, 2]), array([0, 1, 2]))
269 """
270 # create a grid of NSIDE differences where the rows correspond to
271 # ``indices1`` and columns correspond to ``indices2``. ``nside_diff[i][j]``
272 # compares the ith index in ``indices1`` to the jth index in ``indices2``
273 # and calculates the factor that the higher-resolution pixel will need to
274 # be multiplied by to make the pixel comparison. It is less than 1 if the
275 # ith index from ``indices1`` has higher NSIDE (resolution) than the jth
276 # from ``indices2``; 1 if they are the same; and greater than 1 if that
277 # index from ``indices2`` has higher NSIDE. the case of equal resolution
278 # is symmetric, so don't sweat it. Here we are using the fact that pixels
279 # 4p, 4p+1, 4p+2, and 4p+3 in a double resolution NUNIQ skymap correspond
280 # to pixel p to find overlap. uniq2nside implicitly checks whether the
281 # indices are valid NUNIQ indices.
282 import numpy as np
283 nside_factor = (uniq2nside(indices1).reshape((-1, 1)) /
284 uniq2nside(indices2))**-2
285 ind1_highres = np.nonzero(nside_factor < 1)
286 ind2_highres = np.nonzero(nside_factor >= 1)
287 # reuse nside_factor as a temp array for the differences in pixels, which
288 # will be in [0, 1) if the pixels overlap
289 nside_factor[ind1_highres] = (
290 indices1[ind1_highres[0]] * nside_factor[ind1_highres] -
291 indices2[ind1_highres[1]]
292 )
293 nside_factor[ind2_highres] = (
294 indices2[ind2_highres[1]] / nside_factor[ind2_highres] -
295 indices1[ind2_highres[0]]
296 )
297 return np.nonzero(np.floor(nside_factor) == 0)
300def uniq_rasterize(indices, values, density=True, nside=None):
301 """Increase resolution of an NUNIQ HEALPix skymap and return NUNIQ indices
302 and values at the same higher resolution (same NSIDE). Doesn't do any
303 interpolation; just used to make calculations easier.
305 Parameters
306 ----------
307 indices : array
308 Non-overlapping HEALPix NUNIQ indices into the skymap.
309 values : array
310 Pixel values corresponding to the ``indices``.
311 density : bool, optional
312 Whether the pixel values are some sort of density. If True (the
313 default), the pixel values will not be modified (since they don't
314 depend on pixel area. Otherwise, the pixel value will be split evenly
315 between sub-pixels.
316 nside : int, optional
317 The NSIDE value to convert the skymap to. If not provided, the skymap
318 will be converted to the largest NSIDE of the provided pixels.
320 Returns
321 -------
322 raster_indices : array
323 A unique flat array of HEALPix NUNIQ indices into the rasterized
324 skymap, sorted in ascending order.
325 raster_values : array
326 A flat array of pixel values of the rasterized skymap corresponding to
327 ``raster_indices``.
329 Raises
330 ------
331 ValueError
332 If the given ``nside`` value is either an invalid NSIDE value or if it
333 is lower than the highest NSIDE value calculated from ``indices``, or
334 if the given ``indices`` are overlapping.
336 Examples
337 --------
338 Some example inputs with NSIDE of 16, 32, and 64, respectively:
339 >>> import numpy as np
340 >>> indices = np.array([1024, 4100, 44096])
341 >>> values = np.ones_like(indices)
343 Rasterize to the max resolution of NSIDE=64, assuming values represent a
344 density:
345 >>> uniq_rasterize(indices, values)
346 (array([16384, 16385, 16386, 16387, 16388, 16389, 16390, 16391, 16392,
347 16393, 16394, 16395, 16396, 16397, 16398, 16399, 16400, 16401,
348 16402, 16403, 44096]), array([1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]))
350 Rasterize just the first two indices, this time explicitly to NSIDE=64 (vs.
351 the default max NSIDE, 32 in this case), explicitly treating the values NOT
352 as a density (and therefore splitting them amongst subpixels):
353 >>> uniq_rasterize(indices[0:2], values[0:2], density=False, nside=64)
354 (array([16384, 16385, 16386, 16387, 16388, 16389, 16390, 16391, 16392,
355 16393, 16394, 16395, 16396, 16397, 16398, 16399, 16400, 16401,
356 16402, 16403]), array([0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625,
357 0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625, 0.0625,
358 0.25 , 0.25 , 0.25 , 0.25 ]))
359 """
360 import numpy as np
361 indices = indices.flatten()
362 values = values.flatten()
363 nside_in = uniq2nside(indices)
364 if nside is None:
365 nside = np.max(nside_in)
366 else:
367 check_valid_nside(nside)
368 if np.max(nside_in) > nside:
369 raise ValueError("Invalid nside; must be greater than or equal to"
370 " max nside of input skymap.")
371 nside_factor = np.array((nside / nside_in)**2, dtype=int)
372 scaled_values = values if density else values/nside_factor
373 raster_values = np.repeat(scaled_values, nside_factor)
374 raster_indices = np.repeat(nside_factor*indices, nside_factor)
375 # This section is O(L) where L is the max number of LIGO MOC pixels
376 # overlapping with the search area, bounded by ~20k pixels in current LIGO
377 # MOC skymaps (L will be the number of for loops). In practice, it will
378 # perform best when subdividing large pixels, which corresponds with the
379 # typical non-detection case (which will make this method not only faster
380 # asymptotically, but also much faster for the typical background search).
381 # pre-allocate incrementor
382 inc = np.arange(np.max(nside_factor))
383 first_indices = np.roll(nside_factor.cumsum(), shift=1)
384 first_indices[0] = 0
385 for i, ind1 in enumerate(first_indices):
386 raster_indices[ind1:ind1+nside_factor[i]] += inc[:nside_factor[i]]
387 sort = raster_indices.argsort()
388 raster_indices = raster_indices[sort]
389 if not np.all(np.unique(raster_indices) == raster_indices):
390 raise ValueError("Must provide non-overlapping pixels in ``indices``.")
391 return raster_indices, raster_values[sort]
394def upres_nest(mask_nested, mask_nside, nside, mask_sorted=False):
395 """
396 Get ``nested`` indices for a skymap mask at a higher resolution. Output
397 will be sorted.
399 Parameters
400 ----------
401 mask_nested : array-like
402 The set of ``nested`` HEALPix pixels that you'd like to have at higher
403 resolution.
404 mask_nside : int
405 The HEALPix ``nside`` parameter of this mask. Must be a power of 2.
406 nside : int
407 The ``nside`` value you would like for your output. Must be a power of
408 2.
409 mask_sorted : bool, optional
410 Whether the input ``mask_nested`` is already sorted. If it is sorted
411 and is large, specifying this will skip a sort step, potentially saving
412 time.
414 Returns
415 -------
416 highres_mask_nested : np.array
417 The same sky region specified in ``mask_nested`` at ``mask_nside`` but
418 at the resolution specified in ``nside``. This will be a factor of
419 ``(nside/mask_nside)**2`` larger. Will be a new array independent of
420 inputs, so feel free to mutate it.
422 Raises
423 ------
424 ValueError
425 if ``mask_nside`` is larger than ``nside``
426 """
427 import numpy as np
428 check_valid_nside(mask_nside)
429 check_valid_nside(nside)
430 if not mask_sorted:
431 mask = np.sort(mask_nested)
432 else:
433 mask = np.array(mask_nested, copy=True)
434 if mask_nside == nside:
435 return mask
436 elif mask_nside > nside:
437 raise ValueError(f"mask_nside ({mask_nside}) must be smaller than "
438 f"nside ({nside})")
439 factor = nside//mask_nside
440 factor *= factor
441 assert isinstance(factor, int)
442 return (np.array([mask*factor]).T +
443 np.array([np.arange(factor)])).flatten()
446def downres_mask_nest(mask_nested, mask_nside, nside, mask_sorted=False):
447 """
448 Get a lower-resolution set of pixel indices that fully cover the region
449 specified in ``mask_nested``. The returned indices correspond to larger
450 pixels which might contain part of the sky that are not included in the
451 original mask (e.g. if the larger pixel only contains one of the original
452 pixels in ``mask_nested``).
454 Parameters
455 ----------
456 mask_nested : array-like
457 The set of ``nested`` HEALPix pixels that you'd like to have at lower
458 resolution.
459 mask_nside : int
460 The HEALPix ``nside`` parameter of this mask. Must be a power of 2.
461 nside : int
462 The ``nside`` value you would like for your output. Must be a power of
463 2.
464 mask_sorted : bool, optional
465 Whether the input ``mask_nested`` is already sorted. If it is sorted
466 and is large, specifying this will skip a sort step, potentially saving
467 time.
469 Returns
470 -------
471 lowres_mask_nested : np.array
472 The smallest sky region that can be specified at the resolution given
473 by ``nside`` that contains the sky region specified in ``mask_nested``
474 at ``mask_nside``. This will be roughly a factor of
475 ``(nside/mask_nside)**2`` smaller, but do **not** count on this ratio,
476 since the pixels in the returned mask might not have all of their
477 corresponding subpixels in the original ``mask_nested``; in
478 pathological cases, the returned array can be up to the same size (if
479 every pixel in ``mask_nested`` only appears in a single larger pixel in
480 at resolution ``nside``). Will be a new array independent of inputs, so
481 feel free to mutate it.
483 Raises
484 ------
485 ValueError
486 if ``mask_nside`` is smaller than ``nside``
487 """
488 import numpy as np
489 check_valid_nside(mask_nside)
490 check_valid_nside(nside)
491 if not mask_sorted:
492 mask = np.sort(mask_nested)
493 else:
494 mask = np.array(mask_nested, copy=True)
495 if mask_nside == nside:
496 return mask
497 elif mask_nside < nside:
498 raise ValueError(f"mask_nside ({mask_nside}) must be larger than "
499 f"nside ({nside})")
500 factor = mask_nside//nside
501 factor *= factor
502 assert isinstance(factor, int)
503 return np.unique(mask // factor)
506def read_partial_skymap(infile, mask_nested, mask_nside,
507 mask_sorted=False, memmap=True):
508 """
509 Read in pixels from a FITS skymap (or a gzip-compressed FITS skymap) that
510 lie in a specific sky region. Attempts to minimize memory usage by
511 memory-mapping pixels in the input file and only loading those specified in
512 ``mask_nested``.
514 Parameters
515 ----------
516 infile : str
517 A FITS HEALPix skymap (optionally compressed; see note under
518 ``memmap``)
519 mask_nested : array-like
520 HEALPix pixel indices in ``nested`` ordering specifying the region
521 of the skymap that should be loaded
522 mask_nside : int
523 The resolution of the pixel mask
524 mask_nested : bool, optional
525 Whether ``mask_nested`` is already sorted. You can specify
526 ``mask_sorted=True``, potentially saving some sorting time
527 memmap : bool, optional
528 If ``True``, use memory mapping when reading in files. This can VASTLY
529 reduce memory required for high-resolution skymaps. If ``infile`` is
530 gzip-compressed and ``memmap`` is ``True``, then ``infile`` will be
531 decompressed to a temporary file and data will be read from it
532 (necessary to constrain memory usage); for high-resolution skymaps,
533 this can require the availability of several gigabytes of tempfile
534 storage.
536 Returns
537 -------
538 partial_skymap : astropy.table.Table
539 A partial skymap table in ``nested`` ordering. Has two columns:
540 ``INDICES`` and ``PROB``. If the resolution of the original
541 HEALPix skymap file is lower than that of the mask, then any pixels
542 overlapping with those in ``mask_nested`` will be used; this might
543 result in a larger portion of the skymap being used than that
544 specified in ``mask_nested``. The resolution of this skymap will be
545 the resolution of the smallest pixel loaded from the input file (in
546 the case of ``ring`` or ``nested`` ordering, this is just the
547 resolution of the input skymap).
548 """
549 from astropy.table import Table
550 import numpy as np
551 import healpy as hp
552 # from https://stackoverflow.com/a/47080739/3601493
553 with open(infile, 'rb') as test_f:
554 magic_number = binascii.hexlify(test_f.read(2))
555 if memmap and magic_number == b'1f8b':
556 try:
557 with NamedTemporaryFile('wb', suffix='.fits') as tmp:
558 LOGGER.debug(f"zipped skymap {infile} detected, attempting to "
559 f"decompress to tempfile {tmp}")
560 with gzip.open("lvc_skymap.fits.gz", "rb") as infile:
561 buf = infile.read(GZIP_BUFFSIZE)
562 while buf:
563 tmp.write(buf)
564 buf = infile.read(GZIP_BUFFSIZE)
565 return read_partial_skymap(tmp.name, mask_nested, mask_nside,
566 mask_sorted=mask_sorted,
567 memmap=memmap)
568 except OSError as err:
569 if 'gzip' in str(err):
570 LOGGER.error("OSError while attempting to decompress, giving "
571 "up and letting astropy try to read it. error: " +
572 str(err))
573 check_valid_nside(mask_nside)
574 tab = Table.read(infile, memmap=memmap)
575 if tab.meta['ORDERING'] == 'NUNIQ':
576 #TODO: fill out
577 raise NotImplementedError()
578 nside = tab.meta['NSIDE']
579 if nside >= mask_nside:
580 mask = upres_nest(mask_nested, mask_nside, nside,
581 mask_sorted=mask_sorted)
582 else:
583 mask = downres_mask_nest(mask_nested, mask_nside, nside,
584 mask_sorted=mask_sorted)
585 if tab.meta['ORDERING'] == 'NESTED':
586 prob = tab['PROB'][mask].copy()
587 return Table(
588 np.array([mask, prob]).transpose(),
589 names=['INDICES', 'PROB'],
590 meta=tab.meta
591 )
592 if tab.meta['ORDERING'] == 'RING':
593 # we'll use the PROB column as a temporary holder for the RING
594 # pixel indices while we load them in order to preserve sort orders
595 results = Table(
596 np.array([
597 mask,
598 hp.nest2ring(nside, mask),
599 ]).transpose(),
600 names=['INDICES', 'RING'],
601 meta=tab.meta,
602 )
603 results.sort(keys='RING')
604 prob = tab['PROB'][results['RING']].copy()
605 results['RING'] = prob
606 # change the name to PROB now that we've replaced the contents
607 results['RING'].name = 'PROB'
608 results.sort(keys='INDICES')
609 return results
610 raise ValueError(f"Unexpected ORDERING in {infile}: "
611 f"{tab.meta['ORDERING']}")