<commit_before>from __future__ import print_function

import os
import fitsio

import numpy as np
from glob import glob

from astrometry.util.util import wcs_pv2sip_hdr

from tractor.basics import ConstantFitsWcs

from legacypipe.image import LegacySurveyImage, CalibMixin
from legacypipe.cpimage import CPImage
from legacypipe.survey import LegacySurveyData

class MosaicImage(CPImage, CalibMixin):
    '''
    Class for handling images from the Mosaic3 camera processed by the
    NOAO Community Pipeline.
    '''

    # this is defined here for testing purposes (to handle small images)
    splinesky_boxsize = 512

    def __init__(self, survey, t):
        super(MosaicImage, self).__init__(survey, t)
        # convert FWHM into pixel units
        self.fwhm /= self.pixscale

    @classmethod
    def get_bad_expids(self):
        import legacyccds
        fn = os.path.join(os.path.dirname(legacyccds.__file__),
                          'bad_expid_mzls.txt')
        bad_expids = np.loadtxt(fn, dtype=int, usecols=(0,))
        return bad_expids

    @classmethod
    def nominal_zeropoints(self):
        # See legacypipe/ccd_cuts.py and Photometric cuts email 12/21/2016
        return dict(z = 26.20)

    @classmethod
    def photometric_ccds(self, survey, ccds):
        '''
        Returns an index array for the members of the table 'ccds'
        that are photometric.

        This recipe is adapted from the DECam one.
        '''
        # Nominal zeropoints (DECam)
        z0 = self.nominal_zeropoints()
        z0 = np.array([z0[f[0]] for f in ccds.filter])
        good = np.ones(len(ccds), bool)
        n0 = sum(good)
        # See Photometric cuts email 12/21/2016
        # This is our list of cuts to remove non-photometric CCD images
        for name,crit in [
            ('exptime < 30 s', (ccds.exptime < 30)),
            ('ccdnmatch < 20', (ccds.ccdnmatch < 20)),
            ('sky too bright: ccdskycounts >= 150', (ccds.ccdskycounts >= 150)),
            ('abs(zpt - ccdzpt) > 0.1',
             (np.abs(ccds.zpt - ccds.ccdzpt) > 0.1)),
            ('zpt < 0.6 mag of nominal',
             (ccds.zpt < (z0 - 0.6))),
            ('zpt > 0.6 mag of nominal',
             (ccds.zpt > (z0 + 0.6))),
        ]:
            good[crit] = False
            #continue as usual
            n = sum(good)
            print('Flagged', n0-n, 'more non-photometric using criterion:',
                  name)
            n0 = n
        return np.flatnonzero(good)

    @classmethod
    def ccd_cuts(self, survey, ccds):
        ccdcuts = super(MosaicImage, self).ccd_cuts(survey, ccds)

        bits = LegacySurveyData.ccd_cut_bits

        I = self.bad_third_pixel(survey, ccds)
        print(np.sum(I), 'CCDs have bad THIRD_PIXEL')
        ccdcuts[I] += bits['THIRD_PIXEL']

        return ccdcuts

    @classmethod
    def bad_third_pixel(self, survey, ccds):
        '''
        For mosaic this is inconsistent YSHIFT header
        Ensures that ccds are 1/3 pixel interpolated.
        '''
        from astropy.io import fits
        bad = np.zeros(len(ccds), bool)
        # Remove if primary header does NOT have keyword YSHIFT
        rootdir = survey.get_image_dir()
        # The 1/3-pixel shift problem was fixed in hardware on MJD 57674,
        # so only check for problems in data before then.
        I = np.flatnonzero(ccds.mjd_obs < 57674.)
        for i in I:
            fn = ccds.image_filename[i]
            fn = fn.strip()
            fn = os.path.join(rootdir, fn)
            hdulist = fits.open(fn)
            if not 'YSHIFT' in hdulist[0].header:
                print('Did not find YSHIFT in primary header of', fn)
                bad[i] = True
        return bad

    def read_dq(self, **kwargs):
        '''
        Reads the Data Quality (DQ) mask image.
        '''
        print('Reading data quality image', self.dqfn, 'ext', self.hdu)
        dq = self._read_fits(self.dqfn, self.hdu, **kwargs)
        return dq

    def read_invvar(self, clip=True, **kwargs):
        '''
        Reads the inverse-variance (weight) map image.
        '''
        print('Reading weight map image', self.wtfn, 'ext', self.hdu)
        invvar = self._read_fits(self.wtfn, self.hdu, **kwargs)
        return invvar

    def remap_invvar(self, invvar, primhdr, img, dq):
        return self.remap_invvar_shotnoise(invvar, primhdr, img, dq)

    def get_wcs(self):
        '''cpimage.py get_wcs() but wcs comes from interpolated image if this is an
        uninterpolated image'''
        prim= self.read_image_primary_header()
        if 'YSHIFT' in prim.keys():
            # Interpolated image, use its wcs
            hdr = self.read_image_header()
        else:
            # Non-interpolated, use WCS of interpolated instead
            # Temporarily set imgfn to Interpolated image
            imgfn_backup= self.imgfn
            # Change CP*v3 --> CP*v2
            cpdir=os.path.basename(os.path.dirname(imgfn_backup)).replace('v3','v2')
            dirnm= os.path.dirname(os.path.dirname(imgfn_backup))
            i=os.path.basename(imgfn_backup).find('_ooi_')
            searchnm= os.path.basename(imgfn_backup)[:i+5]+'*.fits.fz'
            self.imgfn= np.array( glob(os.path.join(dirnm,cpdir,searchnm)) )
            assert(self.imgfn.size == 1)
            self.imgfn= self.imgfn[0]
            newprim= self.read_image_primary_header()
            from astropy.io import fits
            hdulist = fits.open(self.imgfn)
            print("'YSHIFT' in hdulist[0].header=",'YSHIFT' in hdulist[0].header)
            assert('YSHIFT' in newprim.keys())
            hdr = self.read_image_header()
            self.imgfn= imgfn_backup
            # Continue with wcs using the interpolated hdr
        # First child of MosaicImage is CPImage
        return super(MosaicImage,self).get_wcs(hdr=hdr)

    def get_tractor_wcs(self, wcs, x0, y0,
                        primhdr=None, imghdr=None):
        '''1/3 pixel shift if nont-interpolated image'''
        prim= self.read_image_primary_header()
        if 'YSHIFT' in prim.keys():
            # Use Default wcs class, this is an interpolated image
            return super(MosaicImage, self).get_tractor_wcs(wcs, x0, y0)
        else:
            # IDENTICAL to image.py get_tractor_wcs() except uses OneThirdPixelShiftWcs()
            # Instead of ConstantFitsWcs()
            # class OneThirdPixelShiftWcs is a ConstantFitsWcs class with1/3 pixel function
            twcs= OneThirdPixelShiftWcs(wcs)
            if x0 or y0:
                twcs.setX0Y0(x0,y0)
            return twcs

    #def run_calibs(self, psfex=True, funpack=False, git_version=None,
    #              force=False, **kwargs):
    def run_calibs(self, psfex=True, sky=True, se=False,
                   funpack=False, fcopy=False, use_mask=True,
                   force=False, just_check=False, git_version=None,
                   splinesky=False):
        se = False
        if psfex and os.path.exists(self.psffn) and (not force):
            if self.check_psf(self.psffn):
                psfex = False
        # dependency
        if psfex:
            se = True

        if se and os.path.exists(self.sefn) and (not force):
            if self.check_se_cat(self.sefn):
                se = False
        # dependency
        if se:
            funpack = True

        if sky and (not force) and (
            (os.path.exists(self.skyfn) and not splinesky) or
            (os.path.exists(self.splineskyfn) and splinesky)):
            fn = self.skyfn
            if splinesky:
                fn = self.splineskyfn

            if os.path.exists(fn):
                try:
                    hdr = fitsio.read_header(fn)
                except:
                    print('Failed to read sky file', fn, '-- deleting')
                    os.unlink(fn)
            if os.path.exists(fn):
                print('File', fn, 'exists -- skipping')
                sky = False

        if just_check:
            return (se or psfex or sky)

        todelete = []
        if funpack:
            # The image & mask files to process (funpacked if necessary)
            imgfn,maskfn = self.funpack_files(self.imgfn, self.dqfn, self.hdu, todelete)
        else:
            imgfn,maskfn = self.imgfn,self.dqfn

        if se:
            self.run_se('mzls', imgfn, maskfn)
        if psfex:
            self.run_psfex('mzls')

        if sky:
            self.run_sky('mzls', splinesky=splinesky,\
                         git_version=git_version)

        for fn in todelete:
            os.unlink(fn)

class OneThirdPixelShiftWcs(ConstantFitsWcs):
    def __init__(self,wcs):
        super(OneThirdPixelShiftWcs,self).__init__(wcs)

    def positionToPixel(self, pos, src=None):
        '''
        Converts an :class:`tractor.RaDecPos` to a pixel position.
        Returns: tuple of floats ``(x, y)``
        '''
        x,y = super(OneThirdPixelShiftWcs, self).positionToPixel(pos, src=src)
        # Top half of CCD needs be shifted up by 1./3 pixel
        if (y + self.y0 > 2048):
            #y += 1./3
            y -= 1./3
        return x,y

def main():

    from astrometry.util.fits import fits_table, merge_tables
    from legacypipe.survey import exposure_metadata
    # Fake up a survey-ccds.fits table from MzLS_CP
    from glob import glob
    #fns = glob('/project/projectdirs/cosmo/staging/mosaicz/MZLS_CP/CP20160202/k4m_160203_*oki*')
    fns = glob('/project/projectdirs/cosmo/staging/mosaicz/MZLS_CP/CP20160202/k4m_160203_08*oki*')

    print('Filenames:', fns)
    T = exposure_metadata(fns)

    # HACK
    T.fwhm = T.seeing / 0.262

    # FAKE
    T.ccdnmatch = np.zeros(len(T), np.int32) + 50
    T.zpt = np.zeros(len(T), np.float32) + 26.518
    T.ccdzpt = T.zpt.copy()
    T.ccdraoff = np.zeros(len(T), np.float32)
    T.ccddecoff = np.zeros(len(T), np.float32)

    fmap = {'zd':'z'}
    T.filter = np.array([fmap[f] for f in T.filter])

    T.writeto('mzls-ccds.fits')

    os.system('cp mzls-ccds.fits ~/legacypipe-dir/survey-ccds.fits')
    os.system('gzip -f ~/legacypipe-dir/survey-ccds.fits')

    import sys
    sys.exit(0)

    import logging
    import sys
    from legacypipe.runbrick import run_brick, get_runbrick_kwargs, get_parser

    parser = get_parser()
    opt = parser.parse_args()
    if opt.brick is None and opt.radec is None:
        parser.print_help()
        return -1
    survey, kwargs = get_runbrick_kwargs(opt)
    if kwargs in [-1, 0]:
        return kwargs

    if opt.verbose == 0:
        lvl = logging.INFO
    else:
        lvl = logging.DEBUG
    logging.basicConfig(level=lvl, format='%(message)s', stream=sys.stdout)

    kwargs.update(splinesky=True, pixPsf=True)

    run_brick(opt.brick, survey, **kwargs)

if __name__ == '__main__':
    main()
<commit_msg>mosaic one-third-pixel shift was fixed in hardware<commit_after>from __future__ import print_function

import os
import fitsio

import numpy as np
from glob import glob

from astrometry.util.util import wcs_pv2sip_hdr

from tractor.basics import ConstantFitsWcs

from legacypipe.image import LegacySurveyImage, CalibMixin
from legacypipe.cpimage import CPImage
from legacypipe.survey import LegacySurveyData

