From 51f1e7c5c32531eac0ab7196f10a9b936e464aaa Mon Sep 17 00:00:00 2001 From: davidwalter2 Date: Thu, 26 Mar 2026 13:10:48 -0400 Subject: [PATCH 1/9] Add composite POI model --- rabbit/poi_models/poi_model.py | 30 ++++++++++++++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/rabbit/poi_models/poi_model.py b/rabbit/poi_models/poi_model.py index e011de1..e5a1bbe 100644 --- a/rabbit/poi_models/poi_model.py +++ b/rabbit/poi_models/poi_model.py @@ -52,6 +52,36 @@ def set_poi_default(self, expectSignal, allowNegativePOI=False): self.xpoidefault = tf.sqrt(poidefault) +class CompositePOIModel(POIModel): + """ + multiply different POI models together + """ + + def __init__( + self, + poi_models, + allowNegativePOI=False, + ): + + self.poi_models = poi_models + + self.npoi = sum([m.npoi for m in poi_models]) + + self.pois = np.concatenate([m.pois for m in poi_models]) + + self.allowNegativePOI = allowNegativePOI + + self.is_linear = self.npoi == 0 or self.allowNegativePOI + + def compute(self, poi, full=False): + start = 0 + results = [] + for m in self.poi_models: + results.append(m.compute(poi[start : start + m.npoi], full)) + start += m.npoi + return tf.math.reduce_prod(tf.stack(results, axis=0), axis=0) + + class Ones(POIModel): """ multiply all processes with ones From d0890e4ef502af02d528fa1fe28e5a01274f16fc Mon Sep 17 00:00:00 2001 From: davidwalter2 Date: Thu, 26 Mar 2026 15:49:09 -0400 Subject: [PATCH 2/9] Implement saturated likelihood test for projections --- bin/rabbit_fit.py | 41 ++++++++- rabbit/fitter.py | 162 +++++++++++++++++++-------------- rabbit/poi_models/poi_model.py | 70 +++++++++++++- 3 files changed, 203 insertions(+), 70 deletions(-) diff --git a/bin/rabbit_fit.py b/bin/rabbit_fit.py index 5668e26..d56a2c8 100755 --- a/bin/rabbit_fit.py +++ b/bin/rabbit_fit.py @@ -12,7 +12,9 @@ from rabbit import fitter, inputdata, parsing, workspace from rabbit.mappings import helpers as mh from rabbit.mappings import mapping as mp +from rabbit.mappings import project from rabbit.poi_models import helpers as ph +from rabbit.poi_models import poi_model from rabbit.tfhelpers import edmval_cov from wums import output_tools, logging # isort: skip @@ -248,7 +250,44 @@ def save_hists(args, mappings, fitter, ws, prefit=True, profile=False): ) if aux[-2] is not None: - ws.add_chi2(aux[-2], aux[-1], prefit, mapping) + chi2val = float(aux[-2]) + ndf = int(aux[-1]) + p_val = chi2.sf(chi2val, ndf) + + logger.info("Traditional chi2:") + logger.info(f" ndof: {ndf}") + logger.info(f" chi2/ndf = {round(chi2val)}") + logger.info(rf" p-value: {round(p_val * 100, 2)}%") + + ws.add_chi2(chi2val, ndf, prefit, mapping) + + if not prefit and type(mapping) == project.Project: + # saturated likelihood test + + saturated_model = poi_model.SaturatedProjectModel( + fitter.indata, mapping.channel_info + ) + composite_model = poi_model.CompositePOIModel( + [fitter.poi_model, saturated_model] + ) + + fitter.init_fit_parms( + composite_model, + args.setConstraintMinimum, + unblind=args.unblind, + freeze_parameters=args.freezeParameters, + ) + cb = fitter.minimize() + nllvalreduced = fitter.reduced_nll().numpy() + + ndfsat = saturated_model.npoi + chi2_val = 2.0 * (ws.results["nllvalreduced"] - nllvalreduced) + p_val = chi2.sf(chi2_val, ndfsat) + + logger.info("Saturated chi2:") + logger.info(f" ndof: {ndfsat}") + logger.info(f" 2*deltaNLL: {round(chi2_val, 2)}") + logger.info(rf" p-value: {round(p_val * 100, 2)}%") if args.saveHistsPerProcess and not mapping.skip_per_process: logger.info(f"Save processes histogram for {mapping.key}") diff --git a/rabbit/fitter.py b/rabbit/fitter.py index c6b7e5d..9482111 100644 --- a/rabbit/fitter.py +++ b/rabbit/fitter.py @@ -129,62 +129,21 @@ def __init__( self.chisqFit = options.chisqFit self.covarianceFit = options.covarianceFit - self.prefitUnconstrainedNuisanceUncertainty = ( - options.prefitUnconstrainedNuisanceUncertainty - ) - - self.poi_model = poi_model self.do_blinding = do_blinding - if self.do_blinding: - self._blinding_offsets_poi = tf.Variable( - tf.ones([self.poi_model.npoi], dtype=self.indata.dtype), - trainable=False, - name="offset_poi", - ) - self._blinding_offsets_theta = tf.Variable( - tf.zeros([self.indata.nsyst], dtype=self.indata.dtype), - trainable=False, - name="offset_theta", - ) - self.init_blinding_values(options.unblind) - - self.parms = np.concatenate([self.poi_model.pois, self.indata.systs]) - - # tf tensor containing default constraint minima - theta0default = np.zeros(self.indata.nsyst) - for parm, val in options.setConstraintMinimum: - idx = np.where(self.indata.systs.astype(str) == parm)[0] - if len(idx) != 1: - raise RuntimeError( - f"Expect to find exactly one match for {parm} to set constraint minimum, but found {len(idx)}" - ) - theta0default[idx[0]] = val - - self.theta0default = tf.convert_to_tensor( - theta0default, dtype=self.indata.dtype + self.prefit_unconstrained_nuisance_uncertainty = ( + options.prefitUnconstrainedNuisanceUncertainty ) - # tf variable containing all fit parameters - if self.poi_model.npoi > 0: - xdefault = tf.concat( - [self.poi_model.xpoidefault, self.theta0default], axis=0 - ) - else: - xdefault = self.theta0default - - self.x = tf.Variable(xdefault, trainable=True, name="x") - - # for freezing parameters - self.frozen_params = [] - self.frozen_params_mask = tf.Variable( - tf.zeros_like(self.x, dtype=tf.bool), trainable=False, dtype=tf.bool + # --- fit params + self.init_fit_parms( + poi_model, + options.setConstraintMinimum, + unblind=options.unblind, + freeze_parameters=options.freezeParameters, ) - self.frozen_indices = np.array([]) - self.freeze_params(options.freezeParameters) - - # observed number of events per bin + # --- observed number of events per bin self.nobs = tf.Variable( tf.zeros_like(self.indata.data_obs), trainable=False, name="nobs" ) @@ -211,21 +170,10 @@ def __init__( # provided covariance self.data_cov_inv = self.indata.data_cov_inv - # constraint minima for nuisance parameters - self.theta0 = tf.Variable( - self.theta0default, - trainable=False, - name="theta0", - ) - self.var_theta0 = tf.where( - self.indata.constraintweights == 0.0, - tf.zeros_like(self.indata.constraintweights), - tf.math.reciprocal(self.indata.constraintweights), - ) - # FIXME for now this is needed even if binByBinStat is off because of how it is used in the global impacts # and uncertainty band computations (gradient is allowed to be zero or None and then propagated or skipped only later) + # --- MC stat # global observables for mc stat uncertainty if self.binByBinStatMode == "full": self.beta_shape = self.indata.sumw.shape @@ -308,15 +256,84 @@ def __init__( self.expected_yield(), trainable=False, name="nexpnom" ) + def init_fit_parms( + self, + poi_model, + set_constraint_minimum=[], + unblind=False, + freeze_parameters=[], + ): + self.poi_model = poi_model + + if self.do_blinding: + self._blinding_offsets_poi = tf.Variable( + tf.ones([self.poi_model.npoi], dtype=self.indata.dtype), + trainable=False, + name="offset_poi", + ) + self._blinding_offsets_theta = tf.Variable( + tf.zeros([self.indata.nsyst], dtype=self.indata.dtype), + trainable=False, + name="offset_theta", + ) + self.init_blinding_values(unblind) + + self.parms = np.concatenate([self.poi_model.pois, self.indata.systs]) + + # tf tensor containing default constraint minima + theta0default = np.zeros(self.indata.nsyst) + for parm, val in set_constraint_minimum: + idx = np.where(self.indata.systs.astype(str) == parm)[0] + if len(idx) != 1: + raise RuntimeError( + f"Expect to find exactly one match for {parm} to set constraint minimum, but found {len(idx)}" + ) + theta0default[idx[0]] = val + + self.theta0default = tf.convert_to_tensor( + theta0default, dtype=self.indata.dtype + ) + + # tf variable containing all fit parameters + if self.poi_model.npoi > 0: + xdefault = tf.concat( + [self.poi_model.xpoidefault, self.theta0default], axis=0 + ) + else: + xdefault = self.theta0default + + self.x = tf.Variable(xdefault, trainable=True, name="x") + # parameter covariance matrix self.cov = tf.Variable( self.prefit_covariance( - unconstrained_err=self.prefitUnconstrainedNuisanceUncertainty + unconstrained_err=self.prefit_unconstrained_nuisance_uncertainty ), trainable=False, name="cov", ) + # constraint minima for nuisance parameters + self.theta0 = tf.Variable( + self.theta0default, + trainable=False, + name="theta0", + ) + self.var_theta0 = tf.where( + self.indata.constraintweights == 0.0, + tf.zeros_like(self.indata.constraintweights), + tf.math.reciprocal(self.indata.constraintweights), + ) + + # for freezing parameters + self.frozen_params = [] + self.frozen_params_mask = tf.Variable( + tf.zeros_like(self.x, dtype=tf.bool), trainable=False, dtype=tf.bool + ) + + self.frozen_indices = np.array([]) + self.freeze_params(freeze_parameters) + # determine if problem is linear (ie likelihood is purely quadratic) self.is_linear = ( (self.chisqFit or self.covarianceFit) @@ -326,6 +343,16 @@ def __init__( and ((not self.binByBinStat) or self.binByBinStatType == "normal-additive") ) + # force retrace of @tf.function methods since self.x shape may have changed + for name in dir(type(self)): + val = getattr(type(self), name, None) + if hasattr(val, "python_function"): + setattr( + self, + name, + tf.function(val.python_function.__get__(self, type(self))), + ) + def load_fitresult(self, fitresult_file, fitresult_key): # load results from external fit and set postfit value and covariance elements for common parameters cov_ext = None @@ -549,7 +576,7 @@ def betadefaultassign(self): def defaultassign(self): self.cov.assign( self.prefit_covariance( - unconstrained_err=self.prefitUnconstrainedNuisanceUncertainty + unconstrained_err=self.prefit_unconstrained_nuisance_uncertainty ) ) self.theta0defaultassign() @@ -1028,7 +1055,7 @@ def compute_derivatives(dvars): profile, pdexpdbeta, pd2ldbeta2_pdexpdbeta if pdexpdbeta is not None else None, - self.prefitUnconstrainedNuisanceUncertainty, + self.prefit_unconstrained_nuisance_uncertainty, ) else: impacts = None @@ -1826,7 +1853,7 @@ def expected_events( fun = mapping.compute_flat if inclusive else mapping.compute_flat_per_process - aux = [None] * 4 + aux = [None] * 6 if ( compute_cov or compute_variance @@ -1866,8 +1893,6 @@ def expected_events( mapping.ndf_reduction, profile=profile, ) - logger.info(f"chi2/ndf = {chi2val.numpy()}/{ndf.numpy()}") - aux.append(chi2val) aux.append(ndf) else: @@ -2091,6 +2116,7 @@ def loss_val_grad_hess_beta(self, profile=True): return val, grad, hess def fit(self): + logger.info("Perform iterative fit") def scipy_loss(xval): self.x.assign(xval) diff --git a/rabbit/poi_models/poi_model.py b/rabbit/poi_models/poi_model.py index e5a1bbe..8df4365 100644 --- a/rabbit/poi_models/poi_model.py +++ b/rabbit/poi_models/poi_model.py @@ -1,3 +1,6 @@ +import functools +import itertools + import numpy as np import tensorflow as tf @@ -73,13 +76,17 @@ def __init__( self.is_linear = self.npoi == 0 or self.allowNegativePOI + self.xpoidefault = tf.concat([m.xpoidefault for m in poi_models], axis=0) + def compute(self, poi, full=False): start = 0 results = [] for m in self.poi_models: results.append(m.compute(poi[start : start + m.npoi], full)) start += m.npoi - return tf.math.reduce_prod(tf.stack(results, axis=0), axis=0) + + rnorm = functools.reduce(lambda a, b: a * b, results) + return rnorm class Ones(POIModel): @@ -225,3 +232,64 @@ def compute(self, poi, full=False): rnorm = tf.reshape(rnorm, [1, -1]) return rnorm + + +class SaturatedProjectModel(POIModel): + """ + For computing the saturated test statistic of a projection. + Add one free parameter for each projected bin + """ + + def __init__( + self, indata, channel_info, expectSignal=None, allowNegativePOI=False, **kwargs + ): + self.indata = indata + self.channel_info_mapping = channel_info + + self.npoi = np.sum( + [np.prod([a.size for a in v["axes"]]) for v in channel_info.values()] + ) + + names = [] + for k, v in self.channel_info_mapping.items(): + for idxs in itertools.product(*[range(a.size) for a in v["axes"]]): + label = "_".join(f"{a.name}{i}" for a, i in zip(v["axes"], idxs)) + names.append(f"saturated_{k}_{label}".encode()) + + self.pois = np.array(names) + + self.allowNegativePOI = allowNegativePOI + + self.is_linear = self.npoi == 0 or self.allowNegativePOI + + self.set_poi_default(expectSignal, allowNegativePOI) + + def compute(self, poi, full=False): + start = 0 + rnorms = [] + for k, v in self.indata.channel_info.items(): + if v["masked"] and not full: + continue + shape_input = [a.size for a in v["axes"]] + + irnorm = tf.ones(shape_input, dtype=self.indata.dtype) + if k in self.channel_info_mapping.keys(): + mapping_axes = self.channel_info_mapping[k]["axes"] + shape_mapping = [a.size if a in mapping_axes else 1 for a in v["axes"]] + npoi = np.prod([a.size for a in mapping_axes]) + ipoi = poi[start : start + npoi] + irnorm *= tf.reshape(ipoi, shape_mapping) + start += npoi + + irnorm = tf.reshape( + irnorm, + [ + -1, + ], + ) + rnorms.append(irnorm) + + rnorm = tf.concat(rnorms, axis=0) + rnorm = tf.reshape(rnorm, [-1, 1]) + + return rnorm From a311929506f1e074eed0349d73e0321b9b5fb964 Mon Sep 17 00:00:00 2001 From: davidwalter2 Date: Thu, 26 Mar 2026 17:15:56 -0400 Subject: [PATCH 3/9] Fix on saturated chi squared for projections --- bin/rabbit_fit.py | 23 ++++++++++++++--------- rabbit/fitter.py | 19 +++++++++++++++++++ rabbit/poi_models/poi_model.py | 5 ++++- rabbit/workspace.py | 4 +++- 4 files changed, 40 insertions(+), 11 deletions(-) diff --git a/bin/rabbit_fit.py b/bin/rabbit_fit.py index d56a2c8..44184ef 100755 --- a/bin/rabbit_fit.py +++ b/bin/rabbit_fit.py @@ -1,5 +1,7 @@ #!/usr/bin/env python3 +import copy + import tensorflow as tf tf.config.experimental.enable_op_determinism() @@ -254,7 +256,7 @@ def save_hists(args, mappings, fitter, ws, prefit=True, profile=False): ndf = int(aux[-1]) p_val = chi2.sf(chi2val, ndf) - logger.info("Traditional chi2:") + logger.info("Linear chi2:") logger.info(f" ndof: {ndf}") logger.info(f" chi2/ndf = {round(chi2val)}") logger.info(rf" p-value: {round(p_val * 100, 2)}%") @@ -271,24 +273,27 @@ def save_hists(args, mappings, fitter, ws, prefit=True, profile=False): [fitter.poi_model, saturated_model] ) - fitter.init_fit_parms( + fitter_saturated = copy.deepcopy(fitter) + fitter_saturated.init_fit_parms( composite_model, args.setConstraintMinimum, unblind=args.unblind, freeze_parameters=args.freezeParameters, ) - cb = fitter.minimize() - nllvalreduced = fitter.reduced_nll().numpy() + cb = fitter_saturated.minimize() + nllvalreduced = fitter_saturated.reduced_nll().numpy() - ndfsat = saturated_model.npoi - chi2_val = 2.0 * (ws.results["nllvalreduced"] - nllvalreduced) - p_val = chi2.sf(chi2_val, ndfsat) + ndf = saturated_model.npoi + chi2val = 2.0 * (ws.results["nllvalreduced"] - nllvalreduced) + p_val = chi2.sf(chi2val, ndf) logger.info("Saturated chi2:") - logger.info(f" ndof: {ndfsat}") - logger.info(f" 2*deltaNLL: {round(chi2_val, 2)}") + logger.info(f" ndof: {ndf}") + logger.info(f" 2*deltaNLL: {round(chi2val, 2)}") logger.info(rf" p-value: {round(p_val * 100, 2)}%") + ws.add_chi2(chi2val, ndf, prefit, mapping, saturated=True) + if args.saveHistsPerProcess and not mapping.skip_per_process: logger.info(f"Save processes histogram for {mapping.key}") diff --git a/rabbit/fitter.py b/rabbit/fitter.py index 9482111..f7915e6 100644 --- a/rabbit/fitter.py +++ b/rabbit/fitter.py @@ -353,6 +353,25 @@ def init_fit_parms( tf.function(val.python_function.__get__(self, type(self))), ) + def __deepcopy__(self, memo): + import copy + + # Instance-level tf.function overrides (set by init_fit_parms to force retracing) + # contain FuncGraph objects that cannot be deepcopied. Strip them before copying + # so the copy falls back to the class-level @tf.function methods and retraces. + jit_overrides = { + name + for name in self.__dict__ + if hasattr(getattr(type(self), name, None), "python_function") + } + state = {k: v for k, v in self.__dict__.items() if k not in jit_overrides} + cls = type(self) + obj = cls.__new__(cls) + memo[id(self)] = obj + for k, v in state.items(): + setattr(obj, k, copy.deepcopy(v, memo)) + return obj + def load_fitresult(self, fitresult_file, fitresult_key): # load results from external fit and set postfit value and covariance elements for common parameters cov_ext = None diff --git a/rabbit/poi_models/poi_model.py b/rabbit/poi_models/poi_model.py index 8df4365..2c1cd4c 100644 --- a/rabbit/poi_models/poi_model.py +++ b/rabbit/poi_models/poi_model.py @@ -247,7 +247,10 @@ def __init__( self.channel_info_mapping = channel_info self.npoi = np.sum( - [np.prod([a.size for a in v["axes"]]) for v in channel_info.values()] + [ + np.prod([a.size for a in v["axes"]]) if len(v["axes"]) else 1 + for v in channel_info.values() + ] ) names = [] diff --git a/rabbit/workspace.py b/rabbit/workspace.py index 66417dd..fdf56c0 100644 --- a/rabbit/workspace.py +++ b/rabbit/workspace.py @@ -185,8 +185,10 @@ def add_hist( def add_value(self, value, name, *args, **kwargs): self.dump_obj(value, name, *args, **kwargs) - def add_chi2(self, chi2, ndf, prefit, mapping): + def add_chi2(self, chi2, ndf, prefit, mapping, saturated=False): postfix = "_prefit" if prefit else "" + if saturated: + postfix += "_saturated" self.add_value(int(ndf), "ndf" + postfix, mapping.key) self.add_value(float(chi2), "chi2" + postfix, mapping.key) From 8b0f1da4c89a210545ade770d572bf0b3aae930c Mon Sep 17 00:00:00 2001 From: davidwalter2 Date: Fri, 27 Mar 2026 11:15:18 -0400 Subject: [PATCH 4/9] Add option for computing saturated test on projections; load saturated test if available in plotting script --- bin/rabbit_fit.py | 12 ++++++++- bin/rabbit_plot_hists.py | 57 +++++++++++++++++++++++++++++++--------- 2 files changed, 55 insertions(+), 14 deletions(-) diff --git a/bin/rabbit_fit.py b/bin/rabbit_fit.py index 44184ef..4a8b38d 100755 --- a/bin/rabbit_fit.py +++ b/bin/rabbit_fit.py @@ -141,6 +141,12 @@ def make_parser(): action="store_true", help="save postfit histograms with each noi varied up to down", ) + parser.add_argument( + "--computeSaturatedProjectionTests", + default=False, + action="store_true", + help="Compute the saturated likelihood test for Project mappings", + ) parser.add_argument( "--noChi2", default=False, @@ -263,7 +269,11 @@ def save_hists(args, mappings, fitter, ws, prefit=True, profile=False): ws.add_chi2(chi2val, ndf, prefit, mapping) - if not prefit and type(mapping) == project.Project: + if ( + not prefit + and type(mapping) == project.Project + and args.computeSaturatedProjectionTests + ): # saturated likelihood test saturated_model = poi_model.SaturatedProjectModel( diff --git a/bin/rabbit_plot_hists.py b/bin/rabbit_plot_hists.py index 7a7556f..6ad9d77 100755 --- a/bin/rabbit_plot_hists.py +++ b/bin/rabbit_plot_hists.py @@ -1558,19 +1558,46 @@ def make_plots( ) -def get_chi2(result, no_chi2=True, fittype="postfit"): - chi2_key = f"chi2_prefit" if fittype == "prefit" else "chi2" - ndf_key = f"ndf_prefit" if fittype == "prefit" else "ndf" - if not no_chi2 and fittype == "postfit" and result.get("postfit_profile", False): - # use saturated likelihood test if relevant - chi2 = 2.0 * result["nllvalreduced"] - ndf = result["ndfsat"] - return chi2, ndf, True - elif not no_chi2 and chi2_key in result: - return result[chi2_key], result[ndf_key], False - else: +def get_chi2(result, no_chi2=True, fittype="postfit", chi2type="automatic"): + if no_chi2: return None, None, False + if fittype == "prefit": + if chi2type in [ + "saturated", + ]: + raise RuntimeError("No saturated test statistic in prefit") + elif chi2type in ["automatic", "linear"]: + chi2type = "linear" + chi2_key = "chi2_prefit" + ndf_key = "ndf_prefit" + else: + chi2_key = "chi2" + ndf_key = "ndf" + + if chi2type in ["automatic", "saturated"]: + if ( + result.get("postfit_profile", False) + and "nllvalreduced" in result.keys() + ): + # use saturated likelihood test from actual fit + chi2 = 2.0 * result["nllvalreduced"] + ndf = result["ndfsat"] + return chi2, ndf, True + elif f"{chi2_key}_saturated" in result.keys(): + # use saturated likelihood test from mapping + chi2_key += "_saturated" + ndf_key += "_saturated" + saturated = True + elif chi2type == "automatic": + saturated = False + else: + raise ValueError("No saturated test statistic in postfit results found") + else: + saturated = False + + return result.get(chi2_key, None), result.get(ndf_key, None), saturated + def main(): args = parseArgs() @@ -1675,13 +1702,17 @@ def main(): fitresult if fittype == "postfit" and ( - (instance_key == "BaseMapping" and args.chisq != "linear") - or args.chisq == "saturated" + instance_key + in [ + "BaseMapping", + ] + and args.chisq != "linear" ) else instance ), args.chisq in [" ", "none", None], fittype, + args.chisq, ) for channel, result in instance["channels"].items(): From c037776d58da4d36456c36b54e238adada0ea2ac Mon Sep 17 00:00:00 2001 From: davidwalter2 Date: Fri, 27 Mar 2026 13:14:54 -0400 Subject: [PATCH 5/9] Fix plotting and set covariance for pois in prefit to specified value --- bin/rabbit_plot_hists.py | 4 ++-- rabbit/fitter.py | 5 ++++- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/bin/rabbit_plot_hists.py b/bin/rabbit_plot_hists.py index 6ad9d77..322082c 100755 --- a/bin/rabbit_plot_hists.py +++ b/bin/rabbit_plot_hists.py @@ -1567,10 +1567,10 @@ def get_chi2(result, no_chi2=True, fittype="postfit", chi2type="automatic"): "saturated", ]: raise RuntimeError("No saturated test statistic in prefit") - elif chi2type in ["automatic", "linear"]: - chi2type = "linear" + chi2_key = "chi2_prefit" ndf_key = "ndf_prefit" + saturated = False else: chi2_key = "chi2" ndf_key = "ndf" diff --git a/rabbit/fitter.py b/rabbit/fitter.py index f7915e6..3b654a7 100644 --- a/rabbit/fitter.py +++ b/rabbit/fitter.py @@ -532,7 +532,10 @@ def _default_beta0(self): def prefit_covariance(self, unconstrained_err=0.0): # free parameters are taken to have zero uncertainty for the purposes of prefit uncertainties - var_poi = tf.zeros([self.poi_model.npoi], dtype=self.indata.dtype) + var_poi = ( + tf.ones([self.poi_model.npoi], dtype=self.indata.dtype) + * unconstrained_err**2 + ) # nuisances have their uncertainty taken from the constraint term, but unconstrained nuisances # are set to a placeholder uncertainty (zero by default) for the purposes of prefit uncertainties From 31f4152905f8af7697c49d87164c0110ab7e5e68 Mon Sep 17 00:00:00 2001 From: davidwalter2 Date: Sat, 28 Mar 2026 09:30:14 -0400 Subject: [PATCH 6/9] Store data covariance and add new option '--dataCovariance' to plot data hist --- bin/rabbit_plot_hists.py | 345 ++++++++++++++++++++++++--------------- rabbit/workspace.py | 41 ++++- 2 files changed, 250 insertions(+), 136 deletions(-) diff --git a/bin/rabbit_plot_hists.py b/bin/rabbit_plot_hists.py index 322082c..650e2a7 100755 --- a/bin/rabbit_plot_hists.py +++ b/bin/rabbit_plot_hists.py @@ -418,6 +418,11 @@ def parseArgs(): default=None, help="Label for uncertainty shown in the (ratio) plot", ) + parser.add_argument( + "--dataCovariance", + action="store_true", + help="Use covariance information to plot the data uncertainty", + ) args = parser.parse_args() return args @@ -451,6 +456,7 @@ def make_plot( is_normalized=False, binwnorm=1.0, counts=True, + dataCovariance=False, ): ratio = not args.noLowerPanel and h_data is not None diff = not args.noLowerPanel and args.diff and h_data is not None @@ -559,6 +565,16 @@ def make_plot( xlabel = plot_tools.get_axis_label(config, axes_names, args.xlabel) + # for uncertaity bands + edges = h_data.axes[0].edges + + # need to divide by bin width + binwidth = edges[1:] - edges[:-1] if binwnorm else 1.0 + if h_inclusive.storage_type != hist.storage.Weight: + raise ValueError( + f"Did not find uncertainties in {fittype} hist. Make sure you run rabbit_fit with --computeHistErrors!" + ) + if ratio or diff: if args.ratioToData: rlabel = r"Pred\ " @@ -767,50 +783,129 @@ def make_plot( ) if data: - hep.histplot( - h_data, - yerr=True if counts else h_data.variances() ** 0.5, - histtype=histtype_data, - color="black", - label=args.dataName, - binwnorm=binwnorm, - ax=ax1, - alpha=1.0, - zorder=2, - flow="none", - ) + zorder_data = 1 + if dataCovariance: + nom_data = h_data.values() / binwidth + std_data = np.sqrt(h_data.variances()) / binwidth + hatchstyle_data = None + linecolor_data = "red" # "black" + facecolor_data = "red" # silver" + facecolor_alpha_data = 0.5 + + ax1.fill_between( + edges, + np.append((nom_data + std_data), ((nom_data + std_data))[-1]), + np.append((nom_data - std_data), ((nom_data - std_data))[-1]), + step="post", + facecolor=facecolor_data, + alpha=facecolor_alpha_data, + hatch=hatchstyle_data, + edgecolor="k", + zorder=zorder_data, + linewidth=0.0, + ) - if h_data_stat is not None: - var_stat = h_data_stat.values() ** 2 - h_data_stat = h_data.copy() - h_data_stat.variances()[...] = var_stat + hep.histplot( + h_data, + histtype="step", + color=linecolor_data, + binwnorm=binwnorm, + yerr=False, + ax=ax1, + linestyle="--", + zorder=zorder_data, + flow="none", + ) + extra_handles_upper.append( + plot_tools.LineBandPolygon( + [[0, 0], [1, 0], [1, 1], [0, 1]], + closed=True, + facecolor=facecolor_data, + edgecolor=linecolor_data, + linestyle="--", + alpha=facecolor_alpha_data, + ) + ) + extra_labels_upper.append(args.dataName) + else: hep.histplot( - h_data_stat, - yerr=True if counts else h_data_stat.variances() ** 0.5, + h_data, + yerr=True if counts else h_data.variances() ** 0.5, histtype=histtype_data, color="black", + label=args.dataName, binwnorm=binwnorm, - capsize=2, ax=ax1, alpha=1.0, - zorder=2, + zorder=zorder_data, flow="none", ) + + if h_data_stat is not None: + var_stat = h_data_stat.values() ** 2 + h_data_stat = h_data.copy() + h_data_stat.variances()[...] = var_stat + + hep.histplot( + h_data_stat, + yerr=True if counts else h_data_stat.variances() ** 0.5, + histtype=histtype_data, + color="black", + binwnorm=binwnorm, + capsize=2, + ax=ax1, + alpha=1.0, + zorder=zorder_data, + flow="none", + ) if (args.unfoldedXsec or len(h_stack) == 0) and not args.noPrefit: + zorder_pred = 0 + hatchstyle_pred = None + facecolor_pred = "silver" + facecolor_alpha_pred = 0.5 + pred_label = "Prefit model" if args.unfoldedXsec else args.predName + hep.histplot( h_inclusive, yerr=False, histtype="step", color="black", - label="Prefit model" if args.unfoldedXsec else args.predName, binwnorm=binwnorm, ax=ax1, alpha=1.0, - zorder=2, + zorder=zorder_pred, + label=pred_label if not args.upperPanelUncertaintyBand else None, flow="none", ) + if args.upperPanelUncertaintyBand: + nom = h_inclusive.values() / binwidth + std = np.sqrt(h_inclusive.variances()) / binwidth + + ax1.fill_between( + edges, + np.append((nom + std), ((nom + std))[-1]), + np.append((nom - std), ((nom - std))[-1]), + step="post", + facecolor=facecolor_pred, + alpha=facecolor_alpha_pred, + zorder=zorder_pred, + hatch=hatchstyle_pred, + edgecolor="k", + linewidth=0.0, + ) + + extra_handles_upper.append( + plot_tools.LineBandPolygon( + [[0, 0], [1, 0], [1, 1], [0, 1]], + closed=True, + facecolor=facecolor_pred, + edgecolor="black", + ) + ) + extra_labels_upper.append(pred_label) + if args.ylim is None and binwnorm is None: max_y = np.max(h_inclusive.values() + h_inclusive.variances() ** 0.5) min_y = np.min(h_inclusive.values() - h_inclusive.variances() ** 0.5) @@ -882,88 +977,103 @@ def make_plot( else: cutoff = 0.01 - if args.ratioToData: - h_num = h_inclusive - h_den = h_data - else: - h_num = h_data - h_den = h_inclusive - if diff: - h0 = hh.addHists(h_num, h_num, scale2=-1) - h2 = hh.addHists(h_num, h_den, scale2=-1) - if h_data_stat is not None: - if args.ratioToData: - h2_stat = hh.addHists(h_num, h_data_stat, scale2=-1) - else: - h2_stat = hh.addHists(h_data_stat, h_den, scale2=-1) + r_data = lambda x, y: hh.addHists(x, y, scale2=-1) + r_pred = lambda x, y: hh.addHists(x, y, scale2=-1) + r_stat = lambda x, y: hh.addHists(x, y, scale2=-1) else: - h0 = hh.divideHists( - h_num, - h_num, + r_data = lambda x, y: hh.divideHists( + x, + y, cutoff=1e-8, rel_unc=True, flow=False, by_ax_name=False, ) - h2 = hh.divideHists(h_data, h_den, cutoff=cutoff, rel_unc=True) - if h_data_stat is not None: - h2_stat = hh.divideHists( - h_data_stat, h_den, cutoff=cutoff, rel_unc=True - ) + r_pred = lambda x, y: hh.divideHists(x, y, cutoff=cutoff, rel_unc=True) + r_stat = lambda x, y: hh.divideHists(x, y, cutoff=cutoff, rel_unc=True) + + h_den = h_data if args.ratioToData else h_inclusive + if not dataCovariance: + ax2.axhline(0 if diff else 1, linestyle="--", color="black") + + # prediction hep.histplot( - h0, + r_pred(h_inclusive, h_den), histtype="step", - color="grey", - alpha=0.5, + color="black", yerr=False, ax=ax2, - linewidth=2, + zorder=zorder_pred, flow="none", ) if data: - hep.histplot( - h2, - histtype="errorbar", - color="black", - yerr=True if counts else h2.variances() ** 0.5, - linewidth=2, - ax=ax2, - zorder=2, - flow="none", - ) - if h_data_stat is not None: + h2 = r_data(h_data, h_den) + if dataCovariance: + den_data = 1 if diff else h_den.values() / binwidth + hep.histplot( - h2_stat, + r_pred(h_data, h_den), + histtype="step", + color=linecolor_data, + yerr=False, + ax=ax2, + linestyle="--", + zorder=zorder_data, + flow="none", + ) + + ax2.fill_between( + edges, + np.append( + (nom_data + std_data) / den_data, + ((nom_data + std_data) / den_data)[-1], + ), + np.append( + (nom_data - std_data) / den_data, + ((nom_data - std_data) / den_data)[-1], + ), + step="post", + facecolor=facecolor_data, + alpha=facecolor_alpha_data, + zorder=zorder_data, + hatch=hatchstyle_data, + edgecolor="k", + linewidth=0.0, + ) + + else: + hep.histplot( + h2, histtype="errorbar", color="black", - yerr=True if counts else h2_stat.variances() ** 0.5, + yerr=True if counts else h2.variances() ** 0.5, linewidth=2, - capsize=2, ax=ax2, - zorder=2, + zorder=zorder_data, flow="none", ) - - # for uncertaity bands - edges = h_den.axes[0].edges - - # need to divide by bin width - binwidth = edges[1:] - edges[:-1] if binwnorm else 1.0 - if h_den.storage_type != hist.storage.Weight: - raise ValueError( - f"Did not find uncertainties in {fittype} hist. Make sure you run rabbit_fit with --computeHistErrors!" - ) + if h_data_stat is not None: + h2_stat = r_stat(h_data_stat, h_den) + hep.histplot( + h2_stat, + histtype="errorbar", + color="black", + yerr=True if counts else h2_stat.variances() ** 0.5, + linewidth=2, + capsize=2, + ax=ax2, + zorder=zorder_data, + flow="none", + ) if not args.noUncertainty: nom = h_inclusive.values() / binwidth std = np.sqrt(h_inclusive.variances()) / binwidth + den = 1 if diff else h_den.values() / binwidth - hatchstyle = None - facecolor = "silver" - # label_unc = "Pred. unc." default_unc_label = ( "Normalized model unc." if is_normalized else f"{args.predName} unc." ) @@ -971,62 +1081,19 @@ def make_plot( if args.uncertaintyLabel: label_unc = args.uncertaintyLabel - if diff: - # The difference panel is centered at zero; only the upper-panel - # uncertainty band is centered on the nominal prediction. - lower_hi = std - lower_lo = -std - ax2.fill_between( - edges, - np.append(lower_hi, lower_hi[-1]), - np.append(lower_lo, lower_lo[-1]), - step="post", - facecolor=facecolor, - zorder=0, - hatch=hatchstyle, - edgecolor="k", - linewidth=0.0, - label=label_unc if not args.upperPanelUncertaintyBand else None, - ) - if args.upperPanelUncertaintyBand: - ax1.fill_between( - edges, - np.append((nom + std), ((nom + std))[-1]), - np.append((nom - std), ((nom - std))[-1]), - step="post", - facecolor=facecolor, - zorder=0, - hatch=hatchstyle, - edgecolor="k", - linewidth=0.0, - label=label_unc, - ) - else: - ax2.fill_between( - edges, - np.append((nom + std) / nom, ((nom + std) / nom)[-1]), - np.append((nom - std) / nom, ((nom - std) / nom)[-1]), - step="post", - facecolor=facecolor, - zorder=0, - hatch=hatchstyle, - edgecolor="k", - linewidth=0.0, - label=label_unc if not args.upperPanelUncertaintyBand else None, - ) - if args.upperPanelUncertaintyBand: - ax1.fill_between( - edges, - np.append((nom + std), ((nom + std))[-1]), - np.append((nom - std), ((nom - std))[-1]), - step="post", - facecolor=facecolor, - zorder=0, - hatch=hatchstyle, - edgecolor="k", - linewidth=0.0, - label=label_unc, - ) + ax2.fill_between( + edges, + np.append((nom + std) / den, ((nom + std) / den)[-1]), + np.append((nom - std) / den, ((nom - std) / den)[-1]), + step="post", + facecolor=facecolor_pred, + alpha=facecolor_alpha_pred, + zorder=zorder_pred, + hatch=hatchstyle_pred, + edgecolor="k", + linewidth=0.0, + label=label_unc if not args.upperPanelUncertaintyBand else None, + ) if ( args.showVariations in ["lower", "both"] @@ -1239,9 +1306,9 @@ def make_plot( extra_handles=extra_handles_upper, extra_labels=extra_labels_upper, custom_handlers=( - ["bandfilled"] + ["bandfilled", "lineband"] if any(alpha > 0 for alpha in args.fillVariationsAlphas) - else [] + else ["lineband"] ), extra_text=text_pieces if not args.noExtraText else None, extra_text_loc=None if args.extraTextLoc is None else args.extraTextLoc[:2], @@ -1339,6 +1406,15 @@ def make_plots( else: hist_stack = [] + if args.dataCovariance: + hist_data_cov = result[f"cov_data_obs"].get() + + # from sklearn.decomposition import FactorAnalysis + # fa = FactorAnalysis(n_components=1) + # fa.fit(hist_data_cov) + # D_elements = fa.noise_variance_ + # D = np.diag(D_elements) + axes = [a for a in hist_inclusive.axes] if args.processGrouping is not None and len(hist_stack): @@ -1675,6 +1751,7 @@ def main(): varLabels=varLabels, varColors=varColors, varMarkers=args.varMarkers, + dataCovariance=args.dataCovariance, ) results = fitresult.get("mappings", fitresult.get("physics_models")) diff --git a/rabbit/workspace.py b/rabbit/workspace.py index fdf56c0..3ca3bf5 100644 --- a/rabbit/workspace.py +++ b/rabbit/workspace.py @@ -1,4 +1,5 @@ import os +from copy import deepcopy import h5py import hist @@ -175,9 +176,16 @@ def add_hist( if not isinstance(axes, (list, tuple, np.ndarray)): axes = [axes] if start is not None or stop is not None: - values = values[start:stop] + if values.ndim == 1: + values = values[start:stop] + elif values.ndim == 2: + values = values[start:stop, start:stop] + if variances is not None: - variances = variances[start:stop] + if values.ndim == 1: + variances = variances[start:stop] + elif values.ndim == 2: + variances = variances[start:stop, start:stop] h = self.hist(name, axes, values, variances, label, flow=flow) self.dump_hist(h, mapping_key, channel) @@ -250,6 +258,35 @@ def add_observed_hists( **opts, ) + if data_cov_inv is not None: + axes_x = deepcopy(axes) + axes_y = deepcopy(axes) + for ax, ay in zip(axes_x, axes_y): + ax.__dict__["name"] = f"{ax.name}_x" + ay.__dict__["name"] = f"{ay.name}_y" + + self.add_hist( + "cov_data_obs", + [*axes_x, *axes_y], + cov_data_obs, + label="covariance of observed number of events in data", + **opts, + ) + if nobs_cov_inv is not None: + axes_x = deepcopy(axes) + axes_y = deepcopy(axes) + for ax, ay in zip(axes_x, axes_y): + ax.__dict__["name"] = f"{ax.name}_x" + ay.__dict__["name"] = f"{ay.name}_y" + + self.add_hist( + "cov_nobs_obs", + [*axes_x, *axes_y], + cov_nobs, + label="covariance of observed number of events in data", + **opts, + ) + start = stop return hists_data_obs, hists_nobs From 6d304c6af77755a9a6571b5cf30caebcfe91d73b Mon Sep 17 00:00:00 2001 From: davidwalter2 Date: Sat, 28 Mar 2026 10:17:43 -0400 Subject: [PATCH 7/9] Add support to plot data covariance --- bin/rabbit_plot_cov.py | 78 ++++++++++++++++++++++++++++++------------ 1 file changed, 56 insertions(+), 22 deletions(-) diff --git a/bin/rabbit_plot_cov.py b/bin/rabbit_plot_cov.py index 8bb7389..3af28d8 100755 --- a/bin/rabbit_plot_cov.py +++ b/bin/rabbit_plot_cov.py @@ -57,6 +57,11 @@ def make_parser(): default=["charge", "passIso", "passMT", "cosThetaStarll", "qGen"], help="List of axes where for each bin a separate plot is created", ) + parser.add_argument( + "--dataCovariance", + action="store_true", + help="Use covariance information to plot the data uncertainty", + ) return parser @@ -65,7 +70,7 @@ def plot_matrix( matrix, args, channel=None, - axes=None, + axes_names=None, cmap="coolwarm", config={}, meta=None, @@ -107,8 +112,8 @@ def plot_matrix( **opts, ) if ticklabels is None: - xlabel = plot_tools.get_axis_label(config, axes, args.xlabel, is_bin=True) - ylabel = plot_tools.get_axis_label(config, axes, args.ylabel, is_bin=True) + xlabel = plot_tools.get_axis_label(config, axes_names, args.xlabel, is_bin=True) + ylabel = plot_tools.get_axis_label(config, axes_names, args.ylabel, is_bin=True) ax.set_xlabel(xlabel) ax.set_ylabel(ylabel) @@ -124,11 +129,14 @@ def plot_matrix( ) to_join = [f"hist_{'corr' if args.correlation else 'cov'}"] - to_join.append("prefit" if args.prefit else "postfit") + if args.dataCovariance: + to_join.append("data") + else: + to_join.append("prefit" if args.prefit else "postfit") if channel is not None: to_join.append(channel) - if axes is not None: - to_join.append("_".join(axes)) + if axes_names is not None: + to_join.append("_".join(axes_names)) if suffix is not None: to_join.append(suffix) to_join = [*to_join, args.postfix] @@ -177,7 +185,7 @@ def main(): if args.params is not None: h_cov = fitresult["cov"].get() - axes = h_cov.axes.name + axes_names = h_cov.axes.name if len(args.params) > 0: h_param = fitresult["parms"].get() @@ -203,17 +211,21 @@ def main(): outdir, h_cov, args, - axes=axes, + axes_names=axes_names, config=config, meta=meta, suffix="params", ticklabels=ticklabels, ) - hist_cov_key = f"hist_{'prefit' if args.prefit else 'postfit'}_inclusive_cov" + if args.dataCovariance: + hist_cov_key = "cov_data_obs" + else: + hist_cov_key = f"hist_{'prefit' if args.prefit else 'postfit'}_inclusive_cov" results = fitresult.get("mappings", fitresult.get("physics_models")) for margs in args.mapping: + if margs == []: instance_keys = results.keys() else: @@ -221,14 +233,12 @@ def main(): instance_keys = [k for k in results.keys() if k.startswith(mapping_key)] if len(instance_keys) == 0: raise ValueError( - f"No mapping found under {mapping_key}, available mappings are {results.keys()}" + f"No mapping found under {mapping_key}; available mappings are {results.keys()}" ) for instance_key in instance_keys: instance = results[instance_key] - h_cov = instance[hist_cov_key].get() - suffix = ( instance_key.replace(" ", "_") .replace(".", "p") @@ -238,20 +248,43 @@ def main(): .replace(")", "") ) - plot_matrix( - outdir, - h_cov, - args, - config=config, - meta=meta, - suffix=suffix, - ) + if hist_cov_key in instance.keys(): + h_cov = instance[hist_cov_key].get() + plot_matrix( + outdir, + h_cov, + args, + config=config, + meta=meta, + suffix=suffix, + ) + else: + h_cov = None start = 0 for channel, info in instance["channels"].items(): - channel_hist = info[f"hist_postfit_inclusive"].get() + channel_hist = info["hist_postfit_inclusive"].get() axes = [a for a in channel_hist.axes] - if len(instance.get("channels", {}).keys()) > 1: + axes_names = channel_hist.axes.name + + if h_cov is None: + if hist_cov_key not in info.keys(): + raise ValueError( + f"No key {hist_cov_key}; available; keys are {info.keys()}" + ) + h_cov = info[hist_cov_key].get() + + plot_matrix( + outdir, + h_cov, + args, + channel=channel, + config=config, + meta=meta, + suffix=suffix, + axes_names=axes_names, + ) + elif len(instance.get("channels", {}).keys()) > 1: # plot covariance matrix in each channel nbins = np.prod(channel_hist.shape) stop = int(start + nbins) @@ -268,6 +301,7 @@ def main(): config=config, meta=meta, suffix=suffix, + axes_names=axes_names, ) else: h_cov_channel = h_cov From 85e66c2359bc1a421963e8ce625018cb618a5785 Mon Sep 17 00:00:00 2001 From: davidwalter2 Date: Sat, 28 Mar 2026 14:17:54 -0400 Subject: [PATCH 8/9] Fix slicing --- rabbit/workspace.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/rabbit/workspace.py b/rabbit/workspace.py index 3ca3bf5..1acbfc3 100644 --- a/rabbit/workspace.py +++ b/rabbit/workspace.py @@ -171,21 +171,22 @@ def add_hist( label=None, channel=None, mapping_key=None, + is_matrix=False, flow=False, ): if not isinstance(axes, (list, tuple, np.ndarray)): axes = [axes] if start is not None or stop is not None: - if values.ndim == 1: - values = values[start:stop] - elif values.ndim == 2: + if is_matrix: values = values[start:stop, start:stop] + else: + values = values[start:stop] if variances is not None: - if values.ndim == 1: - variances = variances[start:stop] - elif values.ndim == 2: + if is_matrix: variances = variances[start:stop, start:stop] + else: + variances = variances[start:stop] h = self.hist(name, axes, values, variances, label, flow=flow) self.dump_hist(h, mapping_key, channel) @@ -270,6 +271,7 @@ def add_observed_hists( [*axes_x, *axes_y], cov_data_obs, label="covariance of observed number of events in data", + is_matrix=True, **opts, ) if nobs_cov_inv is not None: @@ -283,6 +285,7 @@ def add_observed_hists( "cov_nobs_obs", [*axes_x, *axes_y], cov_nobs, + is_matrix=True, label="covariance of observed number of events in data", **opts, ) From 07772a0ae70ed35c066af7f824a95bf91bffa5ad Mon Sep 17 00:00:00 2001 From: davidwalter2 Date: Sat, 28 Mar 2026 14:26:22 -0400 Subject: [PATCH 9/9] Fix plotting --- bin/rabbit_plot_hists.py | 19 +++++++++++-------- 1 file changed, 11 insertions(+), 8 deletions(-) diff --git a/bin/rabbit_plot_hists.py b/bin/rabbit_plot_hists.py index 0fe2999..ce6e38d 100755 --- a/bin/rabbit_plot_hists.py +++ b/bin/rabbit_plot_hists.py @@ -412,8 +412,17 @@ def make_plot( xlabel = plot_tools.get_axis_label(config, axes_names, args.xlabel) + # plot prediction first, the data on top + zorder_pred = 0 + zorder_data = 1 + + hatchstyle_pred = None + facecolor_pred = "silver" + facecolor_alpha_pred = 0.5 + pred_label = "Prefit model" if args.unfoldedXsec else args.predName + # for uncertaity bands - edges = h_data.axes[0].edges + edges = h_inclusive.axes[0].edges # need to divide by bin width binwidth = edges[1:] - edges[:-1] if binwnorm else 1.0 @@ -630,7 +639,6 @@ def make_plot( ) if data: - zorder_data = 1 if dataCovariance: nom_data = h_data.values() / binwidth std_data = np.sqrt(h_data.variances()) / binwidth @@ -706,13 +714,8 @@ def make_plot( zorder=zorder_data, flow="none", ) - if (args.unfoldedXsec or len(h_stack) == 0) and not args.noPrefit: - zorder_pred = 0 - hatchstyle_pred = None - facecolor_pred = "silver" - facecolor_alpha_pred = 0.5 - pred_label = "Prefit model" if args.unfoldedXsec else args.predName + if (args.unfoldedXsec or len(h_stack) == 0) and not args.noPrefit: hep.histplot( h_inclusive, yerr=False,