Exemple #1
0
def findLnLike(floats, datas):
    """ SciPy Minimizer.
        This could be moved to waveModel.py IF you decide you
        don't need to look at the trace arrays anymore. """
    global stTrace, enTrace, sloTrace

    # Yeezus.  At this point should you just write a damn class?
    # Unpack parameters.
    dataList, tempList, InterpFn = datas
    data, dataTS, dataE, dataST, loWin, hiWin, dataNoise = dataList
    temp, tempTS, tempE, tempST = tempList
    st = en = slo = 0
    if len(floats) == 1:  # basinhopping case
        slo, st, en = floats[0], dataST, dataE

    if len(floats) == 3:  # minimize case
        st, en, slo = floats

    # Make a trace
    stTrace = np.append(stTrace, st)
    enTrace = np.append(enTrace, en)
    sloTrace = np.append(sloTrace, slo)
    # print st, en, slo

    # Build the model to compare with data
    model, modelTS = wm.MakeModel(dataList,
                                  tempList, [st, en, slo],
                                  fn=InterpFn)

    # Find LL and return
    lnLike = 0.5 * np.sum(
        np.power((data - model) / dataNoise, 2) -
        np.log(1 / np.power(dataNoise, 2)))
    return lnLike
def findLnLike(floats, datas):
    """ SciPy Minimizer.
        This could be moved to waveModel.py IF you decide you
        don't need to look at the trace arrays anymore.

        Also, the version in this file (match-minimizer.py)
        is different because it tries to align the MAX time
        (not the start time) of the waveforms.
    """
    global mtTrace, enTrace, sloTrace

    # Unpack parameters.
    dataList, tempList, InterpFn = datas
    data, dataTS, dataE, dataMax, loWin, hiWin, dataNoise = dataList
    temp, tempTS, tempE, tempMax = tempList

    # Can fix parameters here by setting them back to their input guess values
    mt, en, slo = 0, 0, 0
    if len(floats)==1: # basinhopping case
        slo, mt, en = floats[0], dataMax, dataE
    if len(floats)==3: # minimize case (no fixing)
        mt, en, slo = floats

    # Make a trace
    mtTrace = np.append(mtTrace,mt)
    enTrace = np.append(enTrace,en)
    sloTrace = np.append(sloTrace,slo)
    # print mt, en, slo

    # Build the model to compare with data
    model, modelTS = wm.MakeModel(dataList, tempList, [mt,en,slo], fn=InterpFn)

    # Find LL of data vs. model
    lnLike = 0.5 * np.sum ( np.power((data-model)/dataNoise, 2) - np.log( 1 / np.power(dataNoise,2) ) )
    return lnLike
Exemple #3
0
def findLnLike(floats, datas):
    global stTrace, enTrace, sloTrace

    # Unpack all parameters
    st, en, slo = floats
    dataList, tempList, InterpFn = datas
    data, dataTS, dataE, dataST, loWin, hiWin, dataNoise = dataList
    temp, tempTS, tempE, tempST = tempList

    # Make a trace
    stTrace = np.append(stTrace, st)
    enTrace = np.append(enTrace, en)
    sloTrace = np.append(sloTrace, slo)

    # Build the model to compare with data
    model, modelTS = wm.MakeModel(dataList, tempList, floats, fn=InterpFn)

    # Return NLL
    lnLike = -0.5 * np.sum(
        np.power((data - model) / dataNoise, 2) -
        np.log(1 / np.power(dataNoise, 2)))
    return -1.0 * lnLike
Exemple #4
0
def main():
    """ Get a cool SciPy minimizer working and compare to MCMC. """

    # Load data and template
    npzfile = np.load("./data/optimumInputs.npz")
    rl, tl = npzfile['arr_0'], npzfile['arr_1']
    wave, waveTS, dataE, dataST = rl[0], rl[1], rl[2], rl[3]
    temp, tempTS, tempE, tempST = tl[0], tl[1], tl[2], tl[3]

    # Window the fit around rising edge - start time calculator method
    loWin, hiWin = dataST - 1000, dataST + 4000  # ns
    if loWin < waveTS[0] or hiWin > waveTS[-1]:
        print "Window out of range!  dataST: %.1f  loWin %.1f  hiWin %.1f" % (
            dataST, loWin, hiWin)
    idx = np.where((waveTS >= loWin) & (waveTS <= hiWin))
    data = wave[idx]
    dataTS = waveTS[idx]

    # Pack into lists
    dataNoise = 2.  # just a guess - 1 sigma baseline adc values
    rawList = [wave, waveTS, dataE, dataST]
    dataList = [data, dataTS, dataE, dataST, loWin, hiWin, dataNoise]
    tempList = [temp, tempTS, tempE, tempST]

    # Recreate the guess and the guess's rising edge
    guessFull, guessFullTS = wm.MakeModel(rawList,
                                          tempList, [dataST, dataE, 1.],
                                          opt="full")
    guess, guessTS = wm.MakeModel(dataList,
                                  tempList, [dataST, dataE, 1.],
                                  opt="!fancy")

    # Make an "almost complete" guess - no MCMC
    # st, en, slo = dataST-100, dataE, 5
    # InterpFn = interpolate.interp1d(tempTS, temp, kind="linear", copy="False", assume_sorted="True")
    # model, modelTS = wm.MakeModel(dataList, tempList, [st,en,slo], fn=InterpFn)

    # Fit with MCMC and get best-fit parameters
    numSteps, burnIn = 3000, 1800  # default: 10000, 5000.  fast: 3000, 1800  long test: 20000,10000
    wfModel = wm.TemplateModel(dataList, dataNoise, tempList)
    M = pymc.MCMC(pymc.Model(wfModel))
    M.use_step_method(pymc.Metropolis,
                      M.startTime,
                      proposal_sd=100.,
                      proposal_distribution='Normal')
    M.use_step_method(pymc.Metropolis,
                      M.energy,
                      proposal_sd=1.,
                      proposal_distribution='Normal')
    M.use_step_method(pymc.Metropolis,
                      M.slowness,
                      proposal_sd=100.,
                      proposal_distribution='Normal')
    M.sample(iter=numSteps, verbose=0)
    st = np.median(M.trace("startTime")[:])
    en = np.median(M.trace("energy")[:])
    slo = np.median(M.trace("slowness")[:])
    InterpFn = interpolate.interp1d(tempTS,
                                    temp,
                                    kind="linear",
                                    copy="False",
                                    assume_sorted="True")
    model, modelTS = wm.MakeModel(dataList,
                                  tempList, [st, en, slo],
                                  fn=InterpFn)
    print "MCMC:", st, en, slo

    # Fit with SciPy minimizer
    MakeTracesGlobal(
    )  # creates 3 global arrays: startTrace, enTrace, sloTrace
    floats = [dataST, dataE, 1]
    print "Minimizer guesses:", floats
    datas = [dataList, tempList, InterpFn]
    result = minimize(findLnLike, floats, args=datas, method="Nelder-Mead")
    st, en, slo = result["x"]
    print "Minimizer: %.1f  %.1f  %.1f  Success: %s.  %s" % (
        st, en, slo, result["success"], result["message"])
    minimizer, minimizerTS = wm.MakeModel(dataList,
                                          tempList, [st, en, slo],
                                          fn=InterpFn)

    # plots
    fig = plt.figure(figsize=(11, 7), facecolor='w')
    p1 = plt.subplot2grid((6, 7), (0, 0), colspan=4, rowspan=2)  # original
    p2 = plt.subplot2grid((6, 7), (2, 0), colspan=4, rowspan=3)  # rising edge
    p3 = plt.subplot2grid((6, 7), (0, 4), colspan=3, rowspan=2)  # trace 1
    p4 = plt.subplot2grid((6, 7), (2, 4), colspan=3, rowspan=2,
                          sharex=p3)  # trace 2
    p5 = plt.subplot2grid((6, 7), (4, 4), colspan=3, rowspan=2,
                          sharex=p3)  # trace 3

    # p1 = plt.subplot(211)
    p1.set_title("Energy %.1f keV  Start Time %.0f ns" % (dataE, dataST))
    p1.set_ylabel("ADC [A.U.]", y=0.95, ha='right')
    p1.set_xlabel("Time (ns)", x=0.95, ha='right')
    p1.plot(waveTS, wave, color='blue', alpha=0.8, label='Data WF')
    p1.plot(guessFullTS,
            guessFull,
            color='orange',
            alpha=0.8,
            label='Guess WF')
    p1.axvline(x=dataST, color='green')
    p1.legend(loc=4)

    # p2 = plt.subplot(212)
    p2.plot(dataTS, data, color='blue', label='Data')
    p2.plot(guessTS, guess, color='orange', label='Guess')
    p2.plot(modelTS, model, color='red', linewidth=4, alpha=0.8, label='MCMC')
    p2.plot(minimizerTS,
            minimizer,
            color='cyan',
            linewidth=1,
            label='Nelder-Mead')
    p2.legend(loc=4)

    p3.cla()
    p3.set_title("startTime %.1f  Energy %.2f  Slow %.1f" % (st, en, slo))
    p3.plot(stTrace[1:])
    p3.set_ylabel('startTime')

    p4.cla()
    p4.plot(enTrace[1:])
    p4.set_ylabel('energy')

    p5.cla()
    p5.plot(sloTrace[1:])
    p5.set_ylabel('slowness')

    plt.tight_layout()
    plt.subplots_adjust(hspace=0.35)
    # plt.show(block=False)
    # plt.show()
    plt.savefig("./plots/minimizer-test.pdf")