class MosaicImage(CPImage, CalibMixin):
    '''
    Class for handling images from the Mosaic3 camera processed by the
    NOAO Community Pipeline.
    '''

    mjd_third_pixel_fixed = 57674.

    # this is defined here for testing purposes (to handle small images)
    splinesky_boxsize = 512

    def __init__(self, survey, t):
        super(MosaicImage, self).__init__(survey, t)
        # convert FWHM into pixel units
        self.fwhm /= self.pixscale

    @classmethod
    def get_bad_expids(self):
        import legacyccds
        fn = os.path.join(os.path.dirname(legacyccds.__file__),
                          'bad_expid_mzls.txt')
        bad_expids = np.loadtxt(fn, dtype=int, usecols=(0,))
        return bad_expids

    @classmethod
    def nominal_zeropoints(self):
        # See legacypipe/ccd_cuts.py and Photometric cuts email 12/21/2016
        return dict(z = 26.20)

    @classmethod
    def photometric_ccds(self, survey, ccds):
        '''
        Returns an index array for the members of the table 'ccds'
        that are photometric.

        This recipe is adapted from the DECam one.
        '''
        # Nominal zeropoints (DECam)
        z0 = self.nominal_zeropoints()
        z0 = np.array([z0[f[0]] for f in ccds.filter])
        good = np.ones(len(ccds), bool)
        n0 = sum(good)
        # See Photometric cuts email 12/21/2016
        # This is our list of cuts to remove non-photometric CCD images
        for name,crit in [
            ('exptime < 30 s', (ccds.exptime < 30)),
            ('ccdnmatch < 20', (ccds.ccdnmatch < 20)),
            ('sky too bright: ccdskycounts >= 150', (ccds.ccdskycounts >= 150)),
            ('abs(zpt - ccdzpt) > 0.1',
             (np.abs(ccds.zpt - ccds.ccdzpt) > 0.1)),
            ('zpt < 0.6 mag of nominal',
             (ccds.zpt < (z0 - 0.6))),
            ('zpt > 0.6 mag of nominal',
             (ccds.zpt > (z0 + 0.6))),
        ]:
            good[crit] = False
            #continue as usual
            n = sum(good)
            print('Flagged', n0-n, 'more non-photometric using criterion:',
                  name)
            n0 = n
        return np.flatnonzero(good)

    @classmethod
    def ccd_cuts(self, survey, ccds):
        ccdcuts = super(MosaicImage, self).ccd_cuts(survey, ccds)

        bits = LegacySurveyData.ccd_cut_bits

        I = self.bad_third_pixel(survey, ccds)
        print(np.sum(I), 'CCDs have bad THIRD_PIXEL')
        ccdcuts[I] += bits['THIRD_PIXEL']

        return ccdcuts

    @classmethod
    def bad_third_pixel(self, survey, ccds):
        '''
        For mosaic this is inconsistent YSHIFT header
        Ensures that ccds are 1/3 pixel interpolated.
        '''
        from astropy.io import fits
        bad = np.zeros(len(ccds), bool)
        # Remove if primary header does NOT have keyword YSHIFT
        rootdir = survey.get_image_dir()
        # The 1/3-pixel shift problem was fixed in hardware on MJD 57674,
        # so only check for problems in data before then.
        I = np.flatnonzero(ccds.mjd_obs < MosaicImage.mjd_third_pixel_fixed)
        for i in I:
            fn = ccds.image_filename[i]
            fn = fn.strip()
            fn = os.path.join(rootdir, fn)
            hdulist = fits.open(fn)
            if not 'YSHIFT' in hdulist[0].header:
                print('Did not find YSHIFT in primary header of', fn)
                bad[i] = True
        return bad

    def read_dq(self, **kwargs):
        '''
        Reads the Data Quality (DQ) mask image.
        '''
        print('Reading data quality image', self.dqfn, 'ext', self.hdu)
        dq = self._read_fits(self.dqfn, self.hdu, **kwargs)
        return dq

    def read_invvar(self, clip=True, **kwargs):
        '''
        Reads the inverse-variance (weight) map image.
        '''
        print('Reading weight map image', self.wtfn, 'ext', self.hdu)
        invvar = self._read_fits(self.wtfn, self.hdu, **kwargs)
        return invvar

    def remap_invvar(self, invvar, primhdr, img, dq):
        return self.remap_invvar_shotnoise(invvar, primhdr, img, dq)

    def get_wcs(self):
        '''cpimage.py get_wcs() but wcs comes from interpolated image if this is an
        uninterpolated image'''
        prim= self.read_image_primary_header()
        if 'YSHIFT' in prim.keys() or self.mjdobs > MosaicImage.mjd_third_pixel_fixed:
            # Interpolated image, use its wcs
            hdr = self.read_image_header()
        else:
            # Non-interpolated, use WCS of interpolated instead
            # Temporarily set imgfn to Interpolated image
            imgfn_backup= self.imgfn
            # Change CP*v3 --> CP*v2
            cpdir=os.path.basename(os.path.dirname(imgfn_backup)).replace('v3','v2')
            dirnm= os.path.dirname(os.path.dirname(imgfn_backup))
            i=os.path.basename(imgfn_backup).find('_ooi_')
            searchnm= os.path.basename(imgfn_backup)[:i+5]+'*.fits.fz'
            self.imgfn= np.array( glob(os.path.join(dirnm,cpdir,searchnm)) )
            assert(self.imgfn.size == 1)
            self.imgfn= self.imgfn[0]
            newprim= self.read_image_primary_header()
            from astropy.io import fits
            hdulist = fits.open(self.imgfn)
            print("'YSHIFT' in hdulist[0].header=",'YSHIFT' in hdulist[0].header)
            assert('YSHIFT' in newprim.keys())
            hdr = self.read_image_header()
            self.imgfn= imgfn_backup
            # Continue with wcs using the interpolated hdr
        # First child of MosaicImage is CPImage
        return super(MosaicImage,self).get_wcs(hdr=hdr)

    def get_tractor_wcs(self, wcs, x0, y0,
                        primhdr=None, imghdr=None):
        '''1/3 pixel shift if nont-interpolated image'''
        prim= self.read_image_primary_header()
        if 'YSHIFT' in prim.keys():
            # Use Default wcs class, this is an interpolated image
            return super(MosaicImage, self).get_tractor_wcs(wcs, x0, y0)
        else:
            # IDENTICAL to image.py get_tractor_wcs() except uses OneThirdPixelShiftWcs()
            # Instead of ConstantFitsWcs()
            # class OneThirdPixelShiftWcs is a ConstantFitsWcs class with1/3 pixel function
            twcs= OneThirdPixelShiftWcs(wcs)
            if x0 or y0:
                twcs.setX0Y0(x0,y0)
            return twcs

    #def run_calibs(self, psfex=True, funpack=False, git_version=None,
    #              force=False, **kwargs):
    def run_calibs(self, psfex=True, sky=True, se=False,
                   funpack=False, fcopy=False, use_mask=True,
                   force=False, just_check=False, git_version=None,
                   splinesky=False):
        se = False
        if psfex and os.path.exists(self.psffn) and (not force):
            if self.check_psf(self.psffn):
                psfex = False
        # dependency
        if psfex:
            se = True

        if se and os.path.exists(self.sefn) and (not force):
            if self.check_se_cat(self.sefn):
                se = False
        # dependency
        if se:
            funpack = True

        if sky and (not force) and (
            (os.path.exists(self.skyfn) and not splinesky) or
            (os.path.exists(self.splineskyfn) and splinesky)):
            fn = self.skyfn
            if splinesky:
                fn = self.splineskyfn

            if os.path.exists(fn):
                try:
                    hdr = fitsio.read_header(fn)
                except:
                    print('Failed to read sky file', fn, '-- deleting')
                    os.unlink(fn)
            if os.path.exists(fn):
                print('File', fn, 'exists -- skipping')
                sky = False

        if just_check:
            return (se or psfex or sky)

        todelete = []
        if funpack:
            # The image & mask files to process (funpacked if necessary)
            imgfn,maskfn = self.funpack_files(self.imgfn, self.dqfn, self.hdu, todelete)
        else:
            imgfn,maskfn = self.imgfn,self.dqfn

        if se:
            self.run_se('mzls', imgfn, maskfn)
        if psfex:
            self.run_psfex('mzls')

        if sky:
            self.run_sky('mzls', splinesky=splinesky,\
                         git_version=git_version)

        for fn in todelete:
            os.unlink(fn)

class OneThirdPixelShiftWcs(ConstantFitsWcs):
    def __init__(self,wcs):
        super(OneThirdPixelShiftWcs,self).__init__(wcs)

    def positionToPixel(self, pos, src=None):
        '''
        Converts an :class:`tractor.RaDecPos` to a pixel position.
        Returns: tuple of floats ``(x, y)``
        '''
        x,y = super(OneThirdPixelShiftWcs, self).positionToPixel(pos, src=src)
        # Top half of CCD needs be shifted up by 1./3 pixel
        if (y + self.y0 > 2048):
            #y += 1./3
            y -= 1./3
        return x,y

def main():

    from astrometry.util.fits import fits_table, merge_tables
    from legacypipe.survey import exposure_metadata
    # Fake up a survey-ccds.fits table from MzLS_CP
    from glob import glob
    #fns = glob('/project/projectdirs/cosmo/staging/mosaicz/MZLS_CP/CP20160202/k4m_160203_*oki*')
    fns = glob('/project/projectdirs/cosmo/staging/mosaicz/MZLS_CP/CP20160202/k4m_160203_08*oki*')

    print('Filenames:', fns)
    T = exposure_metadata(fns)

    # HACK
    T.fwhm = T.seeing / 0.262

    # FAKE
    T.ccdnmatch = np.zeros(len(T), np.int32) + 50
    T.zpt = np.zeros(len(T), np.float32) + 26.518
    T.ccdzpt = T.zpt.copy()
    T.ccdraoff = np.zeros(len(T), np.float32)
    T.ccddecoff = np.zeros(len(T), np.float32)

    fmap = {'zd':'z'}
    T.filter = np.array([fmap[f] for f in T.filter])

    T.writeto('mzls-ccds.fits')

    os.system('cp mzls-ccds.fits ~/legacypipe-dir/survey-ccds.fits')
    os.system('gzip -f ~/legacypipe-dir/survey-ccds.fits')

    import sys
    sys.exit(0)

    import logging
    import sys
    from legacypipe.runbrick import run_brick, get_runbrick_kwargs, get_parser

    parser = get_parser()
    opt = parser.parse_args()
    if opt.brick is None and opt.radec is None:
        parser.print_help()
        return -1
    survey, kwargs = get_runbrick_kwargs(opt)
    if kwargs in [-1, 0]:
        return kwargs

    if opt.verbose == 0:
        lvl = logging.INFO
    else:
        lvl = logging.DEBUG
    logging.basicConfig(level=lvl, format='%(message)s', stream=sys.stdout)

    kwargs.update(splinesky=True, pixPsf=True)

    run_brick(opt.brick, survey, **kwargs)

if __name__ == '__main__':
    main()
<|endoftext|>


<commit_before>#! /usr/bin/env python

"""
Find all the PS1 stars in a given DECaLS CCD.
"""

import os
import numpy as np

class HealpixedCatalog(object):
    def __init__(self, fnpattern, nside=32, indexing='ring'):
        '''
        fnpattern: string formatter with key "hp", eg
        'dir/fn-%(hp)05i.fits'
        '''
        self.fnpattern = fnpattern
        self.nside = nside
        self.indexing = indexing

    def healpix_for_radec(self, ra, dec):
        '''
        Returns the healpix number for a given single (scalar) RA,Dec.
        '''
        from astrometry.util.util import radecdegtohealpix, healpix_xy_to_ring, healpix_xy_to_nested
        hpxy = radecdegtohealpix(ra, dec, self.nside)
        if self.indexing == 'ring':
            hp = healpix_xy_to_ring(hpxy, self.nside)
        elif self.indexing == 'nested':
            hp = healpix_xy_to_nested(hpxy, self.nside)
        else:
            # assume XY!
            hp = hpxy
        return hp

    def get_healpix_catalog(self, healpix):
        from astrometry.util.fits import fits_table
        fname = self.fnpattern % dict(hp=healpix)
        print('Reading', fname)
        return fits_table(fname)

    def get_healpix_catalogs(self, healpixes):
        from astrometry.util.fits import merge_tables
        cats = []
        for hp in healpixes:
            cats.append(self.get_healpix_catalog(hp))
        if len(cats) == 1:
            return cats[0]
        return merge_tables(cats)

    def get_catalog_in_wcs(self, wcs, step=100., margin=10):
        # Grid the CCD in pixel space
        W,H = wcs.get_width(), wcs.get_height()
        xx,yy = np.meshgrid(
            np.linspace(1-margin, W+margin, 2+int((W+2*margin)/step)),
            np.linspace(1-margin, H+margin, 2+int((H+2*margin)/step)))
        # Convert to RA,Dec and then to unique healpixes
        ra,dec = wcs.pixelxy2radec(xx.ravel(), yy.ravel())
        healpixes = set()
        for r,d in zip(ra,dec):
            healpixes.add(self.healpix_for_radec(r, d))
        # Read catalog in those healpixes
        cat = self.get_healpix_catalogs(healpixes)
        # Cut to sources actually within the CCD.
        _,xx,yy = wcs.radec2pixelxy(cat.ra, cat.dec)
        cat.x = xx
        cat.y = yy
        onccd = np.flatnonzero((xx >= 1.-margin) * (xx <= W+margin) *
                               (yy >= 1.-margin) * (yy <= H+margin))
        cat.cut(onccd)
        return cat

