def get_yields(comb_bkg=False):
helpers.allow_root()
from . import fit_config
import ROOT
from ROOT import RooFit as RF
shapes.load_shape_class('RooCruijff')
shapes.load_shape_class('RooJohnsonSU')
shapes.load_shape_class('RooBackground')
wsp = fit_config.load_workspace(gcm())
calculator = _metric_base(wsp, comb_bkg)
sig = calculator._get_number_of_signal()
bkg = calculator._get_number_of_background()
return sig, bkg
def get_sweights(do_comb_bkg=False):
helpers.allow_root()
df = gcm().get_data([m(gcm().D0), dtf_dm()])
from . import fit_config
from hep_ml import splot
shapes.load_shape_class('RooCruijff')
shapes.load_shape_class('RooJohnsonSU')
shapes.load_shape_class('RooBackground')
wsp = fit_config.load_workspace(gcm())
sel = selection.get_final_selection()
do_comb_bkg = gcm().mode in config.twotag_modes
df = df[sel]
sig_pdf = wsp.pdf('signal')
rnd_pdf = wsp.pdf('random')
comb_pdf = wsp.pdf('combinatorial')
sig_prob = call_after_set(sig_pdf, wsp, **df)
rnd_prob = call_after_set(rnd_pdf, wsp, **df)
if do_comb_bkg:
comb_prob = call_after_set(comb_pdf, wsp, **df)
if do_comb_bkg:
probs = pd.DataFrame(dict(sig=sig_prob*wsp.var('NSig').getVal(),
rnd=rnd_prob*wsp.var('NSPi').getVal(),
comb=comb_prob*wsp.var('NBkg').getVal()),
index=df.index)
else:
probs = pd.DataFrame(dict(sig=sig_prob*wsp.var('NSig').getVal(),
rnd=rnd_prob*wsp.var('NSPi').getVal()),
index=df.index)
probs = probs.div(probs.sum(axis=1), axis=0)
sweights = splot.compute_sweights(probs)
sweights.index = probs.index
if not do_comb_bkg:
sweights['comb'] = 0.0
return sweights
def get_model_ws_alt(redo=False):
files = [
'/afs/cern.ch/user/c/chasse/public/forDominik/Sig_49_61_WSNR_Smaller.root'
] # NOQA
bcolz_folder = config.bcolz_locations.format('generated_model_ws_alt')
if redo:
try:
shutil.rmtree(bcolz_folder)
except:
pass
helpers.allow_root()
import root_pandas
df = root_pandas.read_root(files, 'events')
# Now rename stuff and fix units to MeV and ns.
# Ugly hardcoded for now.
df.rename(columns={
'c12': vars.cos1(),
'c34': vars.cos2(),
'dtime': vars.ltime(mode_config.D0),
'phi': vars.phi1(),
'm12': vars.m12(),
'm34': vars.m34()
},
inplace=True)
df[vars.m12()] = df[vars.m12()] * 1000.
df[vars.m34()] = df[vars.m34()] * 1000.
df[vars.ltime(mode_config.D0)] = df[vars.ltime(mode_config.D0)] / 1000.
df = df.query('{} > 0.0001725'.format(vars.ltime(mode_config.D0)))
df = df.query('{} < 0.003256'.format(vars.ltime(mode_config.D0)))
bcolz.ctable.fromdataframe(df, rootdir=bcolz_folder)
return df
else:
bc = bcolz.open(bcolz_folder)
return bc.todataframe()
return df
def get_model(redo=False):
files = filelists.Generated.paths
bcolz_folder = config.bcolz_locations.format('generated_model')
if redo:
try:
shutil.rmtree(bcolz_folder)
except:
pass
helpers.allow_root()
import root_pandas
df = root_pandas.read_root(files, 'events')
# Now rename stuff and fix units to MeV and ns.
# Ugly hardcoded for now.
df.rename(columns={
'c12': vars.cos1(),
'c34': vars.cos2(),
'dtime': vars.ltime(mode_config.D0),
'phi': vars.phi1(),
'm12': vars.m12(),
'm34': vars.m34()
},
inplace=True)
df[vars.m12()] = df[vars.m12()] * 1000.
df[vars.m34()] = df[vars.m34()] * 1000.
df[vars.ltime(mode_config.D0)] = df[vars.ltime(mode_config.D0)] / 1000.
df = df.query('{} > 0.0001725'.format(vars.ltime(mode_config.D0)))
df = df.query('{} < 0.003256'.format(vars.ltime(mode_config.D0)))
bcolz.ctable.fromdataframe(df, rootdir=bcolz_folder)
return df
else:
bc = bcolz.open(bcolz_folder)
return bc.todataframe()
return df
df = df.rename(index=str, columns=new_names)
if ctuple is None:
ctuple = bcolz.ctable.fromdataframe(df, rootdir=bcolz_folder)
else:
ctuple.append(df.to_records(index=False))
for f in temp_files:
os.remove(f)
# Loop and delete everything in the datastore that needs to be recached
remove_buffer_for_mode(mode.mode)
if modename == 'WS' and year == 2016:
# As this is the start, hack name of the particle in the mode.
mode.Dstp.name = 'Dstp'
if __name__ == '__main__':
args = parser.create_parser(log)
helpers.allow_root()
if args.polarity == config.magboth:
pols = [config.magup, config.magdown]
else:
pols = [args.polarity]
if args.year == 1516:
years = [2015, 2016]
else:
years = [args.year]
for p, y in product(pols, years):
download(args.mode, p, y, args.fraction, args.test, args.mc,
njobs=args.jobs)