示例#1
0
def make_plots(prefix,
               im,
               tr=None,
               plots=['data', 'model', 'chi'],
               mags=['i'],
               radecroi_in=None):
    import pylab as plt
    import tractor

    if radecroi_in:
        global radecroi
        radecroi = radecroi_in

    if tr is None:
        srcs = get_cfht_catalog(mags=mags)
        tr = tractor.Tractor(tractor.Images(im), srcs)

    mod = tr.getModelImage(im)

    ima = dict(interpolation='nearest',
               origin='lower',
               vmin=im.zr[0],
               vmax=im.zr[1])
    imchi = dict(interpolation='nearest', origin='lower', vmin=-5, vmax=+5)
    if hasattr(im, 'extent'):
        ima.update(extent=im.extent)
        imchi.update(extent=im.extent)

    fignum = 1

    if 'data' in plots:
        plt.figure(fignum)
        fignum += 1
        plt.clf()
        plt.imshow(im.getImage(), **ima)
        plt.gray()
        plt.title(im.name + ': data')
        plt.savefig(prefix + 'data.png')

    if 'model' in plots:
        plt.figure(fignum)
        fignum += 1
        plt.clf()
        plt.imshow(mod, **ima)
        plt.gray()
        plt.title(im.name + ': model')
        plt.savefig(prefix + 'mod.png')

    if 'chi' in plots:
        plt.figure(fignum)
        fignum += 1
        plt.clf()
        plt.imshow((mod - im.getImage()) * im.getInvError(), **imchi)
        plt.gray()
        plt.title(im.name + ': chi')
        plt.savefig(prefix + 'chi.png')
示例#2
0
    def fit_general_gaussian(self,
                             img,
                             sig1,
                             xi,
                             yi,
                             fluxi,
                             psf_r=15,
                             ps=None):
        import tractor
        H, W = img.shape
        ix = int(np.round(xi))
        iy = int(np.round(yi))
        xlo = max(0, ix - psf_r)
        xhi = min(W, ix + psf_r + 1)
        ylo = max(0, iy - psf_r)
        yhi = min(H, iy + psf_r + 1)
        xx, yy = np.meshgrid(np.arange(xlo, xhi), np.arange(ylo, yhi))
        r2 = (xx - xi)**2 + (yy - yi)**2
        keep = (r2 < psf_r**2)
        pix = img[ylo:yhi, xlo:xhi].copy()
        ie = np.zeros_like(pix)
        ie[keep] = 1. / sig1

        psf = tractor.NCircularGaussianPSF([4.], [1.])
        tim = tractor.Image(data=pix, inverr=ie, psf=psf)
        src = tractor.PointSource(tractor.PixPos(xi - xlo, yi - ylo),
                                  tractor.Flux(fluxi))
        tr = tractor.Tractor([tim], [src])

        src.pos.addGaussianPrior('x', 0., 1.)

        doplot = (ps is not None)
        if doplot:
            mod0 = tr.getModelImage(0)

        tim.freezeAllBut('psf')
        psf.freezeAllBut('sigmas')

        # print('Optimizing params:')
        # tr.printThawedParams()

        #print('Parameter step sizes:', tr.getStepSizes())
        optargs = dict(priors=False, shared_params=False)
        for step in range(50):
            dlnp, x, alpha = tr.optimize(**optargs)
            if dlnp == 0:
                break

        # Now fit only the PSF size
        tr.freezeParam('catalog')
        # print('Optimizing params:')
        # tr.printThawedParams()

        for step in range(50):
            dlnp, x, alpha = tr.optimize(**optargs)
            if dlnp == 0:
                break

        # fwhms.append(psf.sigmas[0] * 2.35 * self.pixscale)

        if doplot:
            mod1 = tr.getModelImage(0)
            chi1 = tr.getChiImage(0)

        # Now switch to a non-isotropic PSF
        s = psf.sigmas[0]
        #print('Isotropic fit sigma', s)
        s = np.clip(s, 1., 5.)
        tim.psf = tractor.GaussianMixturePSF(1., 0., 0., s**2, s**2, 0.)

        #print('Optimizing params:')
        #tr.printThawedParams()

        try:
            for step in range(50):
                dlnp, x, alpha = tr.optimize(**optargs)
                #print('PSF:', tim.psf)
                if dlnp == 0:
                    break
        except:
            import traceback
            print('Error during fitting PSF in a focus frame; not to worry')
            traceback.print_exc()
            print(
                '(The above was just an error during fitting one star in a focus frame; not to worry.)'
            )

        # Don't need to re-fit source params because PSF ampl and mean
        # can fit for flux and position.

        if doplot:
            mod2 = tr.getModelImage(0)
            chi2 = tr.getChiImage(0)
            kwa = dict(vmin=-3 * sig1, vmax=50 * sig1, cmap='gray')

            plt.clf()
            plt.subplot(2, 3, 1)
            plt.title('Image')
            dimshow(pix, ticks=False, **kwa)
            plt.subplot(2, 3, 2)
            plt.title('Initial model')
            dimshow(mod0, ticks=False, **kwa)
            plt.subplot(2, 3, 3)
            plt.title('Isotropic model')
            dimshow(mod1, ticks=False, **kwa)
            plt.subplot(2, 3, 4)
            plt.title('Final model')
            dimshow(mod2, ticks=False, **kwa)
            plt.subplot(2, 3, 5)
            plt.title('Isotropic chi')
            dimshow(chi1, vmin=-10, vmax=10, ticks=False)
            plt.subplot(2, 3, 6)
            plt.title('Final chi')
            dimshow(chi2, vmin=-10, vmax=10, ticks=False)
            plt.suptitle('PSF fit')
            ps.savefig()

        return tim.psf.getParams()[-3:]