class ps1cat(HealpixedCatalog):
    ps1band = dict(g=0,r=1,i=2,z=3,Y=4)
    def __init__(self,expnum=None,ccdname=None,ccdwcs=None):
        """Read PS1 or gaia sources for an exposure number + CCD name or CCD WCS

        Args:
            expnum, ccdname: select catalogue with these
            ccdwcs: or select catalogue with this

        """
        self.ps1catdir = os.getenv('PS1CAT_DIR')
        if self.ps1catdir is None:
            raise ValueError('You must have the PS1CAT_DIR environment variable set to point to healpixed PS1 catalogs')
        fnpattern = os.path.join(self.ps1catdir, 'ps1-%(hp)05d.fits')
        super(ps1cat, self).__init__(fnpattern)

        if ccdwcs is None:
            from legacypipe.survey import LegacySurveyData
            survey = LegacySurveyData()
            ccd = survey.find_ccds(expnum=expnum,ccdname=ccdname)[0]
            im = survey.get_image_object(ccd)
            self.ccdwcs = im.get_wcs()
        else:
            self.ccdwcs = ccdwcs

    def get_stars(self,magrange=None,band='r'):
        """Return the set of PS1 or gaia-PS1 matched stars on a given CCD with well-measured grz
        magnitudes. Optionally trim the stars to a desired r-band magnitude
        range.
        """
        cat = self.get_catalog_in_wcs(self.ccdwcs)
        print('Found {} good PS1 stars'.format(len(cat)))
        if magrange is not None:
            keep = np.where((cat.median[:,ps1cat.ps1band[band]]>magrange[0])*
                            (cat.median[:,ps1cat.ps1band[band]]<magrange[1]))[0]
            cat = cat[keep]
            print('Trimming to {} stars with {}=[{},{}]'.
                  format(len(cat),band,magrange[0],magrange[1]))
        return cat

def ps1_to_decam(psmags, band):
    '''
    psmags: 2-d array (Nstars, Nbands)
    band: [grz]
    '''
    # https://desi.lbl.gov/trac/wiki/DecamLegacy/Reductions/Photometric
    g_index = ps1cat.ps1band['g']
    i_index = ps1cat.ps1band['i']
    gmag = psmags[:,g_index]
    imag = psmags[:,i_index]
    gi = gmag - imag

    # Eddie says (2017-Feb-14), in
    # [decam-chatter 4783] DECam grizY color transformations,
    # if ind EQ 0 then coeff = [ 0.00613, -0.01028, -0.03604, -0.00062] ; g
    # if ind EQ 1 then coeff = [ 0.01140, -0.03222,  0.08435, -0.00495] ; r
    # if ind EQ 2 then coeff = [ 0.00829, -0.00566,  0.04171, -0.00904] ; i
    # if ind EQ 3 then coeff = [ 0.00898, -0.02824,  0.07690, -0.02583] ; z
    # if ind EQ 4 then coeff = [ 0.00572, -0.02840,  0.05992, -0.02332] ; Y
    # coeffs = dict(
    #     g = [ 0.00613, -0.01028, -0.03604, -0.00062],
    #     r = [ 0.01140, -0.03222,  0.08435, -0.00495],
    #     i = [ 0.00829, -0.00566,  0.04171, -0.00904],
    #     z = [ 0.00898, -0.02824,  0.07690, -0.02583],
    #     Y = [ 0.00572, -0.02840,  0.05992, -0.02332],)[band]

    # But turns out the coeffs are in the opposite order!
    coeffs = dict(
        g = [ 0.00062,  0.03604, 0.01028, -0.00613 ],
        r = [ 0.00495, -0.08435, 0.03222, -0.01140 ],
        i = [ 0.00904, -0.04171, 0.00566, -0.00829 ],
        z = [ 0.02583, -0.07690, 0.02824, -0.00898 ],
        Y = [ 0.02332, -0.05992, 0.02840, -0.00572 ],)[band]

    # Previously, we used: (with an overall negative)
    # g = [0.0, -0.04709, -0.00084, 0.00340],
    # r = [0.0,  0.09939, -0.04509, 0.01488],
    # z = [0.0,  0.13404, -0.06591, 0.01695])[band]

    colorterm = coeffs[0] + coeffs[1]*gi + coeffs[2]*gi**2 + coeffs[3]*gi**3
    #print('Using DECam ColorTerm')
    return colorterm

def ps1_to_90prime(psmags, band):
    '''
    psmags: 2-d array (Nstars, Nbands)
    band: [gr]

    color terms are taken from:
      https://desi.lbl.gov/trac/wiki/BokLegacy/Photometric
    '''
    g_index = ps1cat.ps1band['g']
    i_index = ps1cat.ps1band['i']
    gmag = psmags[:, g_index]
    imag = psmags[:, i_index]
    gi = gmag - imag
    # 15 Nov 2018
    # https://desi.lbl.gov/trac/wiki/ImagingStandardBandpass
    coeffs = dict(
        g = [-0.00464, -0.08672,  0.00668,  0.00255],
        r = [-0.00110,  0.06875, -0.02480,  0.00855])[band]
        # Old version, from
        # July 22, 2016
        # https://desi.lbl.gov/trac/wiki/BokLegacy/Photometric
        #g = [0.0, +0.06630, +0.00958, -0.00672],
        #r = [0.0, -0.04836, +0.01100, -0.00563])[band]
        #Even older version.
        #g = [0.0, +0.08612, -0.00392, -0.00393],
        #r = [0.0, -0.07831, +0.03304, -0.01027])[band]
        # Note older versions used opposite sign.
    colorterm = -(coeffs[0] + coeffs[1]*gi + coeffs[2]*gi**2 + coeffs[3]*gi**3)
    print('Using 90prime ColorTerm')
    return colorterm

def ps1_to_mosaic(psmags, band):
    '''
    psmags: 2-d array (Nstars, Nbands)
    band: [gr]
    '''
    g_index = ps1cat.ps1band['g']
    i_index = ps1cat.ps1band['i']
    gmag = psmags[:, g_index]
    imag = psmags[:, i_index]
    gi = gmag - imag
    # Average color term for Mosaic3
    # 15 Nov 2018
    # https://desi.lbl.gov/trac/wiki/ImagingStandardBandpass
    coeffs = dict(z = [-0.03664,  0.11084, -0.04477,  0.01223])[band]
    # Old version from
    # https://desi.lbl.gov/trac/wiki/MayallZbandLegacy/CPReductions
    #coeffs = dict(z = [0.0, 0.121315, -0.046082623, 0.011642475])[band]

    colorterm = -(coeffs[0] + coeffs[1]*gi + coeffs[2]*gi**2 + coeffs[3]*gi**3)
    print('Using Mosaic3 ColorTerm')
    return colorterm
<commit_msg>add very preliminary PS1 entries for narrow-band N501 and N673 filters<commit_after>#! /usr/bin/env python
from __future__ import print_function

"""
Find all the PS1 stars in a given DECaLS CCD.
"""

import os
import numpy as np

class HealpixedCatalog(object):
    def __init__(self, fnpattern, nside=32, indexing='ring'):
        '''
        fnpattern: string formatter with key "hp", eg
        'dir/fn-%(hp)05i.fits'
        '''
        self.fnpattern = fnpattern
        self.nside = nside
        self.indexing = indexing

    def healpix_for_radec(self, ra, dec):
        '''
        Returns the healpix number for a given single (scalar) RA,Dec.
        '''
        from astrometry.util.util import radecdegtohealpix, healpix_xy_to_ring, healpix_xy_to_nested
        hpxy = radecdegtohealpix(ra, dec, self.nside)
        if self.indexing == 'ring':
            hp = healpix_xy_to_ring(hpxy, self.nside)
        elif self.indexing == 'nested':
            hp = healpix_xy_to_nested(hpxy, self.nside)
        else:
            # assume XY!
            hp = hpxy
        return hp

    def get_healpix_catalog(self, healpix):
        from astrometry.util.fits import fits_table
        fname = self.fnpattern % dict(hp=healpix)
        print('Reading', fname)
        return fits_table(fname)

    def get_healpix_catalogs(self, healpixes):
        from astrometry.util.fits import merge_tables
        cats = []
        for hp in healpixes:
            cats.append(self.get_healpix_catalog(hp))
        if len(cats) == 1:
            return cats[0]
        return merge_tables(cats)

    def get_catalog_in_wcs(self, wcs, step=100., margin=10):
        # Grid the CCD in pixel space
        W,H = wcs.get_width(), wcs.get_height()
        xx,yy = np.meshgrid(
            np.linspace(1-margin, W+margin, 2+int((W+2*margin)/step)),
            np.linspace(1-margin, H+margin, 2+int((H+2*margin)/step)))
        # Convert to RA,Dec and then to unique healpixes
        ra,dec = wcs.pixelxy2radec(xx.ravel(), yy.ravel())
        healpixes = set()
        for r,d in zip(ra,dec):
            healpixes.add(self.healpix_for_radec(r, d))
        # Read catalog in those healpixes
        cat = self.get_healpix_catalogs(healpixes)
        # Cut to sources actually within the CCD.
        _,xx,yy = wcs.radec2pixelxy(cat.ra, cat.dec)
        cat.x = xx
        cat.y = yy
        onccd = np.flatnonzero((xx >= 1.-margin) * (xx <= W+margin) *
                               (yy >= 1.-margin) * (yy <= H+margin))
        cat.cut(onccd)
        return cat

class ps1cat(HealpixedCatalog):
    ps1band = dict(g=0,r=1,i=2,z=3,Y=4,
                   N501=0,
                   N673=1,
    )
    def __init__(self,expnum=None,ccdname=None,ccdwcs=None):
        """Read PS1 or gaia sources for an exposure number + CCD name or CCD WCS

        Args:
            expnum, ccdname: select catalogue with these
            ccdwcs: or select catalogue with this

        """
        self.ps1catdir = os.getenv('PS1CAT_DIR')
        if self.ps1catdir is None:
            raise ValueError('You must have the PS1CAT_DIR environment variable set to point to healpixed PS1 catalogs')
        fnpattern = os.path.join(self.ps1catdir, 'ps1-%(hp)05d.fits')
        super(ps1cat, self).__init__(fnpattern)

        if ccdwcs is None:
            from legacypipe.survey import LegacySurveyData
            survey = LegacySurveyData()
            ccd = survey.find_ccds(expnum=expnum,ccdname=ccdname)[0]
            im = survey.get_image_object(ccd)
            self.ccdwcs = im.get_wcs()
        else:
            self.ccdwcs = ccdwcs

    def get_stars(self,magrange=None,band='r'):
        """Return the set of PS1 or gaia-PS1 matched stars on a given CCD with well-measured grz
        magnitudes. Optionally trim the stars to a desired r-band magnitude
        range.
        """
        cat = self.get_catalog_in_wcs(self.ccdwcs)
        print('Found {} good PS1 stars'.format(len(cat)))
        if magrange is not None:
            keep = np.where((cat.median[:,ps1cat.ps1band[band]]>magrange[0])*
                            (cat.median[:,ps1cat.ps1band[band]]<magrange[1]))[0]
            cat = cat[keep]
            print('Trimming to {} stars with {}=[{},{}]'.
                  format(len(cat),band,magrange[0],magrange[1]))
        return cat