Exemple #5
0
def main():
    # Load a data and a temp waveform
    npzfile = np.load("./data/tailSlopeInputs.npz")
    dl, tl = npzfile['arr_0'], npzfile['arr_1']
    data, dataTS, dataE, dataRT = dl[0], dl[1], dl[2], dl[3]
    temp, tempTS, tempE, tempRT = tl[0], tl[1], tl[2], tl[3]

    print len(dataTS)

    # Denoise, just for fun
    wp = pywt.WaveletPacket(data=data,
                            wavelet='haar',
                            mode='symmetric',
                            maxlevel=3)
    new_wp = pywt.WaveletPacket(data=None, wavelet='haar', mode='symmetric')
    new_wp['aaa'] = wp['aaa'].data
    waveDenoised = new_wp.reconstruct(update=False)
    # ix = len(newWave)
    # waveDenoised = np.delete(newWave, [ix-1,ix-2,ix-3])

    # Start the tail slope fit ranged based on the calculated start time
    # Plot the template "guess" to illustrate what we would expect to see
    guessFull, guessFullTS = wm.MakeModel(dl,
                                          tl, [dataRT, dataE, 1.],
                                          opt="full")
    idxMax = np.where(guessFull == guessFull.max())  # returns an array/tuple
    idxMax = idxMax[0][0]  # "cast" to int
    idxMax += 100  # move it ahead in case this is a slow pulse
    tail, tailTS = data[idxMax:], dataTS[idxMax:]

    # Exponential fit
    a, b = dataE, 72000
    popt, pcov = curve_fit(wl.tailModelExp, tailTS, tail, p0=[a, b])
    a, b = popt[0], popt[1]

    # Polynomial curve fit
    popt2, pcov2 = curve_fit(wl.tailModelPol, tailTS, tail)
    d, e, f, g = popt2[0], popt2[1], popt2[2], popt2[3]

    # Plots
    fig = plt.figure(figsize=(8, 5), facecolor='w')
    a1 = plt.subplot(111)
    a1.set_ylabel("ADC", y=0.95, ha='right')
    a1.set_xlabel("Time (ns)", x=0.95, ha='right')
    a1.set_title("Energy %.1f  StartTime %.1f  %.1e * exp(-t/%.0f[us]) " %
                 (dataE, dataRT, a, b / 1000))
    a1.plot(dataTS, data, color='blue', alpha=0.7, label='Data')
    a1.plot(dataTS,
            waveDenoised,
            color='red',
            alpha=0.8,
            label='Lvl3 Haar Denoised')
    a1.axvline(x=guessFullTS[np.where(guessFull == guessFull.max())],
               color='black',
               label='Guess Max ADC')
    a1.plot(guessFullTS, guessFull, color='orange', label='Scaled Template WF')
    a1.plot(tailTS,
            wl.tailModelExp(tailTS, *popt),
            color='cyan',
            linewidth=2,
            label='Tail Exp. Fit')
    a1.plot(tailTS,
            wl.tailModelPol(tailTS, *popt2),
            color='green',
            linewidth=2,
            label='Tail Curve Fit')
    a1.legend(loc=4)
    plt.show()
