Esempio n. 1
0
def PatchFitter(
    all_data,
    all_dq,
    ini_psf,
    patch_shape,
    id_start,
    background="linear",
    sequence=["shifts", "psf"],
    tol=1.0e-10,
    eps=1.0e-4,
    gamma=1.0e0,
    ini_shifts=None,
    Nthreads=20,
    floor=None,
    plotfilebase=None,
    gain=None,
    maxiter=np.Inf,
    dumpfilebase=None,
    trim_frac=0.005,
    min_data_frac=0.75,
    core_size=5,
    lower=1.0e-5,
    k=1,
    plot=False,
    clip_parms=None,
    final_clip=[1, 3.0],
    q=1.0,
    clip_shifts=False,
    h=1.4901161193847656e-08,
    Nplot=20,
    small=1.0e-5,
    Nsearch=256,
    search_rate=0.05,
    search_scale=1e-2,
    shift_test_thresh=0.475,
    min_frac=0.5,
    max_nll=1.0e10,
    Nburn=10,
):
    """
    Patch fitting routines for psf inference.
    """
    assert background in [None, "constant", "linear"]
    assert np.mod(patch_shape[0], 2) == 1, "Patch shape[0] must be odd"
    assert np.mod(patch_shape[1], 2) == 1, "Patch shape[1] must be odd"
    assert np.mod(core_size, 2) == 1, "Core size must be odd"
    assert (patch_shape[0] * patch_shape[1]) == all_data.shape[1], "Patch shape does not match data shape"
    kinds = ["shifts", "psf", "evaluate", "plot_data"]
    for i in range(len(sequence)):
        assert sequence[i] in kinds, "sequence not allowed"

    # set parameters
    parms = InferenceParms(
        h,
        k,
        q,
        eps,
        tol,
        gain,
        plot,
        floor,
        gamma,
        all_data.shape[0],
        Nplot,
        small,
        Nsearch,
        id_start,
        max_nll,
        min_frac,
        Nthreads,
        core_size,
        background,
        None,
        patch_shape,
        search_rate,
        plotfilebase,
        search_scale,
        ini_psf.shape,
        shift_test_thresh,
        ini_psf,
    )

    # initialize
    current_psf = ini_psf.copy()
    current_psf /= current_psf.max()
    current_cost = np.inf
    if ini_shifts is not None:
        shifts = ini_shifts
        ref_shifts = ini_shifts.copy()
    else:
        ref_shifts = np.zeros((all_data.shape[0], 2))
    if Nburn is not None:
        current_cost = None
        burn_iter = 0
    else:
        data = all_data
        dq = all_dq
        current_cost = np.inf
        burn_iter = None

    # run
    t0 = time.time()
    while True:
        t = time.time()
        # assign data used during burnin
        if burn_iter is not None:
            burn_iter += 1
            if burn_iter == Nburn:
                data = all_data
                dq = all_dq
                current_cost = np.inf
            else:
                burn_size = np.ceil(1.0 * all_data.shape[0] / Nburn)
                data = all_data[: burn_iter * burn_size]
                dq = all_dq[: burn_iter * burn_size]

        # minimum number of patches, mask initialization
        Nmin = np.ceil(min_data_frac * data.shape[0]).astype(np.int)
        mask = np.arange(data.shape[0], dtype=np.int)
        parms.Ndata = data.shape[0]

        # run a iteration
        for kind in sequence:

            if parms.iter >= maxiter:
                return current_psf

            if kind == "shifts":
                parms.clip_parms = None
                shifts, nll = update_shifts(
                    data[:, parms.core_ind], dq[:, parms.core_ind], current_psf, np.zeros((data.shape[0], 2)), parms
                )
                ref_shifts = shifts.copy()

                print "Shift step 1 done nll, total: ", nll.sum()
                print "Shift step 1 done nll, min: ", nll.min()
                print "Shift step 1 done nll, median: ", np.median(nll)
                print "Shift step 1 done nll, max: ", nll.max()

                if (trim_frac is not None) & (mask.size > Nmin):
                    assert trim_frac > 0.0, "trim_frac must be positive or None"
                    Ntrim = np.ceil(mask.size * trim_frac).astype(np.int)
                    if mask.size - Ntrim < Nmin:
                        Ntrim = mask.size - Nmin

                    # sort and trim the arrays
                    ind = np.sort(np.argsort(nll)[:-Ntrim])
                    dq = dq[ind]
                    data = data[ind]
                    mask = mask[ind]
                    ref_shifts = ref_shifts[ind]
                    parms.Ndata = data.shape[0]
                    parms.data_ids = parms.data_ids[ind]

                    # re-run shifts
                    shifts, nll = update_shifts(
                        data[:, parms.core_ind], dq[:, parms.core_ind], current_psf, ref_shifts, parms
                    )
                else:
                    ind = np.arange(data.shape[0])

                print "Shift step 2 done nll, total: ", nll.sum()
                print "Shift step 2 done nll, min: ", nll.min()
                print "Shift step 2 done nll, median: ", np.median(nll)
                print "Shift step 2 done nll, max: ", nll.max()

                if dumpfilebase is not None:
                    name = dumpfilebase + "_mask_%d.dat" % parms.iter
                    np.savetxt(name, mask, fmt="%d")
                    name = dumpfilebase + "_shifts_%d.dat" % parms.iter
                    np.savetxt(name, shifts)
                    name = dumpfilebase + "_shift_nll_%d.dat" % parms.iter
                    np.savetxt(name, nll)

            if kind == "evaluate":
                parms.return_parms = True
                set_clip_parameters(clip_parms, parms, final_clip)
                nll, fit_parms, masks = evaluate((data, dq, shifts, current_psf, parms, False))
                parms.return_parms = False

            if kind == "psf":
                set_clip_parameters(clip_parms, parms, final_clip)
                if parms.k == 1:
                    new_psf, cost = update_psf_linear(current_psf, data, dq, shifts, nll, fit_parms, masks, parms)
                else:
                    new_psf, cost = update_psf(current_psf, data, dq, shifts, nll, fit_parms, masks, parms)

                if new_psf is not None:
                    psf_plot(
                        ini_psf,
                        np.maximum(parms.small, current_psf),
                        np.maximum(parms.small, new_psf),
                        parms.small,
                        parms,
                    )
                    current_psf = new_psf
                    if dumpfilebase is not None:
                        hdu = pf.PrimaryHDU(current_psf / current_psf.max())
                        hdu.writeto(dumpfilebase + "_psf_%d.fits" % parms.iter, clobber=True)

            if kind == "plot_data":
                if clip_parms is None:
                    parms.clip_parms = [1, np.inf]
                else:
                    try:
                        parms.clip_parms = clip_parms[parms.iter]
                    except:
                        parms.clip_parms = final_clip

                parms.plot_data = True
                foo = evaluate(
                    (data[: parms.Nplot], dq[: parms.Nplot], shifts[: parms.Nplot], current_psf, parms, False)
                )
                parms.plot_data = False

        if current_cost is None:
            tup = (-99.0, cost)
        else:
            tup = (current_cost, cost)
        print "\n\nUsing %d patches" % data.shape[0]
        print "Current cost: %0.2e, new cost %0.2e" % tup
        dt = (time.time() - t) / 3600.0
        dt0 = (time.time() - t0) / 3600.0
        print "Iter %d took %0.2e hrs, total %0.2e hrs\n\n" % (parms.iter, dt, dt0)
        if current_cost is not None:
            # assert cost < current_cost, 'Global cost did not decrease'
            if np.abs((current_cost - cost) / cost) < tol:
                print "Converged at cost %s" % cost
                return current_psf
            else:
                current_cost = cost
        parms.iter += 1