def ps1_to_decam(psmags, band):
    '''
    psmags: 2-d array (Nstars, Nbands)
    band: [grz]
    '''
    # https://desi.lbl.gov/trac/wiki/DecamLegacy/Reductions/Photometric
    g_index = ps1cat.ps1band['g']
    i_index = ps1cat.ps1band['i']
    gmag = psmags[:,g_index]
    imag = psmags[:,i_index]
    gi = gmag - imag

    # Eddie says (2017-Feb-14), in
    # [decam-chatter 4783] DECam grizY color transformations,
    # if ind EQ 0 then coeff = [ 0.00613, -0.01028, -0.03604, -0.00062] ; g
    # if ind EQ 1 then coeff = [ 0.01140, -0.03222,  0.08435, -0.00495] ; r
    # if ind EQ 2 then coeff = [ 0.00829, -0.00566,  0.04171, -0.00904] ; i
    # if ind EQ 3 then coeff = [ 0.00898, -0.02824,  0.07690, -0.02583] ; z
    # if ind EQ 4 then coeff = [ 0.00572, -0.02840,  0.05992, -0.02332] ; Y
    # coeffs = dict(
    #     g = [ 0.00613, -0.01028, -0.03604, -0.00062],
    #     r = [ 0.01140, -0.03222,  0.08435, -0.00495],
    #     i = [ 0.00829, -0.00566,  0.04171, -0.00904],
    #     z = [ 0.00898, -0.02824,  0.07690, -0.02583],
    #     Y = [ 0.00572, -0.02840,  0.05992, -0.02332],)[band]

    # But turns out the coeffs are in the opposite order!
    coeffs = dict(
        g = [ 0.00062,  0.03604, 0.01028, -0.00613 ],
        r = [ 0.00495, -0.08435, 0.03222, -0.01140 ],
        i = [ 0.00904, -0.04171, 0.00566, -0.00829 ],
        z = [ 0.02583, -0.07690, 0.02824, -0.00898 ],
        Y = [ 0.02332, -0.05992, 0.02840, -0.00572 ],

        # FIXME!
        N501 = [ 0.00062,  0.03604, 0.01028, -0.00613 ],
        N673 = [ 0.00495, -0.08435, 0.03222, -0.01140 ],
    )[band]

    # Previously, we used: (with an overall negative)
    # g = [0.0, -0.04709, -0.00084, 0.00340],
    # r = [0.0,  0.09939, -0.04509, 0.01488],
    # z = [0.0,  0.13404, -0.06591, 0.01695])[band]

    colorterm = coeffs[0] + coeffs[1]*gi + coeffs[2]*gi**2 + coeffs[3]*gi**3
    #print('Using DECam ColorTerm')
    return colorterm

def ps1_to_90prime(psmags, band):
    '''
    psmags: 2-d array (Nstars, Nbands)
    band: [gr]

    color terms are taken from:
      https://desi.lbl.gov/trac/wiki/BokLegacy/Photometric
    '''
    g_index = ps1cat.ps1band['g']
    i_index = ps1cat.ps1band['i']
    gmag = psmags[:, g_index]
    imag = psmags[:, i_index]
    gi = gmag - imag
    # 15 Nov 2018
    # https://desi.lbl.gov/trac/wiki/ImagingStandardBandpass
    coeffs = dict(
        g = [-0.00464, -0.08672,  0.00668,  0.00255],
        r = [-0.00110,  0.06875, -0.02480,  0.00855])[band]
        # Old version, from
        # July 22, 2016
        # https://desi.lbl.gov/trac/wiki/BokLegacy/Photometric
        #g = [0.0, +0.06630, +0.00958, -0.00672],
        #r = [0.0, -0.04836, +0.01100, -0.00563])[band]
        #Even older version.
        #g = [0.0, +0.08612, -0.00392, -0.00393],
        #r = [0.0, -0.07831, +0.03304, -0.01027])[band]
        # Note older versions used opposite sign.
    colorterm = -(coeffs[0] + coeffs[1]*gi + coeffs[2]*gi**2 + coeffs[3]*gi**3)
    print('Using 90prime ColorTerm')
    return colorterm

def ps1_to_mosaic(psmags, band):
    '''
    psmags: 2-d array (Nstars, Nbands)
    band: [gr]
    '''
    g_index = ps1cat.ps1band['g']
    i_index = ps1cat.ps1band['i']
    gmag = psmags[:, g_index]
    imag = psmags[:, i_index]
    gi = gmag - imag
    # Average color term for Mosaic3
    # 15 Nov 2018
    # https://desi.lbl.gov/trac/wiki/ImagingStandardBandpass
    coeffs = dict(z = [-0.03664,  0.11084, -0.04477,  0.01223])[band]
    # Old version from
    # https://desi.lbl.gov/trac/wiki/MayallZbandLegacy/CPReductions
    #coeffs = dict(z = [0.0, 0.121315, -0.046082623, 0.011642475])[band]

    colorterm = -(coeffs[0] + coeffs[1]*gi + coeffs[2]*gi**2 + coeffs[3]*gi**3)
    print('Using Mosaic3 ColorTerm')
    return colorterm
<|endoftext|>


<commit_before>'''
>>> candidate_matrices.print_C_stats(10)
0 [(1, 1)]
1 [(0, 1), (1, 1)]
2 [(0, 3), (1, 3)]
3 [(0, 11), (1, 4)]
4 [(0, 33), (1, 7)]
5 [(0, 92), (1, 12)]
6 [(0, 254), (1, 19)]
7 [(0, 682), (1, 32)]
8 [(0, 1818), (1, 52)]
9 [(0, 4810), (1, 85)]
10 [(0, 12677), (1, 139)]

As expected, all zeros and ones.
>>> candidate_matrices.print_D_stats(10)
0 [(0, 1), (1, 1)]
1 [(0, 2), (1, 1)]
2 [(0, 8), (1, 2)]
3 [(0, 21), (1, 3)]
4 [(0, 60), (1, 5)]
5 [(0, 160), (1, 8)]
6 [(0, 429), (1, 13)]
7 [(0, 1134), (1, 21)]
8 [(0, 2992), (1, 34)]
9 [(0, 7865), (1, 55)]
10 [(0, 20648), (1, 89)]

Some negatives, that must be removed. Good outcome for little input.
>>> candidate_matrices.print_product_stats(10)
2 1 [(1, 2)]
3 1 [(0, 3), (1, 3)]
4 1 [(0, 9), (1, 6)]
4 2 [(0, 14), (1, 6)]
5 1 [(0, 31), (1, 9)]
5 2 [(0, 37), (1, 11)]
6 1 [(0, 88), (1, 16)]
6 2 [(0, 112), (1, 18)]
6 3 [(0, 93), (1, 23), (2, 1)]
7 1 [(0, 247), (1, 26)]
7 2 [(0, 305), (1, 31)]
7 3 [(0, 275), (1, 38), (2, 2)]
8 1 [(0, 671), (1, 43)]
8 2 [(0, 831), (1, 53)]
8 3 [(0, 741), (1, 71), (2, 4)]
8 4 [(-2, 1), (-1, 1), (0, 778), (1, 64), (2, 6)]
9 1 [(0, 1799), (1, 71)]
9 2 [(0, 2223), (1, 87)]
9 3 [(0, 2011), (1, 127), (2, 6), (3, 1)]
9 4 [(-2, 2), (-1, 3), (0, 2061), (1, 122), (2, 12)]
10 1 [(0, 4779), (1, 116)]
10 2 [(0, 5905), (1, 147)]
10 3 [(0, 5374), (1, 220), (2, 11), (3, 2)]
10 4 [(-2, 4), (-1, 10), (0, 5519), (1, 228), (2, 22), (3, 2)]
10 5 [(-2, 6), (-1, 7), (0, 5419), (1, 236), (2, 28)]
'''

# For Python2 compatibility
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals
__metaclass__ = type

import re

from .rankmatrices import RankMatrices
from .cdr_matrices import b_empty
from .cdr_matrices import b1
from .cdr_matrices import b2
from .cdr_matrices import rules_factory
from .cdr_matrices import rule_matrices_from_rules
from .cdr_matrices import cdr_print

if b1 != 'C' or b2 != 'D':
    ddt                         # Binary variant not coded.

candidate_2_re = re.compile('(C+)(D+)(.*)')
def candidate_2_split(s):
    mo = candidate_2_re.match(s)
    if mo:
        return mo.groups()

candidate_12_re = re.compile('(D+)(.*)')
def candidate_12_split(s):
    mo = candidate_12_re.match(s)
    if mo:
        return mo.groups()

# Start as clone of basic_matrices.
def candidate_11(word):
    yield b1 + word

def candidate_12(word):
    yield b1 + word
    if 0:
        bits = candidate_12_split(word)
        if bits != None:
            n = len(bits[0])
            for i in range(1, n):
                yield b1 + b2 * (n - i) + b1 * (2*i) + bits[1]

def candidate_2(word):
    yield  b2 + word
    if 0:
        bits = candidate_2_split(word)
        if bits != None:
            yield b1 + bits[0] + bits[1] + b1 + bits[2]

candidate_rules = rules_factory(candidate_11, candidate_12, candidate_2)
candidate_rule_matrices = rule_matrices_from_rules(candidate_rules)
candidate_matrices = RankMatrices(matrices=candidate_rule_matrices)

if __name__ == '__main__':

    import doctest
    print(doctest.testmod())
<commit_msg>Make simplest change thate removes the 8 4 product negative.<commit_after>'''
>>> candidate_matrices.print_C_stats(10)
0 [(1, 1)]
1 [(0, 1), (1, 1)]
2 [(0, 3), (1, 3)]
3 [(0, 11), (1, 4)]
4 [(0, 33), (1, 7)]
5 [(0, 92), (1, 12)]
6 [(0, 254), (1, 19)]
7 [(0, 681), (1, 33)]
8 [(-1, 2), (0, 1816), (1, 52)]
9 [(0, 4809), (1, 86)]
10 [(-1, 2), (0, 12675), (1, 139)]

As expected, all zeros and ones.
>>> candidate_matrices.print_D_stats(10)
0 [(0, 1), (1, 1)]
1 [(0, 2), (1, 1)]
2 [(0, 8), (1, 2)]
3 [(0, 21), (1, 3)]
4 [(0, 60), (1, 5)]
5 [(0, 160), (1, 8)]
6 [(0, 427), (1, 15)]
7 [(0, 1134), (1, 21)]
8 [(0, 2992), (1, 34)]
9 [(0, 7865), (1, 55)]
10 [(0, 20648), (1, 89)]

Some negatives, that must be removed. Good outcome for little input.
>>> candidate_matrices.print_product_stats(10)
2 1 [(1, 2)]
3 1 [(0, 3), (1, 3)]
4 1 [(0, 9), (1, 6)]
4 2 [(0, 14), (1, 6)]
5 1 [(0, 31), (1, 9)]
5 2 [(0, 37), (1, 11)]
6 1 [(0, 88), (1, 16)]
6 2 [(0, 112), (1, 18)]
6 3 [(0, 93), (1, 23), (2, 1)]
7 1 [(0, 247), (1, 26)]
7 2 [(0, 305), (1, 31)]
7 3 [(0, 275), (1, 38), (2, 2)]
8 1 [(0, 668), (1, 46)]
8 2 [(0, 828), (1, 56)]
8 3 [(0, 738), (1, 73), (2, 5)]
8 4 [(0, 780), (1, 64), (2, 6)]
9 1 [(-1, 2), (0, 1799), (1, 69)]
9 2 [(0, 2223), (1, 87)]
9 3 [(0, 2011), (1, 127), (2, 6), (3, 1)]
9 4 [(-2, 2), (-1, 3), (0, 2061), (1, 122), (2, 12)]
10 1 [(0, 4776), (1, 119)]
10 2 [(-1, 2), (0, 5905), (1, 145)]
10 3 [(0, 5368), (1, 225), (2, 12), (3, 2)]
10 4 [(-2, 3), (-1, 9), (0, 5518), (1, 231), (2, 22), (3, 2)]
10 5 [(-2, 5), (-1, 6), (0, 5417), (1, 239), (2, 29)]
'''