Exemple #6
0
def main(argv):
    """ Interactive-fit or 'rapid'-fit waveforms that pass a given TCut.
        BUG: Doesn't always work with a TChain.  Add input files together
        with hadd and use a single TFile.
    """
    scanSpeed = 0.2
    iList = -1
    opt1 = ""
    intMode, batMode = False, False
    if (len(argv) >= 1): opt1 = argv[0]
    if "-i" in (opt1):
        intMode = True
        print "Interactive mode selected."
    if "-b" in (opt1):
        batMode = True
        print "Batch mode selected.  A new file will be created."

    # Load template waveform (created by gen-template.py)
    npzfile = np.load("./data/genTemplateWF.npz")
    temp, tempTS, tempE, tempST = npzfile['arr_0'] + 1, npzfile[
        'arr_1'], npzfile['arr_2'], npzfile['arr_3'] * 10

    # Set cuts
    # theCut = inputFile.Get("cutUsedHere").GetTitle()
    # DS3 "big DC" cut - PRELIMINARY
    theCut = "trapENFCal > 0.8 && gain==0 && mHClean==1 && isGood && !muVeto && !wfDCBits && !isLNFill1 && !isLNFill2 && trapETailMin < 0 && channel!=596 && channel!=676 && channel!=676 && channel!=612 && channel!=1104 && channel!=1200 && channel!=1334 && channel!=1336 && tOffset < 10 && waveS5/TMath::Power(trapENFCal,1/4) < 1200 && (waveS3-waveS2)/trapENFCal > 100 && (waveS3-waveS2)/trapENFCal < 300 && !(channel==692 && (run==16974 || run==16975 || run==16976 || run==16977 || run==16978 || run==16979)) && butterTime < 11000"
    # theCut += " && trapENFCal > 1.5 && trapENFCal < 2.1"
    # theCut += " && trapENFCal < 20 && trapENFCal > 2 && run > 17480"
    # theCut += " && kvorrT/trapENFCal > 2.2 && trapENFCal > 2 && trapENFCal < 10"

    # Set file I/O and create entry lists
    startT = time.clock()
    inFile = "~/project/wavelet-skim/waveletSkimDS3_1.root"
    # inFile = "~/project/wavelet-skim/hadd/waveletSkimDS3.root"
    outFile = "~/project/fit-skim/fitSkimDS3_1.root"
    inputFile = TFile(inFile)
    waveTree = inputFile.Get("skimTree")
    print "Found", waveTree.GetEntries(
    ), "input entries.  Using cut:\n", theCut, "\n"
    waveTree.Draw(">>elist", theCut, "entrylist")
    elist = gDirectory.Get("elist")
    waveTree.SetEntryList(elist)
    nList = elist.GetN()
    print "Found", nList, "entries passing cuts."
    stopT = time.clock()
    print "Data loading time (s): ", (stopT - startT)

    # In batch mode ONLY, create an output file+tree & append new branches
    outputFile = TFile()
    outTree = TTree()
    if batMode:
        outputFile = TFile(outFile, "RECREATE")
        print "Attempting tree copy to", outFile
        outTree = waveTree.CopyTree("")
        outTree.Write()
        print "Wrote", outTree.GetEntries(), "entries."
        cutUsed = TNamed("cutUsedHere", theCut)
        cutUsed.Write()
    fitStart = std.vector("double")()
    fitE = std.vector("double")()
    fitSlo = std.vector("double")()
    fitStartSD = std.vector("double")()
    fitESD = std.vector("double")()
    fitSloSD = std.vector("double")()
    fitChi2NDF = std.vector("double")()
    fitLnLike = std.vector("double")()
    tExp1 = std.vector("double")()
    tExp2 = std.vector("double")()
    tPol0 = std.vector("double")()
    tPol1 = std.vector("double")()
    tPol2 = std.vector("double")()
    tPol3 = std.vector("double")()
    baseAvg = std.vector("double")()
    baseNoise = std.vector("double")()
    bFitStart = outTree.Branch("fitStart", fitStart)
    bFitE = outTree.Branch("fitE", fitE)
    bFitSlo = outTree.Branch("fitSlo", fitSlo)
    bFitStart_sd = outTree.Branch("fitStartSD", fitStartSD)
    bFitE_sd = outTree.Branch("fitESD", fitESD)
    bFitSlo_sd = outTree.Branch("fitSloSD", fitSloSD)
    bFitChi2NDF = outTree.Branch("fitChi2NDF", fitChi2NDF)
    bFitLnLike = outTree.Branch("fitLnLike", fitLnLike)
    bTExp1 = outTree.Branch("tExp1", tExp1)
    bTExp2 = outTree.Branch("tExp2", tExp2)
    bTPol0 = outTree.Branch("tPol0", tPol0)
    bTPol1 = outTree.Branch("tPol1", tPol1)
    bTPol2 = outTree.Branch("tPol2", tPol2)
    bTPol3 = outTree.Branch("tPol3", tPol3)
    bBaseAvg = outTree.Branch("baseAvg", baseAvg)
    bBaseNoise = outTree.Branch("baseNoise", baseNoise)

    # Make a figure
    # with PdfPages('multipage_pdf.pdf') as pdf:
    fig = plt.figure(figsize=(11, 7), facecolor='w')
    p1 = plt.subplot2grid((6, 7), (0, 0), colspan=4, rowspan=2)  # original
    p2 = plt.subplot2grid((6, 7), (2, 0), colspan=4, rowspan=3)  # rising edge
    p3 = plt.subplot2grid((6, 7), (5, 0), colspan=4, rowspan=1)  # residual
    p4 = plt.subplot2grid((6, 7), (0, 4), colspan=3, rowspan=2)  # trace 1
    p5 = plt.subplot2grid((6, 7), (2, 4), colspan=3, rowspan=2,
                          sharex=p4)  # trace 2
    p6 = plt.subplot2grid((6, 7), (4, 4), colspan=3, rowspan=2,
                          sharex=p4)  # trace 3
    if not batMode: plt.show(block=False)

    # Setup multiprocessing
    manager = Manager()
    returnDict = manager.dict()

    # Loop over events
    while (True):
        saveMe = False
        iList += 1
        if intMode == True and iList != 0:
            value = raw_input()
            if value == 'q': break
            if value == 's': saveMe = True
            if value == 'p': iList -= 2  # previous
            if (value.isdigit()): iList = int(value)  # go to entry
        if iList >= elist.GetN(): break

        entry = waveTree.GetEntryNumber(iList)
        waveTree.LoadTree(entry)
        waveTree.GetEntry(entry)
        nChans = waveTree.channel.size()
        numPass = waveTree.Draw("channel", theCut, "GOFF", 1, iList)
        chans = waveTree.GetV1()
        chanList = list(set(int(chans[n]) for n in xrange(numPass)))

        fitStart.assign(nChans, -9999)
        fitE.assign(nChans, -9999)
        fitSlo.assign(nChans, -9999)
        fitStartSD.assign(nChans, -9999)
        fitESD.assign(nChans, -9999)
        fitSloSD.assign(nChans, -9999)
        fitChi2NDF.assign(nChans, -9999)
        fitLnLike.assign(nChans, -9999)
        tExp1.assign(nChans, -9999)
        tExp2.assign(nChans, -9999)
        tPol0.assign(nChans, -9999)
        tPol1.assign(nChans, -9999)
        tPol2.assign(nChans, -9999)
        tPol3.assign(nChans, -9999)
        baseAvg.assign(nChans, -9999)
        baseNoise.assign(nChans, -9999)

        # Loop over hits passing cuts
        hitList = (iH for iH in xrange(nChans)
                   if waveTree.channel.at(iH) in chanList
                   )  # a 'generator expression'
        for iH in hitList:

            # Load waveform for this hit
            run = waveTree.run
            chan = waveTree.channel.at(iH)
            dataE = waveTree.trapENFCal.at(iH)
            dataST = waveTree.butterTime.at(iH)  # replace with blrwfFMR50?
            toe = waveTree.kvorrT.at(iH) / dataE
            print "%d / %d  Run %d  nCh %d  chan %d  trapENF %.1f  t/e %.1f" % (
                iList, nList, run, nChans, chan, dataE, toe)
            signal = wl.processWaveform(waveTree.MGTWaveforms.at(iH),
                                        opt='full')
            waveBLSub = signal.GetWaveBLSub()
            waveFilt = signal.GetWaveFilt()
            waveTS = signal.GetTS()
            dataBaseline, dataNoise = signal.GetBaseNoise()

            # Denoise the data waveform (take only lowest-frequency components)
            wp = pywt.WaveletPacket(data=waveBLSub,
                                    wavelet='haar',
                                    mode='symmetric',
                                    maxlevel=3)
            new_wp = pywt.WaveletPacket(data=None,
                                        wavelet='haar',
                                        mode='symmetric')
            new_wp['aaa'] = wp['aaa'].data
            waveDenoised = new_wp.reconstruct(update=False)

            # Window the fit around rising edge - start time calculator method
            loWin, hiWin = dataST - 1000, dataST + 4000  # ns
            if loWin < waveTS[0] or hiWin > waveTS[-1]:
                print "Window out of range!  dataST: %.1f  loWin %.1f  hiWin %.1f" % (
                    dataST, loWin, hiWin)
            idx = np.where((waveTS >= loWin) & (waveTS <= hiWin))
            data = waveBLSub[idx]
            # data = waveDenoised[idx]
            dataTS = waveTS[idx]

            # Pack into lists
            rawList = [waveBLSub, waveTS, dataE, dataST]
            dataList = [data, dataTS, dataE, dataST, loWin, hiWin]
            tempList = [temp, tempTS, tempE, tempST]

            # Optionally save something to a file
            if saveMe:
                print "Saved entry", iList, iH
                np.savez("./data/tailSlopeInputs.npz", rawList, tempList)

            # Recreate the guess and the guess's rising edge
            guessFull, guessFullTS = wm.MakeModel(dataList,
                                                  tempList,
                                                  [dataST, dataE, 1.],
                                                  opt="full")
            guess, guessTS = wm.MakeModel(dataList,
                                          tempList, [dataST, dataE, 1.],
                                          opt="!fancy")

            # Make an "almost complete" guess - no MCMC
            # st, en, slo = dataST-100, dataE, 5
            # InterpFn = interpolate.interp1d(tempTS, temp, kind="linear", copy="False", assume_sorted="True")
            # model, modelTS = wm.MakeModel(dataList, tempList, [st,en,slo], fn=InterpFn)

            # Fit with MCMC and get best-fit parameters
            numSteps, burnIn = 3000, 1800  # default: 10000, 5000.  fast: 3000, 1800  long test: 20000,10000
            wfModel = wm.TemplateModel(dataList, dataNoise, tempList)
            p = Process(target=RunMCMC,
                        args=(wfModel, numSteps, burnIn, returnDict))
            p.start()
            p.join()
            startTimeTr = returnDict["startTimeTr"]
            energyTr = returnDict["energyTr"]
            slownessTr = returnDict["slownessTr"]
            st = np.median(startTimeTr[burnIn:])
            en = np.median(energyTr[burnIn:])
            slo = np.median(slownessTr[burnIn:])
            InterpFn = interpolate.interp1d(tempTS,
                                            temp,
                                            kind="linear",
                                            copy="False",
                                            assume_sorted="True")
            model, modelTS = wm.MakeModel(dataList,
                                          tempList, [st, en, slo],
                                          fn=InterpFn)

            # Save some extra parameters for the ROOT output
            # Calculate residual, Chi2/NDF, likelihood, etc.
            st_std = np.std(startTimeTr[burnIn:])
            en_std = np.std(energyTr[burnIn:])
            slo_std = np.std(slownessTr[burnIn:])
            residual = model - data
            frac = (np.power(data - model, 2)) / np.abs(model)
            chi2NDF = np.sum(frac) / len(model)
            inv_sigma2 = 1.0 / (dataNoise**2)
            lnLike = -0.5 * (np.sum((data - model)**2 * inv_sigma2 -
                                    np.log(inv_sigma2)))

            # ** Do a separate & simple fit of the tail slope **
            # TODO: Add this to process-waveforms.py
            idxMax = np.where(
                guessFull == guessFull.max())  # returns an array/tuple
            idxMax = idxMax[0][0]  # "cast" to int
            tail, tailTS = waveDenoised[idxMax:], waveTS[idxMax:]
            popt, _ = curve_fit(wl.tailModelPol, tailTS, tail)  # poly fit
            pol0, pol1, pol2, pol3 = popt[0], popt[1], popt[2], popt[3]
            a, b = dataE, 72000
            popt2, _ = curve_fit(wl.tailModelExp, tailTS, tail,
                                 p0=[a, b])  # expo fit
            e1, e2 = popt2[0], popt2[1]

            # Assign values to output vectors and fill branches
            fitStart[iH], fitStartSD[iH] = st, st_std
            fitE[iH], fitESD[iH] = en, en_std
            fitSlo[iH], fitSloSD[iH] = slo, slo_std
            fitChi2NDF[iH] = chi2NDF
            fitLnLike[iH] = lnLike
            tExp1[iH], tExp2[iH] = e1, e2
            tPol0[iH], tPol1[iH], tPol2[iH], tPol3[iH] = pol0, pol1, pol2, pol3
            baseAvg[iH] = dataBaseline
            baseNoise[iH] = dataNoise
            if batMode:
                bFitStart.Fill()
                bFitE.Fill()
                bFitSlo.Fill()
                bFitStart_sd.Fill()
                bFitE_sd.Fill()
                bFitSlo_sd.Fill()
                bFitChi2NDF.Fill()
                bFitLnLike.Fill()
                bTExp1.Fill()
                bTExp2.Fill()
                bTPol0.Fill()
                bTPol1.Fill()
                bTPol2.Fill()
                bTPol3.Fill()
                bBaseAvg.Fill()
                bBaseNoise.Fill()
                if iList % 5000 == 0:
                    outTree.Write("", TObject.kOverwrite)
                    print "%d / %d entries saved (%.2f %% done)." % (
                        iList, nList, 100 * (float(iList) / nList))

            # If not in batch mode, fill the figure
            if batMode: continue
            p1.cla()
            p1.set_ylabel("ADC")
            p1.set_title(
                "Run %d  Channel %d  Entry %d\ntrapENFCal %.1f  T/E %.1f  ST %.1f"
                % (run, chan, iList, dataE, toe, dataST))
            p1.plot(waveTS, waveBLSub, color='blue')
            p1.plot(waveTS, waveDenoised, color='red', alpha=0.8)
            p1.plot(guessFullTS, guessFull, color='orange', linewidth=2)
            p1.axvline(x=dataST, color='green', linewidth=2)
            p1.axvline(x=loWin, color='black')
            p1.axvline(x=hiWin, color='black')
            p1.plot(tailTS,
                    wl.tailModelPol(tailTS, *popt),
                    color='cyan',
                    linewidth=1)  # tail poly fit
            p1.plot(tailTS,
                    wl.tailModelExp(tailTS, *popt2),
                    color='cyan',
                    linewidth=1)  # tail expo fit

            p2.cla()
            p2.plot(dataTS, data, color='blue', label='Data')
            p2.plot(guessTS, guess, color='orange', label='Guess')
            # special: plot the values of the trace after burn-in
            # to see how the model is covering the "money-zone"/rising edge after it's converged.
            # for i in range(burnIn,numSteps):
            # st_tr, en_tr, slo_tr = M.trace('startTime')[i], M.trace('energy')[i], M.trace('slowness')[i]
            # trace, traceTS = wm.MakeModel(dataList, tempList, [st_tr,en_tr,slo_tr], fn=InterpFn)
            # p2.plot(traceTS, trace, color='red',alpha=0.1,linewidth=2)
            p2.plot(modelTS,
                    model,
                    color='red',
                    linewidth=3,
                    alpha=0.7,
                    label='Best Fit')
            p2.legend(loc=4)

            p3.cla()
            p3.set_xlabel("Time [ns]", x=0.95, ha='right')
            p3.set_ylabel("Residual [ADC]")
            p3.plot(modelTS, residual, color='red')
            p3.axhline(y=0, color='blue', alpha=0.3)
            p3.axhline(y=dataNoise, color='blue', alpha=0.3)
            p3.axhline(y=-1.0 * dataNoise, color='blue', alpha=0.3)

            p4.cla()
            minST = tempST - tempTS[-1] + hiWin
            maxST = tempST - tempTS[0] + loWin
            p4.set_title(
                "startTime %.1f  Energy %.2f\nSlow %.1f  chi2/ndf %.1f  Min %d  Max %d"
                % (st, en, slo, chi2NDF, minST, maxST))
            p4.plot(startTimeTr[:])
            p4.set_ylabel('startTime')
            p4.axvline(x=burnIn, color='red', alpha=0.5)

            p5.cla()
            p5.plot(energyTr[:])
            p5.set_ylabel('energy')
            p5.axvline(x=burnIn, color='red', alpha=0.5)

            p6.cla()
            p6.plot(slownessTr[:])
            p6.set_ylabel('slowness')
            p6.axvline(x=burnIn, color='red', alpha=0.5)

            plt.tight_layout()
            plt.subplots_adjust(hspace=0.35)
            plt.pause(scanSpeed)
            # pdf.savefig()

    # End loop over events
    if batMode:
        outTree.Write("", TObject.kOverwrite)
        print "Wrote", outTree.GetBranch(
            "channel").GetEntries(), "entries in the copied tree,"
        print "and wrote", bFitStart.GetEntries(
        ), "entries in the new branches."
