Source code for MPF.examples.exampleHelpers

import ROOT

import random, math
import sys
import os
import glob

import MPF.pyrootHelpers as PH
from MPF.commonHelpers.logger import logger
logger = logger.getChild(__name__)

# Do some dummy plots
def histogramFromList(myList, nbins):
    xmin = min(myList, key=lambda x: x[0])
    xmax = max(myList, key=lambda x: x[0])
    logger.debug(xmin)
    logger.debug(xmax)
    histo = ROOT.TH1D(next(PH.tempNames), '', nbins, float(xmin[0]), float(xmax[0]))
    histo.Sumw2()
    for entry, value in myList:
        histo.Fill(entry, value)
    return histo


def randHist():
    histo = ROOT.TH1D(next(PH.tempNames), '', 10, 0, 1)
    for i in range(10):
        for j in range(i*2):
            histo.Fill(random.random())
            histo.SetDirectory(0)
    return histo

def sampleFunc(func, **kwargs):
    histo = ROOT.TH1D(next(PH.tempNames), '', 10, 0, 10)
    stats = kwargs.get('simulateStats', 0)
    for i in range(10):
        val = func(i)
        if val < 0:
            logger.error('sampling negative values not properly supported')
        if stats > 0:
            for j in range(int(math.ceil(val/stats))):
                histo.Fill(i, stats)
        else:
            histo.Fill(i, val)
    return histo


def loadExampleTreeCode():
    if not ROOT.gROOT.GetClass("TGenPhaseSpace"):
        ROOT.gSystem.Load("libPhysics")
    try:
        ROOT.exampleTree
    except AttributeError:
        cxxFile = os.path.join(os.path.dirname(__file__), "exampleTree.cxx")
        compileFiles = os.path.join(os.path.dirname(__file__), "exampleTree_cxx*")
        # clean up previously existing files
        for filename in glob.glob(compileFiles):
            logger.debug("deleting {}".format(filename))
            os.unlink(filename)
        logger.debug("loading {}".format(cxxFile))

        # i don't know why, but if i put the "+" here there are
        # sometimes problems with the loaded classes for multihistdraw
        # (after there temporary files got deleted)
        # Note: we do atexit deletion again ... anyways - probably don't need the "+" here
        # ROOT.gROOT.ProcessLine(".L {}+".format(cxxFile))

        # so lets do it without "+"
        ROOT.gROOT.ProcessLine(".L {}".format(cxxFile))


def createExampleTrees(filename):
    loadExampleTreeCode()
    exampleTree = ROOT.exampleTree()
    f = ROOT.TFile.Open(filename, "RECREATE")
    exampleTree.makeWTree(80., "w1", 100000)
    exampleTree.makeWTree(100., "w2", 100000)
    exampleTree.makeWTree(150., "w3", 100000)
    exampleTree.makeWTree(2000., "w4", 100000)
    f.Write()
    f.Close()

def createExampleSignalGrid(filename,
                            nevents=1000,
                            m1range=(1000, 2000, 100),
                            m2range=(1000, 2000, 100)):
    loadExampleTreeCode()
    exampleTree = ROOT.exampleTree()
    f = ROOT.TFile.Open(filename, "RECREATE")
    for m1 in range(*m1range):
        for m2 in range(*m2range):
            if m2 >= m1:
                continue
            exampleTree.makeWTree(m1, "w_{}_{}".format(m1, m2), nevents, 0., m2)
    f.Write()
    f.Close()

def parseExampleArgs():
    import argparse
    parser = argparse.ArgumentParser(description='Run the MPF examples')
    parser.add_argument('--log', help='logging level')
    args = parser.parse_args()
    if args.log:
        import logging
        logger.parent.setLevel(getattr(logging, args.log.upper()))