generate_hists.py

Go to the documentation of this file.

"""
  generate_hists.py:
   Make histograms to use for testing the validation system.
  
   Original author: J. Wolcott <jwolcott@fnal.gov>
              Date: September 2016
"""

import os.path
import random


import ROOT

OUTDIR = "/nova/ana/users/jwolcott/mc_valid/new_test"

# three "datasets" will have the same distribution,
# which is broken into three "categories"/components.
# the datasets differ by 'exposure' (number of entries
# sampled from each distribution).
DATASET_NUMEVTS = {
     "test_dataset1": 1e5,
     "test_dataset2": 1e4,
     "test_dataset3": 1e3,
}

MOCK_DATA_DATASETS = [
    "test_dataset3",
]

COMPONENT_DISTRIBUTIONS = [
   "TMath::Gaus(x, 1)",
   "TMath::Landau(x-2, 1)",
   "2*TMath::Student(3*x, 5)",
]
sum_f = ROOT.TF1("fsum", " + ".join(COMPONENT_DISTRIBUTIONS), 0, 5)
integral = sum_f.Integral(0, 5)


COMPONENT_FUNCTIONS = [ ROOT.TF1("%f * f%d" % (1./integral, fnum), distr, 0, 5) for fnum, distr in enumerate(COMPONENT_DISTRIBUTIONS) ]
thresholds = [f.Integral(0, 5)/integral for f in COMPONENT_FUNCTIONS]  # threshold for each distribution for rejection sampling
for i in range(1, len(thresholds)):
    thresholds[i] += thresholds[i-1]

for dataset_name, num_evts in DATASET_NUMEVTS.iteritems():
    outf = ROOT.TFile( os.path.join(OUTDIR, "test_hists_%s.root" % dataset_name), "recreate" )
    assert not outf.IsZombie()
    
    hists = []
    if dataset_name in MOCK_DATA_DATASETS:
        hists.append(ROOT.TH1D("dir1/h_var", "Test ariable;Var (unit);Events", 50, 0, 5))
    else:
        for hist_num in range(len(COMPONENT_FUNCTIONS)):
            hists.append(ROOT.TH1D("dir1/h_var{cat%d}" % hist_num, "Test variable;Var (unit);Events", 50, 0, 5))
    
    for evt_num in xrange(int(num_evts)):
        # pick distribution via rejection sampling
        rnd = random.random()
        for distr_num, threshold in enumerate(thresholds):
            if rnd < threshold or distr_num == len(thresholds)-1:
                break
        
        # now sample from that distribution itself
        val = COMPONENT_FUNCTIONS[distr_num].GetRandom()
        h = hists[distr_num] if len(hists) == len(COMPONENT_DISTRIBUTIONS) else hists[0]
        h.Fill(val)
    
    for h in hists:
        h.Write()
    
    outf.Close()