Exemple #7
0
def main(argv):
    """ Interactive-fit or 'rapid'-fit waveforms that pass a given TCut.
        BUG: Doesn't always work with a TChain.  Add input files together
        with hadd and use a single TFile.
    """
    scanSpeed = 1.
    iList = -1
    opt1 = ""
    intMode = False
    if (len(argv) >= 1): opt1 = argv[0]
    if "-i" in (opt1): intMode = True

    # Set input data and cuts
    inputFile = TFile(
        "~/project/match-skim/waveletSkimDS5_run23920.root")  # calibration
    # inputFile = TFile("~/project/v2-processwfs/waveletSkimDS5_90.root") # noisy BG
    waveTree = inputFile.Get("skimTree")
    theCut = inputFile.Get("cutUsedHere").GetTitle()
    theCut += " && waveS5/trapENFCal < 1200 && trapENFCal > 1.5 && trapENFCal < 3"

    # Print cut and events passing cut
    waveTree.Draw(">>elist", theCut, "entrylist")
    elist = gDirectory.Get("elist")
    waveTree.SetEntryList(elist)
    nList = elist.GetN()
    print "Using cut:\n", theCut, "\n"
    print "Found", waveTree.GetEntries(), "input entries."
    print "Found", nList, "entries passing cuts."

    # Make a figure
    fig = plt.figure(figsize=(15, 10), facecolor='w')
    p1 = plt.subplot2grid((6, 7), (0, 0), colspan=4, rowspan=2)  # original
    p2 = plt.subplot2grid((6, 7), (2, 0), colspan=4, rowspan=3)  # rising edge
    p3 = plt.subplot2grid((6, 7), (5, 0), colspan=4, rowspan=1)  # residual
    p4 = plt.subplot2grid((6, 7), (0, 4), colspan=3, rowspan=2)  # trace 1
    p5 = plt.subplot2grid((6, 7), (2, 4), colspan=3, rowspan=2)  # trace 2
    p6 = plt.subplot2grid((6, 7), (4, 4), colspan=3, rowspan=2)  # trace 3
    plt.show(block=False)

    # Make template(s)
    # npzfile = np.load("./data/genTemplateWF.npz") # gen-template.py
    # temp, tempTS, tempE, tempST = npzfile['arr_0']+1, npzfile['arr_1'], npzfile['arr_2'], npzfile['arr_3']*10

    samp, r, z, tempE, tempST, smooth = 2016, 0, 15, 3938, 1000, 100  # gen-template.py
    # samp, r, z, tempE, tempST, smooth = 5000, 0, 15, 10, 2500, 100  # huge template
    # samp, r, z, tempE, tempST, smooth = 2016, 0, 15, 10, 1000, 100 # regular size template
    # samp, r, z, tempE, tempST, smooth = 2016, 30, 30, 10, 1000, 100 # regular size temp, slower rise
    # samp, r, z, tempE, tempST, smooth = 500, 0, 15, 10, 100, 100  # small template
    temp, tempTS = wl.MakeSiggenWaveform(samp, r, z, tempE, tempST, smooth)
    tempST = tempST * 10  # convert to ns

    # Loop over events
    while True:
        saveMe = False
        iList += 1
        if intMode == True and iList != 0:
            value = raw_input()
            if value == 'q': break
            if value == 's': saveMe = True
            if value == 'p': iList -= 2  # previous
            if (value.isdigit()): iList = int(value)  # go to entry
        if iList >= elist.GetN(): break

        entry = waveTree.GetEntryNumber(iList)
        waveTree.LoadTree(entry)
        waveTree.GetEntry(entry)
        nChans = waveTree.channel.size()
        numPass = waveTree.Draw("channel", theCut, "GOFF", 1, iList)
        chans = waveTree.GetV1()
        chanList = list(set(int(chans[n]) for n in xrange(numPass)))

        # Loop over hits passing cuts
        hitList = (iH for iH in xrange(nChans)
                   if waveTree.channel.at(iH) in chanList
                   )  # a 'generator expression'
        for iH in hitList:

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

            # Load waveform for this hit
            run = waveTree.run
            chan = waveTree.channel.at(iH)
            dataE = waveTree.trapENFCal.at(iH)
            dataST = waveTree.butterTime.at(iH)  # replace with blrwfFMR50?
            toe = waveTree.kvorrT.at(iH) / dataE
            print "%d / %d  Run %d  nCh %d  chan %d  trapENF %.1f  t/e %.1f" % (
                iList, nList, run, nChans, chan, dataE, toe)
            signal = wl.processWaveform(waveTree.MGTWaveforms.at(iH),
                                        opt='full')
            waveBLSub = signal.GetWaveBLSub()
            waveFilt = signal.GetWaveFilt()
            waveTS = signal.GetTS()
            baseAvg, dataNoise = signal.GetBaseNoise()

            # Denoise the data waveform (take only lowest-frequency components)
            wp = pywt.WaveletPacket(data=waveBLSub,
                                    wavelet='haar',
                                    mode='symmetric',
                                    maxlevel=3)
            new_wp = pywt.WaveletPacket(data=None,
                                        wavelet='haar',
                                        mode='symmetric')
            new_wp['aaa'] = wp['aaa'].data
            waveDenoised = new_wp.reconstruct(update=False)

            # Window the fit around rising edge - start time calculator method
            loWin, hiWin = dataST - 1000, dataST + 4000  # ns
            if loWin < waveTS[0] or hiWin > waveTS[-1]:
                print "Window out of range!  dataST: %.1f  loWin %.1f  hiWin %.1f" % (
                    dataST, loWin, hiWin)
            idx = np.where((waveTS >= loWin) & (waveTS <= hiWin))
            data = waveBLSub[idx]
            # data = waveDenoised[idx]
            dataTS = waveTS[idx]

            # Pack into lists
            rawList = [waveBLSub, waveTS, dataE, dataST]
            dataList = [data, dataTS, dataE, dataST, loWin, hiWin, dataNoise]
            tempList = [temp, tempTS, tempE, tempST]

            # Optionally save something to a file
            if saveMe:
                np.savez("./data/tailSlopeInputs.npz", rawList, tempList)

            # Recreate the guess and the guess's rising edge
            guessFull, guessFullTS = wm.MakeModel(dataList,
                                                  tempList,
                                                  [dataST, dataE, 1.],
                                                  opt="full")
            guess, guessTS = wm.MakeModel(dataList,
                                          tempList, [dataST, dataE, 1.],
                                          opt="!fancy")

            # Make an "almost complete" guess - no fitting
            # st, en, slo = dataST-100, dataE, 5
            InterpFn = interpolate.interp1d(tempTS,
                                            temp,
                                            kind="linear",
                                            copy="False",
                                            assume_sorted="True")
            # model, modelTS = wm.MakeModel(dataList, tempList, [st,en,slo], fn=InterpFn)

            # Fit with MCMC and get best-fit parameters
            # numSteps, burnIn = 10000, 5000 # default - 10000, 5000.  try 3000, 2000.  long test: 20000,10000
            # myModel = TemplateModel( dataList, dataNoise, tempList )
            # waveModel = pymc.Model( myModel )
            # M = pymc.MCMC( waveModel )
            # M.use_step_method(pymc.Metropolis, M.startTime, proposal_sd=100., proposal_distribution='Normal')
            # M.use_step_method(pymc.Metropolis, M.energy, proposal_sd=1., proposal_distribution='Normal')
            # M.use_step_method(pymc.Metropolis, M.slowness, proposal_sd=100., proposal_distribution='Normal')
            # M.sample(iter=numSteps, verbose=0) # do the fit
            # st = np.median(M.trace('startTime')[burnIn:])
            # en = np.median(M.trace('energy')[burnIn:])
            # slo = np.median(M.trace('slowness')[burnIn:])
            # InterpFn = interpolate.interp1d(tempTS, temp, kind="linear", copy="False", assume_sorted="True")
            # model, modelTS = wm.MakeModel(dataList, tempList, [st,en,slo], fn=InterpFn)

            # Fit with SciPy minimizer(s) and get best-fit parameters

            # how long does it take?
            start = time.time()

            MakeTracesGlobal(
            )  # creates 3 global arrays: startTrace, enTrace, sloTrace
            datas = [dataList, tempList, InterpFn]

            # floats = [1.]
            # result = op.basinhopping(findLnLike, floats, T=1000, stepsize=15, niter_success=2, minimizer_kwargs={"args":datas})
            # print "slowness:",result["x"],"dataE",dataE,"dataST",dataST
            # floats = [dataST, dataE, result["x"]]
            # MakeTracesGlobal() # resets traces

            floats = [dataST, dataE, 20.]
            result = op.minimize(findLnLike,
                                 floats,
                                 args=datas,
                                 method="Nelder-Mead")
            if not result["success"]: print result["message"]
            st, en, slo = result["x"]

            model, modelTS = wm.MakeModel(dataList,
                                          tempList, [st, en, slo],
                                          fn=InterpFn)

            stop = time.time()
            print "fitting took", stop - start

            # Calculate residual, Chi2/NDF, likelihood, etc.
            residual = model - data
            frac = (np.power(data - model, 2)) / np.abs(model)
            chi2NDF = np.sum(frac) / len(model)
            inv_sigma2 = 1.0 / (dataNoise**2)
            lnLike = -0.5 * (np.sum((data - model)**2 * inv_sigma2 -
                                    np.log(inv_sigma2)))

            # Fill the figure
            p1.cla()
            p1.set_ylabel("ADC")
            p1.set_title(
                "Run %d  Channel %d  Entry %d\ntrapENFCal %.1f  T/E %.1f  ST %.1f"
                % (run, chan, iList, dataE, toe, dataST))
            p1.plot(waveTS, waveBLSub, color='blue')
            p1.plot(waveTS, waveDenoised, color='red', alpha=0.8)
            p1.plot(guessFullTS, guessFull, color='orange', linewidth=2)
            p1.axvline(x=dataST, color='green', linewidth=2)
            p1.axvline(x=loWin, color='black')
            p1.axvline(x=hiWin, color='black')

            p2.cla()
            p2.plot(dataTS, data, color='blue', label='Data')
            p2.plot(guessTS, guess, color='orange', label='Guess')
            p2.plot(modelTS, model, color='red', linewidth=3, label='Best Fit')
            p2.legend(loc=4)

            p3.cla()
            p3.set_xlabel("Time [ns]", x=0.95, ha='right')
            p3.set_ylabel("Residual [ADC]")
            p3.plot(guessTS, residual, color='red')
            p3.axhline(y=0, color='blue', alpha=0.3)
            p3.axhline(y=dataNoise, color='blue', alpha=0.3)
            p3.axhline(y=-1.0 * dataNoise, color='blue', alpha=0.3)

            p4.cla()
            p4.set_title("startTime %.1f  Energy %.2f  Slow %.1f" %
                         (st, en, slo))

            p4.plot(stTrace[1:])
            p4.set_ylabel('startTime')

            p5.cla()
            p5.plot(enTrace[1:])
            p5.set_ylabel('energy')

            p6.cla()
            p6.plot(sloTrace[1:])
            p6.set_ylabel('slowness')

            plt.tight_layout()
            plt.subplots_adjust(hspace=0.35)
            plt.pause(scanSpeed)