# For Python2 compatibility
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals
__metaclass__ = type

import re

from .rankmatrices import RankMatrices
from .cdr_matrices import b_empty
from .cdr_matrices import b1
from .cdr_matrices import b2
from .cdr_matrices import rules_factory
from .cdr_matrices import rule_matrices_from_rules
from .cdr_matrices import cdr_print

if b1 != 'C' or b2 != 'D':
    ddt                         # Binary variant not coded.

candidate_2_re = re.compile('(C+)(D+)(.*)')
def candidate_2_split(s):
    mo = candidate_2_re.match(s)
    if mo:
        return mo.groups()

candidate_12_re = re.compile('(D+)(.*)')
def candidate_12_split(s):
    mo = candidate_12_re.match(s)
    if mo:
        return mo.groups()

# Start as clone of basic_matrices.
def candidate_11(word):
    yield b1 + word

def candidate_12(word):
    yield b1 + word
    if 0:
        bits = candidate_12_split(word)
        if bits != None:
            n = len(bits[0])
            for i in range(1, n):
                yield b1 + b2 * (n - i) + b1 * (2*i) + bits[1]

def candidate_2(word):
    yield  b2 + word
    if word == 'CDCD':
        yield 'CCDCDC'
    elif word == 'CDDC':
        yield 'CCDDCC'
    if 0:
        bits = candidate_2_split(word)
        if bits != None:
            yield b1 + bits[0] + bits[1] + b1 + bits[2]

candidate_rules = rules_factory(candidate_11, candidate_12, candidate_2)
candidate_rule_matrices = rule_matrices_from_rules(candidate_rules)
candidate_matrices = RankMatrices(matrices=candidate_rule_matrices)

if __name__ == '__main__':

    import doctest
    print(doctest.testmod())
<|endoftext|>


<commit_before>
""" boxing - wrapping process with another process so we can run
a process inside and see if it crashes
"""

import py
import os
import sys
import marshal
from py.__.test import config as pytestconfig

PYTESTSTDOUT = "pyteststdout"
PYTESTSTDERR = "pyteststderr"
PYTESTRETVAL = "pytestretval"

import tempfile
import itertools
from StringIO import StringIO

counter = itertools.count().next

class FileBox(object):
    def __init__(self, fun, args=None, kwargs=None, config=None):
        if args is None:
            args = []
        if kwargs is None:
            kwargs = {}
        self.fun = fun
        self.config = config
        assert self.config
        self.args = args
        self.kwargs = kwargs

    def run(self, continuation=False):
        # XXX we should not use py.test.ensuretemp here
        count = counter()
        tempdir = py.test.ensuretemp("box%d" % count)
        self.tempdir = tempdir
        self.PYTESTRETVAL = tempdir.join('retval')
        self.PYTESTSTDOUT = tempdir.join('stdout')
        self.PYTESTSTDERR = tempdir.join('stderr')

        nice_level = self.config.getvalue('dist_nicelevel')
        pid = os.fork()
        if pid:
            if not continuation:
                self.parent(pid)
            else:
                return self.parent, pid
        else:
            try:
                outcome = self.children(nice_level)
            except:
                excinfo = py.code.ExceptionInfo()
                x = open("/tmp/traceback", "w")
                print >>x, "Internal box error"
                for i in excinfo.traceback:
                    print >>x, str(i)[2:-1]
                print >>x, excinfo
                x.close()
                os._exit(1)
            os.close(1)
            os.close(2)
            os._exit(0)
        return pid

    def children(self, nice_level):
        # right now we need to call a function, but first we need to
        # map all IO that might happen
        # make sure sys.stdout points to file descriptor one
        sys.stdout = stdout = self.PYTESTSTDOUT.open('w')
        sys.stdout.flush()
        fdstdout = stdout.fileno()
        if fdstdout != 1:
            os.dup2(fdstdout, 1)
        sys.stderr = stderr = self.PYTESTSTDERR.open('w')
        fdstderr = stderr.fileno()
        if fdstderr != 2:
            os.dup2(fdstderr, 2)
        retvalf = self.PYTESTRETVAL.open("w")
        try:
            if nice_level:
                os.nice(nice_level)
            # with fork() we have duplicated py.test's basetemp
            # directory so we set it manually here.
            # this may be expensive for some test setups,
            # but that is what you get with boxing.
            pytestconfig.basetemp = self.tempdir.join("childbasetemp")
            retval = self.fun(*self.args, **self.kwargs)
            retvalf.write(marshal.dumps(retval))
        finally:
            stdout.close()
            stderr.close()
            retvalf.close()
        os._exit(0)

    def parent(self, pid, waiter=os.waitpid):
        pid, exitstat = waiter(pid, 0)
        self.signal = exitstat & 0x7f
        self.exitstat = exitstat & 0xff00

        if not exitstat:
            retval = self.PYTESTRETVAL.open()
            try:
                retval_data = retval.read()
            finally:
                retval.close()
            self.retval = marshal.loads(retval_data)
        else:
            self.retval = None

        self.stdoutrepr = self.PYTESTSTDOUT.read()
        self.stderrrepr = self.PYTESTSTDERR.read()
        return self.stdoutrepr, self.stderrrepr

Box = FileBox
<commit_msg>[svn r38218] i give up for now, somehow the local change in the FileBox affects --dist testing in ways i can't fully understand at the moment (i would think it only affects --boxed test runs). So we are back to share tempdirs when running boxed (causing failures elswhere).  We care later, i guess.<commit_after>
""" boxing - wrapping process with another process so we can run
a process inside and see if it crashes
"""

import py
import os
import sys
import marshal
from py.__.test import config as pytestconfig

PYTESTSTDOUT = "pyteststdout"
PYTESTSTDERR = "pyteststderr"
PYTESTRETVAL = "pytestretval"

import tempfile
import itertools
from StringIO import StringIO

counter = itertools.count().next

class FileBox(object):
    def __init__(self, fun, args=None, kwargs=None, config=None):
        if args is None:
            args = []
        if kwargs is None:
            kwargs = {}
        self.fun = fun
        self.config = config
        assert self.config
        self.args = args
        self.kwargs = kwargs

    def run(self, continuation=False):
        # XXX we should not use py.test.ensuretemp here
        count = counter()
        tempdir = py.test.ensuretemp("box%d" % count)
        self.tempdir = tempdir
        self.PYTESTRETVAL = tempdir.join('retval')
        self.PYTESTSTDOUT = tempdir.join('stdout')
        self.PYTESTSTDERR = tempdir.join('stderr')

        nice_level = self.config.getvalue('dist_nicelevel')
        pid = os.fork()
        if pid:
            if not continuation:
                self.parent(pid)
            else:
                return self.parent, pid
        else:
            try:
                outcome = self.children(nice_level)
            except:
                excinfo = py.code.ExceptionInfo()
                x = open("/tmp/traceback", "w")
                print >>x, "Internal box error"
                for i in excinfo.traceback:
                    print >>x, str(i)[2:-1]
                print >>x, excinfo
                x.close()
                os._exit(1)
            os.close(1)
            os.close(2)
            os._exit(0)
        return pid

    def children(self, nice_level):
        # right now we need to call a function, but first we need to
        # map all IO that might happen
        # make sure sys.stdout points to file descriptor one
        sys.stdout = stdout = self.PYTESTSTDOUT.open('w')
        sys.stdout.flush()
        fdstdout = stdout.fileno()
        if fdstdout != 1:
            os.dup2(fdstdout, 1)
        sys.stderr = stderr = self.PYTESTSTDERR.open('w')
        fdstderr = stderr.fileno()
        if fdstderr != 2:
            os.dup2(fdstderr, 2)
        retvalf = self.PYTESTRETVAL.open("w")
        try:
            if nice_level:
                os.nice(nice_level)
            # with fork() we have duplicated py.test's basetemp
            # directory so we want to set it manually here.
            # this may be expensive for some test setups,
            # but that is what you get with boxing.
            # XXX but we are called in more than strict boxing
            # mode ("AsyncExecutor") so we can't do the following without
            # inflicting on --dist speed, hum:
            # pytestconfig.basetemp = self.tempdir.join("childbasetemp")
            retval = self.fun(*self.args, **self.kwargs)
            retvalf.write(marshal.dumps(retval))
        finally:
            stdout.close()
            stderr.close()
            retvalf.close()
        os._exit(0)

    def parent(self, pid, waiter=os.waitpid):
        pid, exitstat = waiter(pid, 0)
        self.signal = exitstat & 0x7f
        self.exitstat = exitstat & 0xff00

        if not exitstat:
            retval = self.PYTESTRETVAL.open()
            try:
                retval_data = retval.read()
            finally:
                retval.close()
            self.retval = marshal.loads(retval_data)
        else:
            self.retval = None

        self.stdoutrepr = self.PYTESTSTDOUT.read()
        self.stderrrepr = self.PYTESTSTDERR.read()
        return self.stdoutrepr, self.stderrrepr

Box = FileBox
<|endoftext|>


<commit_before>#!/usr/bin/env python
# encoding: utf-8

# The MIT License (MIT)

# Copyright (c) 2012-2014 CNRS (Hervé BREDIN - http://herve.niderb.fr)

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.

# THE SOFTWARE S PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

from __future__ import unicode_literals

# import sys
# import pandas
# import numpy as np
# from pyannote.core import Segment, Timeline, Annotation, Scores
# from pyannote.core.feature import SlidingWindowFeature
# from pyannote.core import PYANNOTE_URI, PYANNOTE_MODALITY, \
#     PYANNOTE_SEGMENT, PYANNOTE_TRACK, PYANNOTE_LABEL, PYANNOTE_SCORE

from abc import ABCMeta, abstractmethod

class Parser(object):

    __metaclass__ = ABCMeta

    @classmethod
    @abstractmethod
    def file_extensions(cls):
        pass

    @abstractmethod
    def read(self, path, **kwargs):
        pass

    @abstractmethod
    def empty(self, uri=None, modality=None, **kwargs):
        pass

    def __call__(self, uri=None, modality=None, **kwargs):

        match = dict(self._loaded)

        # filter out all annotations
        # but the ones for the requested resource
        if uri is not None:
            match = {(v, m): ann for (v, m), ann in match.iteritems()
                     if v == uri}

        # filter out all remaining annotations
        # but the ones for the requested modality
        if modality is not None:
            match = {(v, m): ann for (v, m), ann in match.iteritems()
                     if m == modality}

        if len(match) == 0:
            A = self.empty(uri=uri, modality=modality, **kwargs)

        elif len(match) == 1:
            A = match.values()[0]

        else:
            msg = 'Found more than one matching annotation: %s'
            raise ValueError(msg % match.keys())

        return A
<commit_msg>fix: add missing properties 'uris' and 'modalities'<commit_after>#!/usr/bin/env python
# encoding: utf-8

# The MIT License (MIT)

# Copyright (c) 2012-2014 CNRS (Hervé BREDIN - http://herve.niderb.fr)

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.

# THE SOFTWARE S PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

from __future__ import unicode_literals

# import sys
# import pandas
# import numpy as np
# from pyannote.core import Segment, Timeline, Annotation, Scores
# from pyannote.core.feature import SlidingWindowFeature
# from pyannote.core import PYANNOTE_URI, PYANNOTE_MODALITY, \
#     PYANNOTE_SEGMENT, PYANNOTE_TRACK, PYANNOTE_LABEL, PYANNOTE_SCORE

from abc import ABCMeta, abstractmethod

class Parser(object):

    __metaclass__ = ABCMeta

    @classmethod
    @abstractmethod
    def file_extensions(cls):
        pass

    @abstractmethod
    def read(self, path, **kwargs):
        pass

    @abstractmethod
    def empty(self, uri=None, modality=None, **kwargs):
        pass

    def __get_uris(self):
        return sorted(set([v for (v, m) in self._loaded]))
    uris = property(fget=__get_uris)
    """"""

    def __get_modalities(self):
        return sorted(set([m for (v, m) in self._loaded]))
    modalities = property(fget=__get_modalities)
    """"""

    def __call__(self, uri=None, modality=None, **kwargs):

        match = dict(self._loaded)

        # filter out all annotations
        # but the ones for the requested resource
        if uri is not None:
            match = {(v, m): ann for (v, m), ann in match.iteritems()
                     if v == uri}

        # filter out all remaining annotations
        # but the ones for the requested modality
        if modality is not None:
            match = {(v, m): ann for (v, m), ann in match.iteritems()
                     if m == modality}

        if len(match) == 0:
            A = self.empty(uri=uri, modality=modality, **kwargs)

        elif len(match) == 1:
            A = match.values()[0]

        else:
            msg = 'Found more than one matching annotation: %s'
            raise ValueError(msg % match.keys())

        return A
