def test_parse_func_var_valid(self):
test_vars = ('JpsiPt', 'JpsiRap', 'chic_Mass', 'someFunkyVariable_34')
test_funcs = ('abs', 'sqrt', 'sin')
for var in test_vars:
self.assertEqual((var, None), mh.parse_func_var(var))
for func in test_funcs:
self.assertEqual((var, getattr(np, func)),
(mh.parse_func_var(func + '(' + var + ')')))
def get_load_vars(self):
"""
Get the variable names and bounds that are necessary to define the
variables in a workspace such that data can be imported from a TTree.
Returns:
load_vars, bounds: List of tuples of strings, where first element is
the name of the variable and the second is the boundaries as
comma separated list
"""
variables = [
(self.fit_var, ",".join([str(v) for v in self.fit_range])),
]
for ivar, var in enumerate(self.bin_cut_vars):
var_name, var_func = parse_func_var(var)
bounds_up = np.max(self.binning[ivar])
# For now only handle the 'abs(var)' case separately
if var_func is not None and 'abs' in var_func.__name__:
bounds_low = -bounds_up
else:
bounds_low = np.min(self.binning[ivar])
variables.append((var_name, "{}, {}".format(bounds_low,
bounds_up)))
return variables
def test_parse_func_var_invalid(self, mock_logger):
res = mh.parse_func_var('some_funky_func(some_funky_var)')
mock_logger.error.assert_called_with(
'Could not find a numpy function named \'some_funky_func\' that '\
'was parsed from \'some_funky_func(some_funky_var)\''
)
self.assertEqual(None, res)
def run_fit(model,
datafile,
outfile,
bin_var,
binning,
massrange,
fix_shape=False,
refit=False,
weights=None):
"""Import data, run fits and store the results"""
if bin_var is None:
logging.fatal('Trying to run a fit without specifying a bin variable. '
'It is also possible that this was not implemented for '
'the model that is trying to be used. Sorry :(')
sys.exit(1)
wsp, costh_bins = create_workspace(model, datafile, bin_var, binning,
massrange, outfile, weights)
model.define_model(wsp)
wsp.Print()
if fix_shape:
# Refitting doesn't make sense if we already fixed the shapes beforehand
refit = False
model.fit(wsp, 'costh_integrated', weighted_fit=weights is not None)
shape_pars = get_shape_params(wsp, 'costh_integrated', model)
fix_params(wsp, [(sp, None) for sp in shape_pars])
bin_var = parse_func_var(bin_var)[0]
do_binned_fits(model, wsp, costh_bins, refit, weights is not None, bin_var)
wsp.writeToFile(outfile)
def get_binning(config):
"""
Read the fit config json and get the bin variable and the binning
"""
binning = parse_binning(config['binning'][1])
bins = zip(binning[:-1], binning[1:])
return parse_func_var(config['bin_vars'][1])[0], bins
def create_workspace(model,
datafile,
binvar,
binning,
massrange,
fitfile,
weights=None):
"""
Create the workspace with the data already imported and the model defined,
also in charge of writing the bin info json file
"""
wsp = r.RooWorkspace('ws_mass_fit')
massrange = [float(v) for v in massrange.split(',')]
# load the data and apply the mass selection of the fitting range immediately
bin_var = parse_func_var(binvar) # necessary for loading
variables = [model.mname, bin_var[0]]
if weights is not None:
variables.append(weights)
data = apply_selections(get_dataframe(datafile, columns=variables),
select_bin(model.mname, *massrange))
costh_bins, costh_means = get_costh_bins(binning, bin_var, data)
create_bin_info_json(fitfile.replace('.root', '_bin_sel_info.json'),
costh_bins, costh_means, bin_var[0], datafile)
# Create the variables in the workspace
try_factory(wsp, '{}[{}, {}]'.format(model.mname, *massrange))
if 'abs' in bin_var[1].__name__:
try_factory(
wsp, '{}[{}, {}]'.format(bin_var[0], -np.max(costh_bins),
np.max(costh_bins)))
else:
try_factory(
wsp, '{}[{}, {}]'.format(bin_var[0], np.min(costh_bins),
np.max(costh_bins)))
dset_vars = r.RooArgSet(get_var(wsp, model.mname),
get_var(wsp, bin_var[0]))
tree = array2tree(data.to_records(index=False))
if weights is not None:
try_factory(wsp, '{}[0, 1e5]'.format(weights))
dataset = r.RooDataSet('full_data', 'full data sample', tree,
dset_vars, '', weights)
else:
dataset = r.RooDataSet('full_data', 'full data sample', tree,
dset_vars)
ws_import(wsp, dataset)
return wsp, costh_bins
def __init__(self, config):
"""
Read all the info from the config and do some computation which are then
needed throughout
"""
# TODO: read the configuration from a json file and make that the
# argument to the function
self.binning = np.array([
parse_binning(config['binning'][0]),
parse_binning(config['binning'][1])
])
# variables as they are defined and used for deciding the binning
# inside ROOT with string cuts
self.bin_cut_vars = config['bin_vars']
# For the internal use and especially for naming the variables it is not
# possible to have functional variable definitions. Simply use whatever
# is the argument of func(var) (i.e. var)
self.bin_vars = [parse_func_var(v)[0] for v in config['bin_vars']]
self.bins = get_bins(self.binning[0], self.binning[1],
self.bin_vars[0], self.bin_vars[1])
self.fit_var = config['fit_variable']
self.fit_range = config['fit_range']
self.expression_strings = config['expression_strings']
self.proto_params = config['proto_parameters']
self.sub_models = config['sub_models']
self.nevent_yields = [m['event_yield'] for m in self.sub_models]
self.components = [(m['name'], m['plot']) for m in self.sub_models]
self.full_model = config['full_model']['name']
# prototype full model expression, where the bin model can be achieved
# with simple replacements
self.proto_fexpr = _full_model_expr(config['full_model']['type'],
self.full_model,
[c[0] for c in self.components],
self.nevent_yields)
self.fix_vars = []
if 'fix_vars' in config:
for var in config['fix_vars']:
self.fix_vars.append((var.keys()[0], var.values()[0]))
self.plot_config = config['plot_config']
def main(args):
"""Main"""
frames = args.frames.split(',')
if args.bin_variable is not None:
bin_var = parse_func_var(args.bin_variable)
if args.binning is None:
logging.error('You have to define a binning for \'{}\' if you want '
'to use it as binning variable'.format(bin_var))
sys.exit(1)
binning = parse_binning(args.binning)
else:
binning = None
bin_var = None
logging.info('Processing gen level file')
gen_hists = create_histograms(get_dataframe(args.genlevelfile), frames,
args.ncosth, args.nphi, bin_var, binning)
logging.info('Processing reco level file')
reco_hists = create_histograms(get_dataframe(args.recolevelfile), frames,
args.ncosth, args.nphi, bin_var, binning,
WEIGHT_F)
logging.debug('Scaling gen level hists by '.format(args.scale_gen))
for hist in gen_hists.values():
hist.Scale(args.scale_gen)
logging.info('calculating acceptance maps')
acc_maps = OrderedDict()
for name, hist in reco_hists.iteritems():
gen_map_n = [n for n in gen_hists if n in name]
if len(gen_map_n) > 1:
logging.warning('Found more than 1 gen level map for \'{}\': {}'
.format(name, gen_map_n))
# Still use just the first, since we should always just have 1
acc_maps[name] = divide(hist, gen_hists[gen_map_n[0]])
logging.debug('storing to output file')
outfile = r.TFile(args.outfile, 'recreate' if args.recreate else 'update')
store_hists(outfile, gen_hists, 'gen_costh_phi', bin_var)
store_hists(outfile, reco_hists, 'reco_costh_phi', bin_var)
store_hists(outfile, acc_maps, 'acc_map_costh_phi', bin_var)
def corr_w(data):
corr = corrmap.eval(_get_var(data, 'costh_' + frame),
_get_var(data, 'phi_' + frame),
_get_var(data, *parse_func_var(var)))
print('Percentage in acceptance: ', np.sum(corr > 0) / float(len(corr)))
return (corr > 0) * corr