Exemple #8
0
def main(argv):
    """ Interactive-fit or 'rapid'-fit waveforms that pass a given TCut.
        BUG: Doesn't always work with a TChain.  Add input files together
        with hadd and use a single TFile.
    """
    scanSpeed = 0.2
    iList = -1
    opt1, opt2 = "", "-"
    intMode = False
    if (len(argv) >= 1): opt1 = argv[0]
    if (len(argv) >= 2): opt2 = argv[1]
    if "-i" in (opt1, opt2):
        intMode = True
        print "Interactive mode selected."

    startT = time.clock()

    # Load template waveforms
    templateFile = TFile("./data/wave-m1cal.root")
    waveTemplates = templateFile.Get("waveTree")
    calDict = {640:28, 672:18, 610:16, 580:19, 582:34, 648:38, 600:21, 578:39, 592:27, 664:55, 626:62, 692:8, 598:22, 690:52, 632:9, 608:7}
    # PlotTemplateWaveforms(waveTemplates,calDict)
    # return

    # Set input data and cuts
    inputFile = TFile("~/project/wavelet-skim/hadd/waveletSkimDS3.root")
    # inputFile = TFile("~/project/wavelet-skim/waveletSkimDS3_13.root")
    waveTree = inputFile.Get("skimTree")
    print "Found",waveTree.GetEntries(),"input entries."

    # theCut = inputFile.Get("cutUsedHere").GetTitle()

    # DS3 big cut
    theCut = "trapENFCal > 0.8 && gain==0 && mHClean==1 && isGood && !muVeto && !wfDCBits && !isLNFill1 && !isLNFill2 && trapETailMin < 0 && channel!=596 && channel!=676 && channel!=676 && channel!=612 && channel!=1104 && channel!=1200 && channel!=1334 && channel!=1336 && tOffset < 10 && waveS5/TMath::Power(trapENFCal,1/4) < 1200 && (waveS3-waveS2)/trapENFCal > 100 && (waveS3-waveS2)/trapENFCal < 300 && !(channel==692 && (run==16974 || run==16975 || run==16976 || run==16977 || run==16978 || run==16979)) && butterTime < 11000"

    # theCut += " && trapENFCal > 2 && trapENFCal < 2.1 && channel==578 && run==16868"  # template fast WF (2kev)
    theCut += " && trapENFCal > 7 && trapENFCal < 10"
    # theCut += " && trapENFCal < 20 && trapENFCal > 2 && run > 17480"
    # theCut += " && kvorrT/trapENFCal > 2.2 && trapENFCal > 2 && trapENFCal < 10"

    # Print cut and events passing cut
    print "Using cut:\n",theCut,"\n"
    waveTree.Draw(">>elist", theCut, "entrylist")
    elist = gDirectory.Get("elist")
    waveTree.SetEntryList(elist)
    nList = elist.GetN()
    print "Found",nList,"entries passing cuts."

    stopT=time.clock()
    print "Data loading time (s): ", (stopT-startT)

    # Make a figure
    fig = plt.figure(figsize=(11,7), facecolor='w')
    p1 = plt.subplot2grid((6,7), (0,0), colspan=4, rowspan=2) # original
    p2 = plt.subplot2grid((6,7), (2,0), colspan=4, rowspan=3) # rising edge
    p3 = plt.subplot2grid((6,7), (5,0), colspan=4, rowspan=1) # residual
    p4 = plt.subplot2grid((6,7), (0,4), colspan=3, rowspan=2 )           # trace 1
    p5 = plt.subplot2grid((6,7), (2,4), colspan=3, rowspan=2, sharex=p4) # trace 2
    p6 = plt.subplot2grid((6,7), (4,4), colspan=3, rowspan=2, sharex=p4) # trace 3
    plt.show(block=False)

    # Loop over events
    while(True):
        saveMe = False
        iList += 1
        if intMode==True and iList != 0:
            value = raw_input()
            if value=='q': break
            if value=='s': saveMe=True
            if value=='p': iList -= 2  # previous
            if (value.isdigit()): iList = int(value) # go to entry
        if iList >= elist.GetN(): break

        entry = waveTree.GetEntryNumber(iList);
        waveTree.LoadTree(entry)
        waveTree.GetEntry(entry)
        nChans = waveTree.channel.size()
        numPass = waveTree.Draw("channel",theCut,"GOFF",1,iList)
        chans = waveTree.GetV1()
        chanList = list(set(int(chans[n]) for n in xrange(numPass)))

        # Loop over hits passing cuts
        hitList = (iH for iH in xrange(nChans) if waveTree.channel.at(iH) in chanList)  # a 'generator expression'
        for iH in hitList:

            # Load waveform for this hit
            run = waveTree.run
            chan = waveTree.channel.at(iH)
            dataE = waveTree.trapENFCal.at(iH)
            dataST = waveTree.butterTime.at(iH) # replace with blrwfFMR50?
            toe = waveTree.kvorrT.at(iH)/dataE
            print "%d / %d  Run %d  nCh %d  chan %d  trapENF %.1f  t/e %.1f" % (iList,nList,run,nChans,chan,dataE,toe)
            signal = wl.processWaveform(waveTree.MGTWaveforms.at(iH),opt='full')
            waveBLSub = signal.GetWaveBLSub()
            waveFilt = signal.GetWaveFilt()
            waveTS = signal.GetTS()
            baseAvg, dataNoise = signal.GetBaseNoise()

            # Denoise the data waveform (take only lowest-frequency components)
            wp = pywt.WaveletPacket(data=waveBLSub, wavelet='haar', mode='symmetric',maxlevel=3)
            new_wp = pywt.WaveletPacket(data=None, wavelet='haar', mode='symmetric')
            new_wp['aaa'] = wp['aaa'].data
            waveDenoised = new_wp.reconstruct(update=False)

            # Load channel template waveform
            try: tempEntry = calDict[chan]
            except:
                print "ain't got channel",chan
                continue
            waveTemplates.GetEntry(calDict[chan])
            template = wl.processWaveform(waveTemplates.event.GetWaveform(waveTemplates.itr),opt='full')
            temp = template.GetWaveBLSub()
            tempTS = template.GetTS()
            tempE = waveTemplates.trapENFCal
            tempST = waveTemplates.blrwfFMR50

            # Window the fit around rising edge - start time calculator method
            loWin, hiWin = dataST - 1000, dataST + 4000 # ns
            if loWin < waveTS[0] or hiWin > waveTS[-1]:
                print "Window out of range!  dataST: %.1f  loWin %.1f  hiWin %.1f" % (dataST,loWin,hiWin)
            idx = np.where((waveTS >= loWin) & (waveTS <= hiWin))
            data = waveBLSub[idx]
            # data = waveDenoised[idx]
            dataTS = waveTS[idx]

            # Pack into lists
            rawList = [waveBLSub, waveTS, dataE, dataST]
            dataList = [data, dataTS, dataE, dataST, loWin, hiWin]
            tempList = [temp, tempTS, tempE, tempST]

            # Optionally save to a file
            if saveMe:
                print "Saved entry",iList,iH
                np.savez("./data/optimumInputs.npz",rawList,tempList)
                # np.savez("./data/rawInputs.npz",rawList)

            # Recreate the guess and the guess's rising edge
            guessFull, guessFullTS = wm.MakeModel(dataList, tempList, [dataST,dataE,1.], opt="full")
            guess, guessTS = wm.MakeModel(dataList, tempList, [dataST,dataE,1.], opt="!fancy")

            # Make an "almost complete" guess - no MCMC
            # st, en, slo = dataST-100, dataE, 5
            # InterpFn = interpolate.interp1d(tempTS, temp, kind="linear", copy="False", assume_sorted="True")
            # model, modelTS = wm.MakeModel(dataList, tempList, [st,en,slo], fn=InterpFn)

            # Fit with MCMC and get best-fit parameters
            numSteps, burnIn = 2000, 1000 # default - 10000, 5000.  try 3000, 2000
            myModel = wm.TemplateModel( dataList, dataNoise, tempList )
            waveModel = pymc.Model( myModel )
            M = pymc.MCMC( waveModel )
            M.use_step_method(pymc.Metropolis, M.startTime, proposal_sd=100., proposal_distribution='Normal')
            M.use_step_method(pymc.Metropolis, M.energy, proposal_sd=1., proposal_distribution='Normal')
            M.use_step_method(pymc.Metropolis, M.slowness, proposal_sd=100., proposal_distribution='Normal')
            M.sample(iter=numSteps, verbose=0) # do the fit
            st = np.median(M.trace('startTime')[burnIn:])
            en = np.median(M.trace('energy')[burnIn:])
            slo = np.median(M.trace('slowness')[burnIn:])
            InterpFn = interpolate.interp1d(tempTS, temp, kind="linear", copy="False", assume_sorted="True")
            model, modelTS = wm.MakeModel(dataList, tempList, [st,en,slo], fn=InterpFn)

            # Calculate residual, Chi2/NDF
            residual = model - data
            frac = (np.power(data - model, 2)) / np.abs(model)
            chi2NDF = np.sum(frac) / len(model)
            print "chi2/NDF:",chi2NDF

            # ** Do a separate & simple fit of the tail slope **
            # TODO: Add this to process-waveforms.py
            idxMax = np.where(guessFull == guessFull.max()) # returns an array/tuple
            idxMax = idxMax[0][0] # "cast" to int
            tail, tailTS = waveDenoised[idxMax:], waveTS[idxMax:]
            popt,_ = curve_fit(wl.tailModelPol, tailTS, tail) # poly fit
            a, b = dataE, 72000
            popt2,_ = curve_fit(wl.tailModelExp, tailTS, tail, p0=[a,b]) # expo fit

            # Fill the figure
            p1.cla()
            p1.set_ylabel("ADC")
            p1.set_title("Run %d  Channel %d  Entry %d\ntrapENFCal %.1f  T/E %.1f  ST %.1f" % (run,chan,iList,dataE,toe,dataST))
            p1.plot(waveTS,waveBLSub,color='blue')
            p1.plot(waveTS,waveDenoised,color='red',alpha=0.8)
            p1.plot(guessFullTS,guessFull,color='orange',linewidth=2)
            p1.axvline(x = dataST, color='green',linewidth=2)
            p1.axvline(x = loWin, color='black')
            p1.axvline(x = hiWin, color='black')
            p1.plot(tailTS, wl.tailModelPol(tailTS, *popt), color='cyan', linewidth=1) # tail poly fit
            p1.plot(tailTS, wl.tailModelExp(tailTS, *popt2), color='cyan', linewidth=1) # tail expo fit

            p2.cla()
            p2.plot(dataTS,data,color='blue',label='Data')
            p2.plot(guessTS,guess,color='orange',label='Guess')
            p2.plot(modelTS,model,color='red',linewidth=3,alpha=0.7,label='Fit')
            p2.legend(loc=4)

            p3.cla()
            p3.set_xlabel("Time [ns]", x=0.95, ha='right')
            p3.set_ylabel("Residual [ADC]")
            p3.plot(modelTS,residual, color='red')
            p3.axhline(y = 0, color='blue', alpha=0.3)
            p3.axhline(y = dataNoise, color='blue', alpha=0.3)
            p3.axhline(y = -1.0*dataNoise, color='blue', alpha=0.3)

            p4.cla()
            minST = tempST - tempTS[-1] + hiWin
            maxST = tempST - tempTS[0] + loWin
            p4.set_title("startTime %.1f  Energy %.2f\nSlow %.1f  chi2/ndf %.1f  Min %d  Max %d" % (st,en,slo,chi2NDF,minST,maxST))
            p4.plot(M.trace('startTime')[:])
            p4.set_ylabel('startTime')

            p5.cla()
            p5.plot(M.trace('energy')[:])
            p5.set_ylabel('energy')

            p6.cla()
            p6.plot(M.trace('slowness')[:])
            p6.set_ylabel('slowness')

            plt.tight_layout()
            plt.subplots_adjust(hspace=0.35)
            plt.pause(scanSpeed)