<|endoftext|>


<commit_before>from pyasn1.type import base, tag

class EndOfOctets(base.AbstractSimpleAsn1Item):
    tagSet = tag.initTagSet(
        tag.Tag(tag.tagClassUniversal, tag.tagFormatSimple, 0x00)
        )
endOfOctets = EndOfOctets()
<commit_msg>set default value (0) to EndOfOctets<commit_after>from pyasn1.type import base, tag

class EndOfOctets(base.AbstractSimpleAsn1Item):
    defaultValue = 0
    tagSet = tag.initTagSet(
        tag.Tag(tag.tagClassUniversal, tag.tagFormatSimple, 0x00)
        )
endOfOctets = EndOfOctets()
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
# vim: set ts=4 et

import random
import re

from plugin import *

class Plugin(BasePlugin):
    @hook
    def choose_trigger(self, msg, args, argstr):
        argstr = argstr.replace(' or ', ',').replace(',or ', ',')
        choices = [x for x in map(str.strip, argstr.split(',')) if x]
        if not choices:
            responses = [
                "Choices choices.. which do I choose from.  Oh right, I don't.",
                "You've got to give me something to work with here!",
                "...",
                "You're supposed to be asking me the questions.  What exactly is your plan?",
                "Look {0}, you've got to give me some options.",
                "No choices?  Empty result set.",
                "\x01ACTION looks at {0} and just walks away.  Slowly.\x01",
            ]
        elif len(choices) == 1:
            responses = [
                "You've forced me to go with {1}.",
                "{1}.  Is there anything else I can help you with, {0}?",
                "I'm not playing this game, {0}.",
                "If you go with {1}, I'm leaving.",
                "{0}, you displease me.",
                "No.",
                "Why would you ask me to choose one thing? WHY?!",
            ]
        else:
            responses = [
                "I suggest you choose {1}.",
                "I'd go with {1}.",
                "Perhaps {1} is the right choice.",
                "Well, I wouldn't go with {1}.",
                "Are you sure you want to go with {1}?",
                "I like {1}.",
                "{1}. Definately {1}.",
            ]

        response = random.choice(responses)
        choice = random.choice(choices) if choices else None
        msg.reply(response.format(msg.source, choice))

    @hook
    def privmsg_command(self, msg):
        if not msg.channel:
            return

        m = re.match('^(?:can|should) (i|we|%s)(\'s|s)?.*\?$' % '|'.join([re.escape(nick) for nick in self.bot.channels[msg.channel]['nicks']]), msg.param[-1], re.I)
        if not m:
            return

        m = re.match(' or ', msg.param[-1], re.I)
        if m:
            return

        responses = [
            "Sure!",
            "Of course.",
            "Why not?",
            "Yes.",
            "Maybe.",
            "Are you sure thats a good idea?",
            "Nah.",
            "Better not.",
            "Nope.",
            "No way!",
        ]

        response = random.choice(responses)
        msg.reply(response)

<commit_msg>Fix misspelling of definitely in choose plugin<commit_after># -*- coding: utf-8 -*-
# vim: set ts=4 et

import random
import re

from plugin import *

class Plugin(BasePlugin):
    @hook
    def choose_trigger(self, msg, args, argstr):
        argstr = argstr.replace(' or ', ',').replace(',or ', ',')
        choices = [x for x in map(str.strip, argstr.split(',')) if x]
        if not choices:
            responses = [
                "Choices choices.. which do I choose from.  Oh right, I don't.",
                "You've got to give me something to work with here!",
                "...",
                "You're supposed to be asking me the questions.  What exactly is your plan?",
                "Look {0}, you've got to give me some options.",
                "No choices?  Empty result set.",
                "\x01ACTION looks at {0} and just walks away.  Slowly.\x01",
            ]
        elif len(choices) == 1:
            responses = [
                "You've forced me to go with {1}.",
                "{1}.  Is there anything else I can help you with, {0}?",
                "I'm not playing this game, {0}.",
                "If you go with {1}, I'm leaving.",
                "{0}, you displease me.",
                "No.",
                "Why would you ask me to choose one thing? WHY?!",
            ]
        else:
            responses = [
                "I suggest you choose {1}.",
                "I'd go with {1}.",
                "Perhaps {1} is the right choice.",
                "Well, I wouldn't go with {1}.",
                "Are you sure you want to go with {1}?",
                "I like {1}.",
                "{1}. Definitely {1}.",
            ]

        response = random.choice(responses)
        choice = random.choice(choices) if choices else None
        msg.reply(response.format(msg.source, choice))

    @hook
    def privmsg_command(self, msg):
        if not msg.channel:
            return

        m = re.match('^(?:can|should) (i|we|%s)(\'s|s)?.*\?$' % '|'.join([re.escape(nick) for nick in self.bot.channels[msg.channel]['nicks']]), msg.param[-1], re.I)
        if not m:
            return

        m = re.match(' or ', msg.param[-1], re.I)
        if m:
            return

        responses = [
            "Sure!",
            "Of course.",
            "Why not?",
            "Yes.",
            "Maybe.",
            "Are you sure thats a good idea?",
            "Nah.",
            "Better not.",
            "Nope.",
            "No way!",
        ]

        response = random.choice(responses)
        msg.reply(response)

<|endoftext|>


<commit_before>
import os
import json
import logging
from pybythec import utils

log = logging.getLogger('pybythec')

class BuildStatus:
  '''
    member variables:
    name: name of target
    path: path to write status.json file to
    status: failed, built, up to date, or locked
    description: what happened
  '''

  def __init__(self, name, path = ''):
    self.name = name
    self.path = path
    self.status = 'failed'
    self.description = ''

  # def readFromFile(self, buildPath):
  def readFromFile(self, libSrcDir, buildDir, buildType, compiler, binaryFormat):
    '''
      buildPath (in): where to read the status.json file from
    '''

    buildPath = '{0}/{1}/{2}/{3}/{4}'.format(libSrcDir, buildDir, buildType, compiler, binaryFormat)
    if not os.path.exists(buildPath): # try without the buildDir being hidden
      buildPath = '{0}/{1}/{2}/{3}/{4}'.format(libSrcDir, buildDir.lstrip('.'), buildType, compiler, binaryFormat)

    contents = utils.loadJsonFile(buildPath + '/status.json')
    if not contents:
      self.description = 'couldn\'t find contents in ' + buildPath
      log.error('couldn\'t find contents in ' + buildPath)
      return
    if 'status' in contents:
      self.status = contents['status']
    else:
      self.description = 'couldn\'t find the build status in ' + buildPath
      log.error(self.description)
    if 'description' in contents:
      self.description = contents['description']
    else:
      self.description = buildPath + ' doesn\'t contain a description'
      log.warning(self.description)

  def writeInfo(self, status, msg):
    log.info(msg)
    self.status = status
    self.description = msg
    self._writeToFile()

  def writeError(self, msg):
    log.error(msg)
    self.description = msg
    self._writeToFile()

  def _writeToFile(self):
    if not os.path.exists(self.path):
      return
    with open(self.path + '/status.json', 'w') as f:
      json.dump({'status': self.status, 'description': self.description}, f, indent = 4)
<commit_msg>minor cleanup<commit_after>
import os
import json
import logging
from pybythec import utils

log = logging.getLogger('pybythec')

class BuildStatus:
  '''
    member variables:
    name: name of target
    path: path to write status.json file to
    status: failed, built, up to date, or locked
    description: what happened
  '''

  def __init__(self, name, path = ''):
    self.name = name
    self.path = path
    self.status = 'failed'
    self.description = ''

  def readFromFile(self, libSrcDir, buildDir, buildType, compiler, binaryFormat):
    '''
      buildPath (in): where to read the status.json file from
    '''

    buildPath = '{0}/{1}/{2}/{3}/{4}'.format(libSrcDir, buildDir, buildType, compiler, binaryFormat)
    if not os.path.exists(buildPath): # try the non-hidden version
      buildPath = '{0}/{1}/{2}/{3}/{4}'.format(libSrcDir, buildDir.lstrip('.'), buildType, compiler, binaryFormat)

    contents = utils.loadJsonFile(buildPath + '/status.json')
    if not contents:
      self.description = 'couldn\'t find contents in ' + buildPath
      log.error('couldn\'t find contents in ' + buildPath)
      return
    if 'status' in contents:
      self.status = contents['status']
    else:
      self.description = 'couldn\'t find the build status in ' + buildPath
      log.error(self.description)
    if 'description' in contents:
      self.description = contents['description']
    else:
      self.description = buildPath + ' doesn\'t contain a description'
      log.warning(self.description)

  def writeInfo(self, status, msg):
    log.info(msg)
    self.status = status
    self.description = msg
    self._writeToFile()

  def writeError(self, msg):
    log.error(msg)
    self.description = msg
    self._writeToFile()

  def _writeToFile(self):
    if not os.path.exists(self.path):
      return
    with open(self.path + '/status.json', 'w') as f:
      json.dump({'status': self.status, 'description': self.description}, f, indent = 4)
<|endoftext|>


<commit_before>"""
The utils module handles helper routines for equilibrium calculation.
"""
from __future__ import division
import pycalphad.variables as v
from pycalphad.core.halton import halton
from pycalphad.core.constants import MIN_SITE_FRACTION
from sympy.utilities.lambdify import lambdify
from sympy.printing.lambdarepr import LambdaPrinter
import numpy as np
import operator
import functools
import itertools
import collections
from importlib import import_module

_NUMBA = None

try:
    _NUMBA = import_module('numba')

    @_NUMBA.vectorize
    def nbwhere(condition, x, y):
        if condition:
            return x
        else:
            return y
except ImportError:
    pass

try:
    # Only available in numpy 1.10 and newer
    from numpy import broadcast_to
except ImportError:
    def broadcast_to(arr, shape):
        "Broadcast an array to a desired shape. Returns a view."
        return np.broadcast_arrays(arr, np.empty(shape, dtype=np.bool))[0]