示例#3
0
def main():

    # In LSST meas-deblend (on lsst6):
    # python examples/suprimePlot.py --data ~dstn/lsst/ACT-data -v 126969 -c 5 --data-range -100 300 --roi 0 500 0 500 --psf psf.fits --image img.fits --sources srcs.fits

    from optparse import OptionParser
    import sys

    parser = OptionParser(usage=('%prog <img> <psf> <srcs>'))
    parser.add_option('-v',
                      '--verbose',
                      dest='verbose',
                      action='count',
                      default=0,
                      help='Make more verbose')
    opt, args = parser.parse_args()

    if len(args) != 3:
        parser.print_help()
        sys.exit(-1)

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

    imgfn, psffn, srcfn = args

    pimg = pyfits.open(imgfn)
    if len(pimg) != 4:
        print('Image must have 3 extensions')
        sys.exit(-1)
    img = pimg[1].data
    mask = pimg[2].data
    maskhdr = pimg[2].header
    var = pimg[3].data
    del pimg

    print('var', var.shape)
    #print var
    print('mask', mask.shape)
    #print mask
    print('img', img.shape)
    #print img

    mask = mask.astype(np.int16)
    for bit in range(16):
        on = ((mask & (1 << bit)) != 0)
        print('Bit', bit, 'has', np.sum(on), 'pixels set')
    '''
	MP_BAD  =                    0
	Bit 0 has 2500 pixels set
	MP_SAT  =                    1
	Bit 1 has 5771 pixels set
	MP_INTRP=                    2
	Bit 2 has 11269 pixels set
	MP_CR   =                    3
	Bit 3 has 136 pixels set
	MP_EDGE =                    4
	Bit 4 has 11856 pixels set
	HIERARCH MP_DETECTED =       5
	Bit 5 has 37032 pixels set
	'''

    print('Mask header:', maskhdr)
    maskplanes = {}
    print('Mask planes:')
    for card in maskhdr.ascardlist():
        if not card.key.startswith('MP_'):
            continue
        print(card.value, card.key)
        maskplanes[card.key[3:]] = card.value

    print('Variance range:', var.min(), var.max())

    print('Image median:', np.median(img.ravel()))

    invvar = 1. / var
    invvar[var == 0] = 0.
    invvar[var < 0] = 0.

    sig = np.sqrt(np.median(var))
    H, W = img.shape
    for k, v in maskplanes.items():
        plt.clf()

        I = ((mask & (1 << v)) != 0)
        rgb = np.zeros((H, W, 3))
        clipimg = np.clip((img - (-3. * sig)) / (13. * sig), 0, 1)
        cimg = clipimg.copy()
        cimg[I] = 1
        rgb[:, :, 0] = cimg
        cimg = clipimg.copy()
        cimg[I] = 0
        rgb[:, :, 1] = cimg
        rgb[:, :, 2] = cimg
        plt.imshow(rgb, interpolation='nearest', origin='lower')
        plt.title(k)
        plt.savefig('mask-%s.png' % k.lower())

    badmask = sum([(1 << maskplanes[k])
                   for k in ['BAD', 'SAT', 'INTRP', 'CR']])
    # HACK -- left EDGE sucks
    badmask += (1 << maskplanes['EDGE'])
    #badmask = (1 << 0) | (1 << 1) | (1 << 2) | (1 << 3)
    #badmask |= (1 << 4)
    print('Masking out: 0x%x' % badmask)
    invvar[(mask & badmask) != 0] = 0.

    assert (all(np.isfinite(img.ravel())))
    assert (all(np.isfinite(invvar.ravel())))

    psf = pyfits.open(psffn)[0].data
    print('psf', psf.shape)
    psf /= psf.sum()

    from tractor.emfit import em_fit_2d
    from tractor.fitpsf import em_init_params

    # Create Gaussian mixture model PSF approximation.
    S = psf.shape[0]
    # number of Gaussian components
    K = 3
    w, mu, sig = em_init_params(K, None, None, None)
    II = psf.copy()
    II /= II.sum()
    # HIDEOUS HACK
    II = np.maximum(II, 0)
    print('Multi-Gaussian PSF fit...')
    xm, ym = -(S / 2), -(S / 2)
    em_fit_2d(II, xm, ym, w, mu, sig)
    print('w,mu,sig', w, mu, sig)
    mypsf = tractor.GaussianMixturePSF(w, mu, sig)

    P = mypsf.getPointSourcePatch(S / 2, S / 2)
    mn, mx = psf.min(), psf.max()
    ima = dict(interpolation='nearest', origin='lower', vmin=mn, vmax=mx)
    plt.clf()
    plt.subplot(1, 2, 1)
    plt.imshow(psf, **ima)
    plt.subplot(1, 2, 2)
    pimg = np.zeros_like(psf)
    P.addTo(pimg)
    plt.imshow(pimg, **ima)
    plt.savefig('psf.png')

    sig = np.sqrt(np.median(var))

    plt.clf()
    plt.hist(img.ravel(), 100, range=(-3. * sig, 3. * sig))
    plt.savefig('imghist.png')

    srcs = fits_table(srcfn)
    print('Initial:', len(srcs), 'sources')
    # Trim sources with x=0 or y=0
    srcs = srcs[(srcs.x != 0) * (srcs.y != 0)]
    print('Trim on x,y:', len(srcs), 'sources left')
    # Zero out nans & infs
    for c in ['theta', 'a', 'b']:
        I = np.logical_not(np.isfinite(srcs.get(c)))
        srcs.get(c)[I] = 0.
    # Set sources with flux=NaN to something more sensible...
    I = np.logical_not(np.isfinite(srcs.flux))
    srcs.flux[I] = 1.
    # Sort sources by flux.
    srcs = srcs[np.argsort(-srcs.flux)]

    # Trim sources that are way outside the image.
    margin = 8. * np.maximum(srcs.a, srcs.b)
    H, W = img.shape
    srcs = srcs[(srcs.x > -margin) * (srcs.y > -margin) *
                (srcs.x < (W + margin) * (srcs.y < (H + margin)))]
    print('Trim out-of-bounds:', len(srcs), 'sources left')

    wcs = tractor.FitsWcs(Sip(imgfn, 1))
    #wcs = tractor.NullWCS()

    timg = tractor.Image(data=img,
                         invvar=invvar,
                         psf=mypsf,
                         wcs=wcs,
                         sky=tractor.ConstantSky(0.),
                         photocal=tractor.NullPhotoCal(),
                         name='image')

    inverr = timg.getInvError()
    assert (all(np.isfinite(inverr.ravel())))

    tsrcs = []
    for s in srcs:
        #pos = tractor.PixPos(s.x, s.y)
        pos = tractor.RaDecPos(s.ra, s.dec)
        if s.a > 0 and s.b > 0:
            eflux = tractor.Flux(s.flux / 2.)
            dflux = tractor.Flux(s.flux / 2.)
            re, ab, phi = s.a, s.b / s.a, 90. - s.theta
            eshape = gal.GalaxyShape(re, ab, phi)
            dshape = gal.GalaxyShape(re, ab, phi)
            print('Fluxes', eflux, dflux)
            tsrc = gal.CompositeGalaxy(pos, eflux, eshape, dflux, dshape)
        else:
            flux = tractor.Flux(s.flux)
            print('Flux', flux)
            tsrc = tractor.PointSource(pos, flux)
        tsrcs.append(tsrc)

    chug = tractor.Tractor([timg])
    for src in tsrcs:
        if chug.getModelPatch(timg, src) is None:
            print('Dropping non-overlapping source:', src)
            continue
        chug.addSource(src)
    print('Kept a total of', len(chug.catalog), 'sources')

    ima = dict(interpolation='nearest',
               origin='lower',
               vmin=-3. * sig,
               vmax=10. * sig)
    chia = dict(interpolation='nearest', origin='lower', vmin=-5., vmax=5.)

    plt.clf()
    plt.imshow(img, **ima)
    plt.colorbar()
    plt.savefig('img.png')

    plt.clf()
    plt.imshow(invvar, interpolation='nearest', origin='lower')
    plt.colorbar()
    plt.savefig('invvar.png')

    mod = chug.getModelImages()[0]
    plt.clf()
    plt.imshow(mod, **ima)
    plt.colorbar()
    plt.savefig('mod-0.png')

    chi = chug.getChiImage(0)
    plt.clf()
    plt.imshow(chi, **chia)
    plt.colorbar()
    plt.savefig('chi-0.png')

    for step in range(5):
        cat = chug.getCatalog()
        for src in cat:
            if chug.getModelPatch(timg, src) is None:
                print('Dropping non-overlapping source:', src)
                chug.removeSource(src)
        print('Kept a total of', len(chug.catalog), 'sources')

        #cat = chug.getCatalog()
        #for i,src in enumerate([]):
        #for i,src in enumerate(chug.getCatalog()):
        #for i in range(len(cat)):
        i = 0
        while i < len(cat):
            src = cat[i]

            #print 'Step', i
            #for j,s in enumerate(cat):
            #	x,y = timg.getWcs().positionToPixel(s, s.getPosition())
            #	print '  ',
            #	if j == i:
            #		print '*',
            #	print '(%6.1f, %6.1f)'%(x,y), s

            print('Optimizing source', i, 'of', len(cat))

            x, y = timg.getWcs().positionToPixel(src.getPosition(), src)
            print('(%6.1f, %6.1f)' % (x, y), src)
            # pre = src.getModelPatch(timg)

            s1 = str(src)
            print('src1 ', s1)
            dlnp1, X, a = chug.optimizeCatalogFluxes(srcs=[src])
            s2 = str(src)
            dlnp2, X, a = chug.optimizeCatalogAtFixedComplexityStep(srcs=[src],
                                                                    sky=False)
            s3 = str(src)

            #post = src.getModelPatch(timg)

            print('src1 ', s1)
            print('src2 ', s2)
            print('src3 ', s3)
            print('dlnp', dlnp1, dlnp2)

            if chug.getModelPatch(timg, src) is None:
                print('After optimizing, no overlap!')
                print('Removing source', src)
                chug.removeSource(src)
                i -= 1
            i += 1

            # plt.clf()
            # plt.subplot(2,2,1)
            # img = timg.getImage()
            # (x0,x1,y0,y1) = pre.getExtent()
            # plt.imshow(img, **ima)
            # ax = plt.axis()
            # plt.plot([x0,x0,x1,x1,x0], [y0,y1,y1,y0,y0], 'k-', lw=2)
            # plt.axis(ax)
            # plt.subplot(2,2,3)
            # plt.imshow(pre.getImage(), **ima)
            # plt.subplot(2,2,4)
            # plt.imshow(post.getImage(), **ima)
            # plt.savefig('prepost-s%i-s%03i.png' % (step, i))
            #
            # mod = chug.getModelImages()[0]
            # plt.clf()
            # plt.imshow(mod, **ima)
            # plt.colorbar()
            # plt.savefig('mod-s%i-s%03i.png' % (step, i))
            # chi = chug.getChiImage(0)
            # plt.clf()
            # plt.imshow(chi, **chia)
            # plt.colorbar()
            # plt.savefig('chi-s%i-s%03i.png' % (step, i))

        #dlnp,x,a = chug.optimizeCatalogFluxes()
        #print 'fluxes: dlnp', dlnp
        #dlnp,x,a = chug.optimizeCatalogAtFixedComplexityStep()
        #print 'opt: dlnp', dlnp

        mod = chug.getModelImages()[0]
        plt.clf()
        plt.imshow(mod, **ima)
        plt.colorbar()
        plt.savefig('mod-%i.png' % (step + 1))

        chi = chug.getChiImage(0)
        plt.clf()
        plt.imshow(chi, **chia)
        plt.colorbar()
        plt.savefig('chi-%i.png' % (step + 1))

    return

    for step in range(5):
        chug.optimizeCatalogFluxes()
        mod = chug.getModelImages()[0]
        plt.clf()
        plt.imshow(mod, **ima)
        plt.colorbar()
        plt.savefig('mod-s%i.png' % step)

        chi = chug.getChiImage(0)
        plt.clf()
        plt.imshow(chi, **chia)
        plt.colorbar()
        plt.savefig('chi-s%i.png' % step)