Esempio n. 2
0
def learn_psf(
    data,
    dq,
    initial_psf,
    clip_parms,
    noise_parms,
    plotfilebase,
    kernel_parms,
    patch_shape,
    knn=32,
    min_patch_frac=0.75,
    core_size=5,
    nll_tol=1.0e-5,
    k=3,
    q=1.0,
    Nplot=20,
    plot=False,
    flann_precision=0.99,
    final_clip=[1, 3.0],
    background="constant",
    Nthreads=20,
    max_iter=20,
    max_nll=1.0e10,
    shift_test_thresh=0.475,
):
    """
    Inference routine for learning a psf model via scaled data and a 
    kernel basis.
    """
    assert background in [None, "constant", "linear"]
    assert np.mod(patch_shape[0], 2) == 1, "Patch shape[0] must be odd"
    assert np.mod(patch_shape[1], 2) == 1, "Patch shape[1] must be odd"
    assert np.mod(core_size, 2) == 1, "Core size must be odd"
    assert (patch_shape[0] * patch_shape[1]) == data.shape[1], "Patch shape does not match data shape"

    # bundle parameters to be passed to other functions
    parms = InferenceParms(
        k,
        q,
        knn,
        plot,
        data.shape[0],
        Nplot,
        nll_tol,
        max_nll,
        Nthreads,
        core_size,
        background,
        clip_parms,
        patch_shape,
        noise_parms,
        kernel_parms,
        plotfilebase,
        min_patch_frac,
        flann_precision,
        initial_psf.shape,
        shift_test_thresh,
    )

    # initialize
    print "Initialized with %d patches\n" % data.shape[0]
    initial_psf /= initial_psf.max()
    psf_model = initial_psf.copy()
    cost = np.Inf

    # Run through data, reject patches that are bad/crowded.
    parms.clip_parms = None
    shifts, nll = update_shifts(data[:, parms.core_ind], dq[:, parms.core_ind], psf_model, parms)
    set_clip_parameters(clip_parms, parms, final_clip)
    fit_parms, fit_vars, nll, masks = fit_patches(data, dq, shifts, psf_model, parms)
    nll = np.sum(nll, axis=1)
    ind = nll < parms.max_nll
    shifts = shifts[ind]
    data = data[ind]
    dq = dq[ind]
    parms.data_ids = data[ind]
    print "%d patches are ok under the initial model\n" % data.shape[0]
    print "Initial NLL is %0.6e" % np.sum(nll)

    # Build a new psf
    psf_model = psf_builder(data, masks, shifts, fit_parms, fit_vars, parms)

    for blah in range(12):
        parms.clip_parms = None
        shifts, nll = update_shifts(data[:, parms.core_ind], dq[:, parms.core_ind], psf_model, parms)
        print blah, nll.sum(), shifts[0], shifts[-1]
        set_clip_parameters(clip_parms, parms, final_clip)
        fit_parms, fit_vars, nll, masks = fit_patches(data, dq, shifts, psf_model, parms)
        nll = np.sum(nll, axis=1)
        ind = nll < parms.max_nll
        shifts = shifts[ind]
        data = data[ind]
        dq = dq[ind]
        parms.data_ids = data[ind]
        print "New NLL is %0.6e" % np.sum(nll)
        f = pl.figure(figsize=(16, 8))
        pl.subplot(121)
        pl.imshow(np.abs(psf_model), interpolation="nearest", origin="lower", norm=LogNorm(vmin=1.0e-6, vmax=1.0))
        psf_model = psf_builder(data, masks, shifts, fit_parms, fit_vars, parms)
        pl.colorbar(shrink=0.7)
        pl.subplot(122)
        pl.imshow(np.abs(psf_model), interpolation="nearest", origin="lower", norm=LogNorm(vmin=1.0e-6, vmax=1.0))
        pl.colorbar(shrink=0.7)
        print psf_model[50, 50]
        print psf_model.max()
        f.savefig("../../plots/foo.png")
    assert 0