class NumPyPrinter(LambdaPrinter):
    """
    Numpy printer which handles vectorized piecewise functions,
    logical operators, etc.
    """
    _default_settings = {
        "order": "none",
        "full_prec": "auto",
    }

    def _print_seq(self, seq, delimiter=', '):
        "General sequence printer: converts to tuple"
        # Print tuples here instead of lists because numba supports
        #     tuples in nopython mode.
        return '({},)'.format(delimiter.join(self._print(item) for item in seq))

    #def _print_Integer(self, expr):
    #    # Upcast to work around an integer overflow bug
    #    return 'asarray({0}, dtype=float)'.format(expr)

    def _print_MatrixBase(self, expr):
        return "%s(%s)" % ('array', self._print(expr.tolist()))

    _print_SparseMatrix = \
        _print_MutableSparseMatrix = \
        _print_ImmutableSparseMatrix = \
        _print_Matrix = \
        _print_DenseMatrix = \
        _print_MutableDenseMatrix = \
        _print_ImmutableMatrix = \
        _print_ImmutableDenseMatrix = \
        _print_MatrixBase

    def _print_MatMul(self, expr):
        "Matrix multiplication printer"
        return '({0})'.format(').dot('.join(self._print(i) for i in expr.args))

    def _print_Piecewise(self, expr):
        "Piecewise function printer"
        exprs = [self._print(arg.expr) for arg in expr.args]
        conds = [self._print(arg.cond) for arg in expr.args]
        if (len(conds) == 1) and exprs[0] == '0':
            return '0'
        # We expect exprs and conds to be in order here
        # First condition to evaluate to True is returned
        # Default result: float zero
        result = 'where({0}, {1}, 0.)'.format(conds[-1], exprs[-1])
        for expr, cond in reversed(list(zip(exprs[:-1], conds[:-1]))):
            result = 'where({0}, {1}, {2})'.format(cond, expr, result)

        return result

    def _print_And(self, expr):
        "Logical And printer"
        # We have to override LambdaPrinter because it uses Python 'and' keyword.
        # If LambdaPrinter didn't define it, we could use StrPrinter's
        # version of the function and add 'logical_and' to NUMPY_TRANSLATIONS.
        if len(expr.args) == 1:
            return self._print(expr.args[0])
        result = 'logical_and({0}, {1})'.format(self._print(expr.args[-2]), self._print(expr.args[-1]))
        for arg in reversed(expr.args[:-2]):
            result = 'logical_and({0}, {1})'.format(self._print(arg), result)
        return  result

    def _print_Or(self, expr):
        "Logical Or printer"
        # We have to override LambdaPrinter because it uses Python 'or' keyword.
        # If LambdaPrinter didn't define it, we could use StrPrinter's
        # version of the function and add 'logical_or' to NUMPY_TRANSLATIONS.
        if len(expr.args) == 1:
            return self._print(expr.args[0])
        result = 'logical_or({0}, {1})'.format(self._print(expr.args[-2]), self._print(expr.args[-1]))
        for arg in reversed(expr.args[:-2]):
            result = 'logical_or({0}, {1})'.format(self._print(arg), result)
        return  result

    def _print_Not(self, expr):
        "Logical Not printer"
        # We have to override LambdaPrinter because it uses Python 'not' keyword.
        # If LambdaPrinter didn't define it, we would still have to define our
        #     own because StrPrinter doesn't define it.
        if len(expr.args) == 1:
            return self._print(expr.args[0])
        result = 'logical_not({0}, {1})'.format(self._print(expr.args[-2]), self._print(expr.args[-1]))
        for arg in reversed(expr.args[:-2]):
            result = 'logical_not({0}, {1})'.format(self._print(arg), result)
        return  result

def point_sample(comp_count, pdof=10):
    """
    Sample 'pdof * (sum(comp_count) - len(comp_count))' points in
    composition space for the sublattice configuration specified
    by 'comp_count'. Points are sampled quasi-randomly from a Halton sequence.
    A Halton sequence is like a uniform random distribution, but the
    result will always be the same for a given 'comp_count' and 'pdof'.
    Note: For systems with only one component, only one point will be
    returned, regardless of 'pdof'. This is because the degrees of freedom
    are zero for that case.

    Parameters
    ----------
    comp_count : list
        Number of components in each sublattice.
    pdof : int
        Number of points to sample per degree of freedom.

    Returns
    -------
    ndarray of generated points satisfying the mass balance.

    Examples
    --------
    >>> comps = [8,1] # 8 components in sublattice 1; only 1 in sublattice 2
    >>> pts = point_sample(comps, pdof=20) # 7 d.o.f, returns a 140x7 ndarray
    """
    # Generate Halton sequence with appropriate dimensions and size
    pts = halton(sum(comp_count),
                 pdof * (sum(comp_count) - len(comp_count)), scramble=True)
    # Convert low-discrepancy sequence to normalized exponential
    # This will be uniformly distributed over the simplices
    pts = -np.log(pts)
    cur_idx = 0
    for ctx in comp_count:
        end_idx = cur_idx + ctx
        pts[:, cur_idx:end_idx] /= pts[:, cur_idx:end_idx].sum(axis=1)[:, None]
        cur_idx = end_idx

    if len(pts) == 0:
        pts = np.atleast_2d([1] * len(comp_count))
    return pts

def make_callable(model, variables, mode=None):
    """
    Take a SymPy object and create a callable function.

    Parameters
    ----------
    model, SymPy object
        Abstract representation of function
    variables, list
        Input variables, ordered in the way the return function will expect
    mode, ['numpy', 'numba', 'sympy'], optional
        Method to use when 'compiling' the function. SymPy mode is
        slow and should only be used for debugging. If Numba is installed,
        it can offer speed-ups when calling the energy function many
        times on multi-core CPUs.

    Returns
    -------
    Function that takes arguments in the same order as 'variables'
    and returns the energy according to 'model'.

    Examples
    --------
    None yet.
    """
    energy = None
    if mode is None:
        # no mode specified; use numba if available, otherwise numpy
        if _NUMBA:
            mode = 'numba'
        else:
            mode = 'numpy'

    if mode == 'sympy':
        energy = lambda *vs: model.subs(zip(variables, vs)).evalf()
    elif mode == 'numpy':
        energy = lambdify(tuple(variables), model, dummify=True,
                          modules=[{'where': np.where, 'Abs': np.abs}, 'numpy'], printer=NumPyPrinter)
    elif mode == 'numba':
        variables = tuple(variables)
        varsig = 'float64({})'.format(','.join(['float64'] * len(variables)))
        energy = lambdify(variables, model, dummify=True,
                          modules=[{'where': nbwhere, 'Abs': np.abs}, 'numpy'], printer=NumPyPrinter)
        # target=parallel seems to incur too much overhead on small arrays
        energy = _NUMBA.vectorize([varsig], nopython=True, target='cpu')(energy)
    else:
        energy = lambdify(tuple(variables), model, dummify=True,
                          modules=mode)

    return energy

def sizeof_fmt(num, suffix='B'):
    """
    Human-readable string for a number of bytes.
    http://stackoverflow.com/questions/1094841/reusable-library-to-get-human-readable-version-of-file-size
    """
    for unit in ['', 'K', 'M', 'G', 'T', 'P', 'E', 'Z']:
        if abs(num) < 1000.0:
            return "%3.1f%s%s" % (num, unit, suffix)
        num /= 1000.0
    return "%.1f%s%s" % (num, 'Y', suffix)

def unpack_condition(tup):
    """
    Convert a condition to a list of values.

    Notes
    -----
    Rules for keys of conditions dicts:
    (1) If it's numeric, treat as a point value
    (2) If it's a tuple with one element, treat as a point value
    (3) If it's a tuple with two elements, treat as lower/upper limits and guess a step size.
    (4) If it's a tuple with three elements, treat as lower/upper/step
    (5) If it's a list, ndarray or other non-tuple ordered iterable, use those values directly.

    """
    if isinstance(tup, tuple):
        if len(tup) == 1:
            return [float(tup[0])]
        elif len(tup) == 2:
            return np.arange(tup[0], tup[1], dtype=np.float)
        elif len(tup) == 3:
            return np.arange(tup[0], tup[1], tup[2], dtype=np.float)
        else:
            raise ValueError('Condition tuple is length {}'.format(len(tup)))
    elif isinstance(tup, collections.Iterable):
        return [float(x) for x in tup]
    else:
        return [float(tup)]

def unpack_phases(phases):
    "Convert a phases list/dict into a sorted list."
    active_phases = None
    if isinstance(phases, (list, tuple, set)):
        active_phases = sorted(phases)
    elif isinstance(phases, dict):
        active_phases = sorted([phn for phn, status in phases.items() \
            if status == 'entered'])
    elif type(phases) is str:
        active_phases = [phases]
    return active_phases

def generate_dof(phase, active_comps):
    """
    Accept a Phase object and a set() of the active components.
    Return a tuple of variable names and the sublattice degrees of freedom.
    """
    variables = []
    sublattice_dof = []
    for idx, sublattice in enumerate(phase.constituents):
        dof = 0
        for component in sorted(set(sublattice).intersection(active_comps)):
            variables.append(v.SiteFraction(phase.name.upper(), idx, component))
            dof += 1
        sublattice_dof.append(dof)
    return variables, sublattice_dof

def endmember_matrix(dof, vacancy_indices=None):
    """
    Accept the number of components in each sublattice.
    Return a matrix corresponding to the compositions of all endmembers.

    Parameters
    ----------
    dof : list of int
        Number of components in each sublattice.
    vacancy_indices, list of int, optional
        If vacancies are present in every sublattice, specify their indices
        in each sublattice to ensure the "pure vacancy" endmember is excluded.

    Examples
    --------
    Sublattice configuration like: `(AL, NI, VA):(AL, NI, VA):(VA)`
    >>> endmember_matrix([3,3,1], vacancy_indices=[2, 2, 0])
    """
    total_endmembers = functools.reduce(operator.mul, dof, 1)
    res_matrix = np.empty((total_endmembers, sum(dof)), dtype=np.float)
    dof_arrays = [np.eye(d).tolist() for d in dof]
    row_idx = 0
    for row in itertools.product(*dof_arrays):
        res_matrix[row_idx, :] = np.concatenate(row, axis=0)
        row_idx += 1
    if vacancy_indices is not None and len(vacancy_indices) == len(dof):
        dof_adj = np.array([sum(dof[0:i]) for i in range(len(dof))])
        indices = np.array(vacancy_indices) + dof_adj
        row_idx_to_delete = np.where(np.all(res_matrix[:, indices] == 1,
                                            axis=1))
        res_matrix = np.delete(res_matrix, (row_idx_to_delete), axis=0)
    # Adjust site fractions to the numerical limit
    cur_idx = 0
    res_matrix[res_matrix == 0] = MIN_SITE_FRACTION
    for ctx in dof:
        end_idx = cur_idx + ctx
        res_matrix[:, cur_idx:end_idx] /= res_matrix[:, cur_idx:end_idx].sum(axis=1)[:, None]
        cur_idx = end_idx
    return res_matrix

def unpack_kwarg(kwarg_obj, default_arg=None):
    """
    Keyword arguments in pycalphad can be passed as a constant value, a
    dict of phase names and values, or a list containing both of these. If
    the latter, then the dict is checked first; if the phase of interest is not
    there, then the constant value is used.

    This function is a way to construct defaultdicts out of keyword arguments.

    Parameters
    ----------
    kwarg_obj : dict, iterable, or None
        Argument to unpack into a defaultdict
    default_arg : object
        Default value to use if iterable isn't specified

    Returns
    -------
    defaultdict for the keyword argument of interest

    Examples
    --------
    >>> test_func = lambda **kwargs: print(unpack_kwarg('opt'))
    >>> test_func(opt=100)
    >>> test_func(opt={'FCC_A1': 50, 'BCC_B2': 10})
    >>> test_func(opt=[{'FCC_A1': 30}, 200])
    >>> test_func()
    >>> test_func2 = lambda **kwargs: print(unpack_kwarg('opt', default_arg=1))
    >>> test_func2()
    """
    new_dict = collections.defaultdict(lambda: default_arg)

    if isinstance(kwarg_obj, collections.Mapping):
        new_dict.update(kwarg_obj)
    # kwarg_obj is a list containing a dict and a default
    # For now at least, we don't treat ndarrays the same as other iterables
    # ndarrays are assumed to be numeric arrays containing "default values", so don't match here
    elif isinstance(kwarg_obj, collections.Iterable) and not isinstance(kwarg_obj, np.ndarray):
        for element in kwarg_obj:
            if isinstance(element, collections.Mapping):
                new_dict.update(element)
            else:
                # element=element syntax to silence var-from-loop warning
                new_dict = collections.defaultdict(
                    lambda element=element: element, new_dict)
    elif kwarg_obj is None:
        pass
    else:
        new_dict = collections.defaultdict(lambda: kwarg_obj)

    return new_dict

<commit_msg>MAINT: Remove numba support<commit_after>"""
The utils module handles helper routines for equilibrium calculation.
"""
from __future__ import division
import pycalphad.variables as v
from pycalphad.core.halton import halton
from pycalphad.core.constants import MIN_SITE_FRACTION
from sympy.utilities.lambdify import lambdify
from sympy.printing.lambdarepr import LambdaPrinter
import numpy as np
import operator
import functools
import itertools
import collections

try:
    # Only available in numpy 1.10 and newer
    from numpy import broadcast_to
except ImportError:
    def broadcast_to(arr, shape):
        "Broadcast an array to a desired shape. Returns a view."
        return np.broadcast_arrays(arr, np.empty(shape, dtype=np.bool))[0]

