Useful snippets for computing correlation functions

corrsnips

#!/usr/bin/env python

import sys
import pdb
import numpy as np
import math

try:
    from esutil import coords
    import esutil
except:
    print "You need to install esutils!"
    sys.exit(1)


def count_pairs(ra1, dec1, ra2, dec2, theta_bins):
    """Fast pair counting routine.

    Arguments:
        ra1: RA of first sample, in degrees
        dec1: DEC of first sample, in degrees
        ra2: RA of second sample, in degrees
        dec2: DEC of second sample, in degrees
        theta_bins: np.array of theta bins

    Returns:
        pairs: Pair counts
        r1: Left side bin edges
        r2: Right side bin edges
    """
    depth = 10
    h = esutil.htm.HTM(depth)
    ntheta = len(theta_bins)
    r1, r2, pairs = h.bincount(
        min(theta_bins), max(theta_bins), ntheta, ra1, dec1,
        ra2, dec2)
    print 'Pair results: ', pairs
    return pairs, r1, r2


def define_random(spec, photo, nrand):
    """Generates a random catalog in the same range as an input catalog.

    Arguments: 
        spec, photo: np.recarrays of spectroscopic and photometric catalogs
        nrand: integer that is the desired number of random points

    Output:
        rand: random catalog (recarray with 'RA' and 'DEC')
    """
    ra_min = np.min(np.array(np.min(spec['RA']), np.min(photo['RA'])))
    ra_max = np.max(np.array(np.max(spec['RA']), np.max(photo['RA'])))
    dec_min = np.min(np.array(np.min(spec['DEC']), np.min(photo['DEC'])))
    dec_max = np.max(np.array(np.max(spec['DEC']), np.max(photo['DEC'])))

    print '[+] Making a random catalog'
    rarand, decrand = coords.randsphere(
        nrand, ra_range=[ra_min, ra_max],
        dec_range=[dec_min, dec_max], system='eq')
   
    rand = np.recarray((nrand,), dtype=[('RA', float), ('DEC', float)])
    rand['RA'] = rarand
    rand['DEC'] = decrand
    return rand