def main():
    """
   NAME
     LensTractor.py

   PURPOSE
     Run the Tractor on a deck of single object cutout images.
     Read in an image and its weight map, guess the PSF, put an object at the
     centre image of the field and then optimize the catalog and PSF.

   COMMENTS

   FLAGS
     -h --help        Print this message
     -v --verbose     Verbose operation
     -s --sample      Sample the posterior PDF instead of optimizing
     -x --no-plots    Do not plot progress
     -l --lens        Only fit lens model
     -n --nebula      Only fit nebula model

   INPUTS
     *.fits           Deck of postcard images

   OPTIONAL INPUTS

   OUTPUTS
     stdout                       Useful information
     *.png                        Plots in png format
     
     To be implemented:
       lenstractor_progress.log     Logged output
       lenstractor_results.txt      Model comparison results
       lenstractor_lens.cat         Lens model parameters, including lightcurves
       lenstractor_nebula.cat       Nebula model parameters, including lightcurves

   EXAMPLES

     python LensTractor.py -x examples/H1413+117_10x10arcsec_55*fits > examples/H1413+117_10x10arcsec_lenstractor.log
   
   DEPENDENCIES
     * The Tractor     astrometry.net/svn/trunk/projects/tractor
     * emcee           github.com/danfm/emcee
     * astrometry.net  astrometry.net/svn/trunk/util

   BUGS
     - PSFs not being optimized correctly - permafrost?

   HISTORY
     2012-07-06       First predicted Lens images Marshall/Hogg (Oxford/NYU)
   """

    # --------------------------------------------------------------------

    from optparse import OptionParser
    import sys

    # Set available options:
    parser = OptionParser(usage=('%prog *.fits'))
    # Verbosity:
    parser.add_option('-v',
                      '--verbose',
                      dest='verbose',
                      action='count',
                      default=False,
                      help='Make more verbose')
    vb = True  # for usual outputs.
    # Sampling:
    parser.add_option('-s',
                      '--sample',
                      dest='MCMC',
                      action='count',
                      default=False,
                      help='Sample posterior PDF')
    # Plotting:
    parser.add_option('-x',
                      '--no-plots',
                      dest='noplots',
                      action='count',
                      default=False,
                      help='Skip plotting')
    # Lens model only:
    parser.add_option('-l',
                      '--lens',
                      dest='lens',
                      action='count',
                      default=False,
                      help='Fit lens model')
    # Nebula model only:
    parser.add_option('-n',
                      '--nebula',
                      dest='nebula',
                      action='count',
                      default=False,
                      help='Fit nebula model')

    # Read in options and arguments - note only sci and wht images are supplied:
    opt, args = parser.parse_args()

    if len(args) < 2:
        parser.print_help()
        sys.exit(-1)

    # The rest of the command line is assumed to be a list of files:
    inputfiles = args

    # Workflow:
    if opt.lens:
        models = ['lens']
    elif opt.nebula:
        models = ['nebula']
    else:
        models = ['nebula', 'lens']
    BIC = dict(zip(models, np.zeros(len(models))))
    # NB. default operation is to fit both and compare.
    # Do nebula first: PSF and sky then roll over into lens.

    if vb:
        print "* * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * *"
        print "                               LensTractor "
        print "    Fitting", models, " models to a deck of FITS postcards"
        print "* * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * *"

    # -------------------------------------------------------------------------

    # Organise the deck of inputfiles into scifiles and varfiles:
    scifiles, varfiles = lenstractor.Riffle(inputfiles, vb=vb)

    # Read into Tractor Image objects, and see what filters we have:
    images, total_mags, bands = lenstractor.Deal(scifiles,
                                                 varfiles,
                                                 SURVEY='PS1',
                                                 vb=vb)

    # -------------------------------------------------------------------------
    # Generic items needed to initialize the Tractor's catalog.

    # Get rough idea of object position from wcs of first image- works
    # well if all images are the same size and well registered!
    wcs = images[0].wcs
    NX, NY = np.shape(images[0].data)
    if vb: print "Generic initial position ", NX, NY, "(pixels)"

    # m0 = 15.0
    # magnitudes = m0*np.ones(len(bandnames))
    # MAGIC initial magnitude. This should be estimated from
    # the total flux in each (background-subtracted) image...

    bandnames = np.unique(bands)
    magnitudes = np.zeros(len(bandnames))
    for i, bandname in enumerate(bandnames):
        index = np.where(bands == bandname)
        magnitudes[i] = np.median(total_mags[index])
    if vb: print "Generic initial SED ", dict(zip(bandnames, magnitudes))

    # -------------------------------------------------------------------------

    # Loop over models:

    for model in models:

        if vb:
            print "* * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * *"
            print "Initializing model: " + model

        # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -

        if model == 'nebula':

            # srcs = [InitializeNebula(wcs,bandnames)]

            # Nebula - a flexible galaxy plus four point sources,
            #   5 sources in total. Start them off with equal magnitudes:

            x, y = 0.5 * NX, 0.5 * NY
            fudge = 2
            equalmagnitudes = magnitudes + 2.5 * np.log10(5 * fudge)
            mags = tractor.Mags(order=bandnames,
                                **dict(zip(bandnames, equalmagnitudes)))

            # Exponential galaxy...
            galpos = wcs.pixelToPosition(x, y)
            mg = tractor.Mags(order=bandnames,
                              **dict(zip(bandnames, equalmagnitudes)))
            re = 0.5  # arcsec
            q = 1.0  # axis ratio
            theta = 0.0  # degrees
            galshape = tractor.sdss_galaxy.GalaxyShape(re, q, theta)
            nebulousgalaxy = tractor.sdss_galaxy.ExpGalaxy(
                galpos, mags.copy(), galshape)
            if vb: print nebulousgalaxy

            # ...plus 4 independent point sources, arranged in a small cross:
            e = 3.0  # pixels
            srcs = [
                nebulousgalaxy,
                tractor.PointSource(wcs.pixelToPosition(x + e, y),
                                    mags.copy()),
                tractor.PointSource(wcs.pixelToPosition(x - e, y),
                                    mags.copy()),
                tractor.PointSource(wcs.pixelToPosition(x, y + e),
                                    mags.copy()),
                tractor.PointSource(wcs.pixelToPosition(x, y - e), mags.copy())
            ]

            # Old setup - just 4 point sources:
            # srcs = [tractor.PointSource(wcs.pixelToPosition(x+e,y),mags),
            #         tractor.PointSource(wcs.pixelToPosition(x-e,y),mags),
            #         tractor.PointSource(wcs.pixelToPosition(x,y+e),mags),
            #         tractor.PointSource(wcs.pixelToPosition(x,y-e),mags)]

        # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -

        elif model == 'lens':

            # srcs = [InitializeLens(wcs,bandnames)]

            # Source to be lensed:
            xs, ys = 0.5 * NX, 0.5 * NY
            unlensedmagnitudes = magnitudes + 2.5 * np.log10(40.0)
            ms = tractor.Mags(order=bandnames,
                              **dict(zip(bandnames, unlensedmagnitudes)))
            if vb: print ms
            sourcepos = wcs.pixelToPosition(xs, ys)
            if vb: print sourcepos

            pointsource = tractor.PointSource(sourcepos, ms)
            if vb: print pointsource

            # Lens mass:
            thetaE = lenstractor.EinsteinRadius(0.75)  # arcsec
            if vb: print thetaE
            gamma = 0.2  # to make quad
            phi = 0.0  # deg
            xshear = lenstractor.ExternalShear(gamma, phi)
            if vb: print xshear

            # Lens light:
            x, y = 0.5 * NX, 0.5 * NY
            lenspos = wcs.pixelToPosition(x, y)
            if vb: print lenspos
            halfmagnitudes = magnitudes + 2.5 * np.log10(2.0)
            md = tractor.Mags(order=bandnames,
                              **dict(zip(bandnames, halfmagnitudes)))
            if vb: print md
            re = 1.0  # arcsec
            q = 1.0  # axis ratio
            theta = 0.0  # degrees
            galshape = tractor.sdss_galaxy.GalaxyShape(re, q, theta)
            if vb: print galshape

            lensgalaxy = lenstractor.LensGalaxy(lenspos, md, galshape, thetaE,
                                                xshear)
            if vb: print lensgalaxy

            srcs = [lenstractor.PointSourceLens(lensgalaxy, pointsource)]

        # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -

        if vb:
            print "Model =", srcs
            print " "

        # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -

        # Set up logging to the terminal by The Tractor:
        if opt.verbose:
            lvl = logging.DEBUG
        else:
            lvl = logging.INFO
        logging.basicConfig(level=lvl, format='%(message)s', stream=sys.stdout)

        # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -

        # Start a tractor, and let it make a catalog one src at a time.
        # Pass in a copy of the image list, so that PSF etc are
        # initialised correctly for each model.
        chug = tractor.Tractor(list(images))
        for src in srcs:
            chug.addSource(src)

        # Plot initial state:
        lenstractor.Plot_state(chug, model + '_progress_initial')

        # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -

        if not opt.MCMC:
            # Optimize the model parameters:

            if model == 'nebula':
                Nrounds = 2
                Nsteps_optimizing_catalog = 20
                Nsteps_optimizing_PSFs = 10
            elif model == 'lens':
                Nrounds = 3
                Nsteps_optimizing_catalog = 7
                Nsteps_optimizing_PSFs = 3

            if vb:
                print "Optimizing model:"
                print "   - no. of iterations per round to be spent on catalog: ", Nsteps_optimizing_catalog
                print "   - no. of iterations per round to be spent on PSFs: ", Nsteps_optimizing_PSFs
                print "   - no. of rounds: ", Nrounds

            k = 0
            for round in range(Nrounds):

                print "Fitting " + model + ": seconds out, round", round

                # Freeze the PSF, sky and photocal, leaving the sources:
                chug.thawParam('catalog')
                for image in chug.getImages():
                    image.thawParams('sky')
                    image.freezeParams('photocal', 'wcs', 'psf')
                print "Fitting " + model + ": Catalog parameters to be optimized are:", chug.getParamNames(
                )
                print "Fitting " + model + ": Initial values are:", chug.getParams(
                )
                print "Fitting " + model + ": Step sizes:", chug.getStepSizes()

                # Optimize sources with initial PSF:
                for i in range(Nsteps_optimizing_catalog):
                    dlnp, X, a = chug.optimize()
                    if not opt.noplots:
                        lenstractor.Plot_state(
                            chug, model +
                            '_progress_optimizing_step-%02d_catalog' % k)
                    print "Fitting " + model + ": at step", k, "parameter values are:", chug.getParams(
                    )
                    k += 1

                # Freeze the sources and thaw the psfs:
                chug.freezeParam('catalog')
                for image in chug.getImages():
                    image.thawParams('psf')
                    image.freezeParams('photocal', 'wcs', 'sky')
                print "Fitting PSF: After thawing, zeroth PSF = ", chug.getImage(
                    0).psf
                print "Fitting PSF: PSF parameters to be optimized are:", chug.getParamNames(
                )
                print "Fitting PSF: Initial values are:", chug.getParams()
                print "Fitting PSF: Step sizes:", chug.getStepSizes()

                # Optimize everything that is not frozen:
                for i in range(Nsteps_optimizing_PSFs):
                    dlnp, X, a = chug.optimize()
                    if not opt.noplots:
                        lenstractor.Plot_state(
                            chug, model +
                            '_progress_optimizing_step-%02d_catalog' % k)
                    print "Fitting PSF: at step", k, "parameter values are:", chug.getParams(
                    )
                    k += 1
                print "Fitting PSF: After optimizing, zeroth PSF = ", chug.getImage(
                    0).psf

            # BUG: PSF not being optimized correctly - missing derivatives?

            lenstractor.Plot_state(chug,
                                   model + '_progress_optimizing_zcomplete')

        # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -

        elif opt.MCMC:
            # MCMC sample the model parameters.

            if vb:
                print "Sampling model parameters with emcee:"

            # Freeze the sky and photocal, leaving the PSFs and sources:
            for image in chug.getImages():
                image.freezeParams('photocal', 'wcs')

            # Get the thawed parameters:
            p0 = np.array(chug.getParams())
            print 'Tractor parameters:'
            for i, parname in enumerate(chug.getParamNames()):
                print '  ', parname, '=', p0[i]
            ndim = len(p0)
            print 'Number of parameter space dimensions: ', ndim

            # Make an emcee sampler that uses our tractor to compute its logprob:
            nw = 8 * ndim
            sampler = emcee.EnsembleSampler(nw,
                                            ndim,
                                            chug,
                                            live_dangerously=True)

            # Start the walkers off near the initialisation point -
            # We need it to be ~1 pixel in position, and not too much
            # flux restrction...

            if model == 'lens':
                # The following gets us 0.2" in dec:
                psteps = np.zeros_like(p0) + 0.00004
                # This could be optimized, to allow more initial freedom in eg flux.

            elif model == 'nebula':
                # Good first guess should be some fraction of the optimization step sizes:
                psteps = 0.2 * np.array(chug.getStepSizes())

            print "Initial size (in each dimension) of sample ball = ", psteps

            pp = emcee.EnsembleSampler.sampleBall(p0, psteps, nw)
            rstate = None
            lnp = None

            # Take a few steps - memory leaks fast! (~10Mb per sec)
            for step in range(1, 4):

                print 'Run MCMC step set:', step
                t0 = tractor.Time()
                pp, lnp, rstate = sampler.run_mcmc(pp,
                                                   5,
                                                   lnprob0=lnp,
                                                   rstate0=rstate)

                print 'Mean acceptance fraction after', sampler.iterations, 'iterations =', np.mean(
                    sampler.acceptance_fraction)
                t_mcmc = (tractor.Time() - t0)
                print 'Runtime:', t_mcmc

                # Find the current posterior means:
                pbar = np.mean(pp, axis=0)
                print "Mean parameters: ", pbar, np.mean(lnp)

                # Find the current best sample:
                maxlnp = np.max(lnp)
                best = np.where(lnp == maxlnp)
                pbest = np.ravel(pp[best, :])
                print "Best parameters: ", pbest, maxlnp

                if not opt.noplots:
                    chug.setParams(pbest)
                    lenstractor.Plot_state(
                        chug, model + '_progress_sampling_step-%02d' % step)

            print 'Best lnprob:', np.max(lnp)
            # print 'dlnprobs:', ', '.join(['%.1f' % d for d in lnp - np.max(lnp)])
            # print 'MCMC took', t_mcmc, 'sec'

            # Take the last best sample and call it a result:
            chug.setParams(pbest)

            lenstractor.Plot_state(chug,
                                   model + '_progress_sampling_zcomplete')

        if vb:
            print "* * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * *"

        # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -

        # Collect statistics about this model's fit:

        chisq = -2.0 * chug.getLogLikelihood()
        chug.thawParam('catalog')
        for image in chug.getImages():
            image.thawParams('sky', 'psf')
            image.freezeParams('photocal', 'wcs')
        K = len(chug.getParams())
        N = chug.getNdata()
        BIC[model] = chisq + K * np.log(1.0 * N)
        print "Fitting " + model + ": chisq, K, N, BIC =", chisq, K, N, BIC[
            model]

    # -------------------------------------------------------------------------

    if len(models) == 2:
        # Compare models and report:
        print "BIC = ", BIC
        print "Fitting result: Bayes factor in favour of nebula is exp[", -0.5 * (
            BIC['nebula'] - BIC['lens']), "]"
        print "* * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * *"

    # -------------------------------------------------------------------------

    print "Tractor stopping."

    return
