class Catalog(DataObject):
    """
    A Catalog is an extension of a DataObject.  It can be decomposed into
    individual Entry objects.  This class is used to describe catalogs of all
    sorts, including galaxy catalogs, cluster catalogs, and other ndarrays.
    """

    entry_class = Entry

    @property
    def size(self):
        """
        Return the size of the Catalog.
        """
        return self._ndarray.size

    def append(self, append_cat):
        """
        Append a number of rows to the catalog, in-place.

        Parameters
        ----------
        append_cat: `redmapper.Catalog` or `np.ndarray`
           Catalog to append
        """
        if isinstance(append_cat, Catalog):
            self._ndarray = np.append(self._ndarray, append_cat._ndarray)
        else:
            self._ndarray = np.append(self._ndarray, append_cat)

    def extend(self, n_new):
        """
        Extend catalog with zero-filled rows, in-place.

        Parameters
        ----------
        n_new: `int`
           Number of new rows to append
        """
        temp = np.zeros(n_new, dtype=self._ndarray.dtype)
        self._ndarray = np.append(self._ndarray, temp)

    def __len__(self): return self.size

    def __getitem__(self, key):
        if isinstance(key, int):
            return self.entry_class(self._ndarray.__getitem__(key))
        return type(self)(self._ndarray.__getitem__(key))

    def __setitem__(self, key, val):
        self._ndarray.__setitem__(key, val) class Entry(DataObject):
    """
    An Entry is an extension of a DataObject.  It is intended to be used as a
    single Entry in a Catalog object.
    """
    def __init__(self, array):
        """
        Instantiate an Entry.

        Parameters
        ----------
        array: `np.ndarray` of length 1
        """
        if array.size != 1:
            raise ValueError("Input array must have length one.")
        # If this is an array of length 1, we want it to be a scalar-ish
        if len(array.shape) == 0:
            super(Entry, self).__init__(array)
        else:
            super(Entry, self).__init__(array[0])

    # @classmethod
    # def from_dict(cls, dict): pass

    def add_fields(self, newdtype):
        """
        Add fields to an Entry.

        Parameters
        ----------
        newdtype: data-type
           `np.dtype` description
        """
        array = np.zeros(self._ndarray.size, newdtype)
        self._lower_array(array)
        self._ndarray = self._merge_arrays([self._ndarray, array])[0]

    def __getattr__(self, attr):
        try:
            return self._ndarray[attr.lower()]
        except:
            return object.__getattribute__(self, attr) TOTAL_SQDEG = 4 * 180**2 / np.pi SEC_PER_DEG = 3600 def astro_to_sphere(ra, dec):
    """
    Convert astronomical ra/dec to healpix theta, phi coordinates.

    Parameters
    ----------
    ra: `np.array`
       Float array of right ascension
    dec: `np.array`
       Float array of declination

    Returns
    -------
    theta: `np.array`
       Float array of healpix theta
    phi: `np.array`
       Float array of healpix phi
    """
    return np.radians(90.0-dec), np.radians(ra) def calc_theta_i(mag, mag_err, maxmag, limmag):
    """
    Calculate the luminosity function smooth cutoff function, theta_i.

    Parameters
    ----------
    mag: `np.array`
       Magnitude values to compute theta_i
    mag_err: `np.array`
       Magnitude errors to compute theta_i
    maxmag: `float`
       Maximum possible magnitude in the catalog (hard cutoff)
    limmag: `float`
       Limiting magnitude to compute smooth cutoff

    Returns
    -------
    theta_i: `np.array`
       Float array of theta_i(mag, mag_err)
    """
    theta_i = np.ones(mag.size)
    eff_lim = np.clip(maxmag, 0, limmag)
    dmag = eff_lim - mag
    calc, = np.where(dmag < 5.0 * mag_err)
    if calc.size > 0:
        theta_i[calc] = 0.5 + 0.5 * erf(dmag[calc] / (np.sqrt(2) * mag_err[calc]))
    hi, = np.where(mag > limmag)
    if hi.size > 0:
        theta_i[hi] = 0.0
    return theta_i def apply_errormodels(maskgals, mag_in, b=None, err_ratio=1.0, fluxmode=False,
    nonoise=False, inlup=False, lnscat=None, sigma0=0.0):
    """
    Apply error models to a set of magnitudes.

    Parameters
    ----------
    maskgals: `redmapper.Catalog`
       Catalog of maskgals which have individual limmag, nsig values.
    mag_in: `np.array`
       Float array of raw magnitudes
    b: `np.array`, optional
       Luptitude softening parameters.  Default is None (not luptitudes)
    err_ratio: `float`, optional
       Error scaling ratio (for testing/adjustments).  Default is 1.0.
    fluxmode: `bool`, optional
       Return fluxes rather than magnitudes/luptitudes.  Default is False.
    nonoise: `bool`, optional
       Do not apply noise to output fluxes?  Default is False.
    inlup: `bool`, optional
       Input magnitudes are actually luptitudes.  Default is False.
       If True, b array must be supplied.
    lnscat: `float`, optional
       Apply additional ln(scatter) term (increasing errors).
       Default is None (no ln(scatter) term).
    sigma0: `float`, optional
       Additional noise floor term to apply.  Default is 0.0 (no floor).

    Returns
    -------
    mag: `np.array`
       Array of noised magnitudes/luptitudes/fluxes
    mag_err: `np.array`
       Array of magnitude/luptitude/flux errors
    """
    f1lim = 10.**((maskgals.limmag - maskgals.zp[0])/(-2.5))
    fsky1 = (((f1lim**2.) * maskgals.exptime)/(maskgals.nsig[0]**2.) - f1lim)
    fsky1 = np.clip(fsky1, 0.001, None)

    if inlup:
        bnmgy = b*1e9
        tflux = maskgals.exptime*2.0*bnmgy*np.sinh(-np.log(b)-0.4*np.log(10.0)*mag_in)
    else:
        tflux = maskgals.exptime*10.**((mag_in - maskgals.zp[0])/(-2.5))

    noise = err_ratio*np.sqrt(fsky1*maskgals.exptime + tflux)

    if sigma0 > 0.0:
        noise = np.sqrt(noise**2. + ((np.log(10.)/2.5) * sigma0 * tflux)**2.)

    if lnscat is not None:
        noise = np.exp(np.log(noise) + lnscat * np.random.normal(size=noise.size))

    if nonoise:
        flux = tflux
    else:
        flux = tflux + noise*random.standard_normal(mag_in.size)

    if fluxmode:
        mag = flux/maskgals.exptime
        mag_err = noise/maskgals.exptime
    else:
        if b is not None:
            bnmgy = b*1e9

            flux_new = flux/maskgals.exptime
            noise_new = noise/maskgals.exptime

            mag = 2.5*np.log10(1.0/b) - np.arcsinh(0.5*flux_new/bnmgy)/(0.4*np.log(10.0))
            mag_err = 2.5*noise_new/(2.0*bnmgy*np.log(10.0)*np.sqrt(1.0+(0.5*flux_new/bnmgy)**2.0))
        else:
            # Want to suppress warnings here (because we check for infinites below)
            with warnings.catch_warnings():
                warnings.filterwarnings("ignore")
                mag = maskgals.zp[0]-2.5*np.log10(flux/maskgals.exptime)
                mag_err = (2.5/np.log(10.0))*(noise/flux)

            bad, = np.where(np.isfinite(mag) == False)
            mag[bad] = 99.0
            mag_err[bad] = 99.0

    return mag, mag_err def make_lockfile(lockfile, block=False, maxtry=300, waittime=2):
    """
    Make a lockfile with atomic linking.

    Parameters
    ----------
    lockfile: `str`
       Name of lockfile to create
    block: `bool`, optional
       Block execution until lockfile is created?  Default is False.
    maxtry: `int`, optional
       Maximum number of tries to create lockfile (if blocking).  Default is 300.
    waittime: `float`, optional
       Wait time in seconds between tries (if blocking).  Default is 2.

    Returns
    -------
    locked: `bool`
       True if lockfile created, False if not.
    """

    import os
    import tempfile
    import time

    tf = tempfile.mkstemp(prefix='', dir=os.path.dirname(lockfile))
    tempfilename = tf[1]
    os.close(tf[0])

    if not block:
        # Non-blocking
        try:
            os.link(tempfilename, lockfile)
            os.unlink(tempfilename)
            # We successfully got the lock
            return True
        except:
            os.unlink(tempfilename)
            # We did not get the lock; return anyway
            return False
    else:
        # Blocking ... wait to get the lock up to maxtry

        ctr = 0
        locked = False
        while (ctr < maxtry) and (not locked):
            try:
                os.link(tempfilename, lockfile)
                os.unlink(tempfilename)
                locked = True
            except:
                ctr += 1
                time.sleep(waittime)

        if not locked:
            os.unlink(tempfilename)

        return locked def sample_from_pdf(f, ran, step, nsamp, **kwargs):
    """
    Sample from a PDF described by a function f.

    Parameters
    ----------
    f: `function`
       PDF function that can be called
    ran: `np.arange` or `list`
       Two-element range over which to sample.
    step: `float`
       Step size for interpolation.
    nsamp: `int`
       Number of samples from pdf
    **kwargs: `dict`
       Extra arguments to call f()

    Returns
    -------
    samples: `np.array`
       Float array of samples from PDF.
    """

    x = np.arange(ran[0], ran[1], step)
    pdf = f(x, **kwargs)
    pdf /= np.sum(pdf)
    cdf = np.cumsum(pdf, dtype=np.float64)
    cdfi = (cdf * x.size).astype(np.int32)

    rand = (np.random.uniform(size=nsamp) * x.size).astype(np.int32)

    samples = np.zeros(nsamp)
    for i in range(nsamp):
        test, = np.where(cdfi >= rand[i])
        samples[i] = x[test[0]]

    return samples def chisq_pdf(data, k):
    """
    Compute the chi-squared probability density function for an array of values
    and given number of degrees of freedom.  Implementing Equation 7 of
    Rykoff++2014.

    Parameters
    ----------
    data: `np.array`
       Float array of values
    k: `float`
       Number of effective degrees of freedom.

    Returns
    -------
    pdf: `np.array`
       Float array of pdf values
    """
    normalization = 1./(2**(k/2.) * special.gamma(k/2.))
    return normalization * data**((k/2.)-1) * np.exp(-data/2.) def schechter_pdf(x, alpha=-1.0, mstar=0.0):
    """
    Compute an unnormalized Schechter function pdf.

    Parameters
    ----------
    x: `np.array`
       Float array of x (magnitude) values
    alpha: `float`, optional
       Faint-end slope of the Schechter function.  Default is -1.0.
    mstar: `float`, optional
       Pivot magnitude mstar.  Default is 0.0.

    Returns
    -------
    pdf: `np.array`
       Float array of unnormalized Schechter pdf values.
    """
    pdf = 10.**(0.4*(alpha + 1.0) * (mstar - x)) * np.exp(-10.**(0.4 * (mstar - x)))
    return pdf def nfw_pdf(x, rscale=0.15, corer=0.1, radfactor=False):
    """
    Compute an unnormalized projected NFW pdf.

    Parameters
    ----------
    x: `np.array`
       Float array of x (radius) values
    rscale: `float`, optional
       NFW radial scaling.  Default is 0.15 Mpc.
    corer: `float`, optional
       Maximum radius to substitute a flat core.  Default is 0.1 Mpc.
    radfactor: `bool`, optional
       Multiply by 2*pi*r factor?  Default is False.

    Returns
    -------
    pdf: `np.array`
       Unnormalized projected NFW pdf.
    """

    xscale = x / rscale
    corex = corer / rscale

    sigx = np.zeros(x.size)

    low, = np.where(xscale < corex)
    mid, = np.where((xscale >= corex) & (xscale <= 0.999))
    high, = np.where((xscale >= 1.001) & (xscale < 10.0 / rscale))
    other, = np.where((xscale > 0.999) & (xscale < 1.001))

    if low.size > 0:
        arg = np.sqrt((1. - corex)/(1. + corex))
        pre = 2./(np.sqrt(1. - corex**2))
        front = 1./(corex**2. - 1)
        sigx[low] = front * (1. - pre*0.5*np.log((1. + arg)/(1. - arg)))

    if mid.size > 0:
        arg = np.sqrt((1. - xscale[mid])/(1. + xscale[mid]))
        pre = 2./(np.sqrt(1. - xscale[mid]**2.))
        front = 1./(xscale[mid]**2. - 1.)
        sigx[mid] = front * (1. - pre*0.5*np.log((1. + arg)/(1. - arg)))

    if high.size > 0:
        arg = np.sqrt((xscale[high] - 1.)/(xscale[high] + 1.))
        pre = 2./(np.sqrt(xscale[high]**2 - 1.))
        front = 1./(xscale[high]**2 - 1)
        sigx[high] = front * (1. - pre*np.arctan(arg))

    if other.size > 0:
        xlo, xhi = 0.999, 1.001
        arglo, arghi = np.sqrt((1 - xlo)/(1 + xlo)), np.sqrt((xhi - 1)/(xhi + 1))
        prelo, prehi = 2./np.sqrt(1.-xlo**2), 2./np.sqrt(xhi**2 - 1)
        frontlo, fronthi = 1./(xlo**2 - 1), 1./(xhi**2 - 1)
        testlo = frontlo * (1. - prelo*0.5*np.log((1 + arglo)/(1 - arglo)))
        testhi = fronthi * (1. - prehi*np.arctan(arghi))
        sigx[other] = (testlo + testhi)/2.

    if radfactor:
        sigx *= (2. * np.pi * x)

    return sigx def get_healsparse_subpix_indices(subpix_nside, subpix_hpix, subpix_border, coverage_nside):
    """
    Retrieve the coverage pixels that intersect the region, with a border.

    Parameters
    ----------
    subpix_nside: `int`
       Nside for the subregion
    subpix_hpix: `list`
       Pixel numbers for the subregion (ring format).
    subpix_border: `float`
       Border radius to cover outside subpix_hpix
    coverage_nside: `int`
       Nside of the healsparse coverage map
    """

    if len(subpix_hpix) == 0:
        raise RuntimeError("subpix_hpix cannot be an empty list.")
    if subpix_border > 0.0 and len(subpix_hpix) > 1:
        raise NotImplementedError("Cannot read multiple subpixels with a border.")

    # First, we need to know which pixel(s) from nside_coverage are covered by
    # subpix_hpix

    if subpix_nside == 0:
        # Special case for the full sky
        return None
    elif subpix_nside == coverage_nside:
        covpix = hp.ring2nest(subpix_nside, subpix_hpix)
    elif subpix_nside > coverage_nside:
        # what pixels are these contained in?
        theta, phi = hp.pix2ang(subpix_nside, subpix_hpix, nest=False)
        covpix = hp.ang2pix(coverage_nside, theta, phi, nest=True)
    else:
        # This is subpix_nside < coverage_nside
        # what coverage pixels are contained in subpix_hpix?
        subpix_hpix_nest = hp.ring2nest(subpix_nside, subpix_hpix)
        bit_shift = 2 * int(np.round(np.log(coverage_nside / subpix_nside) / np.log(2)))
        n_pix_per_pix = 2**bit_shift
        covpix = np.zeros(len(subpix_hpix_nest) * n_pix_per_pix, dtype=np.int64)
        for i, hpix in enumerate(subpix_hpix_nest):
            covpix[i * n_pix_per_pix: (i + 1) * n_pix_per_pix] = np.left_shift(hpix, bit_shift) + np.arange(n_pix_per_pix)

    # And now if we have a border...(only if the length is 1)
    if subpix_border > 0.0:
        nside_testing = max([coverage_nside * 4, subpix_nside * 4])
        boundaries = hp.boundaries(subpix_nside, subpix_hpix[0], step=nside_testing/subpix_nside)

        extrapix = np.zeros(0, dtype=np.int64)

        # These are pixels that touch the boundary
        for i in range(boundaries.shape[1]):
            pixint = hp.query_disc(nside_testing, boundaries[:, i],
                                   np.radians(subpix_border), inclusive=True, fact=8)
            extrapix = np.append(extrapix, pixint)

        extrapix = np.unique(extrapix)
        theta, phi = hp.pix2ang(nside_testing, extrapix)
        covpix = np.unique(np.append(covpix, hp.ang2pix(coverage_nside, theta, phi, nest=True)))

    return covpix
