from k3pi_utilities.debugging import call_debug
from k3pi_cpp import t_gen_phasespaced
from k3pi_config import config
from k3pi_config.modes.D0ToKpipipi_RS import D0ToKpipipi_RS as mode_config
from k3pi_utilities import variables as vars
import pandas as pd
import numpy as np
import shutil
import bcolz
from k3pi_utilities import logger
log = logger.get_logger('phsp_generator')
@call_debug
def generate(nevents=10000000):
"""Generates and returns phase-space MC for a 4 body decay of K3pi.
Decay is generated using TGenPhaseSpace and a decaytime is added
from an exponential with lifetime 0.00040995 ns.
Generated sample is cached based on the size and only redone if the number
changes of no cached results are found.
"""
bcolz_folder = config.bcolz_locations.format('generated_phasespace')
try:
bc = bcolz.open(bcolz_folder)
if bc.size == nevents:
log.info('Returning stored PHSP')
return bc.todataframe()
from k3pi_config.modes import gcm, opposite_mode
from analysis import add_variables, selection
import numpy as np
from hep_ml.commonutils import train_test_split
from k3pi_config import config
from k3pi_utilities import logger
log = logger.get_logger('bdt_data')
def prep_data_for_sklearn(**kwargs):
if kwargs.get('comb_data', False):
features = [f.functor(f.particle) for f in gcm().comb_bkg_bdt_vars]
else:
features = [f.functor(f.particle) for f in gcm().rand_spi_bdt_vars]
spectators = [f.functor(f.particle) for f in gcm().spectator_vars]
kwargs.update({'sklearn': True})
data = get_bdt_data(**kwargs)
train, test = train_test_split(data, random_state=43)
return (train, test, train['labels'].astype(np.bool),
test['labels'].astype(np.bool)), features, spectators
def just_the_labels(sw=False, comb_data=False):
data = get_bdt_data(sw=sw, sklearn=True, comb_data=comb_data)
train, test = train_test_split(data, random_state=43)
return train[['labels', 'weights']], test[['labels', 'weights']]
from k3pi_config.modes.D0ToKpipipi_RS import D0ToKpipipi_RS
from k3pi_config.modes.D0ToKpipipi_WS import D0ToKpipipi_WS
from k3pi_config.modes.D0ToKpipipi_2tag_RS import D0ToKpipipi_2tag_RS
from k3pi_config.modes.D0ToKpipipi_2tag_WS import D0ToKpipipi_2tag_WS
from k3pi_config.modes.mode_base import ModeConfig
import sys
from contextlib import contextmanager
from k3pi_utilities import logger
log = logger.get_logger('modes')
_current_mode = None
_previous_modes = []
_thismodule = sys.modules[__name__]
@contextmanager
def MODE(polarity, year, mode, mc=None):
global _current_mode
global _previous_modes
if not hasattr(_thismodule, 'D0ToKpipipi_' + mode):
raise ModeConfig('No mode {} -> {}'.format(mode,
'D0ToKpipipi_' + mode))
if _current_mode is not None:
_previous_modes.append(_current_mode)
_current_mode = getattr(_thismodule, 'D0ToKpipipi_' + mode)(polarity, year,
mc)
log.debug('Entering {} {} {} mc={}'.format(mode, polarity, year, mc))
yield _current_mode
log.debug('Leaving {} {} {} mc={}'.format(mode, polarity, year, mc))
try:
from k3pi_utilities.debugging import call_debug
from k3pi_config import config, filelists
from k3pi_config.modes.D0ToKpipipi_RS import D0ToKpipipi_RS as mode_config
from k3pi_utilities import variables as vars
from analysis import selection, add_variables, final_selection
import matplotlib.pyplot as plt
from k3pi_utilities import parser, helpers
from matplotlib.backends.backend_pdf import PdfPages
from k3pi_plotting import comparison
from k3pi_config.modes import MODE, gcm
import shutil
import bcolz
from k3pi_utilities import logger
log = logger.get_logger('model_converter')
@call_debug
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')
from decorator import decorate
from k3pi_utilities.logger import get_logger
log = get_logger('call_debug')
def call_debug(function):
def wrapper(f, *args, **kwargs):
log.debug('Calling {}.'.format(function.__name__))
return f(*args, **kwargs)
return decorate(function, wrapper)
from k3pi_utilities import buffer, selective_load
from hep_ml import uboost, gradientboosting as ugb, losses # NOQA
from k3pi_utilities import variables as vars
from k3pi_config.modes import MODE, gcm
from k3pi_utilities import bdt_utils
from analysis import bdt_data
from k3pi_utilities import tex_compile
from k3pi_utilities.helpers import add_separation_page
# Because this stuff all lots of memory and I am too lazy to properly make
# sure all functions correctly delete their stuff.
from k3pi_utilities.processify import processify
from k3pi_utilities.logger import get_logger
log = get_logger('bdt_studies')
@processify
def train_bdts(sw=False, comb_bkg=False):
log.info('Training BDTs for {} {} {}'.format(gcm().mode,
gcm().polarity,
gcm().year))
(train, test, train_lbl,
test_lbl), features, spectators = bdt_data.prep_data_for_sklearn(
sw=sw, same_weight=True, comb_data=comb_bkg) # NOQA
uniform_features = [vars.ltime(gcm().D0)]
n_estimators = 400
classifiers = {}
from k3pi_config.modes import MODE, gcm, opposite_mode
from k3pi_utilities import parser, logger, PlotConfig
from analysis import final_selection, add_variables, selection
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.backends.backend_pdf import PdfPages
from k3pi_utilities import variables as vars
from scipy.special import erf
from k3pi_plotting import comparison
from k3pi_plotting import utils as plot_utils
log = logger.get_logger('phsp_comparison')
def phsp_comparison_plots():
"""Plots the mode sidebands and the opposite mode signal region phsp
distributions. Only really meaningful if executed for the WS events.
Opposite mode is plotted as solid, with the uncertainty propagated to
the mode error plot.
"""
# Beside phase space, also plot D0 momentum and flight distance
extra_vars = [
gcm().ltime_var,
PlotConfig(vars.pt,
gcm().D0, (100, 0, 15000)),
PlotConfig(vars.vdchi2,
gcm().D0, (100, 0, 10), np.log,
r'$\ln(\text{{{}}})$'), # NOQA
]
# opposite_mode
from sklearn.metrics import roc_curve, roc_auc_score
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import to_rgba
import palettable
from k3pi_utilities import variables as vars
from k3pi_utilities import helpers
from k3pi_config import config
from k3pi_config.modes import gcm
from k3pi_utilities.logger import get_logger
log = get_logger('k3pi_plotting/bdt')
from k3pi_plotting import utils as plot_utils
def roc(test, bdt, labels, weights):
# Get the probabilities for class 1 (signal)
probs = bdt.predict_proba(test).transpose()[1]
x, y, _ = roc_curve(labels, probs, sample_weight=weights)
# Make physical values so the plot looks nicer
x[x < 0] = 0.
x[x > 1] = 1.
y[y < 0] = 0.
y[y > 1] = 1.
auc = roc_auc_score(labels, probs, sample_weight=weights)
return x, y, auc
def plot_roc(ax, name, colour, test, bdt, labels, weights):
x, y, auc = roc(test, bdt, labels, weights)
from collections import defaultdict
from k3pi_config import config, modes
from k3pi_utilities.logger import get_logger
from inspect import getargspec
import pandas as pd
from pandas import HDFStore
log = get_logger('buffer_load')
accumulated_per_mode = defaultdict(lambda: set())
def remove_buffer_for_function(func):
"""This will trigger the deletion of buffered values matching the function
so they are automatically refreshed when needed."""
with pd.get_store(config.data_store) as store:
for a in store:
if 'Cached' not in a:
continue
if func.__name__ in a:
log.info('Removing {} from store'.format(a))
store.remove(a)
def remove_buffer_for_mode(modename):
"""Cleans out the cache for everything belonging to a certain mode."""
with pd.get_store(config.data_store) as store:
for a in store:
if 'Cached' not in a:
continue
if modename not in a:
from math import floor, log10
import pickle
import random
import shutil
import string
from math import ceil
import sys
from scipy.stats import gaussian_kde
from scipy.optimize import fmin
from collections import OrderedDict
import numpy
import matplotlib.pyplot as plt
from k3pi_utilities import logger
log = logger.get_logger('helpers')
# List of prefixes for branch names
# Can't place in config due to circular dependency
numbered_branch_prefixes = ['First', 'Second', 'Third', 'Fourth', 'Fifth']
def allow_root():
"""
Removes the ban on ROOT (if previously set)
Returns true if successful, false otherwise
:returns: Boolean
"""
if 'ROOT' in sys.modules:
import tempfile
from k3pi_utilities import helpers, logger
log = logger.get_logger('fit_config')
WS_DMASS_NAME = 'dmass_var'
WS_MASS_NAME = 'mass_var'
def get_delta_mass_var(workspace):
"""Return the RooRealVar of the fitted delta mass variable."""
return workspace.var(workspace.obj(WS_DMASS_NAME).GetName())
def get_mass_var(workspace):
"""Return the RooRealVar of the fitted delta mass variable."""
return workspace.var(workspace.obj(WS_MASS_NAME).GetName())
def workspace_import(workspace, obj, name='', cmdargs=None):
import ROOT
"""Import object obj in to the RooWorkspace.
Import as normal if name is given, else use hack to fix ROOT 6 bug.
https://sft.its.cern.ch/jira/browse/ROOT-6785
Keyword arguments:
workspace -- RooWorkspace to import obj into
obj -- Object to import in to the workspace
name -- Optional name to give the object (default: None)
cmdargs -- Optional list of RooCmdArg objects to forward (default: [])
"""
if not cmdargs:
from matplotlib.backends.backend_pdf import PdfPages
from itertools import permutations
from k3pi_config.modes import MODE, gcm
from analysis.model_generated import get_model
from k3pi_utilities import bdt_utils, buffer
from k3pi_utilities.helpers import dump, load
from k3pi_plotting import comparison
from k3pi_plotting import utils as plot_utils
from k3pi_utilities import logger
import numpy as np
import pandas as pd
from hep_ml.reweight import GBReweighter
log = logger.get_logger('efficiencies')
@buffer.buffer_load
def get_efficiency():
"""Returns or first trains the BDT efficiency."""
extra_vars = [
gcm().ltime_var
]
all_vars = gcm().phsp_vars + extra_vars
columns = [v.var for v in all_vars if 'phi' not in v.var]
columns += ['cosphi', 'sinphi']
log.info('Getting efficiencies for {}'.format(', '.join(columns)))
# Current mode stuff
data = gcm().get_data([f.var for f in extra_vars])
from k3pi_config import config
import numpy as np
from k3pi_utilities import logger
log = logger.get_logger('metrics')
class _metric_base(object):
def __init__(self, wsp, comb_bkg=False):
self.wsp = wsp
self._comb_bkg = comb_bkg
# wsp.var('D0_M').setRange(
# 'signal', config.PDG_MASSES['D0'] - 20.,
# config.PDG_MASSES['D0'] + 20.)
wsp.var('delta_m_dtf').setRange(
'signal',
config.PDG_MASSES['delta'] - 0.5,
config.PDG_MASSES['delta'] + 0.5)
self._signal = wsp.pdf('signal')
self._random = wsp.pdf('random')
self._nsig = wsp.var('NSig')
self._nrnd = wsp.var('NSPi')
if comb_bkg:
self._ncmb = wsp.var('NBkg')
self._comb = wsp.pdf('combinatorial')
def _get_number_of_signal(self):
import ROOT.RooFit as RF
_isig = self._signal.createIntegral(
self.wsp.set('datavars'), RF.NormSet(self.wsp.set('datavars')),
RF.Range("signal"))
log.info('Signal yield: {}'.format(self._nsig.getVal()))
from k3pi_utilities import parser, logger
from matplotlib.backends.backend_pdf import PdfPages
from analysis import extended_selection, selection, misid_selection
from itertools import combinations
from k3pi_utilities import variables as vars
from k3pi_utilities import buffer
import matplotlib.pyplot as plt
from k3pi_cpp import compute_delta_angle
import seaborn.apionly as sns
from k3pi_utilities.helpers import add_separation_page
import numpy as np
import pandas as pd
from k3pi_config import config
log = logger.get_logger('candidate_sorting')
@buffer.buffer_load
def remove_right_sign_candidates():
"""Remove wrong sign D0 candidates which are combined and end up
in the signal window in the right sign sample"""
# Get the necessary information from the current mode
year = gcm().year
polarity = gcm().polarity
polarity = gcm().polarity
if gcm().mode not in config.twotag_modes:
rs, ws = 'RS', 'WS'
else:
rs, ws = '2tag_RS', '2tag_WS'
with MODE(polarity, year, rs):