class NumPyPrinter(LambdaPrinter):
    """
    Numpy printer which handles vectorized piecewise functions,
    logical operators, etc.
    """
    _default_settings = {
        "order": "none",
        "full_prec": "auto",
    }

    def _print_seq(self, seq, delimiter=', '):
        "General sequence printer: converts to tuple"
        # Print tuples here instead of lists because numba supports
        #     tuples in nopython mode.
        return '({},)'.format(delimiter.join(self._print(item) for item in seq))

    #def _print_Integer(self, expr):
    #    # Upcast to work around an integer overflow bug
    #    return 'asarray({0}, dtype=float)'.format(expr)

    def _print_MatrixBase(self, expr):
        return "%s(%s)" % ('array', self._print(expr.tolist()))

    _print_SparseMatrix = \
        _print_MutableSparseMatrix = \
        _print_ImmutableSparseMatrix = \
        _print_Matrix = \
        _print_DenseMatrix = \
        _print_MutableDenseMatrix = \
        _print_ImmutableMatrix = \
        _print_ImmutableDenseMatrix = \
        _print_MatrixBase

    def _print_MatMul(self, expr):
        "Matrix multiplication printer"
        return '({0})'.format(').dot('.join(self._print(i) for i in expr.args))

    def _print_Piecewise(self, expr):
        "Piecewise function printer"
        exprs = [self._print(arg.expr) for arg in expr.args]
        conds = [self._print(arg.cond) for arg in expr.args]
        if (len(conds) == 1) and exprs[0] == '0':
            return '0'
        # We expect exprs and conds to be in order here
        # First condition to evaluate to True is returned
        # Default result: float zero
        result = 'where({0}, {1}, 0.)'.format(conds[-1], exprs[-1])
        for expr, cond in reversed(list(zip(exprs[:-1], conds[:-1]))):
            result = 'where({0}, {1}, {2})'.format(cond, expr, result)

        return result

    def _print_And(self, expr):
        "Logical And printer"
        # We have to override LambdaPrinter because it uses Python 'and' keyword.
        # If LambdaPrinter didn't define it, we could use StrPrinter's
        # version of the function and add 'logical_and' to NUMPY_TRANSLATIONS.
        if len(expr.args) == 1:
            return self._print(expr.args[0])
        result = 'logical_and({0}, {1})'.format(self._print(expr.args[-2]), self._print(expr.args[-1]))
        for arg in reversed(expr.args[:-2]):
            result = 'logical_and({0}, {1})'.format(self._print(arg), result)
        return  result

    def _print_Or(self, expr):
        "Logical Or printer"
        # We have to override LambdaPrinter because it uses Python 'or' keyword.
        # If LambdaPrinter didn't define it, we could use StrPrinter's
        # version of the function and add 'logical_or' to NUMPY_TRANSLATIONS.
        if len(expr.args) == 1:
            return self._print(expr.args[0])
        result = 'logical_or({0}, {1})'.format(self._print(expr.args[-2]), self._print(expr.args[-1]))
        for arg in reversed(expr.args[:-2]):
            result = 'logical_or({0}, {1})'.format(self._print(arg), result)
        return  result

    def _print_Not(self, expr):
        "Logical Not printer"
        # We have to override LambdaPrinter because it uses Python 'not' keyword.
        # If LambdaPrinter didn't define it, we would still have to define our
        #     own because StrPrinter doesn't define it.
        if len(expr.args) == 1:
            return self._print(expr.args[0])
        result = 'logical_not({0}, {1})'.format(self._print(expr.args[-2]), self._print(expr.args[-1]))
        for arg in reversed(expr.args[:-2]):
            result = 'logical_not({0}, {1})'.format(self._print(arg), result)
        return  result

def point_sample(comp_count, pdof=10):
    """
    Sample 'pdof * (sum(comp_count) - len(comp_count))' points in
    composition space for the sublattice configuration specified
    by 'comp_count'. Points are sampled quasi-randomly from a Halton sequence.
    A Halton sequence is like a uniform random distribution, but the
    result will always be the same for a given 'comp_count' and 'pdof'.
    Note: For systems with only one component, only one point will be
    returned, regardless of 'pdof'. This is because the degrees of freedom
    are zero for that case.

    Parameters
    ----------
    comp_count : list
        Number of components in each sublattice.
    pdof : int
        Number of points to sample per degree of freedom.

    Returns
    -------
    ndarray of generated points satisfying the mass balance.

    Examples
    --------
    >>> comps = [8,1] # 8 components in sublattice 1; only 1 in sublattice 2
    >>> pts = point_sample(comps, pdof=20) # 7 d.o.f, returns a 140x7 ndarray
    """
    # Generate Halton sequence with appropriate dimensions and size
    pts = halton(sum(comp_count),
                 pdof * (sum(comp_count) - len(comp_count)), scramble=True)
    # Convert low-discrepancy sequence to normalized exponential
    # This will be uniformly distributed over the simplices
    pts = -np.log(pts)
    cur_idx = 0
    for ctx in comp_count:
        end_idx = cur_idx + ctx
        pts[:, cur_idx:end_idx] /= pts[:, cur_idx:end_idx].sum(axis=1)[:, None]
        cur_idx = end_idx

    if len(pts) == 0:
        pts = np.atleast_2d([1] * len(comp_count))
    return pts

def make_callable(model, variables, mode=None):
    """
    Take a SymPy object and create a callable function.

    Parameters
    ----------
    model, SymPy object
        Abstract representation of function
    variables, list
        Input variables, ordered in the way the return function will expect
    mode, ['numpy', 'numba', 'sympy'], optional
        Method to use when 'compiling' the function. SymPy mode is
        slow and should only be used for debugging. If Numba is installed,
        it can offer speed-ups when calling the energy function many
        times on multi-core CPUs.

    Returns
    -------
    Function that takes arguments in the same order as 'variables'
    and returns the energy according to 'model'.

    Examples
    --------
    None yet.
    """
    energy = None
    if mode is None:
        mode = 'numpy'

    if mode == 'sympy':
        energy = lambda *vs: model.subs(zip(variables, vs)).evalf()
    elif mode == 'numpy':
        energy = lambdify(tuple(variables), model, dummify=True,
                          modules=[{'where': np.where, 'Abs': np.abs}, 'numpy'], printer=NumPyPrinter)
    else:
        energy = lambdify(tuple(variables), model, dummify=True,
                          modules=mode)

    return energy

def sizeof_fmt(num, suffix='B'):
    """
    Human-readable string for a number of bytes.
    http://stackoverflow.com/questions/1094841/reusable-library-to-get-human-readable-version-of-file-size
    """
    for unit in ['', 'K', 'M', 'G', 'T', 'P', 'E', 'Z']:
        if abs(num) < 1000.0:
            return "%3.1f%s%s" % (num, unit, suffix)
        num /= 1000.0
    return "%.1f%s%s" % (num, 'Y', suffix)

def unpack_condition(tup):
    """
    Convert a condition to a list of values.

    Notes
    -----
    Rules for keys of conditions dicts:
    (1) If it's numeric, treat as a point value
    (2) If it's a tuple with one element, treat as a point value
    (3) If it's a tuple with two elements, treat as lower/upper limits and guess a step size.
    (4) If it's a tuple with three elements, treat as lower/upper/step
    (5) If it's a list, ndarray or other non-tuple ordered iterable, use those values directly.

    """
    if isinstance(tup, tuple):
        if len(tup) == 1:
            return [float(tup[0])]
        elif len(tup) == 2:
            return np.arange(tup[0], tup[1], dtype=np.float)
        elif len(tup) == 3:
            return np.arange(tup[0], tup[1], tup[2], dtype=np.float)
        else:
            raise ValueError('Condition tuple is length {}'.format(len(tup)))
    elif isinstance(tup, collections.Iterable):
        return [float(x) for x in tup]
    else:
        return [float(tup)]

def unpack_phases(phases):
    "Convert a phases list/dict into a sorted list."
    active_phases = None
    if isinstance(phases, (list, tuple, set)):
        active_phases = sorted(phases)
    elif isinstance(phases, dict):
        active_phases = sorted([phn for phn, status in phases.items() \
            if status == 'entered'])
    elif type(phases) is str:
        active_phases = [phases]
    return active_phases

def generate_dof(phase, active_comps):
    """
    Accept a Phase object and a set() of the active components.
    Return a tuple of variable names and the sublattice degrees of freedom.
    """
    variables = []
    sublattice_dof = []
    for idx, sublattice in enumerate(phase.constituents):
        dof = 0
        for component in sorted(set(sublattice).intersection(active_comps)):
            variables.append(v.SiteFraction(phase.name.upper(), idx, component))
            dof += 1
        sublattice_dof.append(dof)
    return variables, sublattice_dof

def endmember_matrix(dof, vacancy_indices=None):
    """
    Accept the number of components in each sublattice.
    Return a matrix corresponding to the compositions of all endmembers.

    Parameters
    ----------
    dof : list of int
        Number of components in each sublattice.
    vacancy_indices, list of int, optional
        If vacancies are present in every sublattice, specify their indices
        in each sublattice to ensure the "pure vacancy" endmember is excluded.

    Examples
    --------
    Sublattice configuration like: `(AL, NI, VA):(AL, NI, VA):(VA)`
    >>> endmember_matrix([3,3,1], vacancy_indices=[2, 2, 0])
    """
    total_endmembers = functools.reduce(operator.mul, dof, 1)
    res_matrix = np.empty((total_endmembers, sum(dof)), dtype=np.float)
    dof_arrays = [np.eye(d).tolist() for d in dof]
    row_idx = 0
    for row in itertools.product(*dof_arrays):
        res_matrix[row_idx, :] = np.concatenate(row, axis=0)
        row_idx += 1
    if vacancy_indices is not None and len(vacancy_indices) == len(dof):
        dof_adj = np.array([sum(dof[0:i]) for i in range(len(dof))])
        indices = np.array(vacancy_indices) + dof_adj
        row_idx_to_delete = np.where(np.all(res_matrix[:, indices] == 1,
                                            axis=1))
        res_matrix = np.delete(res_matrix, (row_idx_to_delete), axis=0)
    # Adjust site fractions to the numerical limit
    cur_idx = 0
    res_matrix[res_matrix == 0] = MIN_SITE_FRACTION
    for ctx in dof:
        end_idx = cur_idx + ctx
        res_matrix[:, cur_idx:end_idx] /= res_matrix[:, cur_idx:end_idx].sum(axis=1)[:, None]
        cur_idx = end_idx
    return res_matrix

def unpack_kwarg(kwarg_obj, default_arg=None):
    """
    Keyword arguments in pycalphad can be passed as a constant value, a
    dict of phase names and values, or a list containing both of these. If
    the latter, then the dict is checked first; if the phase of interest is not
    there, then the constant value is used.

    This function is a way to construct defaultdicts out of keyword arguments.

    Parameters
    ----------
    kwarg_obj : dict, iterable, or None
        Argument to unpack into a defaultdict
    default_arg : object
        Default value to use if iterable isn't specified

    Returns
    -------
    defaultdict for the keyword argument of interest

    Examples
    --------
    >>> test_func = lambda **kwargs: print(unpack_kwarg('opt'))
    >>> test_func(opt=100)
    >>> test_func(opt={'FCC_A1': 50, 'BCC_B2': 10})
    >>> test_func(opt=[{'FCC_A1': 30}, 200])
    >>> test_func()
    >>> test_func2 = lambda **kwargs: print(unpack_kwarg('opt', default_arg=1))
    >>> test_func2()
    """
    new_dict = collections.defaultdict(lambda: default_arg)

    if isinstance(kwarg_obj, collections.Mapping):
        new_dict.update(kwarg_obj)
    # kwarg_obj is a list containing a dict and a default
    # For now at least, we don't treat ndarrays the same as other iterables
    # ndarrays are assumed to be numeric arrays containing "default values", so don't match here
    elif isinstance(kwarg_obj, collections.Iterable) and not isinstance(kwarg_obj, np.ndarray):
        for element in kwarg_obj:
            if isinstance(element, collections.Mapping):
                new_dict.update(element)
            else:
                # element=element syntax to silence var-from-loop warning
                new_dict = collections.defaultdict(
                    lambda element=element: element, new_dict)
    elif kwarg_obj is None:
        pass
    else:
        new_dict = collections.defaultdict(lambda: kwarg_obj)

    return new_dict

<|endoftext|>