Added the Boubert & Everall (2019) selection function, with tests

selectionfunctions/boubert_everall_2019.py

@@ -0,0 +1,241 @@
+#!/usr/bin/env python
+#
+# boubert_everall_2019.py
+# Reads the Gaia DR2 selection function from Boubert & Everall (2019, submitted).
+#
+# Copyright (C) 2019  Douglas Boubert & Andrew Everall
+#
+# This program is free software; you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation; either version 2 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License along
+# with this program; if not, write to the Free Software Foundation, Inc.,
+# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
+#
+
+from __future__ import print_function, division
+
+import os
+import h5py
+import numpy as np
+
+import astropy.coordinates as coordinates
+import astropy.units as units
+import h5py
+import healpy as hp
+from scipy import interpolate, special
+
+from .std_paths import *
+from .map_base import SelectionFunction, ensure_flat_icrs, coord2healpix
+from .source_base import ensure_gaia_g
+from . import fetch_utils
+
+from time import time
+
+
+class BoubertEverall2019Query(SelectionFunction):
+    """
+    Queries the Gaia DR2 selection function (Boubert & Everall, 2019).
+    """
+
+    def __init__(self, map_fname=None, version='modelAB', crowding=False, bounds=True):
+        """
+        Args:
+            map_fname (Optional[:obj:`str`]): Filename of the BoubertEverall2019 selection function. Defaults to
+                :obj:`None`, meaning that the default location is used.
+            version (Optional[:obj:`str`]): The selection function version to download. Valid versions
+                are :obj:`'modelT'` and :obj:`'modelAB'`
+                Defaults to :obj:`'modelT'`.
+            crowding (Optional[:obj:`bool`]): Whether or not the selection function includes crowding.
+                Defaults to :obj:`'False'`.
+            bounds (Optional[:obj:`bool`]): Whether or not the selection function is bounded to 1.7 < G < 21.5.
+                Defaults to :obj:`'True'`.
+        """
+
+        if map_fname is None:
+            map_fname = os.path.join(data_dir(), 'boubert_everall_2019', 'boubert_everall_2019.h5')
+
+        t_start = time()
+
+        with h5py.File(map_fname, 'r') as f:
+            # Load auxilliary data
+            print('Loading auxilliary data ...')
+            self._nside = 1024
+            self._g_grid = f['g_grid'][...]
+            self._n_field = f['n_field'][...]
+            self._crowding = crowding
+            self._bounds = bounds
+            if crowding == True:
+                self._log10_rho_grid = f['log10_rho_grid'][...]
+                self._rho_field= f['density_field'][...]
+
+            t_auxilliary = time()
+
+            # Load selection function
+            print('Loading selection function ...')
+            if version == 'modelT':
+                if crowding == True:
+                    self._theta = f['t_theta_percentiles'][:,:,2]
+                else:
+                    self._theta = f['t_theta_percentiles'][0,:,2]
+            elif version == 'modelAB':
+                if crowding == True:
+                    self._alpha = f['ab_alpha_percentiles'][:,:,2]
+                    self._beta = f['ab_beta_percentiles'][:,:,2]
+                else:
+                    self._alpha = f['ab_alpha_percentiles'][0,:,2]
+                    self._beta = f['ab_beta_percentiles'][0,:,2]
+
+            if bounds == True:
+                self._g_min = 1.7
+                self._g_max = 21.5
+            else:
+                self._g_min = -np.inf
+                self._g_max = np.inf
+
+            t_sf = time()
+
+        # Create interpolator
+        print('Creating selection function interpolator...')
+        if version == 'modelT':
+            if crowding == True:
+                self._theta_interpolator = interpolate.RectBivariateSpline(self._log10_rho_grid,self._g_grid,self._theta)
+                self._interpolator = lambda _log10_rho, _g : self._theta_interpolator(_log10_rho, _g, grid = False)
+            else:
+                self._theta_interpolator = interpolate.interp1d(self._g_grid,self._theta,kind='linear',fill_value='extrapolate')
+                self._interpolator = lambda _g : self._theta_interpolator(_g)
+        elif version == 'modelAB':
+            if crowding == True:
+                self._alpha_interpolator = interpolate.RectBivariateSpline(self._log10_rho_grid,self._g_grid,self._alpha)
+                self._beta_interpolator = interpolate.RectBivariateSpline(self._log10_rho_grid,self._g_grid,self._beta)
+                self._interpolator = lambda _log10_rho, _g : (self._alpha_interpolator(_log10_rho, _g, grid = False),self._beta_interpolator(_log10_rho, _g, grid = False))
+            else:
+                self._alpha_interpolator = interpolate.interp1d(self._g_grid,self._alpha,kind='linear',fill_value='extrapolate')
+                self._beta_interpolator = interpolate.interp1d(self._g_grid,self._beta,kind='linear',fill_value='extrapolate')
+                self._interpolator = lambda _g : (self._alpha_interpolator(_g),self._beta_interpolator(_g))
+
+        t_interpolator = time()
+
+        t_finish = time()
+
+        print('t = {:.3f} s'.format(t_finish - t_start))
+        print('  auxilliary: {: >7.3f} s'.format(t_auxilliary-t_start))
+        print('          sf: {: >7.3f} s'.format(t_sf-t_auxilliary))
+        print('interpolator: {: >7.3f} s'.format(t_interpolator-t_sf))
+
+    def _selection_function(self,_n,_parameters):
+        if len(_parameters) == 1:
+
+            # This must be Model T, _parameters = (theta)
+            _t = _parameters
+
+            # 0 < theta < 1, make it so!
+            _t[_t<0.0] = 1e-6
+            _t[_t>1.0] = 1-1e-6
+
+            _result = special.betainc(5,_n-4,_t)
+
+        elif len(_parameters) == 2:
+
+            # This must be Model AB, _parameters = (alpha,beta)
+            _a, _b = _parameters
+
+            # 0.01 < alpha,beta < 1000, make it so!
+            _a[_a<1e-2] = 1e-2
+            _a[_a>1e+2] = 1e+2
+            _b[_b<1e-2] = 1e-2
+            _b[_b>1e+2] = 1e+2
+
+            _result = np.ones(_a.shape)
+            for _m in range(1,5)[::-1]:
+                _result = 1 + _result*((_n-_m+1)/_m)*(_a+_m-1)/(_b+_n-_m)
+            _result = 1- _result*special.beta(_a,_b+_n)/special.beta(_a,_b)
+
+        return _result
+
+
+    @ensure_flat_icrs
+    @ensure_gaia_g
+    def query(self, coords):
+        """
+        Returns the selection function at the requested coordinates.
+
+        Args:
+            coords (:obj:`astropy.coordinates.SkyCoord`): The coordinates to query.
+
+        Returns:
+            Selection function at the specified coordinates, as a fraction.
+
+        """
+
+        # Convert coordinates to healpix indices
+        hpxidx = coord2healpix(coords, 'icrs', self._nside, nest=True)
+
+        # Calculate the number of observations of each source
+        n = self._n_field[hpxidx]
+
+        # Extract Gaia G magnitude
+        G = coords.photometry['gaia_g']
+
+        if self._crowding == True:
+
+            # Calculate the local density field at each source
+            log10_rho = self._log10_rho_field[hpxidx]
+
+            # Calculate parameters
+            sf_parameters = self._interpolator(log10_rho,G)
+
+        else:
+
+            # Calculate parameters
+            sf_parameters = self._interpolator(G)
+
+        # Evaluate selection function
+        selection_function = self._selection_function(n,sf_parameters)
+
+        if self._bounds == True:
+            _outside_bounds = np.where( (G<self._g_min) | (G>self._g_max) )
+            selection_function[_outside_bounds] = 0.0
+
+        return selection_function
+
+
+def fetch():
+    """
+    Downloads the specified version of the Bayestar dust map.
+
+    Args:
+        version (Optional[:obj:`str`]): The map version to download. Valid versions are
+            :obj:`'bayestar2019'` (Green, Schlafly, Finkbeiner et al. 2019),
+            :obj:`'bayestar2017'` (Green, Schlafly, Finkbeiner et al. 2018) and
+            :obj:`'bayestar2015'` (Green, Schlafly, Finkbeiner et al. 2015). Defaults
+            to :obj:`'bayestar2019'`.
+
+    Raises:
+        :obj:`ValueError`: The requested version of the map does not exist.
+
+        :obj:`DownloadError`: Either no matching file was found under the given DOI, or
+            the MD5 sum of the file was not as expected.
+
+        :obj:`requests.exceptions.HTTPError`: The given DOI does not exist, or there
+            was a problem connecting to the Dataverse.
+    """
+
+    doi = '10.7910/DVN/PDFOVC'
+
+    requirements = {'filename': 'boubert_everall_2019.h5'}
+
+    local_fname = os.path.join(data_dir(), 'boubert_everall_2019', 'boubert_everall_2019.h5')
+
+    # Download the data
+    fetch_utils.dataverse_download_doi(
+        doi,
+        local_fname,
+        file_requirements=requirements)