def main(argv):
    """ Get a sweet minimizer + time domain match filter working. Also graduate. """
    scanSpeed = 1.
    iList = -1
    opt1 = ""
    intMode = False
    if (len(argv) >= 1): opt1 = argv[0]
    if "-i" in (opt1): intMode = True

    # Set input data and cuts
    inputFile = TFile("~/project/match-skim/waveletSkimDS5_run23920.root") # calibration
    # inputFile = TFile("~/project/v2-processwfs/waveletSkimDS5_90.root") # noisy BG
    waveTree = inputFile.Get("skimTree")
    theCut = inputFile.Get("cutUsedHere").GetTitle()
    # theCut += " && waveS5/trapENFCal < 1200 && trapENFCal < 10"
    theCut += " && trapENFCal < 30"

    # Print cut and events passing cut
    waveTree.Draw(">>elist", theCut, "entrylist")
    elist = gDirectory.Get("elist")
    waveTree.SetEntryList(elist)
    nList = elist.GetN()
    print "Using cut:\n",theCut,"\n"
    print "Found",waveTree.GetEntries(),"input entries."
    print "Found",nList,"entries passing cuts."

    # Make a figure
    fig = plt.figure(figsize=(15,10), facecolor='w')
    p1 = plt.subplot2grid((6,7), (0,0), colspan=4, rowspan=3) # data & fit
    p2 = plt.subplot2grid((6,7), (3,0), colspan=4, rowspan=3) # matched filter
    p3 = plt.subplot2grid((6,7), (0,4), colspan=3, rowspan=2) # trace 1
    p4 = plt.subplot2grid((6,7), (2,4), colspan=3, rowspan=2) # trace 2
    p5 = plt.subplot2grid((6,7), (4,4), colspan=3, rowspan=2) # trace 3
    plt.show(block=False)

    # Make template(s)
    tSamp, tR, tZ, tAmp, tST, tSlo = 5000, 0, 15, 100, 2500, 10
    tOrig, tOrigTS = wl.MakeSiggenWaveform(tSamp,tR,tZ,tAmp,tST,tSlo)

    # Loop over events
    while True:
        saveMe = False
        iList += 1
        if intMode==True and iList != 0:
            value = raw_input()
            if value=='q': break
            if value=='s': saveMe=True
            if value=='p': iList -= 2  # previous
            if (value.isdigit()): iList = int(value) # go to entry
        if iList >= elist.GetN(): break

        entry = waveTree.GetEntryNumber(iList);
        waveTree.LoadTree(entry)
        waveTree.GetEntry(entry)
        nChans = waveTree.channel.size()
        numPass = waveTree.Draw("channel",theCut,"GOFF",1,iList)
        chans = waveTree.GetV1()
        chanList = list(set(int(chans[n]) for n in xrange(numPass)))

        # Loop over hits passing cuts
        hitList = (iH for iH in xrange(nChans) if waveTree.channel.at(iH) in chanList)  # a 'generator expression'
        for iH in hitList:

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

            # Load data
            run = waveTree.run
            chan = waveTree.channel.at(iH)
            dataE = waveTree.trapENFCal.at(iH)
            dataENM = waveTree.trapENM.at(iH)
            # dataMax = waveTree.trapENMSample.at(iH)*10. - 4000
            dataMax = waveTree.rt90.at(iH)
            signal = wl.processWaveform(waveTree.MGTWaveforms.at(iH))
            data = signal.GetWaveBLSub()
            dataTS = signal.GetTS()
            dataBL, dataNoise = signal.GetBaseNoise()
            print "%d / %d  Run %d  nCh %d  chan %d  trapENF %.1f  len %d" % (iList,nList,run,nChans,chan,dataE,len(data))

            # Load and scale template
            temp, tempTS = tOrig, tOrigTS
            temp = temp * (dataENM / tAmp)         # scale by amplitudes (not energies)
            tempMax = np.argmax(temp) * 10         # convert to ns
            tempTS = tempTS - (tempMax - dataMax)  # align @ max of rising edge
            tempMax = tempTS[np.argmax(temp)]      # get the new max TS after the shifting
            tempE = dataE # set 'calibrated' energy of template equal to data's trapENFCal.

            # Lowpass or bandpass filter the data
            # NOTE: This is tough.  There might be a feature due to rising edges, between 3e6 and 5e6 Hz.
            #       But that's only from data. Doing it w/ siggen templates yielded squat.
            # B,A = sig.butter(2, [3e6/(1e8/2),5e6/(1e8/2)], btype='bandpass')
            B,A = sig.butter(2,1e6/(1e8/2),btype='lowpass') # nice lowpass filter
            data_LowPass = sig.lfilter(B, A, data)
            B1,A1 = sig.butter(2, [1e5/(1e8/2),1e6/(1e8/2)], btype='bandpass') # attempt at bandpass
            data_BanPass = sig.lfilter(B1, A1, data)


            # -- Run minimizers

            # Set window and parameters
            loWin, hiWin = dataTS[0], dataTS[-1]
            mt, en, slo = dataMax-5, dataE+1, 20.  # guesses
            # print "guesses:        mt %.0f  en %.3f  slo %.2f" % (mt, en, slo)

            # pack into lists
            InterpFn = interpolate.interp1d(tempTS, temp, kind="linear", copy="False", assume_sorted="True")
            dataList = [data, dataTS, dataE, dataMax, loWin, hiWin, dataNoise]
            tempList = [temp, tempTS, tempE, tempMax]
            datas = [dataList, tempList, InterpFn]
            floats = [mt,en,slo]
            guess, guessTS = wm.MakeModel(dataList, tempList, floats, fn=InterpFn)

            start = time.clock()
            MakeTracesGlobal()

            fitFail = False # you should output this to ROOT

            # 0. brute - trying to make a quick initial guess for the slowness parameter
            # FIXME: this keeps returning the lowest allowed value for slowness.
            # bruteArgs = [en, datas]
            # ranges = (slice(dataMax-500,dataMax+300,100), slice(1,100,10))
            # resbrute = op.brute(bruteLnLike, ranges, args=(bruteArgs,), full_output=False, finish=None, disp=False)
            # mt, slo = resbrute[0], resbrute[1]
            # print "brute results:  mt %.0f  en %.3f  slo %.2f" % (mt, en, slo)
            # floats = [mt, en, slo]

            # 1. nelder-mead (~0.5 sec).  nobody does it better.  works almost like magic.
            result = op.minimize(findLnLike, floats, args=datas, method="Nelder-Mead")#,options={"maxiter":10})
            if not result["success"]:
                print "fit 'fail': ", result["message"]
                fitFail = True
            mt, en, slo = result["x"]
            # print "nelder results: mt %.0f  en %.3f  slo %.2f" % (mt, en, slo)
            nelder, nelderTS = wm.MakeModel(dataList, tempList, [mt,en,slo], fn=InterpFn)
            nelderStop = time.clock()

            # 2. powell "polish step" (~0.1 sec).  Sometimes screws up Nelder's good results!
            # floats = [mt,en,slo]
            # powell = op.minimize(findLnLike, floats, args=datas, method="Powell")
            # if not powell["success"]:
            #     fitFail= True
            #     print powell["message"]
            # mt, en, slo = powell["x"]
            # # print "powell results: mt %.0f  en %.3f  slo %.2f" % (mt, en, slo)
            # model, modelTS = wm.MakeModel(dataList, tempList, [mt,en,slo], fn=InterpFn)
            #
            # powellStop = time.clock()
            # print "Time: nelder %.3f  powell %.3f  total %.3f" % (nelderStop-start, powellStop-nelderStop, powellStop-start)

            # take absolute values for parameters
            mt, en, slo = abs(mt), abs(en), abs(slo)

            model, modelTS = wm.MakeModel(dataList, tempList, [mt,en,slo], fn=InterpFn)


            # compute a trap energy off the model (need pinghan's calibration constants)
            trap = np.zeros(1000)
            trap = np.append(trap,wl.trapezoidalFilter(model))
            trapTS = modelTS
            trapMax = np.amax(trap)



            # -- Time domain match filter

            match, matchTS = wm.MakeModel(dataList, tempList, [mt,en,slo], opt="nowindow")
            match = np.flip(match,0)

            # line up the max of the match with the max of the best-fit signal
            modelMaxTime = modelTS[np.argmax(model)]
            matchMaxTime = matchTS[np.argmax(match)]
            matchTS = matchTS + (modelMaxTime - matchMaxTime)

            idx = np.where((matchTS >= dataTS[0]-5) & (matchTS <= dataTS[-1]+5))
            match, matchTS = match[idx], matchTS[idx]

            # make sure match is same length as data, then compute convolution and parameters
            matchMax, matchWidth, matchArea = 0, 9999, 0
            smoothMF = np.zeros(len(dataTS))
            if len(match)!=len(data):
                print "array mismatch: len match %i  len data %i " % (len(match),len(data))
            else:
                # I can't decide which of these is better.
                # smoothMF = gaussian_filter(match * data, sigma=float(5))
                smoothMF = gaussian_filter(match * data_LowPass,sigma=float(5))

                # computer AWESOME match filter parameters
                matchMax = np.amax(smoothMF)
                idx = np.where(smoothMF > matchMax/2.)
                matchWidth = matchTS[idx][-1] - matchTS[idx][0]
                matchArea = np.sum(match[idx]) / float(len(match[idx]))


            # Plot data and template
            p1.cla()
            p1.plot(dataTS,data,'b',label='data')
            idx2 = np.where((tempTS >= dataTS[0]-5) & (tempTS <= dataTS[-1]+5))
            p1.plot(tempTS[idx2],temp[idx2],'r',label='template')
            # p1.plot(dataTS, data_LowPass, label='lowpass')
            # p1.plot(dataTS, data_BanPass, label='bandpass')
            p1.plot(trapTS,trap,color='k',label='bestfit-trap')
            p1.plot(nelderTS,nelder,color='cyan',label='bestfit')
            # p1.plot(modelTS,model,color='magenta',label='powells')
            p1.set_xlabel('Time (s)')
            p1.set_ylabel('Voltage (arb)')
            p1.set_title("Run %i  Ch %i  E %.2f  ENM %.2f  trapMax %.2f" % (run,chan,dataE,dataENM,trapMax))
            p1.legend(loc=4)

            p2.cla()
            p2.plot(dataTS,data,'b',alpha=0.8,label='data')
            # p2.plot(guessTS,guess,'y',alpha=0.8,label='guess')
            p2.plot(matchTS,match,'k',label='match')
            p2.plot(modelTS,model,color='orange',label='model')

            p2.plot(dataTS,data_LowPass,'r',label='lowpass')

            if len(match)==len(data): p2.plot(matchTS,smoothMF,'g',label='smoothMF')

            p2.axvline(matchTS[idx][-1],color='green',alpha=0.7)
            p2.axvline(matchTS[idx][0],color='green',alpha=0.7)

            p2.set_title("mMax %.2f  mWidth %.2f  mAreaNorm %.2f  mM/mW %.4f" % (matchMax,matchWidth,matchArea,matchMax/matchWidth))
            p2.legend(loc=4)


            # plot traces
            p3.cla()
            p3.set_title("maxTime %.1f  Energy %.2f  Slow %.1f" % (mt,en,slo))
            p3.plot(mtTrace[1:])
            p3.set_ylabel('maxTime')
            p4.cla()
            p4.plot(enTrace[1:])
            p4.set_ylabel('energy')
            p5.cla()
            p5.plot(sloTrace[1:])
            p5.set_ylabel('slowness')

            plt.tight_layout()
            plt.pause(scanSpeed)