示例#5
0
    def run(self, ps=None, focus=False, momentsize=5, n_fwhm=100):
        import pylab as plt
        from astrometry.util.plotutils import dimshow, plothist
        from legacyanalysis.ps1cat import ps1cat
        import photutils
        import tractor

        fn = self.fn
        ext = self.ext
        pixsc = self.pixscale

        F = fitsio.FITS(fn)
        primhdr = F[0].read_header()
        self.primhdr = primhdr
        img, hdr = self.read_raw(F, ext)
        self.hdr = hdr

        # pre sky-sub
        mn, mx = np.percentile(img.ravel(), [25, 98])
        self.imgkwa = dict(vmin=mn, vmax=mx, cmap='gray')

        if self.debug and ps is not None:
            plt.clf()
            dimshow(img, **self.imgkwa)
            plt.title('Raw image')
            ps.savefig()

            M = 200
            plt.clf()
            plt.subplot(2, 2, 1)
            dimshow(img[-M:, :M], ticks=False, **self.imgkwa)
            plt.subplot(2, 2, 2)
            dimshow(img[-M:, -M:], ticks=False, **self.imgkwa)
            plt.subplot(2, 2, 3)
            dimshow(img[:M, :M], ticks=False, **self.imgkwa)
            plt.subplot(2, 2, 4)
            dimshow(img[:M, -M:], ticks=False, **self.imgkwa)
            plt.suptitle('Raw image corners')
            ps.savefig()

        img, trim_x0, trim_y0 = self.trim_edges(img)

        fullH, fullW = img.shape

        if self.debug and ps is not None:
            plt.clf()
            dimshow(img, **self.imgkwa)
            plt.title('Trimmed image')
            ps.savefig()

            M = 200
            plt.clf()
            plt.subplot(2, 2, 1)
            dimshow(img[-M:, :M], ticks=False, **self.imgkwa)
            plt.subplot(2, 2, 2)
            dimshow(img[-M:, -M:], ticks=False, **self.imgkwa)
            plt.subplot(2, 2, 3)
            dimshow(img[:M, :M], ticks=False, **self.imgkwa)
            plt.subplot(2, 2, 4)
            dimshow(img[:M, -M:], ticks=False, **self.imgkwa)
            plt.suptitle('Trimmed corners')
            ps.savefig()

        band = self.get_band(primhdr)
        exptime = primhdr['EXPTIME']
        airmass = primhdr['AIRMASS']
        print('Band', band, 'Exptime', exptime, 'Airmass', airmass)

        zp0 = self.nom.zeropoint(band, ext=self.ext)
        sky0 = self.nom.sky(band)
        kx = self.nom.fiducial_exptime(band).k_co

        # Find the sky value and noise level
        sky, sig1 = self.get_sky_and_sigma(img)

        sky1 = np.median(sky)
        skybr = -2.5 * np.log10(sky1 / pixsc / pixsc / exptime) + zp0
        print('Sky brightness: %8.3f mag/arcsec^2' % skybr)
        print('Fiducial:       %8.3f mag/arcsec^2' % sky0)

        img -= sky

        self.remove_sky_gradients(img)

        # Post sky-sub
        mn, mx = np.percentile(img.ravel(), [25, 98])
        self.imgkwa = dict(vmin=mn, vmax=mx, cmap='gray')

        if ps is not None:
            plt.clf()
            dimshow(img, **self.imgkwa)
            plt.title('Sky-sub image: %s-%s' % (os.path.basename(fn).replace(
                '.fits', '').replace('.fz', ''), ext))
            plt.colorbar()
            ps.savefig()

        # Read WCS header and compute boresight
        wcs = self.get_wcs(hdr)
        ra_ccd, dec_ccd = wcs.pixelxy2radec((fullW + 1) / 2., (fullH + 1) / 2.)

        # Detect stars
        psfsig = self.nominal_fwhm / 2.35
        detsn = self.detection_map(img, sig1, psfsig, ps)

        slices = self.detect_sources(detsn, self.det_thresh, ps)
        print(len(slices), 'sources detected')
        if len(slices) < 20:
            slices = self.detect_sources(detsn, 10., ps)
            print(len(slices), 'sources detected')
        ndetected = len(slices)

        camera = primhdr.get('INSTRUME', '').strip().lower()
        # -> "decam" / "mosaic3"
        meas = dict(band=band,
                    airmass=airmass,
                    skybright=skybr,
                    pixscale=pixsc,
                    primhdr=primhdr,
                    hdr=hdr,
                    wcs=wcs,
                    ra_ccd=ra_ccd,
                    dec_ccd=dec_ccd,
                    extension=ext,
                    camera=camera,
                    ndetected=ndetected)

        if ndetected == 0:
            print('NO SOURCES DETECTED')
            return meas

        xx, yy = [], []
        fx, fy = [], []
        mx2, my2, mxy = [], [], []
        wmx2, wmy2, wmxy = [], [], []
        # "Peak" region to centroid
        P = momentsize
        H, W = img.shape

        for i, slc in enumerate(slices):
            y0 = slc[0].start
            x0 = slc[1].start
            subimg = detsn[slc]
            imax = np.argmax(subimg)
            y, x = np.unravel_index(imax, subimg.shape)
            if (x0 + x) < P or (x0 + x) > W - 1 - P or (y0 + y) < P or (
                    y0 + y) > H - 1 - P:
                #print('Skipping edge peak', x0+x, y0+y)
                continue
            xx.append(x0 + x)
            yy.append(y0 + y)
            pkarea = detsn[y0 + y - P:y0 + y + P + 1,
                           x0 + x - P:x0 + x + P + 1]

            from scipy.ndimage.measurements import center_of_mass
            cy, cx = center_of_mass(pkarea)
            #print('Center of mass', cx,cy)
            fx.append(x0 + x - P + cx)
            fy.append(y0 + y - P + cy)
            #print('x,y', x0+x, y0+y, 'vs centroid', x0+x-P+cx, y0+y-P+cy)

            ### HACK -- measure source ellipticity
            # go back to the image (not detection map)
            #subimg = img[slc]
            subimg = img[y0 + y - P:y0 + y + P + 1,
                         x0 + x - P:x0 + x + P + 1].copy()
            subimg /= subimg.sum()
            ph, pw = subimg.shape
            px, py = np.meshgrid(np.arange(pw), np.arange(ph))
            mx2.append(np.sum(subimg * (px - cx)**2))
            my2.append(np.sum(subimg * (py - cy)**2))
            mxy.append(np.sum(subimg * (px - cx) * (py - cy)))
            # Gaussian windowed version
            s = 1.
            wimg = subimg * np.exp(-0.5 * ((px - cx)**2 + (py - cy)**2) / s**2)
            wimg /= np.sum(wimg)
            wmx2.append(np.sum(wimg * (px - cx)**2))
            wmy2.append(np.sum(wimg * (py - cy)**2))
            wmxy.append(np.sum(wimg * (px - cx) * (py - cy)))

        mx2 = np.array(mx2)
        my2 = np.array(my2)
        mxy = np.array(mxy)
        wmx2 = np.array(wmx2)
        wmy2 = np.array(wmy2)
        wmxy = np.array(wmxy)

        # semi-major/minor axes and position angle
        theta = np.rad2deg(np.arctan2(2 * mxy, mx2 - my2) / 2.)
        theta = np.abs(theta) * np.sign(mxy)
        s = np.sqrt(((mx2 - my2) / 2.)**2 + mxy**2)
        a = np.sqrt((mx2 + my2) / 2. + s)
        b = np.sqrt((mx2 + my2) / 2. - s)
        ell = 1. - b / a

        wtheta = np.rad2deg(np.arctan2(2 * wmxy, wmx2 - wmy2) / 2.)
        wtheta = np.abs(wtheta) * np.sign(wmxy)
        ws = np.sqrt(((wmx2 - wmy2) / 2.)**2 + wmxy**2)
        wa = np.sqrt((wmx2 + wmy2) / 2. + ws)
        wb = np.sqrt((wmx2 + wmy2) / 2. - ws)
        well = 1. - wb / wa

        fx = np.array(fx)
        fy = np.array(fy)
        xx = np.array(xx)
        yy = np.array(yy)

        if ps is not None:

            plt.clf()
            dimshow(detsn, vmin=-3, vmax=50, cmap='gray')
            ax = plt.axis()
            plt.plot(fx, fy, 'go', mec='g', mfc='none', ms=10)
            plt.colorbar()
            plt.title('Detected sources')
            plt.axis(ax)
            ps.savefig()

            # show centroids too
            # plt.plot(xx, yy, 'go', mec='g', mfc='none', ms=8)
            # plt.axis(ax)
            # ps.savefig()

        # if ps is not None:
        #     plt.clf()
        #     plt.subplot(2,1,1)
        #     mx = np.percentile(np.append(mx2,my2), 99)
        #     ha = dict(histtype='step', range=(0,mx), bins=50)
        #     plt.hist(mx2, color='b', label='mx2', **ha)
        #     plt.hist(my2, color='r', label='my2', **ha)
        #     plt.hist(mxy, color='g', label='mxy', **ha)
        #     plt.legend()
        #     plt.xlim(0,mx)
        #     plt.subplot(2,1,2)
        #     mx = np.percentile(np.append(wmx2,wmy2), 99)
        #     ha = dict(histtype='step', range=(0,mx), bins=50, lw=3, alpha=0.3)
        #     plt.hist(wmx2, color='b', label='wx2', **ha)
        #     plt.hist(wmy2, color='r', label='wy2', **ha)
        #     plt.hist(wmxy, color='g', label='wxy', **ha)
        #     plt.legend()
        #     plt.xlim(0,mx)
        #     plt.suptitle('Source moments')
        #     ps.savefig()
        #
        #     #mx = np.percentile(np.abs(np.append(mxy,wmxy)), 99)
        #     plt.clf()
        #     plt.subplot(2,1,1)
        #     ha = dict(histtype='step', range=(0,1), bins=50)
        #     plt.hist(ell, color='g', label='ell', **ha)
        #     plt.hist(well, color='g', lw=3, alpha=0.3, label='windowed ell', **ha)
        #     plt.legend()
        #     plt.subplot(2,1,2)
        #     ha = dict(histtype='step', range=(-90,90), bins=50)
        #     plt.hist(theta, color='g', label='theta', **ha)
        #     plt.hist(wtheta, color='g', lw=3, alpha=0.3,
        #              label='windowed theta', **ha)
        #     plt.xlim(-90,90)
        #     plt.legend()
        #     plt.suptitle('Source ellipticities & angles')
        #     ps.savefig()

        # Cut down to stars whose centroids are within 1 pixel of their peaks...
        #keep = (np.hypot(fx - xx, fy - yy) < 2)
        #print(sum(keep), 'of', len(keep), 'stars have centroids within 2 of peaks')
        #print('mean dx', np.mean(fx-xx), 'dy', np.mean(fy-yy), 'pixels')
        #assert(float(sum(keep)) / len(keep) > 0.9)
        #fx = fx[keep]
        #fy = fy[keep]

        apxy = np.vstack((fx, fy)).T
        ap = []
        aprad_pix = self.aprad / pixsc
        aper = photutils.CircularAperture(apxy, aprad_pix)
        p = photutils.aperture_photometry(img, aper)
        apflux = p.field('aperture_sum')

        # Manual aperture photometry to get clipped means in sky annulus
        sky_inner_r, sky_outer_r = [r / pixsc for r in self.skyrad]
        sky = []
        for xi, yi in zip(fx, fy):
            ix = int(np.round(xi))
            iy = int(np.round(yi))
            skyR = int(np.ceil(sky_outer_r))
            xlo = max(0, ix - skyR)
            xhi = min(W, ix + skyR + 1)
            ylo = max(0, iy - skyR)
            yhi = min(H, iy + skyR + 1)
            xx, yy = np.meshgrid(np.arange(xlo, xhi), np.arange(ylo, yhi))
            r2 = (xx - xi)**2 + (yy - yi)**2
            inannulus = ((r2 >= sky_inner_r**2) * (r2 < sky_outer_r**2))
            skypix = img[ylo:yhi, xlo:xhi][inannulus]
            #print('ylo,yhi, xlo,xhi', ylo,yhi, xlo,xhi, 'img subshape', img[ylo:yhi, xlo:xhi].shape, 'inann shape', inannulus.shape)
            s, nil = sensible_sigmaclip(skypix)
            sky.append(s)
        sky = np.array(sky)

        apflux2 = apflux - sky * (np.pi * aprad_pix**2)
        good = (apflux2 > 0) * (apflux > 0)
        apflux = apflux[good]
        apflux2 = apflux2[good]
        fx = fx[good]
        fy = fy[good]

        # Read in the PS1 catalog, and keep those within 0.25 deg of CCD center
        # and those with main sequence colors
        pscat = ps1cat(ccdwcs=wcs)
        stars = pscat.get_stars()
        #print('Got PS1 stars:', len(stars))

        # we add the color term later
        ps1band = ps1cat.ps1band[band]
        stars.mag = stars.median[:, ps1band]

        ok, px, py = wcs.radec2pixelxy(stars.ra, stars.dec)
        px -= 1
        py -= 1

        if ps is not None:
            #kwa = dict(vmin=-3*sig1, vmax=50*sig1, cmap='gray')
            # Add to the 'detected sources' plot
            # mn,mx = np.percentile(img.ravel(), [50,99])
            # kwa = dict(vmin=mn, vmax=mx, cmap='gray')
            # plt.clf()
            # dimshow(img, **kwa)
            ax = plt.axis()
            #plt.plot(fx, fy, 'go', mec='g', mfc='none', ms=10)
            K = np.argsort(stars.mag)
            plt.plot(px[K[:10]] - trim_x0,
                     py[K[:10]] - trim_y0,
                     'o',
                     mec='m',
                     mfc='none',
                     ms=12,
                     mew=2)
            plt.plot(px[K[10:]] - trim_x0,
                     py[K[10:]] - trim_y0,
                     'o',
                     mec='m',
                     mfc='none',
                     ms=8)
            plt.axis(ax)
            plt.title('PS1 stars')
            #plt.colorbar()
            ps.savefig()

        # we trimmed the image before running detection; re-add that margin
        fullx = fx + trim_x0
        fully = fy + trim_y0

        # Match PS1 to our detections, find offset
        radius = self.maxshift / pixsc

        I, J, dx, dy = self.match_ps1_stars(px, py, fullx, fully, radius,
                                            stars)
        print(len(I), 'spatial matches with large radius', self.maxshift,
              'arcsec,', radius, 'pix')

        bins = 2 * int(np.ceil(radius))
        #print('Histogramming with', bins, 'bins')
        histo, xe, ye = np.histogram2d(dx,
                                       dy,
                                       bins=bins,
                                       range=((-radius, radius), (-radius,
                                                                  radius)))
        # smooth histogram before finding peak -- fuzzy matching
        from scipy.ndimage.filters import gaussian_filter
        histo = gaussian_filter(histo, 1.)
        histo = histo.T
        mx = np.argmax(histo)
        my, mx = np.unravel_index(mx, histo.shape)
        shiftx = (xe[mx] + xe[mx + 1]) / 2.
        shifty = (ye[my] + ye[my + 1]) / 2.

        if ps is not None:
            plt.clf()
            plothist(dx, dy, range=((-radius, radius), (-radius, radius)))
            plt.xlabel('dx (pixels)')
            plt.ylabel('dy (pixels)')
            plt.title('Offsets to PS1 stars')
            ax = plt.axis()
            plt.axhline(0, color='b')
            plt.axvline(0, color='b')
            plt.plot(shiftx, shifty, 'o', mec='m', mfc='none', ms=15, mew=3)
            plt.axis(ax)
            ps.savefig()

        # Refine with smaller search radius
        radius2 = 3. / pixsc
        I, J, dx, dy = self.match_ps1_stars(px, py, fullx + shiftx,
                                            fully + shifty, radius2, stars)
        print(len(J), 'matches to PS1 with small radius', 3, 'arcsec')
        shiftx2 = np.median(dx)
        shifty2 = np.median(dy)
        #print('Stage-1 shift', shiftx, shifty)
        #print('Stage-2 shift', shiftx2, shifty2)
        sx = shiftx + shiftx2
        sy = shifty + shifty2
        print('Astrometric shift (%.0f, %.0f) pixels' % (sx, sy))

        if self.debug and ps is not None:
            plt.clf()
            plothist(dx, dy, range=((-radius2, radius2), (-radius2, radius2)))
            plt.xlabel('dx (pixels)')
            plt.ylabel('dy (pixels)')
            plt.title('Offsets to PS1 stars')
            ax = plt.axis()
            plt.axhline(0, color='b')
            plt.axvline(0, color='b')
            plt.plot(shiftx2, shifty2, 'o', mec='m', mfc='none', ms=15, mew=3)
            plt.axis(ax)
            ps.savefig()

        if ps is not None:
            mn, mx = np.percentile(img.ravel(), [50, 99])
            kwa2 = dict(vmin=mn, vmax=mx, cmap='gray')
            plt.clf()
            dimshow(img, **kwa2)
            ax = plt.axis()
            plt.plot(fx[J], fy[J], 'go', mec='g', mfc='none', ms=10, mew=2)
            plt.plot(px[I] - sx - trim_x0,
                     py[I] - sy - trim_y0,
                     'm+',
                     ms=10,
                     mew=2)
            plt.axis(ax)
            plt.title('Matched PS1 stars')
            plt.colorbar()
            ps.savefig()

            plt.clf()
            dimshow(img, **kwa2)
            ax = plt.axis()
            plt.plot(fx[J], fy[J], 'go', mec='g', mfc='none', ms=10, mew=2)
            K = np.argsort(stars.mag)
            plt.plot(px[K[:10]] - sx - trim_x0,
                     py[K[:10]] - sy - trim_y0,
                     'o',
                     mec='m',
                     mfc='none',
                     ms=12,
                     mew=2)
            plt.plot(px[K[10:]] - sx - trim_x0,
                     py[K[10:]] - sy - trim_y0,
                     'o',
                     mec='m',
                     mfc='none',
                     ms=8,
                     mew=2)
            plt.axis(ax)
            plt.title('All PS1 stars')
            plt.colorbar()
            ps.savefig()

        # Now cut to just *stars* with good colors
        stars.gicolor = stars.median[:, 0] - stars.median[:, 2]
        keep = (stars.gicolor > 0.4) * (stars.gicolor < 2.7)
        stars.cut(keep)
        if len(stars) == 0:
            print('No overlap or too few stars in PS1')
            return None
        px = px[keep]
        py = py[keep]
        # Re-match
        I, J, dx, dy = self.match_ps1_stars(px, py, fullx + sx, fully + sy,
                                            radius2, stars)
        print('Cut to', len(stars), 'PS1 stars with good colors; matched',
              len(I))

        nmatched = len(I)

        meas.update(dx=sx, dy=sy, nmatched=nmatched)

        if focus:
            meas.update(img=img,
                        hdr=hdr,
                        primhdr=primhdr,
                        fx=fx,
                        fy=fy,
                        px=px - trim_x0 - sx,
                        py=py - trim_y0 - sy,
                        sig1=sig1,
                        stars=stars,
                        moments=(mx2, my2, mxy, theta, a, b, ell),
                        wmoments=(wmx2, wmy2, wmxy, wtheta, wa, wb, well),
                        apflux=apflux,
                        apflux2=apflux2)
            return meas

        #print('Mean astrometric shift (arcsec): delta-ra=', -np.mean(dy)*0.263, 'delta-dec=', np.mean(dx)*0.263)

        # Compute photometric offset compared to PS1
        # as the PS1 minus observed mags
        colorterm = self.colorterm_ps1_to_observed(stars.median, band)
        stars.mag += colorterm
        ps1mag = stars.mag[I]

        if False and ps is not None:
            plt.clf()
            plt.semilogy(ps1mag, apflux2[J], 'b.')
            plt.xlabel('PS1 mag')
            plt.ylabel('DECam ap flux (with sky sub)')
            ps.savefig()

            plt.clf()
            plt.semilogy(ps1mag, apflux[J], 'b.')
            plt.xlabel('PS1 mag')
            plt.ylabel('DECam ap flux (no sky sub)')
            ps.savefig()

        apmag2 = -2.5 * np.log10(apflux2) + zp0 + 2.5 * np.log10(exptime)
        apmag = -2.5 * np.log10(apflux) + zp0 + 2.5 * np.log10(exptime)

        if ps is not None:
            plt.clf()
            plt.plot(ps1mag, apmag[J], 'b.', label='No sky sub')
            plt.plot(ps1mag, apmag2[J], 'r.', label='Sky sub')
            # ax = plt.axis()
            # mn = min(ax[0], ax[2])
            # mx = max(ax[1], ax[3])
            # plt.plot([mn,mx], [mn,mx], 'k-', alpha=0.1)
            # plt.axis(ax)
            plt.xlabel('PS1 mag')
            plt.ylabel('DECam ap mag')
            plt.legend(loc='upper left')
            plt.title('Zeropoint')
            ps.savefig()

        dm = ps1mag - apmag[J]
        dmag, dsig = sensible_sigmaclip(dm, nsigma=2.5)
        print('Mag offset: %8.3f' % dmag)
        print('Scatter:    %8.3f' % dsig)

        if not np.isfinite(dmag) or not np.isfinite(dsig):
            print('FAILED TO GET ZEROPOINT!')
            meas.update(zp=None)
            return meas

        from scipy.stats import sigmaclip
        goodpix, lo, hi = sigmaclip(dm, low=3, high=3)
        dmagmed = np.median(goodpix)
        print(len(goodpix), 'stars used for zeropoint median')
        print('Using median zeropoint:')
        zp_med = zp0 + dmagmed
        trans_med = 10.**(-0.4 * (zp0 - zp_med - kx * (airmass - 1.)))
        print('Zeropoint %6.3f' % zp_med)
        print('Transparency: %.3f' % trans_med)

        dm = ps1mag - apmag2[J]
        dmag2, dsig2 = sensible_sigmaclip(dm, nsigma=2.5)
        #print('Sky-sub mag offset', dmag2)
        #print('Scatter', dsig2)

        if ps is not None:
            plt.clf()
            plt.plot(ps1mag,
                     apmag[J] + dmag - ps1mag,
                     'b.',
                     label='No sky sub')
            plt.plot(ps1mag, apmag2[J] + dmag2 - ps1mag, 'r.', label='Sky sub')
            plt.xlabel('PS1 mag')
            plt.ylabel('DECam ap mag - PS1 mag')
            plt.legend(loc='upper left')
            plt.ylim(-0.25, 0.25)
            plt.axhline(0, color='k', alpha=0.25)
            plt.title('Zeropoint')
            ps.savefig()

        zp_obs = zp0 + dmag
        transparency = 10.**(-0.4 * (zp0 - zp_obs - kx * (airmass - 1.)))
        meas.update(zp=zp_obs, transparency=transparency)

        print('Zeropoint %6.3f' % zp_obs)
        print('Fiducial  %6.3f' % zp0)
        print('Transparency: %.3f' % transparency)

        # print('Using sky-subtracted values:')
        # zp_sky = zp0 + dmag2
        # trans_sky = 10.**(-0.4 * (zp0 - zp_sky - kx * (airmass - 1.)))
        # print('Zeropoint %6.3f' % zp_sky)
        # print('Transparency: %.3f' % trans_sky)

        fwhms = []
        psf_r = 15
        if n_fwhm not in [0, None]:
            Jf = J[:n_fwhm]

        for i, (xi, yi, fluxi) in enumerate(zip(fx[Jf], fy[Jf], apflux[Jf])):
            #print('Fitting source', i, 'of', len(Jf))
            ix = int(np.round(xi))
            iy = int(np.round(yi))
            xlo = max(0, ix - psf_r)
            xhi = min(W, ix + psf_r + 1)
            ylo = max(0, iy - psf_r)
            yhi = min(H, iy + psf_r + 1)
            xx, yy = np.meshgrid(np.arange(xlo, xhi), np.arange(ylo, yhi))
            r2 = (xx - xi)**2 + (yy - yi)**2
            keep = (r2 < psf_r**2)
            pix = img[ylo:yhi, xlo:xhi].copy()
            ie = np.zeros_like(pix)
            ie[keep] = 1. / sig1
            #print('fitting source at', ix,iy)
            #print('number of active pixels:', np.sum(ie > 0), 'shape', ie.shape)

            psf = tractor.NCircularGaussianPSF([4.], [1.])
            tim = tractor.Image(data=pix, inverr=ie, psf=psf)
            src = tractor.PointSource(tractor.PixPos(xi - xlo, yi - ylo),
                                      tractor.Flux(fluxi))
            tr = tractor.Tractor([tim], [src])

            #print('Posterior before prior:', tr.getLogProb())
            src.pos.addGaussianPrior('x', 0., 1.)
            #print('Posterior after prior:', tr.getLogProb())

            doplot = (i < 5) * (ps is not None)
            if doplot:
                mod0 = tr.getModelImage(0)

            tim.freezeAllBut('psf')
            psf.freezeAllBut('sigmas')

            # print('Optimizing params:')
            # tr.printThawedParams()

            #print('Parameter step sizes:', tr.getStepSizes())
            optargs = dict(priors=False, shared_params=False)
            for step in range(50):
                dlnp, x, alpha = tr.optimize(**optargs)
                #print('dlnp', dlnp)
                #print('src', src)
                #print('psf', psf)
                if dlnp == 0:
                    break
            # Now fit only the PSF size
            tr.freezeParam('catalog')
            # print('Optimizing params:')
            # tr.printThawedParams()

            for step in range(50):
                dlnp, x, alpha = tr.optimize(**optargs)
                #print('dlnp', dlnp)
                #print('src', src)
                #print('psf', psf)
                if dlnp == 0:
                    break

            fwhms.append(psf.sigmas[0] * 2.35 * pixsc)

            if doplot:
                mod1 = tr.getModelImage(0)
                chi1 = tr.getChiImage(0)

                plt.clf()
                plt.subplot(2, 2, 1)
                plt.title('Image')
                dimshow(pix, **self.imgkwa)
                plt.subplot(2, 2, 2)
                plt.title('Initial model')
                dimshow(mod0, **self.imgkwa)
                plt.subplot(2, 2, 3)
                plt.title('Final model')
                dimshow(mod1, **self.imgkwa)
                plt.subplot(2, 2, 4)
                plt.title('Final chi')
                dimshow(chi1, vmin=-10, vmax=10)
                plt.suptitle('PSF fit')
                ps.savefig()

        fwhms = np.array(fwhms)
        fwhm = np.median(fwhms)
        print('Median FWHM: %.3f' % np.median(fwhms))
        meas.update(seeing=fwhm)

        if False and ps is not None:
            lo, hi = np.percentile(fwhms, [5, 95])
            lo -= 0.1
            hi += 0.1
            plt.clf()
            plt.hist(fwhms, 25, range=(lo, hi), histtype='step', color='b')
            plt.xlabel('FWHM (arcsec)')
            ps.savefig()

        if ps is not None:
            plt.clf()
            for i, (xi, yi) in enumerate(zip(fx[J], fy[J])[:50]):
                ix = int(np.round(xi))
                iy = int(np.round(yi))
                xlo = max(0, ix - psf_r)
                xhi = min(W, ix + psf_r + 1)
                ylo = max(0, iy - psf_r)
                yhi = min(H, iy + psf_r + 1)
                pix = img[ylo:yhi, xlo:xhi]

                slc = pix[iy - ylo, :].copy()
                slc /= np.sum(slc)
                p1 = plt.plot(slc, 'b-', alpha=0.2)
                slc = pix[:, ix - xlo].copy()
                slc /= np.sum(slc)
                p2 = plt.plot(slc, 'r-', alpha=0.2)
                ph, pw = pix.shape
                cx, cy = pw / 2, ph / 2
                if i == 0:
                    xx = np.linspace(0, pw, 300)
                    dx = xx[1] - xx[0]
                    sig = fwhm / pixsc / 2.35
                    yy = np.exp(-0.5 * (xx - cx)**2 / sig**2)  # * np.sum(pix)
                    yy /= (np.sum(yy) * dx)
                    p3 = plt.plot(xx, yy, 'k-', zorder=20)
            #plt.ylim(-0.2, 1.0)
            plt.legend([p1[0], p2[0], p3[0]],
                       ['image slice (y)', 'image slice (x)', 'fit'])
            plt.title('PSF fit')
            ps.savefig()

        return meas