selectionfunctions/map_base.py

@@ -34,6 +34,7 @@
 
 from . import json_serializers
 from . import sfexceptions
+from .source_base import SourceCoord
 
 # import time
 
@@ -75,7 +76,7 @@ def coord2healpix(coords, frame, nside, nest=True):
     elif hasattr(c, 'w'):
         return hp.pixelfunc.vec2pix(nside, c.w.kpc, c.u.kpc, c.v.kpc, nest=nest)
     else:
-        raise dustexceptions.CoordFrameError(
+        raise sfexceptions.CoordFrameError(
             'No method to transform from coordinate frame "{}" to HEALPix.'.format(
                 frame))
 
@@ -158,6 +159,111 @@ def _wrapper_func(self, coords, **kwargs):
     return _wrapper_func
 
 
+def equ_to_shape(equ, shape):
+    ra = np.reshape(equ.ra.deg, shape)*units.deg
+    dec = np.reshape(equ.dec.deg, shape)*units.deg
+
+    has_dist = hasattr(equ.distance, 'kpc')
+    d = np.reshape(equ.distance.kpc, shape)*units.kpc if has_dist else None
+
+    has_photometry = equ.photometry is not None
+    if has_photometry:
+        photometry = {k:np.reshape(v, shape) for k,v in equ.photometry.items()}
+    else:
+        photometry = None
+
+    has_photometry_error = equ.photometry_error is not None
+    if has_photometry_error:
+        photometry_error = {k:np.reshape(v, shape) for k,v in equ.photometry_error.items()}
+    else:
+        photometry_error = None
+
+    return SourceCoord(ra, dec, distance=d, photometry=photometry, photometry_error=photometry_error, frame='icrs')
+
+def ensure_flat_icrs(f):
+    """
+    A decorator for class methods of the form
+
+    .. code-block:: python
+
+        Class.method(self, coords, **kwargs)
+
+    where ``coords`` is an :obj:`astropy.coordinates.SkyCoord` object.
+
+    The decorator ensures that the ``coords`` that gets passed to
+    ``Class.method`` is a flat array of Equatorial coordinates. It also reshapes
+    the output of ``Class.method`` to have the same shape (possibly scalar) as
+    the input ``coords``. If the output of ``Class.method`` is a tuple or list
+    (instead of an array), each element in the output is reshaped instead.
+
+    Args:
+        f (class method): A function with the signature
+            ``(self, coords, **kwargs)``, where ``coords`` is a :obj:`SkyCoord`
+            object containing an array.
+
+    Returns:
+        A function that takes :obj:`SkyCoord` input with any shape (including
+        scalar).
+    """
+
+    @wraps(f)
+    def _wrapper_func(self, coords, **kwargs):
+        # t0 = time.time()
+
+        if coords.frame.name != 'icrs':
+            equ = coords.transform_to('icrs')
+        else:
+            equ = coords
+
+        # t1 = time.time()
+
+        is_array = not coords.isscalar
+        if is_array:
+            orig_shape = coords.shape
+            shape_flat = (np.prod(orig_shape),)
+            # print 'Original shape: {}'.format(orig_shape)
+            # print 'Flattened shape: {}'.format(shape_flat)
+            equ = equ_to_shape(equ, shape_flat)
+        else:
+            equ = equ_to_shape(equ, (1,))
+
+        # t2 = time.time()
+
+        out = f(self, equ, **kwargs)
+
+        # t3 = time.time()
+
+        if is_array:
+            if isinstance(out, list) or isinstance(out, tuple):
+                # Apply to each array in output list
+                for o in out:
+                    o.shape = orig_shape + o.shape[1:]
+            else:   # Only one array in output
+                out.shape = orig_shape + out.shape[1:]
+        else:
+            if isinstance(out, list) or isinstance(out, tuple):
+                out = list(out)
+
+                # Apply to each array in output list
+                for k,o in enumerate(out):
+                    out[k] = o[0]
+            else:   # Only one array in output
+                out = out[0]
+
+        # t4 = time.time()
+
+        # print('')
+        # print('time inside ensure_flat_galactic: {:.4f} s'.format(t4-t0))
+        # print('{: >7.4f} s : {: >6.4f} s : transform_to("galactic")'.format(t1-t0, t1-t0))
+        # print('{: >7.4f} s : {: >6.4f} s : reshape coordinates'.format(t2-t0, t2-t1))
+        # print('{: >7.4f} s : {: >6.4f} s : execute query'.format(t3-t0, t3-t2))
+        # print('{: >7.4f} s : {: >6.4f} s : reshape output'.format(t4-t0, t4-t3))
+        # print('')
+
+        return out
+
+    return _wrapper_func
+
 def gal_to_shape(gal, shape):
     l = np.reshape(gal.l.deg, shape)*units.deg
     b = np.reshape(gal.b.deg, shape)*units.deg
@@ -167,7 +273,6 @@ def gal_to_shape(gal, shape):
 
     return coordinates.SkyCoord(l, b, distance=d, frame='galactic')
 
-
 def ensure_flat_galactic(f):
     """
     A decorator for class methods of the form