sample positions uniformly over HEALPix pixel (#344)

Author: ntessore

docs/position/examples/uniform_in_pixel.yml

@@ -0,0 +1,8 @@
+# healpix map information
+nside: 2
+
+# sample 500 galaxies in healpix pixel 5
+positions: !skypy.position.uniform_in_pixel
+  nside: $nside
+  ipix: 5
+  size: 500

docs/position/index.rst

@@ -33,6 +33,32 @@ use the `~skypy.position.uniform_around` function:
     plt.grid()
 
 
+.. _skypy.position.uniform_in_pixel:
+
+A more complicated uniform distribution is that over a given healpix pixel,
+which is sampled by the `~skypy.position.uniform_in_pixel` function:
+
+.. literalinclude:: examples/uniform_in_pixel.yml
+   :language: yaml
+
+.. plot::
+   :include-source: false
+
+    from skypy.pipeline import Pipeline
+
+    pipeline = Pipeline.read('examples/uniform_in_pixel.yml')
+    pipeline.execute()
+
+    coords = pipeline['positions']
+    ra, dec = coords.ra.wrap_at('180d').radian, coords.dec.radian
+
+    import matplotlib.pyplot as plt
+    plt.figure(figsize=(8,4.2))
+    plt.subplot(111, projection="aitoff")
+    plt.plot(ra, dec, '.', alpha=0.2)
+    plt.grid()
+
+
 Reference/API
 =============
 

setup.cfg

@@ -36,10 +36,12 @@ test =
 all =
     camb
     classy @ git+https://github.com/skypyproject/[email protected]#egg=classy
+    healpy
     skypy-data @ https://github.com/skypyproject/skypy-data/archive/master.tar.gz
     specutils
 docs =
     sphinx-astropy
+    healpy
     matplotlib
 
 [options.package_data]

skypy/position/__init__.py

@@ -2,8 +2,9 @@
 
 __all__ = []
 
-from ._uniform import uniform_around
+from ._uniform import uniform_around, uniform_in_pixel
 
 __all__ += [
     'uniform_around',
+    'uniform_in_pixel',
 ]

skypy/position/_uniform.py

@@ -2,9 +2,9 @@
 
 import numpy as np
 from astropy import units
+from astropy.coordinates import SkyCoord
 
 TWO_PI = 2*np.pi
-PI_HALF = np.pi/2
 
 
 @units.quantity_input(area=units.sr)
@@ -45,3 +45,75 @@ def uniform_around(centre, area, size):
 
     # construct random sky coordinates around centre
     return centre.directional_offset_by(phi, theta)
+
+
+def uniform_in_pixel(nside, ipix, size, nest=False):
+    '''Uniform distribution of points over healpix pixel.
+
+    Draws randomly distributed points from the healpix pixel `ipix` for a map
+    with a given `nside` parameter.
+
+    Parameters
+    ----------
+    nside : int
+        Healpix map `nside` parameter.
+    ipix : int
+        Healpix map pixel index.
+    size : int
+        Number of points to draw.
+    nest : bool, optional
+        If True assume ``NESTED`` pixel ordering, otherwise ``RING`` pixel
+        ordering. Default is ``RING`` pixel ordering.
+
+    Returns
+    -------
+    coords : `~astropy.coordinates.SkyCoord`
+        Randomly distributed points over the healpix pixel.
+
+    Warnings
+    --------
+    This function requires the ``healpy`` package.
+
+    Examples
+    --------
+    See :ref:`User Documentation <skypy.position.uniform_in_pixel>`.
+
+    '''
+
+    from healpy import pix2ang, max_pixrad, nside2pixarea, ang2pix
+
+    # get the centre of the healpix pixel as a SkyCoord
+    centre_lon, centre_lat = pix2ang(nside, ipix, nest=nest, lonlat=True)
+    centre = SkyCoord(centre_lon, centre_lat, unit=units.deg)
+
+    # get the maximum radius of a healpix pixel in radian
+    r = max_pixrad(nside)
+
+    # use that radius as the aperture of a spherical area in steradian
+    area = TWO_PI*(1 - np.cos(r))*units.sr
+
+    # oversampling factor = 1/(probability of the draw)
+    over = area.value/nside2pixarea(nside)
+
+    # the array of longitudes and latitudes of the sample
+    lon, lat = np.empty(0), np.empty(0)
+
+    # rejection sampling over irregularly shaped healpix pixels
+    miss = size
+    while miss > 0:
+        # get the coordinates in a circular aperture around centre
+        sample = uniform_around(centre, area, int(np.ceil(miss*over)))
+
+        # get longitude and latitude of the sample
+        sample_lon, sample_lat = sample.ra.deg, sample.dec.deg
+
+        # accept those positions that are inside the correct pixel
+        accept = ipix == ang2pix(nside, sample_lon, sample_lat, nest=nest, lonlat=True)
+
+        # store the new positions
+        lon = np.append(lon, np.extract(accept, sample_lon))
+        lat = np.append(lat, np.extract(accept, sample_lat))
+        miss = size - len(lon)
+
+    # construct the coordinates
+    return SkyCoord(lon[:size], lat[:size], unit=units.deg)

skypy/position/tests/test_healpy.py

@@ -0,0 +1,48 @@
+import numpy as np
+import pytest
+
+healpy = pytest.importorskip('healpy')
+
+
+def test_uniform_in_pixel_correctness():
+    from skypy.position import uniform_in_pixel
+
+    # 768 healpix pixels
+    nside = 8
+    npix = healpy.nside2npix(nside)
+
+    ipix = np.random.randint(npix)
+
+    size = 1000
+    pos = uniform_in_pixel(nside, ipix, size=size)
+    assert len(pos) == size
+
+    lon, lat = pos.ra.deg, pos.dec.deg
+    apix = healpy.ang2pix(nside, lon, lat, lonlat=True)
+    np.testing.assert_array_equal(apix, ipix)
+
+
+@pytest.mark.flaky
+def test_uniform_in_pixel_distribution():
+    from scipy.stats import kstest
+
+    from skypy.position import uniform_in_pixel
+
+    # 48 healpix pixels
+    nside = 2
+    npix = healpy.nside2npix(nside)
+
+    # sample entire sky with 20 positions per pixel
+    size = 20
+    theta, phi = np.empty(npix*size), np.empty(npix*size)
+    for ipix in range(npix):
+        pos = uniform_in_pixel(nside, ipix, size=size)
+        assert len(pos) == size
+        theta[ipix*size:(ipix+1)*size] = np.pi/2 - pos.dec.rad
+        phi[ipix*size:(ipix+1)*size] = pos.ra.rad
+
+    # test for uniform distribution
+    D, p = kstest(theta, lambda t: (1 - np.cos(t))/2)
+    assert p > 0.01, 'D = {}, p = {}'.format(D, p)
+    D, p = kstest(phi, 'uniform', args=(0, 2*np.pi))
+    assert p > 0.01, 'D = {}, p = {}'.format(D, p)