示例#6
0
    ra, dec = 40., 10.
    psf_sigma = 1.4  # pixels
    v = psf_sigma**2

    ps = pixscale / 3600.
    wcs = Tan(ra, dec, W / 2. + 0.5, H / 2. + 0.5, -ps, 0., 0., ps, float(W),
              float(H))

    tim = tractor.Image(data=np.zeros((H, W), np.float32),
                        inverr=np.ones((H, W), np.float32),
                        psf=tractor.GaussianMixturePSF(1., 0., 0., v, v, 0.),
                        wcs=tractor.ConstantFitsWcs(wcs))
    src = RexGalaxy(tractor.RaDecPos(ra, dec), tractor.Flux(100.),
                    LogRadius(0.))

    tr = tractor.Tractor([tim], [src])
    mod = tr.getModelImage(0)

    plt.clf()
    plt.imshow(mod, interpolation='nearest', origin='lower')
    plt.savefig('rex.png')

    # add noise with std 1.
    noisy = mod + np.random.normal(size=mod.shape)
    # make that the tim's data
    tim.data = noisy

    # reset the source params
    src.brightness.setParams([1.])

    tr.freezeParam('images')
示例#7
0
def main():
    """
   NAME
     LensTractor.py

   PURPOSE
     Run the Tractor on a deck of single object cutout images.
     Read in an image and its weight map, guess the PSF, put an object at the
     centre image of the field and then optimize the catalog and PSF.

   COMMENTS
     The idea is to identify good lens candidates by principled model 
     selection: two well-defined models competing against each other, given 
     multi-epoch imaging data. The Nebula model (1 extended source, plus
     N=1,2,3 or 4 point sources, with sub-models denoted by "NebulaN") is very
     flexible, so should be better at fitting the data in general than the
     Lens model (1 extended source, plus 1 background point source). However,
     when the Lens provides a good fit, it does so at lower cost (fewer
     parameters), so should win by Bayesian information criteria (we use BIC
     as a cheap proxy for evidence ratio).
     
     The default workflow is as follows:
     
       Is the target a lens?
       
       * Try Nebula1
       
       * Try Nebula2
       
           if Nebula1 beats Nebula2: 
             Return NO
           else:
             Nebula = Nebula2
       
       * Try Nebula4
       
           if Nebula4 beats Nebula2: 
             Nebula = Nebula4

                      
       * Try Lens (via Nebula)
           if Lens beats Nebula: 
             Return YES
           else:
             Return NO

      Initialisation of Lens via Nebula depends on the point source 
      multiplicity of the final Nebula model: what we do with three point
      image positions will be different from what we do with 4 point image
      positions, particularly with regard to the deflector centroid. 
      
      Open questions:
      
      Does it make sense to dogmatically associate the extended object with
      the deflector?
         YES: detection of a deflector galaxy will be needed for a convincing
         candidate anyway.
         NO: using the extended object to model a high S/N merging image
         system should not be punished
      
      How are we going to interpret the point image positions if we do not
      have an estimated deflector position?
      

   FLAGS
     -h --help        Print this message
     -v --verbose     Verbose operation
     -s --sample      Sample the posterior PDF instead of optimizing
     -x --no-plots    Do not plot progress
     -l --lens        Only fit lens model, initialized from scratch
     -n --nebula      Only fit nebula model, initialized from scratch

   INPUTS
     *.fits           Deck of postcard images

   OPTIONAL INPUTS
     --optimization-rounds        Nr   Number of rounds of optimization [2]
     --optimization-steps-catalog Nc   Number of steps per round spent
                                        optimizing source catalog [10]
     --optimization-steps-psf     Np   Number of steps per round spent
                                        optimizing source catalog [2]

   OUTPUTS
     stdout                       Useful information
     *.png                        Plots in png format
     
     To be implemented:
       lenstractor_progress.log     Logged output
       lenstractor_results.txt      Model comparison results
       lenstractor_lens.cat         Lens model parameters, including lightcurves
       lenstractor_nebula.cat       Nebula model parameters, including lightcurves

   EXAMPLES

     python LensTractor.py -x examples/H1413+117_10x10arcsec_55*fits > examples/H1413+117_10x10arcsec_lenstractor.log
   
   DEPENDENCIES
     * The Tractor     astrometry.net/svn/trunk/projects/tractor
     * emcee           github.com/danfm/emcee
     * astrometry.net  astrometry.net/svn/trunk/util

   BUGS

   HISTORY
     2012-07-06       First predicted Lens images Marshall/Hogg (Oxford/NYU)
     2013-01-23       Sequential Nebula 1,2,4 scheme implemented
   """

    # --------------------------------------------------------------------

    from argparse import ArgumentParser
    import sys

    # Set available options:
    parser = ArgumentParser()
    # List of files:
    parser.add_argument('inputfiles', metavar='N', nargs='+')
    # Verbosity:
    parser.add_argument('-v',
                        '--verbose',
                        dest='verbose',
                        action='store_true',
                        default=False,
                        help='Make more verbose')
    # Sampling:
    parser.add_argument('-s',
                        '--sample',
                        dest='MCMC',
                        action='store_true',
                        default=False,
                        help='Sample posterior PDF')
    # Plotting:
    parser.add_argument('-x',
                        '--no-plots',
                        dest='noplots',
                        action='store_true',
                        default=False,
                        help='Skip plotting')
    # Lens model only:
    parser.add_argument('-l',
                        '--lens',
                        dest='lens',
                        action='store_true',
                        default=False,
                        help='Fit lens model')
    # Nebula model only:
    parser.add_argument('-n',
                        '--nebula',
                        dest='nebula',
                        action='store_true',
                        default=False,
                        help='Fit nebula model')
    # optimization workflow:
    parser.add_argument('--optimization-rounds',
                        dest='Nr',
                        type=int,
                        default=3,
                        help='No. of optimization rounds')
    parser.add_argument(
        '--optimization-steps-catalog',
        dest='Nc',
        type=int,
        default=10,
        help='No. of optimization steps spent on source catalog')
    parser.add_argument('--optimization-steps-psf',
                        dest='Np',
                        type=int,
                        default=2,
                        help='No. of optimization steps spent on PSFs')
    parser.add_argument('--survey',
                        dest='survey',
                        type=str,
                        default="PS1",
                        help="Survey, either PS1 or KIDS")

    # Read in options and arguments - note only sci and wht images are supplied:
    args = parser.parse_args()

    if len(args.inputfiles) < 2:
        parser.print_help()
        sys.exit(-1)

    vb = args.verbose

    # Workflow:
    if args.lens:
        models = ['lens']
    elif args.nebula:
        models = [
            'nebula1',
            'nebula2',
            'nebula4',
        ]
    else:
        models = ['nebula1', 'nebula2', 'nebula4', 'lens']
    BIC = dict(zip(models, np.zeros(len(models))))
    # NB. default operation is to fit both and compare.

    if vb:
        print "* * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * *"
        print "                               LensTractor "
        print "    Fitting", models, " models to a deck of FITS postcards"
        print "* * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * *"

    # -------------------------------------------------------------------------

    # Organise the deck of inputfiles into scifiles and varfiles:
    scifiles, varfiles = lenstractor.Riffle(args.inputfiles, vb=vb)

    # Read into Tractor Image objects, and see what filters we have:
    images, total_mags, bands = lenstractor.Deal(scifiles,
                                                 varfiles,
                                                 SURVEY=args.survey,
                                                 vb=vb)

    # -------------------------------------------------------------------------
    # Generic items needed to initialize the Tractor's catalog.

    # Get rough idea of object position from wcs of first image- works
    # well if all images are the same size and well registered!
    wcs = images[0].wcs
    NX, NY = np.shape(images[0].data)
    if vb: print "Generic initial position ", NX, NY, "(pixels)"

    # m0 = 15.0
    # magnitudes = m0*np.ones(len(bandnames))
    # MAGIC initial magnitude. This should be estimated from
    # the total flux in each (background-subtracted) image...

    bandnames = np.unique(bands)
    magnitudes = np.zeros(len(bandnames))
    for i, bandname in enumerate(bandnames):
        index = np.where(bands == bandname)
        magnitudes[i] = np.median(total_mags[index])
    if vb: print "Generic initial SED ", dict(zip(bandnames, magnitudes))

    # -------------------------------------------------------------------------

    # Loop over models:

    for model in models:

        if vb:
            print "* * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * *"
            print "Initializing model: " + model

        # Figure out what type of model this is:
        modeltype = model[0:6]

        #   -   -   -   -   -   -   -   -   -   -   -   -   -   -   -   -   - -

        if modeltype == 'nebula':

            # Nebula - a flexible galaxy plus K point sources,

            # How many point sources?
            K = int(model[6:7])

            # The first nebula model we run has 1 point source and one
            # galaxy, initialised sensibly but randomly.
            # All subsequent models just have extra random point sources.

            # Start both sources off with equal magnitudes:
            x, y = 0.5 * NX, 0.5 * NY
            fudge = 2
            equalmagnitudes = magnitudes + 2.5 * np.log10(5 * fudge)
            mags = tractor.Mags(order=bandnames,
                                **dict(zip(bandnames, equalmagnitudes)))

            # Add an exponential galaxy:
            galpos = wcs.pixelToPosition(x, y)
            mg = tractor.Mags(order=bandnames,
                              **dict(zip(bandnames, equalmagnitudes)))
            re = 0.5  # arcsec
            q = 1.0  # axis ratio
            theta = 0.0  # degrees
            galshape = tractor.sdss_galaxy.GalaxyShape(re, q, theta)
            nebulousgalaxy = tractor.sdss_galaxy.ExpGalaxy(
                galpos, mags.copy(), galshape)
            if vb: print nebulousgalaxy
            srcs = [nebulousgalaxy]

            for i in range(K):
                # Add a point source with random position near nebula centre:
                e = 2.0  # pixels
                dx, dy = e * np.random.randn(2)
                star = tractor.PointSource(wcs.pixelToPosition(x + dx, y + dy),
                                           mags.copy())
                if vb: print star
                srcs.append(star)

        #   -   -   -   -   -   -   -   -   -   -   -   -   -   -   -   -   - -

        # Original Nebula4 initialisation went as follows:
        #   tractor.PointSource(wcs.pixelToPosition(x+e,y),mags.copy()),
        #   tractor.PointSource(wcs.pixelToPosition(x-e,y),mags.copy()),
        #   tractor.PointSource(wcs.pixelToPosition(x,y+e),mags.copy()),
        #   tractor.PointSource(wcs.pixelToPosition(x,y-e),mags.copy())]

        #   -   -   -   -   -   -   -   -   -   -   -   -   -   -   -   -   - -

        elif modeltype == 'lens':

            # If nebula has been run, use the best nebula model (by BIC)
            # to initialise the lens model. If it hasn't, do something
            # sensible.

            # Source to be lensed:
            xs, ys = 0.5 * NX, 0.5 * NY
            # Tiny random offset (pixels):
            e = 0.5
            dx, dy = e * np.random.randn(2)
            xs, ys = xs + dx, ys + dy
            unlensedmagnitudes = magnitudes + 2.5 * np.log10(40.0)
            ms = tractor.Mags(order=bandnames,
                              **dict(zip(bandnames, unlensedmagnitudes)))
            if vb: print ms
            sourcepos = wcs.pixelToPosition(xs, ys)
            if vb: print sourcepos

            pointsource = tractor.PointSource(sourcepos, ms)
            if vb: print pointsource

            # Lens mass:
            thetaE = lenstractor.EinsteinRadius(0.75)  # arcsec
            if vb: print thetaE
            gamma = 0.2  # to make quad
            phi = 0.0  # deg
            xshear = lenstractor.ExternalShear(gamma, phi)
            if vb: print xshear

            # Lens light:
            x, y = 0.5 * NX, 0.5 * NY
            lenspos = wcs.pixelToPosition(x, y)
            if vb: print lenspos
            lensmagnitudes = magnitudes + 2.5 * np.log10(10.0)
            md = tractor.Mags(order=bandnames,
                              **dict(zip(bandnames, lensmagnitudes)))
            if vb: print md
            re = 0.5  # arcsec
            q = 0.8  # axis ratio
            theta = 90.0  # degrees
            galshape = tractor.sdss_galaxy.GalaxyShape(re, q, theta)
            if vb: print galshape

            lensgalaxy = lenstractor.LensGalaxy(lenspos, md, galshape, thetaE,
                                                xshear)
            if vb: print lensgalaxy

            srcs = [lenstractor.PointSourceLens(lensgalaxy, pointsource)]

        #   -   -   -   -   -   -   -   -   -   -   -   -   -   -   -   -   - -

        if vb:
            print "Initialization complete."
            print "Model =", srcs
            print " "

        # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -

        # Set up logging to the terminal by The Tractor:
        if vb:
            lvl = logging.DEBUG
        else:
            lvl = logging.INFO
        logging.basicConfig(level=lvl, format='%(message)s', stream=sys.stdout)

        # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -

        # Start a tractor, and let it make a catalog one src at a time.
        # Pass in a copy of the image list, so that PSF etc are
        # initialised correctly for each model.
        chug = tractor.Tractor(list(images))
        for src in srcs:
            chug.addSource(src)

        # Plot initial state:
        lenstractor.Plot_state(chug,
                               model + '_progress_initial',
                               SURVEY=args.survey)

        # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -

        # Do the fit - either by maximizing the posterior PDF ("optimizing")
        # or by exploring the posterior PDF ("MCMC").

        if not args.MCMC:

            # Pre-op! Cycle over N initialisations, trying 5 steps each. Then
            # Optimize from the best of these.

            # Optimize the model parameters:

            if modeltype == 'nebula':
                Nrounds = args.Nr
                Nsteps_optimizing_catalog = args.Nc
                Nsteps_optimizing_PSFs = args.Np
            elif modeltype == 'lens':
                Nrounds = args.Nr
                Nsteps_optimizing_catalog = args.Nc
                Nsteps_optimizing_PSFs = args.Np

            if vb:
                print "Optimizing model:"
                print "   - no. of iterations per round to be spent on catalog: ", Nsteps_optimizing_catalog
                print "   - no. of iterations per round to be spent on PSFs: ", Nsteps_optimizing_PSFs
                print "   - no. of rounds: ", Nrounds

            k = 0
            for round in range(Nrounds):

                print "Fitting " + model + ": seconds out, round", round

                # Freeze the PSF, sky and photocal, leaving the sources:
                print "Thawing catalog..."
                chug.thawParam('catalog')
                for image in chug.getImages():
                    print "Thawing sky..."
                    image.thawParams('sky')
                    print "Freezing photocal, WCS, PSF..."
                    image.freezeParams('photocal', 'wcs', 'psf')
                print "Fitting " + model + ": Catalog parameters to be optimized are:", chug.getParamNames(
                )
                print "Fitting " + model + ": Initial values are:", chug.getParams(
                )
                print "Fitting " + model + ": Step sizes:", chug.getStepSizes()

                # Optimize sources with initial PSF:
                for i in range(Nsteps_optimizing_catalog):
                    dlnp, X, a = chug.optimize(damp=3)
                    # print "Fitting "+model+": at step",k,"parameter values are:",chug.getParams()
                    print "Progress: k,dlnp = ", k, dlnp
                    print ""
                    print "Catalog:", chug.getParams()
                    if dlnp == 0:
                        print "Converged? Exiting..."
                        # Although this only leaves *this* loop...
                        break
                    k += 1
                if not args.noplots:
                    lenstractor.Plot_state(
                        chug,
                        model + '_progress_optimizing_step-%02d_catalog' % k,
                        SURVEY=args.survey)

                # Freeze the sources and sky and thaw the psfs:
                print "Freezing catalog..."
                chug.freezeParam('catalog')
                for image in chug.getImages():
                    print "Thawing PSF..."
                    image.thawParams('psf')
                    print "Freezing photocal, WCS, sky..."
                    image.freezeParams('photocal', 'wcs', 'sky')
                print "Fitting PSF: After thawing, zeroth PSF = ", chug.getImage(
                    0).psf
                print "Fitting PSF: PSF parameters to be optimized are:", chug.getParamNames(
                )
                print "Fitting PSF: Initial values are:", chug.getParams()
                print "Fitting PSF: Step sizes:", chug.getStepSizes()

                # Optimize everything that is not frozen:
                for i in range(Nsteps_optimizing_PSFs):
                    dlnp, X, a = chug.optimize()
                    print "Fitting PSF: at step", k, "parameter values are:", chug.getParams(
                    )
                    k += 1
                print "Fitting PSF: After optimizing, zeroth PSF = ", chug.getImage(
                    0).psf
                if not args.noplots:
                    lenstractor.Plot_state(
                        chug,
                        model + '_progress_optimizing_step-%02d_catalog' % k,
                        SURVEY=args.survey)

            # BUG: PSF not being optimized correctly - missing derivatives?

            lenstractor.Plot_state(chug,
                                   model + '_progress_optimizing_zcomplete',
                                   SURVEY=args.survey)

        #   -   -   -   -   -   -   -   -   -   -   -   -   -   -   -   -   - -

        elif args.MCMC:
            # MCMC sample the model parameters.

            if vb:
                print "Sampling model parameters with emcee:"

            # Freeze the wcs and photocal, leaving the PSFs, sky and sources:
            for image in chug.getImages():
                image.freezeParams('photocal', 'wcs')
                # Temporary expt:
                image.freezeParams('psf')

            # Get the thawed parameters:
            p0 = np.array(chug.getParams())
            print 'Tractor parameters:'
            for i, parname in enumerate(chug.getParamNames()):
                print '  ', parname, '=', p0[i]
            ndim = len(p0)
            print 'Number of parameter space dimensions: ', ndim

            # Make an emcee sampler that uses our tractor to compute its logprob:
            nw = 8 * ndim
            sampler = emcee.EnsembleSampler(nw, ndim, chug, threads=4)

            # Start the walkers off near the initialisation point -
            # We need it to be ~1 pixel in position, and not too much
            # flux restrction...

            if model == 'lens':
                # The following gets us 0.2" in dec:
                psteps = np.zeros_like(p0) + 0.00004
                # This could be optimized, to allow more initial freedom in eg flux.

            elif model == 'nebula':
                # Good first guess should be some fraction of the optimization step sizes:
                psteps = 0.2 * np.array(chug.getStepSizes())

            # BUG - nebula+lens workflow not yet enabled!

            print "Initial size (in each dimension) of sample ball = ", psteps

            pp = emcee.EnsembleSampler.sampleBall(p0, psteps, nw)
            rstate = None
            lnp = None

            # Take a few steps - memory leaks fast! (~10Mb per sec)
            for step in range(10):

                print 'EMCEE: Run MCMC step set:', step
                t0 = tractor.Time()
                pp, lnp, rstate = sampler.run_mcmc(pp,
                                                   50,
                                                   lnprob0=lnp,
                                                   rstate0=rstate)

                print 'EMCEE: Mean acceptance fraction after', sampler.iterations, 'iterations =', np.mean(
                    sampler.acceptance_fraction)
                t_mcmc = (tractor.Time() - t0)
                print 'EMCEE: Runtime:', t_mcmc

                # Find the current posterior means:
                # pbar = np.mean(pp,axis=0)
                # print "Mean parameters: ",pbar,np.mean(lnp)

                # Find the current best sample:
                maxlnp = np.max(lnp)
                best = np.where(lnp == maxlnp)
                pbest = np.ravel(pp[best, :])
                print "EMCEE: Best parameters: ", maxlnp, pbest
                chug.setParams(pbest)
                chisq = -2.0 * chug.getLogLikelihood()
                print "EMCEE: chisq at Best pt: ", chisq
                if not args.noplots:
                    chug.setParams(pbest)
                    lenstractor.Plot_state(
                        chug,
                        model + '_progress_sampling_step-%02d' % step,
                        SURVEY=args.survey)

            # Take the last best sample and call it a result:
            chug.setParams(pbest)

            print 'EMCEE: Best lnprob, chisq:', maxlnp, chisq
            # print 'dlnprobs:', ', '.join(['%.1f' % d for d in lnp - np.max(lnp)])
            # print 'MCMC took', t_mcmc, 'sec'

            # Make the final plot:
            lenstractor.Plot_state(chug,
                                   model + '_progress_sampling_zcomplete',
                                   SURVEY=args.survey)

        #   -   -   -   -   -   -   -   -   -   -   -   -   -   -   -   -   - -

        if vb:
            print "Fit complete."
            print "* * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * *"

        # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -

        # Collect statistics about this model's fit:

        chisq = -2.0 * chug.getLogLikelihood()
        chug.thawParam('catalog')
        for image in chug.getImages():
            image.thawParams('sky', 'psf')
            image.freezeParams('photocal', 'wcs')
        K = len(chug.getParams())
        N = chug.getNdata()
        BIC[model] = chisq + K * np.log(1.0 * N)
        print model + " results: chisq, K, N, BIC =", chisq, K, N, BIC[model]

    # -------------------------------------------------------------------------

    # Make some decision about the nature of this system.

    #    if len(models) > 1:
    #    # Compare models and report:
    #        print "BIC = ",BIC
    #        print "Hypothesis test result: Bayes factor in favour of nebula is exp[",-0.5*(BIC['nebula'] - BIC['lens']),"]"
    #        print "* * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * *"

    # -------------------------------------------------------------------------

    print "Tractor stopping."

    return