import esutil
import fitsio
import healpy as hp
import numpy as np
import os
import scipy.integrate
import healsparse
from scipy.special import erf
from .catalog import Catalog,Entry
from .utilities import TOTAL_SQDEG, SEC_PER_DEG, astro_to_sphere, calc_theta_i, apply_errormodels
from .utilities import make_lockfile, sample_from_pdf, chisq_pdf, schechter_pdf, nfw_pdf
from .utilities import get_healsparse_subpix_indices

        # Read the maskgals
        # These are going to be *all* the maskgals, but we only operate on a subset
        # at a time
        self.maskgals_all = Catalog.from_fits_file(maskgalfile)

    def select_maskgals_sample(self, maskgal_index=None):
        """
        Select a subset of maskgals by sampling.

        This will set self.maskgals to the subset in question.

        Parameters
        ----------
        maskgal_index: `int`, optional
           Pre-selected index to sample from (for reproducibility).
           Default is None (select randomly).
        """

        if maskgal_index is None:
            maskgal_index = np.random.choice(self.config.maskgal_nsamples)

        self.maskgals = self.maskgals_all[maskgal_index * self.config.maskgal_ngals:
                                              (maskgal_index + 1) * self.config.maskgal_ngals]

        return maskgal_index

    def gen_maskgals(self, maskgalfile):
        """
        Method to generate the maskgal monte carlo galaxies.

        Parameters
        ----------
        maskgalfile: `str`
           Name of maskgal file to generate.
        """

        minrad = np.clip(np.floor(10.*self.config.percolation_r0 * (3./100.)**self.config.percolation_beta) / 10., None, 0.5)
        maxrad = np.ceil(10.*self.config.percolation_r0 * (300./100.)**self.config.percolation_beta) / 10.

        nradbins = np.ceil((maxrad - minrad) / self.config.maskgal_rad_stepsize).astype(np.int32) + 1
        radbins = np.arange(nradbins, dtype=np.float32) * self.config.maskgal_rad_stepsize + minrad

        nmag = self.config.nmag
        ncol = nmag - 1

        ngals = self.config.maskgal_ngals * self.config.maskgal_nsamples

        maskgals = Catalog.zeros(ngals, dtype=[('r', 'f4'),
                                               ('phi', 'f4'),
                                               ('x', 'f4'),
                                               ('y', 'f4'),
                                               ('r_uniform', 'f4'),
                                               ('x_uniform', 'f4'),
                                               ('y_uniform', 'f4'),
                                               ('m', 'f4'),
                                               ('refmag', 'f4'),
                                               ('refmag_obs', 'f4'),
                                               ('refmag_obs_err', 'f4'),
                                               ('chisq', 'f4'),
                                               ('cwt', 'f4'),
                                               ('chisq_pdf', 'f4'),
                                               ('nfw', 'f4'),
                                               ('dzred', 'f4'),
                                               ('zwt', 'f4'),
                                               ('lumwt', 'f4'),
                                               ('lum_pdf', 'f4'),
                                               ('limmag', 'f4'),
                                               ('limmag_dered', 'f4'),
                                               ('exptime', 'f4'),
                                               ('m50', 'f4'),
                                               ('eff', 'f4'),
                                               ('w', 'f4'),
                                               ('theta_r', 'f4', nradbins),
                                               ('mark', bool),
                                               ('radbins', 'f4', nradbins),
                                               ('nin', 'f4', nradbins),
                                               ('nin_orig', 'f4', nradbins),
                                               ('zp', 'f4'),
                                               ('ebv', 'f4'),
                                               ('extinction', 'f4'),
                                               ('nsig', 'f4')])

        maskgals['radbins'] = np.tile(radbins, maskgals.size).reshape(maskgals.size, nradbins)

        # Generate chisq
        maskgals.chisq = sample_from_pdf(chisq_pdf, [0.0, self.config.chisq_max],
                                         self.config.chisq_max / 10000.,
                                         maskgals.size, k=ncol)
        # Generate mstar
        maskgals.m = sample_from_pdf(schechter_pdf,
                                     [-2.5*np.log10(10.0),
                                       -2.5*np.log10(self.config.lval_reference) + self.config.maskgal_dmag_extra],
                                     0.002, maskgals.size,
                                     alpha=self.config.calib_lumfunc_alpha, mstar=0.0)
        # Generate nfw(r)
        maskgals.r = sample_from_pdf(nfw_pdf,
                                     [0.001, maxrad],
                                     0.001, maskgals.size, radfactor=True)

        # Generate phi
        maskgals.phi = 2. * np.pi * np.random.random(size=maskgals.size)

        # Precompute x/y
        maskgals.x = maskgals.r * np.cos(maskgals.phi)
        maskgals.y = maskgals.r * np.sin(maskgals.phi)

        # And uniform x/y
        maskgals.r_uniform = self.config.bkg_local_annuli[1] * np.sqrt(np.random.uniform(size=maskgals.size))
        theta_new = np.random.uniform(size=maskgals.size)*2*np.pi
        maskgals.x_uniform = maskgals.r_uniform*np.cos(theta_new)
        maskgals.y_uniform = maskgals.r_uniform*np.sin(theta_new)

        # Compute weights to go with these values

        # Chisq weight
        maskgals.cwt = chisq_pdf(maskgals.chisq, ncol)
        maskgals.chisq_pdf = maskgals.cwt

        # Nfw weight

        maskgals.nfw = nfw_pdf(maskgals.r, radfactor=True)