示例#8
0
文件: coadd.py 项目: dstndstn/euclid
    print(name)
    print()

    # We're simulating a single isolated point source
    src = tractor.PointSource(tractor.PixPos(W // 2, H // 2),
                              tractor.Flux(trueflux))

    # Produce tractor image objects with noise model, PSF model, etc
    tims = []
    for noise, psf_size in imageset:
        tim = tractor.Image(np.zeros((H, W), np.float32),
                            inverr=np.ones((H, W), np.float32) * (1. / noise),
                            psf=tractor.NCircularGaussianPSF([psf_size], [1.]),
                            photocal=tractor.LinearPhotoCal(1.))
        # Create noiseless model image (simulated image)
        tr = tractor.Tractor([tim], [src])
        mod = tr.getModelImage(0)
        tim.data = mod
        tims.append(tim)

    # First we'll run without any noise added to the images, to get estimates of ideal performance.

    # Run the Simultaneous Fitting method
    tr = tractor.Tractor(tims, [src])
    src.brightness.setParams([0.])
    # Freeze the source position -- only fit for flux
    src.freezeParam('pos')
    # Freeze the image calibration parameters (PSF model, sky background, photometric calibration, etc)
    tr.freezeParam('images')
    phot = tr.optimize_forced_photometry(variance=True)
    simult_err = 1. / np.sqrt(phot.IV[0])
示例#9
0
        tim = tractor.Image(data=np.zeros((H, W), np.float32),
                            inverr=np.ones((H, W), np.float32),
                            psf=psf,
                            wcs=wcs,
                            photcal=photcal)

        ##  _tr            = tractor.Tractor([tim], [src])
        ##  mod            = _tr.getModelImage(0)

        tim.data = tim.data + noise.data  ##  + mod.data
        tims.append(tim)

    cat = tractor.Catalog(src)

    ##
    tr = tractor.Tractor(tims, cat)

    # Evaluate likelihood.
    lnp = tr.getLogProb()
    print('Logprob:', lnp)

    for nm, val in zip(tr.getParamNames(), tr.getParams()):
        print('  ', nm, val)

    exit(0)

    mod = tr.getModelImage(0)

    np.savetxt('output/rex_noiseless.txt', mod)

    # Reset the source params.