From 9cf622779a41fe849d6917b80ea95b1488ad1955 Mon Sep 17 00:00:00 2001 From: davidwalter2 Date: Sat, 11 Apr 2026 09:31:52 -0400 Subject: [PATCH 01/10] Rename 'poi model' to 'param model' and add POI/nuisance configurability Rename the poi_models module to param_models throughout (directory, file, classes, attributes, CLI args). Each ParamModel now declares npoi (true POIs), nnui (model nuisances), and nparams = npoi + nnui, enabling future models (e.g. ABCD) where parameters are nuisances rather than POIs. All existing models set nnui=0 so behaviour is unchanged. Co-Authored-By: Claude Sonnet 4.6 --- .github/workflows/main.yml | 6 +- CLAUDE.md | 21 ++- bin/rabbit_fit.py | 22 +-- bin/rabbit_limit.py | 22 +-- rabbit/fitter.py | 86 +++++---- rabbit/impacts/global_impacts.py | 26 +-- rabbit/impacts/nonprofiled_impacts.py | 10 +- rabbit/impacts/traditional_impacts.py | 32 ++-- .../{poi_models => param_models}/__init__.py | 0 .../{poi_models => param_models}/helpers.py | 8 +- .../param_model.py} | 168 ++++++++++-------- rabbit/parsing.py | 8 +- rabbit/workspace.py | 2 +- 13 files changed, 235 insertions(+), 176 deletions(-) rename rabbit/{poi_models => param_models}/__init__.py (100%) rename rabbit/{poi_models => param_models}/helpers.py (65%) rename rabbit/{poi_models/poi_model.py => param_models/param_model.py} (55%) diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index e03717bc..63e6416f 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -277,7 +277,7 @@ jobs: - name: linearized tensor: test_tensor_symmetric.hdf5 fit_args: >- - --binByBinStatType normal-additive --allowNegativePOI + --binByBinStatType normal-additive --allowNegativeParam - name: bb_full_normal_additive tensor: test_tensor.hdf5 fit_args: >- @@ -400,12 +400,12 @@ jobs: - name: bsm limits run: >- rabbit_limit.py $RABBIT_OUTDIR/test_tensor_bsm.hdf5 -o $RABBIT_OUTDIR/ --postfix bsm - -t 0 --unblind --expectSignal bsm1 0 --expectSignal bsm2 1 --allowNegativePOI --asymptoticLimits bsm1 --levels 0.05 + -t 0 --unblind --expectSignal bsm1 0 --expectSignal bsm2 1 --allowNegativeParam --asymptoticLimits bsm1 --levels 0.05 - name: bsm limits on channel run: >- rabbit_limit.py $RABBIT_OUTDIR/test_tensor_bsm.hdf5 -o $RABBIT_OUTDIR/ --postfix bsm - -t 0 --unblind --poiModel Mixture bsm1 sig --expectSignal bsm1_sig_mixing 0 --asymptoticLimits bsm1_sig_mixing --levels 0.05 --allowNegativePOI + -t 0 --unblind --paramModel Mixture bsm1 sig --expectSignal bsm1_sig_mixing 0 --asymptoticLimits bsm1_sig_mixing --levels 0.05 --allowNegativeParam plotting: runs-on: [self-hosted, linux, x64] diff --git a/CLAUDE.md b/CLAUDE.md index 63e2b857..aa5c7694 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -65,26 +65,31 @@ The workflow has three stages: **write input tensor → run fit → post-process Supports dense and sparse tensor representations, symmetric/asymmetric systematics, and log-normal or normal systematic types. ### 2. Fit: `rabbit/fitter.py` -`Fitter` takes a `FitInputData` object (loaded from HDF5 by `rabbit/inputdata.py`), a `POIModel`, and options. It builds a differentiable negative log-likelihood using TensorFlow and minimizes it via SciPy. Results are written through a `Workspace`. +`Fitter` takes a `FitInputData` object (loaded from HDF5 by `rabbit/inputdata.py`), a `ParamModel`, and options. It builds a differentiable negative log-likelihood using TensorFlow and minimizes it via SciPy. Results are written through a `Workspace`. ### 3. Output: `rabbit/workspace.py` `Workspace` collects post-fit histograms, covariance matrices, impacts, and likelihood scans into an HDF5 output file using `hist.Hist` objects (via the `wums` library). -### POI Models: `rabbit/poi_models/poi_model.py` -Base class `POIModel` defines `compute(poi)` which returns a `[1, nproc]` tensor scaling signal processes. Built-in models: `Mu` (default), `Ones`, `Mixture`. Custom models are loaded by providing a dotted Python path (e.g. `--poiModel mymod.MyModel`); the module must be on `$PYTHONPATH` with an `__init__.py`. +### Param Models: `rabbit/param_models/param_model.py` +Base class `ParamModel` defines `compute(param)` which returns a `[1, nproc]` tensor scaling process yields. Each model declares: +- `nparams`: total parameters (POIs + model nuisances) +- `npoi`: true parameters of interest (first `npoi` entries; reported as POIs in outputs) +- `nnui`: model nuisance parameters (`nparams - npoi`; reported as nuisances in outputs) + +Built-in models: `Mu` (default, one POI per signal process), `Ones` (no parameters), `Mixture` (mixing POIs). Custom models are loaded by providing a dotted Python path (e.g. `--paramModel mymod.MyModel`); the module must be on `$PYTHONPATH` with an `__init__.py`. The deprecated `--poiModel` alias is still accepted. ### Mappings: `rabbit/mappings/` -Base class `Mapping` in `mapping.py` defines `compute_flat(params, observables)`, which is a differentiable transformation of the flat bin vector. The framework propagates uncertainties through it via automatic differentiation (`tf.GradientTape`). Built-in mappings (`Select`, `Project`, `Normalize`, `Ratio`, `Normratio`) live in `project.py` and `ratio.py`. Custom mappings follow the same pattern as POI models. +Base class `Mapping` in `mapping.py` defines `compute_flat(params, observables)`, which is a differentiable transformation of the flat bin vector. The framework propagates uncertainties through it via automatic differentiation (`tf.GradientTape`). Built-in mappings (`Select`, `Project`, `Normalize`, `Ratio`, `Normratio`) live in `project.py` and `ratio.py`. Custom mappings follow the same pattern as param models. ### Bin scripts: `bin/` Entry points registered in `pyproject.toml`. The main one is `rabbit_fit.py`; others are diagnostic/plotting scripts. All use `rabbit/parsing.py` for shared CLI arguments. ## Custom Extensions -Custom mappings and POI models must: -1. Subclass `Mapping` or `POIModel` respectively -2. Implement `compute_flat` (mapping) or `compute` (POI model) as TF-differentiable functions +Custom mappings and param models must: +1. Subclass `Mapping` or `ParamModel` respectively +2. Implement `compute_flat` (mapping) or `compute` (param model) as TF-differentiable functions 3. Be importable from `$PYTHONPATH` (directory needs `__init__.py`) -4. Be referenced with dotted module path: `-m my_package.MyMapping` or `--poiModel my_package.MyModel` +4. Be referenced with dotted module path: `-m my_package.MyMapping` or `--paramModel my_package.MyModel` See `tests/my_custom_mapping.py` and `tests/my_custom_model.py` for examples. diff --git a/bin/rabbit_fit.py b/bin/rabbit_fit.py index dcf8b1a1..58d953b3 100755 --- a/bin/rabbit_fit.py +++ b/bin/rabbit_fit.py @@ -15,8 +15,8 @@ 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.param_models import helpers as ph +from rabbit.param_models import param_model from rabbit.regularization import helpers as rh from rabbit.regularization.lcurve import l_curve_optimize_tau, l_curve_scan_tau from rabbit.tfhelpers import edmval_cov @@ -302,11 +302,11 @@ def save_hists(args, mappings, fitter, ws, prefit=True, profile=False): ): # saturated likelihood test - saturated_model = poi_model.SaturatedProjectModel( + saturated_model = param_model.SaturatedProjectModel( fitter.indata, mapping.channel_info ) - composite_model = poi_model.CompositePOIModel( - [fitter.poi_model, saturated_model] + composite_model = param_model.CompositeParamModel( + [fitter.param_model, saturated_model] ) fitter_saturated = copy.deepcopy(fitter) @@ -449,7 +449,7 @@ def fit(args, fitter, ws, dofit=True): nllvalreduced = fitter.reduced_nll().numpy() ndfsat = ( - tf.size(fitter.nobs) - fitter.poi_model.npoi - fitter.indata.nsystnoconstraint + tf.size(fitter.nobs) - fitter.param_model.npoi - fitter.indata.nsystnoconstraint ).numpy() chi2_val = 2.0 * nllvalreduced @@ -574,12 +574,12 @@ def main(): indata = inputdata.FitInputData(args.filename, args.pseudoData) - margs = args.poiModel - poi_model = ph.load_model(margs[0], indata, *margs[1:], **vars(args)) + margs = args.paramModel + param_model = ph.load_model(margs[0], indata, *margs[1:], **vars(args)) ifitter = fitter.Fitter( indata, - poi_model, + param_model, args, do_blinding=any(blinded_fits), globalImpactsFromJVP=not args.globalImpactsDisableJVP, @@ -615,8 +615,8 @@ def main(): "meta_info": output_tools.make_meta_info_dict(args=args), "meta_info_input": ifitter.indata.metadata, "procs": ifitter.indata.procs, - "pois": ifitter.poi_model.pois, - "nois": ifitter.parms[ifitter.poi_model.npoi :][indata.noiidxs], + "pois": ifitter.param_model.params[: ifitter.param_model.npoi], + "nois": ifitter.parms[ifitter.param_model.nparams :][indata.noiidxs], } with workspace.Workspace( diff --git a/bin/rabbit_limit.py b/bin/rabbit_limit.py index e1f5ad24..d76cc0c4 100755 --- a/bin/rabbit_limit.py +++ b/bin/rabbit_limit.py @@ -13,7 +13,7 @@ from rabbit import asymptotic_limits, fitter, inputdata, parsing, workspace from rabbit.mappings import helpers as mh -from rabbit.poi_models import helpers as ph +from rabbit.param_models import helpers as ph from rabbit.tfhelpers import edmval_cov from wums import output_tools, logging # isort: skip @@ -136,14 +136,16 @@ def do_asymptotic_limits(args, fitter, ws, data_values, mapping=None, fit_data=F fun = None idx = np.where(fitter.parms.astype(str) == key)[0][0] - if key in fitter.poi_model.pois.astype(str): + if key in fitter.param_model.params[: fitter.param_model.npoi].astype(str): is_poi = True xbest = fitter_asimov.get_poi()[idx] xobs = fitter.get_poi()[idx] elif key in fitter.parms.astype(str): is_poi = False - xbest = fitter_asimov.get_theta()[idx - fitter_asimov.poi_model.npoi] - xobs = fitter.get_theta()[idx - fitter.poi_model.npoi] + xbest = fitter_asimov.get_theta()[ + idx - fitter_asimov.param_model.nparams + ] + xobs = fitter.get_theta()[idx - fitter.param_model.nparams] xerr = fitter_asimov.cov[idx, idx] ** 0.5 @@ -156,7 +158,7 @@ def do_asymptotic_limits(args, fitter, ws, data_values, mapping=None, fit_data=F fitter.cov.assign(cov) xobs_err = fitter.cov[idx, idx] ** 0.5 - if is_poi and not args.allowNegativePOI: + if is_poi and not args.allowNegativeParam: xerr = 2 * xerr * xbest**0.5 xobs_err = 2 * xobs_err * xobs**0.5 @@ -262,10 +264,10 @@ def main(): indata = inputdata.FitInputData(args.filename, args.pseudoData) - margs = args.poiModel - poi_model = ph.load_model(margs[0], indata, *margs[1:], **vars(args)) + margs = args.paramModel + param_model = ph.load_model(margs[0], indata, *margs[1:], **vars(args)) - ifitter = fitter.Fitter(indata, poi_model, args, do_blinding=any(blinded_fits)) + ifitter = fitter.Fitter(indata, param_model, args, do_blinding=any(blinded_fits)) # mapping for observables and parameters if len(args.mapping) > 0: @@ -282,8 +284,8 @@ def main(): "meta_info": output_tools.make_meta_info_dict(args=args), "meta_info_input": ifitter.indata.metadata, "procs": ifitter.indata.procs, - "pois": ifitter.poi_model.pois, - "nois": ifitter.parms[ifitter.poi_model.npoi :][indata.noiidxs], + "pois": ifitter.param_model.params[: ifitter.param_model.npoi], + "nois": ifitter.parms[ifitter.param_model.nparams :][indata.noiidxs], } with workspace.Workspace( diff --git a/rabbit/fitter.py b/rabbit/fitter.py index 5e3c29f5..4930c499 100644 --- a/rabbit/fitter.py +++ b/rabbit/fitter.py @@ -81,7 +81,7 @@ class Fitter: valid_systematic_types = ["log_normal", "normal"] def __init__( - self, indata, poi_model, options, globalImpactsFromJVP=True, do_blinding=False + self, indata, param_model, options, globalImpactsFromJVP=True, do_blinding=False ): self.indata = indata @@ -137,7 +137,7 @@ def __init__( # --- fit params self.init_fit_parms( - poi_model, + param_model, options.setConstraintMinimum, unblind=options.unblind, freeze_parameters=options.freezeParameters, @@ -258,16 +258,16 @@ def __init__( def init_fit_parms( self, - poi_model, + param_model, set_constraint_minimum=[], unblind=False, freeze_parameters=[], ): - self.poi_model = poi_model + self.param_model = param_model if self.do_blinding: self._blinding_offsets_poi = tf.Variable( - tf.ones([self.poi_model.npoi], dtype=self.indata.dtype), + tf.ones([self.param_model.npoi], dtype=self.indata.dtype), trainable=False, name="offset_poi", ) @@ -278,7 +278,7 @@ def init_fit_parms( ) self.init_blinding_values(unblind) - self.parms = np.concatenate([self.poi_model.pois, self.indata.systs]) + self.parms = np.concatenate([self.param_model.params, self.indata.systs]) # tf tensor containing default constraint minima theta0default = np.zeros(self.indata.nsyst) @@ -295,9 +295,9 @@ def init_fit_parms( ) # tf variable containing all fit parameters - if self.poi_model.npoi > 0: + if self.param_model.nparams > 0: xdefault = tf.concat( - [self.poi_model.xpoidefault, self.theta0default], axis=0 + [self.param_model.xparamdefault, self.theta0default], axis=0 ) else: xdefault = self.theta0default @@ -342,7 +342,7 @@ def init_fit_parms( # determine if problem is linear (ie likelihood is purely quadratic) self.is_linear = ( (self.chisqFit or self.covarianceFit) - and self.poi_model.is_linear + and self.param_model.is_linear and self.indata.symmetric_tensor and self.indata.systematic_type == "normal" and ((not self.binByBinStat) or self.binByBinStatType == "normal-additive") @@ -446,7 +446,7 @@ def init_blinding_values(self, unblind_parameter_expressions=[]): unblind_parameters = match_regexp_params( unblind_parameter_expressions, [ - *self.poi_model.pois, + *self.param_model.params[: self.param_model.npoi], *[self.indata.systs[i] for i in self.indata.noiidxs], ], ) @@ -484,9 +484,9 @@ def deterministic_random_from_string(s, mean=0.0, std=5.0): self._blinding_values_theta[i] = value # add offset to pois - self._blinding_values_poi = np.ones(self.poi_model.npoi, dtype=np.float64) - for i in range(self.poi_model.npoi): - param = self.poi_model.pois[i] + self._blinding_values_poi = np.ones(self.param_model.npoi, dtype=np.float64) + for i in range(self.param_model.npoi): + param = self.param_model.params[i] if param in unblind_parameters: continue logger.debug(f"Blind signal strength modifier for {param}") @@ -501,18 +501,17 @@ def set_blinding_offsets(self, blind=True): self._blinding_offsets_theta.assign(self._blinding_values_theta) else: self._blinding_offsets_poi.assign( - np.ones(self.poi_model.npoi, dtype=np.float64) + np.ones(self.param_model.npoi, dtype=np.float64) ) self._blinding_offsets_theta.assign( np.zeros(self.indata.nsyst, dtype=np.float64) ) def get_theta(self): - theta = self.x[self.poi_model.npoi : self.poi_model.npoi + self.indata.nsyst] + start = self.param_model.nparams + theta = self.x[start : start + self.indata.nsyst] theta = tf.where( - self.frozen_params_mask[ - self.poi_model.npoi : self.poi_model.npoi + self.indata.nsyst - ], + self.frozen_params_mask[start : start + self.indata.nsyst], tf.stop_gradient(theta), theta, ) @@ -521,14 +520,19 @@ def get_theta(self): else: return theta + def get_model_nui(self): + npoi = self.param_model.npoi + nnui = self.param_model.nnui + return self.x[npoi : npoi + nnui] + def get_poi(self): - xpoi = self.x[: self.poi_model.npoi] - if self.poi_model.allowNegativePOI: + xpoi = self.x[: self.param_model.npoi] + if self.param_model.allowNegativeParam: poi = xpoi else: poi = tf.square(xpoi) poi = tf.where( - self.frozen_params_mask[: self.poi_model.npoi], tf.stop_gradient(poi), poi + self.frozen_params_mask[: self.param_model.npoi], tf.stop_gradient(poi), poi ) if self.do_blinding: return poi * self._blinding_offsets_poi @@ -536,7 +540,9 @@ def get_poi(self): return poi def get_x(self): - return tf.concat([self.get_poi(), self.get_theta()], axis=0) + return tf.concat( + [self.get_poi(), self.get_model_nui(), self.get_theta()], axis=0 + ) def _default_beta0(self): if self.binByBinStatType in ["gamma", "normal-multiplicative"]: @@ -546,8 +552,8 @@ 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.ones([self.poi_model.npoi], dtype=self.indata.dtype) + var_params = ( + tf.ones([self.param_model.nparams], dtype=self.indata.dtype) * unconstrained_err**2 ) @@ -559,7 +565,7 @@ def prefit_covariance(self, unconstrained_err=0.0): tf.math.reciprocal(self.indata.constraintweights), ) - invhessianprefit = tf.linalg.diag(tf.concat([var_poi, var_theta], axis=0)) + invhessianprefit = tf.linalg.diag(tf.concat([var_params, var_theta], axis=0)) return invhessianprefit @tf.function @@ -598,10 +604,12 @@ def theta0defaultassign(self): self.theta0.assign(self.theta0default) def xdefaultassign(self): - if self.poi_model.npoi == 0: + if self.param_model.nparams == 0: self.x.assign(self.theta0) else: - self.x.assign(tf.concat([self.poi_model.xpoidefault, self.theta0], axis=0)) + self.x.assign( + tf.concat([self.param_model.xparamdefault, self.theta0], axis=0) + ) def beta0defaultassign(self): self.set_beta0(self._default_beta0()) @@ -630,7 +638,7 @@ def defaultassign(self): def bayesassign(self): # FIXME use theta0 as the mean and constraintweight to scale the width - if self.poi_model.npoi == 0: + if self.param_model.nparams == 0: self.x.assign( self.theta0 + tf.random.normal(shape=self.theta0.shape, dtype=self.theta0.dtype) @@ -639,7 +647,7 @@ def bayesassign(self): self.x.assign( tf.concat( [ - self.poi_model.xpoidefault, + self.param_model.xparamdefault, self.theta0 + tf.random.normal( shape=self.theta0.shape, dtype=self.theta0.dtype @@ -835,7 +843,11 @@ def edmval_cov(self, grad, hess): @tf.function def impacts_parms(self, hess): - nstat = self.poi_model.npoi + self.indata.nsystnoconstraint + nstat = ( + self.param_model.npoi + + self.param_model.nnui + + self.indata.nsystnoconstraint + ) hess_stat = hess[:nstat, :nstat] cov_stat = tf.linalg.inv(hess_stat) @@ -852,7 +864,7 @@ def impacts_parms(self, hess): self.cov, cov_stat, cov_stat_no_bbb, - self.poi_model.npoi, + self.param_model.npoi, self.indata.noiidxs, self.indata.systgroupidxs, ) @@ -868,7 +880,7 @@ def global_impacts_parms(self): self._compute_yields_with_beta, self._compute_lbeta, self._compute_lc, - self.poi_model.npoi, + self.param_model.npoi, self.indata.noiidxs, self.indata.systgroupidxs, self.binByBinStat, @@ -892,7 +904,7 @@ def gaussian_global_impacts_parms(self): if self.binByBinStatType in ["normal-additive"] or not self.binByBinStat else 1.0 / self.kstat ), - self.poi_model.npoi, + self.param_model.npoi, self.indata.noiidxs, self.binByBinStat, self.binByBinStatMode, @@ -912,7 +924,7 @@ def nonprofiled_impacts_parms(self, unconstrained_err=1.0): self.indata.constraintweights, self.indata.systgroups, self.indata.systgroupidxs, - self.poi_model.npoi, + self.param_model.nparams, self.minimize, self.diagnostics, self.loss_val_grad_hess, @@ -1089,7 +1101,7 @@ def compute_derivatives(dvars): self._compute_yields_with_beta, self._compute_lbeta, self._compute_lc, - self.poi_model.npoi, + self.param_model.npoi, self.indata.systgroupidxs, self.binByBinStat, self.binByBinStatMode, @@ -1195,9 +1207,11 @@ def _expected_variations( def _compute_yields_noBBB(self, full=True): # full: compute yields inclduing masked channels poi = self.get_poi() + model_nui = self.get_model_nui() theta = self.get_theta() - rnorm = self.poi_model.compute(poi, full) + all_params = tf.concat([poi, model_nui], axis=0) + rnorm = self.param_model.compute(all_params, full) normcentral = None if self.indata.symmetric_tensor: diff --git a/rabbit/impacts/global_impacts.py b/rabbit/impacts/global_impacts.py index cc134810..98e638a7 100644 --- a/rabbit/impacts/global_impacts.py +++ b/rabbit/impacts/global_impacts.py @@ -16,9 +16,9 @@ import tensorflow as tf -def _gather_poi_noi_vector(v, noiidxs, npoi=0): - v_poi = v[:npoi] - v_noi = tf.gather(v[npoi:], noiidxs) +def _gather_poi_noi_vector(v, noiidxs, nsignal_params=0): + v_poi = v[:nsignal_params] + v_noi = tf.gather(v[nsignal_params:], noiidxs) return tf.concat([v_poi, v_noi], axis=0) @@ -231,7 +231,7 @@ def global_impacts_parms( compute_yields_with_beta_fn, compute_lbeta_fn, compute_lc_fn, - npoi, + nsignal_params, noiidxs, systgroupidxs, bin_by_bin_stat, @@ -239,8 +239,8 @@ def global_impacts_parms( global_impacts_from_jvp, cov, ): - idxs_poi = tf.range(npoi, dtype=tf.int64) - idxs_noi = tf.constant(npoi + noiidxs, dtype=tf.int64) + idxs_poi = tf.range(nsignal_params, dtype=tf.int64) + idxs_noi = tf.constant(nsignal_params + noiidxs, dtype=tf.int64) idxsout = tf.concat([idxs_poi, idxs_noi], axis=0) dexpdx = tf.one_hot(idxsout, depth=cov.shape[0], dtype=cov.dtype) @@ -265,7 +265,7 @@ def global_impacts_parms( ) impacts_x0 = _compute_global_impacts_x0(x, compute_lc_fn, cov_dexpdx) - impacts_theta0 = tf.transpose(impacts_x0[npoi:]) + impacts_theta0 = tf.transpose(impacts_x0[nsignal_params:]) impacts_theta0_sq = tf.square(impacts_theta0) var_theta0 = tf.reduce_sum(impacts_theta0_sq, axis=-1) @@ -293,7 +293,7 @@ def global_impacts_obs( compute_yields_with_beta_fn, compute_lbeta_fn, compute_lc_fn, - npoi, + nsignal_params, systgroupidxs, bin_by_bin_stat, bin_by_bin_stat_mode, @@ -353,7 +353,7 @@ def global_impacts_obs( ) impacts_x0 = _compute_global_impacts_x0(x, compute_lc_fn, cov_dexpdx) - impacts_theta0 = tf.transpose(impacts_x0[npoi:]) + impacts_theta0 = tf.transpose(impacts_x0[nsignal_params:]) impacts_theta0_sq = tf.square(impacts_theta0) var_theta0 = tf.reduce_sum(impacts_theta0_sq, axis=-1) @@ -444,7 +444,7 @@ def gaussian_global_impacts_parms( vartheta0, varnobs, varbeta0, - npoi, + nsignal_params, noiidxs, bin_by_bin_stat, bin_by_bin_stat_mode, @@ -453,9 +453,9 @@ def gaussian_global_impacts_parms( data_cov_inv=None, ): # compute impacts for pois and nois - dxdtheta0 = _gather_poi_noi_vector(dxdtheta0, noiidxs, npoi) - dxdnobs = _gather_poi_noi_vector(dxdnobs, noiidxs, npoi) - dxdbeta0 = _gather_poi_noi_vector(dxdbeta0, noiidxs, npoi) + dxdtheta0 = _gather_poi_noi_vector(dxdtheta0, noiidxs, nsignal_params) + dxdnobs = _gather_poi_noi_vector(dxdnobs, noiidxs, nsignal_params) + dxdbeta0 = _gather_poi_noi_vector(dxdbeta0, noiidxs, nsignal_params) return _gaussian_global_impacts( dxdtheta0, diff --git a/rabbit/impacts/nonprofiled_impacts.py b/rabbit/impacts/nonprofiled_impacts.py index 04355d47..89fc379e 100644 --- a/rabbit/impacts/nonprofiled_impacts.py +++ b/rabbit/impacts/nonprofiled_impacts.py @@ -28,7 +28,7 @@ def nonprofiled_impacts_parms( constraintweights, systgroups, systgroupidxs, - npoi, + nparams, minimize_fn, diagnostics=False, loss_val_grad_hess_fn=None, @@ -43,7 +43,7 @@ def nonprofiled_impacts_parms( constraintweights: constraint weights for each nuisance parameter. systgroups: systematic group names. systgroupidxs: per-group lists of nuisance parameter indices. - npoi: number of POIs. + nparams: total number of model parameters (nparams + nnui); offset from x index to theta0 index. minimize_fn: callable that runs the fit (no arguments). diagnostics: if True, log EDM after each minimization (requires loss_val_grad_hess_fn). loss_val_grad_hess_fn: callable returning (val, grad, hess); used only when diagnostics=True. @@ -65,9 +65,9 @@ def nonprofiled_impacts_parms( logger.info(f"Now at parameter {frozen_params[i]}") for j, sign in enumerate((1, -1)): - variation = sign * err_theta[idx - npoi] + theta0_tmp[idx - npoi] + variation = sign * err_theta[idx - nparams] + theta0_tmp[idx - nparams] # vary the non-profiled parameter - theta0[idx - npoi].assign(variation) + theta0[idx - nparams].assign(variation) x[idx].assign( variation ) # this should not be needed but should accelerate the minimization @@ -83,7 +83,7 @@ def nonprofiled_impacts_parms( x.assign(x_tmp) # back to original value - theta0[idx - npoi].assign(theta0_tmp[idx - npoi]) + theta0[idx - nparams].assign(theta0_tmp[idx - nparams]) impact_group_names = [] impact_groups = [] diff --git a/rabbit/impacts/traditional_impacts.py b/rabbit/impacts/traditional_impacts.py index 5ab46436..f02cdd18 100644 --- a/rabbit/impacts/traditional_impacts.py +++ b/rabbit/impacts/traditional_impacts.py @@ -7,8 +7,8 @@ import tensorflow as tf -def _compute_impact_group(cov, v, idxs, npoi=0): - cov_reduced = tf.gather(cov[npoi:, npoi:], idxs, axis=0) +def _compute_impact_group(cov, v, idxs, nsignal_params=0): + cov_reduced = tf.gather(cov[nsignal_params:, nsignal_params:], idxs, axis=0) cov_reduced = tf.gather(cov_reduced, idxs, axis=1) v_reduced = tf.gather(v, idxs, axis=1) invC_v = tf.linalg.solve(cov_reduced, tf.transpose(v_reduced)) @@ -16,50 +16,58 @@ def _compute_impact_group(cov, v, idxs, npoi=0): return tf.sqrt(v_invC_v) -def _gather_poi_noi_vector(v, noiidxs, npoi=0): - v_poi = v[:npoi] +def _gather_poi_noi_vector(v, noiidxs, nsignal_params=0): + v_poi = v[:nsignal_params] # protection for constained NOIs, set them to 0 - mask = (noiidxs >= 0) & (noiidxs < tf.shape(v[npoi:])[0]) + mask = (noiidxs >= 0) & (noiidxs < tf.shape(v[nsignal_params:])[0]) safe_idxs = tf.where(mask, noiidxs, 0) mask = tf.cast(mask, v.dtype) mask = tf.reshape( mask, tf.concat([tf.shape(mask), tf.ones(tf.rank(v) - 1, dtype=tf.int32)], axis=0), ) - v_noi = tf.gather(v[npoi:], safe_idxs) * mask + v_noi = tf.gather(v[nsignal_params:], safe_idxs) * mask v_gathered = tf.concat([v_poi, v_noi], axis=0) return v_gathered -def impacts_parms(cov, cov_stat, cov_stat_no_bbb, npoi=0, noiidxs=[], systgroupidxs=[]): +def impacts_parms( + cov, cov_stat, cov_stat_no_bbb, nsignal_params=0, noiidxs=[], systgroupidxs=[] +): """ Gaussian approximation """ # impact for poi at index i in covariance matrix from nuisance with index j is C_ij/sqrt(C_jj) = /sqrt() - v = _gather_poi_noi_vector(cov, noiidxs, npoi) + v = _gather_poi_noi_vector(cov, noiidxs, nsignal_params) impacts = v / tf.reshape(tf.sqrt(tf.linalg.diag_part(cov)), [1, -1]) if cov_stat_no_bbb is not None: # impact bin-by-bin stat impacts_data_stat = tf.sqrt(tf.linalg.diag_part(cov_stat_no_bbb)) - impacts_data_stat = _gather_poi_noi_vector(impacts_data_stat, noiidxs, npoi) + impacts_data_stat = _gather_poi_noi_vector( + impacts_data_stat, noiidxs, nsignal_params + ) impacts_data_stat = tf.reshape(impacts_data_stat, (-1, 1)) impacts_bbb_sq = tf.linalg.diag_part(cov_stat - cov_stat_no_bbb) - impacts_bbb_sq = _gather_poi_noi_vector(impacts_bbb_sq, noiidxs, npoi) + impacts_bbb_sq = _gather_poi_noi_vector(impacts_bbb_sq, noiidxs, nsignal_params) impacts_bbb = tf.sqrt(tf.nn.relu(impacts_bbb_sq)) # max(0,x) impacts_bbb = tf.reshape(impacts_bbb, (-1, 1)) impacts_grouped = tf.concat([impacts_data_stat, impacts_bbb], axis=1) else: impacts_data_stat = tf.sqrt(tf.linalg.diag_part(cov_stat)) - impacts_data_stat = _gather_poi_noi_vector(impacts_data_stat, noiidxs, npoi) + impacts_data_stat = _gather_poi_noi_vector( + impacts_data_stat, noiidxs, nsignal_params + ) impacts_data_stat = tf.reshape(impacts_data_stat, (-1, 1)) impacts_grouped = impacts_data_stat if len(systgroupidxs): impacts_grouped_syst = tf.map_fn( - lambda idxs: _compute_impact_group(cov, v[:, npoi:], idxs, npoi), + lambda idxs: _compute_impact_group( + cov, v[:, nsignal_params:], idxs, nsignal_params + ), tf.ragged.constant(systgroupidxs, dtype=tf.int32), fn_output_signature=tf.TensorSpec( shape=(impacts.shape[0],), dtype=tf.float64 diff --git a/rabbit/poi_models/__init__.py b/rabbit/param_models/__init__.py similarity index 100% rename from rabbit/poi_models/__init__.py rename to rabbit/param_models/__init__.py diff --git a/rabbit/poi_models/helpers.py b/rabbit/param_models/helpers.py similarity index 65% rename from rabbit/poi_models/helpers.py rename to rabbit/param_models/helpers.py index 83836f3a..3eddb8d1 100644 --- a/rabbit/poi_models/helpers.py +++ b/rabbit/param_models/helpers.py @@ -2,14 +2,14 @@ # dictionary with class name and the corresponding filename where it is defined baseline_models = { - "Ones": "poi_model", - "Mu": "poi_model", - "Mixture": "poi_model", + "Ones": "param_model", + "Mu": "param_model", + "Mixture": "param_model", } def load_model(class_name, indata, *args, **kwargs): model = common.load_class_from_module( - class_name, baseline_models, base_dir="rabbit.poi_models" + class_name, baseline_models, base_dir="rabbit.param_models" ) return model.parse_args(indata, *args, **kwargs) diff --git a/rabbit/poi_models/poi_model.py b/rabbit/param_models/param_model.py similarity index 55% rename from rabbit/poi_models/poi_model.py rename to rabbit/param_models/param_model.py index 2c1cd4ca..d6dd59d2 100644 --- a/rabbit/poi_models/poi_model.py +++ b/rabbit/param_models/param_model.py @@ -5,131 +5,150 @@ import tensorflow as tf -class POIModel: +class ParamModel: def __init__(self, indata, *args, **kwargs): self.indata = indata - # # a POI model must set these attribues - # self.npoi = # number of parameters of interest (POIs) - # self.pois = # list of names for the POIs - # self.xpoidefault = # default values for the POIs - # self.is_linear = # define if the model is linear in the POIs - # self.allowNegativePOI = # define if the POI can be negative or not + # # a param model must set these attributes + # self.nparams = # total number of parameters (npoi + nnui) + # self.npoi = # number of true parameters of interest (POIs), reported as POIs in outputs + # self.nnui = # number of model nuisance parameters (= nparams - npoi) + # self.params = # list of names for all parameters (POIs first, then model nuisances) + # self.xparamdefault = # default values for all parameters (length nparams) + # self.is_linear = # define if the model is linear in the parameters + # self.allowNegativeParam = # define if the POI parameters can be negative or not - # class function to parse strings as given by the argparse input e.g. --poiModel ... + # class function to parse strings as given by the argparse input e.g. --paramModel ... @classmethod def parse_args(cls, indata, *args, **kwargs): return cls(indata, *args, **kwargs) - def compute(self, poi, full=False): + def compute(self, param, full=False): """ Compute an array for the rate per process - :param params: 1D tensor of explicit parameters in the fit + :param param: 1D tensor of explicit parameters in the fit (length nparams) :return 2D tensor to be multiplied with [proc,bin] tensor """ - def set_poi_default(self, expectSignal, allowNegativePOI=False): + def set_param_default(self, expectSignal, allowNegativeParam=False): """ - Set default POI values, used by different POI models + Set default parameter values, used by different param models. + Only the first npoi entries (true POIs) support the squaring transform; + model nuisance parameters (nnui entries) are always stored directly. """ - poidefault = tf.ones([self.npoi], dtype=self.indata.dtype) + paramdefault = tf.ones([self.nparams], dtype=self.indata.dtype) if expectSignal is not None: indices = [] updates = [] for signal, value in expectSignal: - if signal.encode() not in self.pois: + if signal.encode() not in self.params: raise ValueError( - f"{signal.encode()} not in list of POIs: {self.pois}" + f"{signal.encode()} not in list of params: {self.params}" ) - idx = np.where(np.isin(self.pois, signal.encode()))[0][0] + idx = np.where(np.isin(self.params, signal.encode()))[0][0] indices.append([idx]) updates.append(float(value)) - poidefault = tf.tensor_scatter_nd_update(poidefault, indices, updates) + paramdefault = tf.tensor_scatter_nd_update(paramdefault, indices, updates) - if allowNegativePOI: - self.xpoidefault = poidefault + # squaring transform applies only to the npoi true POI entries + poi_part = paramdefault[: self.npoi] + nui_part = paramdefault[self.npoi :] + + if allowNegativeParam: + xpoi_part = poi_part else: - self.xpoidefault = tf.sqrt(poidefault) + xpoi_part = tf.sqrt(poi_part) + + self.xparamdefault = tf.concat([xpoi_part, nui_part], axis=0) -class CompositePOIModel(POIModel): +class CompositeParamModel(ParamModel): """ - multiply different POI models together + multiply different param models together """ def __init__( self, - poi_models, - allowNegativePOI=False, + param_models, + allowNegativeParam=False, ): - self.poi_models = poi_models + self.param_models = param_models - self.npoi = sum([m.npoi for m in poi_models]) + self.nparams = sum([m.nparams for m in param_models]) + self.npoi = sum([m.npoi for m in param_models]) + self.nnui = sum([m.nnui for m in param_models]) - self.pois = np.concatenate([m.pois for m in poi_models]) + self.params = np.concatenate([m.params for m in param_models]) - self.allowNegativePOI = allowNegativePOI + self.allowNegativeParam = allowNegativeParam - self.is_linear = self.npoi == 0 or self.allowNegativePOI + self.is_linear = self.nparams == 0 or self.allowNegativeParam - self.xpoidefault = tf.concat([m.xpoidefault for m in poi_models], axis=0) + self.xparamdefault = tf.concat([m.xparamdefault for m in param_models], axis=0) - def compute(self, poi, full=False): + def compute(self, param, full=False): start = 0 results = [] - for m in self.poi_models: - results.append(m.compute(poi[start : start + m.npoi], full)) - start += m.npoi + for m in self.param_models: + results.append(m.compute(param[start : start + m.nparams], full)) + start += m.nparams rnorm = functools.reduce(lambda a, b: a * b, results) return rnorm -class Ones(POIModel): +class Ones(ParamModel): """ multiply all processes with ones """ def __init__(self, indata, **kwargs): self.indata = indata + self.nparams = 0 self.npoi = 0 - self.pois = np.array([]) - self.poidefault = tf.zeros([], dtype=self.indata.dtype) + self.nnui = 0 + self.params = np.array([]) + self.xparamdefault = tf.zeros([0], dtype=self.indata.dtype) - self.allowNegativePOI = False + self.allowNegativeParam = False self.is_linear = True - def compute(self, poi, full=False): + def compute(self, param, full=False): rnorm = tf.ones(self.indata.nproc, dtype=self.indata.dtype) rnorm = tf.reshape(rnorm, [1, -1]) return rnorm -class Mu(POIModel): +class Mu(ParamModel): """ multiply unconstrained parameter to signal processes, and ones otherwise """ - def __init__(self, indata, expectSignal=None, allowNegativePOI=False, **kwargs): + def __init__(self, indata, expectSignal=None, allowNegativeParam=False, **kwargs): self.indata = indata - self.npoi = self.indata.nsignals + self.nparams = self.indata.nsignals + self.npoi = self.nparams + self.nnui = 0 - self.pois = np.array([s for s in self.indata.signals]) + self.params = np.array([s for s in self.indata.signals]) - self.allowNegativePOI = allowNegativePOI + self.allowNegativeParam = allowNegativeParam - self.is_linear = self.npoi == 0 or self.allowNegativePOI + self.is_linear = self.nparams == 0 or self.allowNegativeParam - self.set_poi_default(expectSignal, allowNegativePOI) + self.set_param_default(expectSignal, allowNegativeParam) - def compute(self, poi, full=False): + def compute(self, param, full=False): rnorm = tf.concat( - [poi, tf.ones([self.indata.nproc - poi.shape[0]], dtype=self.indata.dtype)], + [ + param, + tf.ones([self.indata.nproc - param.shape[0]], dtype=self.indata.dtype), + ], axis=0, ) @@ -137,7 +156,7 @@ def compute(self, poi, full=False): return rnorm -class Mixture(POIModel): +class Mixture(ParamModel): """ Based on unconstrained parameters x_i multiply `primary` process by x_i @@ -150,7 +169,7 @@ def __init__( primary_processes, complementary_processes, expectSignal=None, - allowNegativePOI=False, + allowNegativeParam=False, **kwargs, ): self.indata = indata @@ -185,8 +204,10 @@ def __init__( )[0] self.all_idx = np.concatenate([self.primary_idxs, self.complementary_idxs]) - self.npoi = len(primary_processes) - self.pois = np.array( + self.nparams = len(primary_processes) + self.npoi = self.nparams + self.nnui = 0 + self.params = np.array( [ f"{p}_{c}_mixing".encode() for p, c in zip( @@ -195,16 +216,16 @@ def __init__( ] ) - self.allowNegativePOI = allowNegativePOI + self.allowNegativeParam = allowNegativeParam self.is_linear = False - self.set_poi_default(expectSignal, allowNegativePOI) + self.set_param_default(expectSignal, allowNegativeParam) @classmethod def parse_args(cls, indata, *args, **kwargs): """ parsing the input arguments into the constructor, is has to be called as - --poiModel Mixture ,,... ,,... + --paramModel Mixture ,,... ,,... to introduce a mixing parameter for proc_0 with proc_a, and proc_1 with proc_b, etc. """ @@ -218,10 +239,10 @@ def parse_args(cls, indata, *args, **kwargs): return cls(indata, primaries, complementaries, **kwargs) - def compute(self, poi, full=False): + def compute(self, param, full=False): - ones = tf.ones(self.npoi, dtype=self.indata.dtype) - updates = tf.concat([ones * poi, ones * (1 - poi)], axis=0) + ones = tf.ones(self.nparams, dtype=self.indata.dtype) + updates = tf.concat([ones * param, ones * (1 - param)], axis=0) # Single scatter update rnorm = tf.tensor_scatter_nd_update( @@ -234,24 +255,31 @@ def compute(self, poi, full=False): return rnorm -class SaturatedProjectModel(POIModel): +class SaturatedProjectModel(ParamModel): """ 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, + channel_info, + expectSignal=None, + allowNegativeParam=False, + **kwargs, ): self.indata = indata self.channel_info_mapping = channel_info - self.npoi = np.sum( + self.nparams = np.sum( [ np.prod([a.size for a in v["axes"]]) if len(v["axes"]) else 1 for v in channel_info.values() ] ) + self.npoi = self.nparams + self.nnui = 0 names = [] for k, v in self.channel_info_mapping.items(): @@ -259,15 +287,15 @@ def __init__( 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.params = np.array(names) - self.allowNegativePOI = allowNegativePOI + self.allowNegativeParam = allowNegativeParam - self.is_linear = self.npoi == 0 or self.allowNegativePOI + self.is_linear = self.nparams == 0 or self.allowNegativeParam - self.set_poi_default(expectSignal, allowNegativePOI) + self.set_param_default(expectSignal, allowNegativeParam) - def compute(self, poi, full=False): + def compute(self, param, full=False): start = 0 rnorms = [] for k, v in self.indata.channel_info.items(): @@ -279,10 +307,10 @@ def compute(self, poi, full=False): 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 + n_mapping_params = np.prod([a.size for a in mapping_axes]) + iparam = param[start : start + n_mapping_params] + irnorm *= tf.reshape(iparam, shape_mapping) + start += n_mapping_params irnorm = tf.reshape( irnorm, diff --git a/rabbit/parsing.py b/rabbit/parsing.py index 1945ece7..96c63fab 100644 --- a/rabbit/parsing.py +++ b/rabbit/parsing.py @@ -307,11 +307,12 @@ def common_parser(): help="Specify tuple with signal name and rate multiplier for signal expectation (used for fit starting values and for toys). E.g. '--expectSignal BSM 0.0 --expectSignal SM 1.0'", ) parser.add_argument( - "--allowNegativePOI", + "--allowNegativeParam", default=False, action="store_true", help="allow signal strengths to be negative (otherwise constrained to be non-negative)", ) + parser.add_argument( "--noBinByBinStat", default=False, @@ -331,11 +332,12 @@ def common_parser(): help="Barlow-Beeston mode bin-by-bin statistical uncertainties", ) parser.add_argument( - "--poiModel", + "--paramModel", default=["Mu"], nargs="+", - help="Specify POI model to be used to introduce non standard parameterization", + help="Specify param model to be used to introduce non standard parameterization", ) + parser.add_argument( "-m", "--mapping", diff --git a/rabbit/workspace.py b/rabbit/workspace.py index a29910be..f9e67cb6 100644 --- a/rabbit/workspace.py +++ b/rabbit/workspace.py @@ -50,7 +50,7 @@ def __init__(self, outdir, outname, fitter, postfix=None): self.results = {} self.parms = fitter.parms - self.npoi = fitter.poi_model.npoi + self.npoi = fitter.param_model.npoi self.noiidxs = fitter.indata.noiidxs # some information for the impact histograms From e6ceedf010dca2c3165fa3e083a4a8d66ba28276 Mon Sep 17 00:00:00 2001 From: davidwalter2 Date: Sat, 11 Apr 2026 12:50:54 -0400 Subject: [PATCH 02/10] Add ABCD background estimation param model Implements the ABCD method as a new ParamModel subclass with npoi=0, nnui=3*n_bins. Free parameters a_i, b_i, c_i scale regions A, B, C per bin; region D is predicted as D_i = a_i*c_i/b_i times an MC correction factor. Positivity is enforced via squaring inside compute(). CLI syntax supports optional per-axis bin selections. Co-Authored-By: Claude Sonnet 4.6 --- CLAUDE.md | 8 +- rabbit/param_models/abcd_model.py | 288 ++++++++++++++++++++++++++++++ rabbit/param_models/helpers.py | 1 + tests/make_abcd_tensor.py | 95 ++++++++++ 4 files changed, 391 insertions(+), 1 deletion(-) create mode 100644 rabbit/param_models/abcd_model.py create mode 100644 tests/make_abcd_tensor.py diff --git a/CLAUDE.md b/CLAUDE.md index aa5c7694..30604d5f 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -76,7 +76,13 @@ Base class `ParamModel` defines `compute(param)` which returns a `[1, nproc]` te - `npoi`: true parameters of interest (first `npoi` entries; reported as POIs in outputs) - `nnui`: model nuisance parameters (`nparams - npoi`; reported as nuisances in outputs) -Built-in models: `Mu` (default, one POI per signal process), `Ones` (no parameters), `Mixture` (mixing POIs). Custom models are loaded by providing a dotted Python path (e.g. `--paramModel mymod.MyModel`); the module must be on `$PYTHONPATH` with an `__init__.py`. The deprecated `--poiModel` alias is still accepted. +Built-in models: +- `Mu` (default): one POI per signal process +- `Ones`: no parameters (all yields fixed to MC) +- `Mixture`: mixing POIs between pairs of processes +- `ABCD`: data-driven background estimation using four regions; `npoi=0`, `nnui=3*n_bins`. CLI: `--paramModel ABCD [ax:val ...] [ax:val ...] [ax:val ...] [ax:val ...]` where `ax:val` pairs optionally select a single bin along a named axis (e.g. `iso:0`). Regions A, B, C have free parameters; D is predicted as `a*c/b` times an MC correction factor. + +Custom models are loaded by providing a dotted Python path (e.g. `--paramModel mymod.MyModel`); the module must be on `$PYTHONPATH` with an `__init__.py`. ### Mappings: `rabbit/mappings/` Base class `Mapping` in `mapping.py` defines `compute_flat(params, observables)`, which is a differentiable transformation of the flat bin vector. The framework propagates uncertainties through it via automatic differentiation (`tf.GradientTape`). Built-in mappings (`Select`, `Project`, `Normalize`, `Ratio`, `Normratio`) live in `project.py` and `ratio.py`. Custom mappings follow the same pattern as param models. diff --git a/rabbit/param_models/abcd_model.py b/rabbit/param_models/abcd_model.py new file mode 100644 index 00000000..13bfd7ad --- /dev/null +++ b/rabbit/param_models/abcd_model.py @@ -0,0 +1,288 @@ +""" +ABCD background estimation param model. + +The ABCD method estimates a background process using four regions defined by +two independent discriminating variables x and y. The independence condition +A/B = C/D is enforced by construction: free parameters a_i, b_i, c_i scale +regions A, B, C per bin, and the signal region D is predicted as +D_i = a_i * c_i / b_i (up to an MC correction factor). +""" + +import itertools + +import numpy as np +import tensorflow as tf + +from rabbit.param_models.param_model import ParamModel + + +def _get_global_indices(indata, channel_dict): + """Convert one region's channel dict to a numpy array of global flat bin indices. + + Args: + indata: FitInputData instance (channel_info must already have start/stop set). + channel_dict: dict with exactly one entry {ch_name: {axis_name: bin_idx, ...}}. + An empty axis selection dict means all bins of the channel are taken. + + Returns: + numpy array of global flat bin indices for the selected bins. + """ + if len(channel_dict) != 1: + raise ValueError( + f"Each ABCD region must specify exactly one channel, got {list(channel_dict.keys())}" + ) + ch_name, axis_sel = next(iter(channel_dict.items())) + + if ch_name not in indata.channel_info: + raise ValueError( + f"Channel '{ch_name}' not found. Available channels: {list(indata.channel_info.keys())}" + ) + + info = indata.channel_info[ch_name] + axes = info["axes"] + start = info["start"] + + shape = [a.size for a in axes] + all_indices = np.arange(int(np.prod(shape))).reshape(shape) + + slices = [slice(None)] * len(axes) + for ax_name, ax_idx in axis_sel.items(): + ax_pos = next((i for i, a in enumerate(axes) if a.name == ax_name), None) + if ax_pos is None: + raise ValueError( + f"Axis '{ax_name}' not found in channel '{ch_name}'. " + f"Available axes: {[a.name for a in axes]}" + ) + slices[ax_pos] = ax_idx + + local_indices = all_indices[tuple(slices)].flatten() + return start + local_indices + + +def _build_param_names(process, ch_name, axes, axis_sel, shape): + """Build parameter name strings for each selected bin of a region. + + Returns a list of byte strings like b'process_chname_ax0name0_ax1name2'. + """ + # Full index grid over all axes + all_axis_ranges = [range(a.size) for a in axes] + names = [] + for idxs in itertools.product(*all_axis_ranges): + # Apply axis selection: skip bins not matching the fixed selections + skip = False + for ax_name, ax_idx in axis_sel.items(): + ax_pos = next(i for i, a in enumerate(axes) if a.name == ax_name) + if idxs[ax_pos] != ax_idx: + skip = True + break + if skip: + continue + label = "_".join(f"{a.name}{i}" for a, i in zip(axes, idxs)) + names.append(f"{process}_{ch_name}_{label}".encode()) + return names + + +class ABCD(ParamModel): + """ + ABCD background estimation model. + + Defines free parameters a_i, b_i, c_i for regions A, B, C per bin. + Region D (signal region) is derived: D_i = a_i * c_i / b_i * mc_factor_D_i, + where mc_factor_D_i = N_A^MC * N_C^MC / (N_B^MC * N_D^MC) accounts for + non-unity MC templates. + + Parameters are pure model nuisances (npoi=0, nnui=3*n_bins). Positivity + is enforced inside compute() via tf.square() on the raw fit variables. + + CLI syntax: + --paramModel ABCD [ax:val ...] [ax:val ...] \\ + [ax:val ...] [ax:val ...] + + Python constructor: + ABCD(indata, "nonprompt", + channel_A={"ch_fakes": {"iso": 0}}, + channel_B={"ch_fakes": {"iso": 1}}, + channel_C={"ch_C": {}}, + channel_D={"ch_D": {}}) + """ + + def __init__( + self, + indata, + abcd_process, + channel_A, + channel_B, + channel_C, + channel_D, + **kwargs, + ): + """ + Args: + indata: FitInputData instance. + abcd_process: name of the background process to apply ABCD to (str). + channel_A/B/C/D: dicts {ch_name: {axis_name: bin_idx, ...}} defining + each ABCD region. An empty inner dict selects all bins of the channel. + """ + self.indata = indata + + # Validate process + proc_name = ( + abcd_process.encode() if isinstance(abcd_process, str) else abcd_process + ) + if proc_name not in indata.procs: + raise ValueError( + f"Process '{abcd_process}' not found. Available: {list(indata.procs)}" + ) + proc_idx = int(np.where(np.array(indata.procs) == proc_name)[0][0]) + + # Get global flat indices for each region + idx_A = _get_global_indices(indata, channel_A) + idx_B = _get_global_indices(indata, channel_B) + idx_C = _get_global_indices(indata, channel_C) + idx_D = _get_global_indices(indata, channel_D) + + n = len(idx_A) + if not (len(idx_B) == len(idx_C) == len(idx_D) == n): + raise ValueError( + f"All ABCD regions must have the same number of bins, got " + f"A={len(idx_A)}, B={len(idx_B)}, C={len(idx_C)}, D={len(idx_D)}" + ) + self.n_bins = n + + # Extract MC templates for the ABCD process + norm_dense = tf.sparse.to_dense(indata.norm) if indata.sparse else indata.norm + norm_proc = tf.cast(norm_dense[:, proc_idx], dtype=indata.dtype) + + mc_A = tf.gather(norm_proc, idx_A) + mc_B = tf.gather(norm_proc, idx_B) + mc_C = tf.gather(norm_proc, idx_C) + mc_D = tf.gather(norm_proc, idx_D) + + # MC correction factor D = A*C / (B*D_MC); protect against zeros + mc_B_safe = tf.where(mc_B == 0.0, tf.ones_like(mc_B), mc_B) + mc_D_safe = tf.where(mc_D == 0.0, tf.ones_like(mc_D), mc_D) + self.mc_factor_D = mc_A * mc_C / (mc_B_safe * mc_D_safe) + + # Scatter indices for compute(): stored as Python lists so they become + # tf.constant at @tf.function trace time (not retraced per call) + self._idx = { + "A": [[int(i), proc_idx] for i in idx_A], + "B": [[int(i), proc_idx] for i in idx_B], + "C": [[int(i), proc_idx] for i in idx_C], + "D": [[int(i), proc_idx] for i in idx_D], + } + + # Per-region activity flags for non-full mode + # (masked channels have bin indices >= indata.nbins) + self._active_nonfull = { + r: bool(all(int(i) < indata.nbins for i in idxs)) + for r, idxs in zip("ABCD", [idx_A, idx_B, idx_C, idx_D]) + } + + # Build parameter names: a_i for A, b_i for B, c_i for C + names = [] + for ch_dict, region_label in [ + (channel_A, "A"), + (channel_B, "B"), + (channel_C, "C"), + ]: + ch_name, axis_sel = next(iter(ch_dict.items())) + ch_axes = indata.channel_info[ch_name]["axes"] + ch_shape = [a.size for a in ch_axes] + names.extend( + _build_param_names(abcd_process, ch_name, ch_axes, axis_sel, ch_shape) + ) + + # Model attributes + self.nparams = 3 * n + self.npoi = 0 + self.nnui = 3 * n + self.params = np.array(names) + self.is_linear = False + self.allowNegativeParam = False + self.xparamdefault = tf.ones([3 * n], dtype=indata.dtype) + + @classmethod + def parse_args(cls, indata, *args, **kwargs): + """Parse CLI arguments for ABCDModel. + + Syntax: + --paramModel ABCD [ax:val ...] [ax:val ...] \\ + [ax:val ...] [ax:val ...] + + Each channel name is followed by zero or more axis:value pairs. + Axis values are integers. + """ + if len(args) < 5: + raise ValueError( + "ABCDModel expects: process ch_A [ax:val ...] ch_B [ax:val ...] " + "ch_C [ax:val ...] ch_D [ax:val ...]" + ) + tokens = list(args) + process = tokens.pop(0) + + def _parse_one_region(tokens, region_name): + if not tokens or ":" in tokens[0]: + raise ValueError( + f"Expected channel name for region {region_name}, " + f"got '{tokens[0] if tokens else ''}'" + ) + ch_name = tokens.pop(0) + axis_sel = {} + while tokens and ":" in tokens[0]: + ax, val = tokens.pop(0).split(":", 1) + axis_sel[ax] = int(val) + return {ch_name: axis_sel} + + channel_A = _parse_one_region(tokens, "A") + channel_B = _parse_one_region(tokens, "B") + channel_C = _parse_one_region(tokens, "C") + channel_D = _parse_one_region(tokens, "D") + + if tokens: + raise ValueError(f"Unexpected extra arguments after channel_D: {tokens}") + + return cls( + indata, process, channel_A, channel_B, channel_C, channel_D, **kwargs + ) + + def compute(self, param, full=False): + """Compute per-bin, per-process yield scale factors. + + Args: + param: 1D tensor of length 3*n_bins (raw fit variables for a, b, c). + full: if True, return [nbinsfull, nproc]; else [nbins, nproc]. + + Returns: + 2D tensor [nbins(_full), nproc] of yield scale factors. + All entries are 1 except for the ABCD process bins which are set to + a_i, b_i, c_i (regions A, B, C) and d_i = a_i*c_i/b_i*mc_factor (D). + """ + n = self.n_bins + # Enforce positivity via squaring; raw param=1 → physical value=1 + a = tf.square(param[:n]) + b = tf.square(param[n : 2 * n]) + c = tf.square(param[2 * n :]) + + b_safe = tf.where(b == 0.0, tf.ones_like(b) * 1e-10, b) + d = a * c / b_safe * self.mc_factor_D + + nbins = self.indata.nbinsfull if full else self.indata.nbins + rnorm = tf.ones([nbins, self.indata.nproc], dtype=self.indata.dtype) + + # Python-level control flow resolved at @tf.function trace time + indices = [] + updates = [] + for region, vals in [("A", a), ("B", b), ("C", c), ("D", d)]: + active = True if full else self._active_nonfull[region] + if active: + indices.extend(self._idx[region]) + updates.append(vals) + + if indices: + rnorm = tf.tensor_scatter_nd_update( + rnorm, + tf.constant(indices, dtype=tf.int64), + tf.concat(updates, axis=0), + ) + return rnorm diff --git a/rabbit/param_models/helpers.py b/rabbit/param_models/helpers.py index 3eddb8d1..db004c37 100644 --- a/rabbit/param_models/helpers.py +++ b/rabbit/param_models/helpers.py @@ -5,6 +5,7 @@ "Ones": "param_model", "Mu": "param_model", "Mixture": "param_model", + "ABCD": "abcd_model", } diff --git a/tests/make_abcd_tensor.py b/tests/make_abcd_tensor.py new file mode 100644 index 00000000..635c964b --- /dev/null +++ b/tests/make_abcd_tensor.py @@ -0,0 +1,95 @@ +""" +Create a test HDF5 tensor with four ABCD regions for testing ABCDModel. + +Layout: + ch_A: pt histogram — fake rate measurement region A (low iso, low mt) + ch_B: pt histogram — fake rate measurement region B (high iso, low mt) + ch_C: pt histogram — application region C (low iso, high mt) + ch_D: pt histogram — signal region D (high iso, high mt) + +The ABCD method predicts D = C * A / B. +A signal process lives in ch_C and ch_D. +A nonprompt background lives in all four channels. +""" + +import argparse +import os + +import hist +import numpy as np + +from rabbit import tensorwriter + +parser = argparse.ArgumentParser() +parser.add_argument("-o", "--output", default="./", help="output directory") +parser.add_argument("--outname", default="test_abcd_tensor", help="output file name") +args = parser.parse_args() + +rng = np.random.default_rng(42) + +n_pt = 5 # number of pt bins +ax_pt = hist.axis.Regular(n_pt, 0, 50, name="pt") + + +def make_hist(n_events, mean_weight): + """Histogram filled with uniform pt values and given mean weight.""" + h = hist.Hist(ax_pt, storage=hist.storage.Weight()) + pt_vals = rng.uniform(0, 50, n_events) + weights = rng.normal(mean_weight, 0.1 * mean_weight, n_events) + h.fill(pt=pt_vals, weight=weights) + return h + + +def make_data(sig_h, bkg_h): + """Poisson-fluctuated data from signal + background templates.""" + h = hist.Hist(ax_pt, storage=hist.storage.Double()) + sig_vals = sig_h.values() if sig_h is not None else 0.0 + expected = sig_vals + bkg_h.values() + h.view()[:] = rng.poisson(np.maximum(expected, 0)) + return h + + +# Signal process: only in ch_C and ch_D +h_sig_C = make_hist(5000, 1.0) +h_sig_D = make_hist(5000, 1.0) + +# Nonprompt background: present in all four channels +# Set up so the ABCD relation A*C/B ≈ D holds approximately per bin +h_np_A = make_hist(8000, 2.0) +h_np_B = make_hist(8000, 4.0) # B ~ 2x A +h_np_C = make_hist(6000, 1.5) +h_np_D = make_hist(3000, 0.75) # D ~ A*C/B per bin + +# Poisson-fluctuated data +h_data_A = make_data(None, h_np_A) +h_data_B = make_data(None, h_np_B) +h_data_C = make_data(h_sig_C, h_np_C) +h_data_D = make_data(h_sig_D, h_np_D) + +# --- Build tensor --- +writer = tensorwriter.TensorWriter(sparse=False) + +writer.add_channel(h_sig_C.axes, "ch_A") +writer.add_channel(h_sig_C.axes, "ch_B") +writer.add_channel(h_sig_C.axes, "ch_C") +writer.add_channel(h_sig_C.axes, "ch_D") + +# Data +writer.add_data(h_data_A, "ch_A") +writer.add_data(h_data_B, "ch_B") +writer.add_data(h_data_C, "ch_C") +writer.add_data(h_data_D, "ch_D") + +# Signal process (only in C and D) +writer.add_process(h_sig_C, "signal", "ch_C", signal=True) +writer.add_process(h_sig_D, "signal", "ch_D", signal=True) + +# Nonprompt background (all four channels) +writer.add_process(h_np_A, "nonprompt", "ch_A", signal=False) +writer.add_process(h_np_B, "nonprompt", "ch_B", signal=False) +writer.add_process(h_np_C, "nonprompt", "ch_C", signal=False) +writer.add_process(h_np_D, "nonprompt", "ch_D", signal=False) + +outdir = os.path.join(args.output, args.outname) +writer.write(outdir) +print(f"Written to {outdir}/rabbit_input.hdf5") From 64a0a137bc4d20b860f7a27016a0669a5ab2dfb6 Mon Sep 17 00:00:00 2001 From: davidwalter2 Date: Sat, 11 Apr 2026 13:20:43 -0400 Subject: [PATCH 03/10] Support multiple --paramModel via CompositeParamModel --paramModel now uses action="append" so it can be specified multiple times. Each invocation defines one model spec; when multiple are given they are automatically combined via CompositeParamModel. Single-model usage and the default (Mu) are fully backward compatible. Co-Authored-By: Claude Sonnet 4.6 --- bin/rabbit_fit.py | 4 ++-- bin/rabbit_limit.py | 4 ++-- rabbit/param_models/helpers.py | 18 ++++++++++++++++++ rabbit/parsing.py | 7 +++++-- 4 files changed, 27 insertions(+), 6 deletions(-) diff --git a/bin/rabbit_fit.py b/bin/rabbit_fit.py index 58d953b3..a7dd740c 100755 --- a/bin/rabbit_fit.py +++ b/bin/rabbit_fit.py @@ -574,8 +574,8 @@ def main(): indata = inputdata.FitInputData(args.filename, args.pseudoData) - margs = args.paramModel - param_model = ph.load_model(margs[0], indata, *margs[1:], **vars(args)) + model_specs = args.paramModel or [["Mu"]] + param_model = ph.load_models(model_specs, indata, **vars(args)) ifitter = fitter.Fitter( indata, diff --git a/bin/rabbit_limit.py b/bin/rabbit_limit.py index d76cc0c4..26383ccf 100755 --- a/bin/rabbit_limit.py +++ b/bin/rabbit_limit.py @@ -264,8 +264,8 @@ def main(): indata = inputdata.FitInputData(args.filename, args.pseudoData) - margs = args.paramModel - param_model = ph.load_model(margs[0], indata, *margs[1:], **vars(args)) + model_specs = args.paramModel or [["Mu"]] + param_model = ph.load_models(model_specs, indata, **vars(args)) ifitter = fitter.Fitter(indata, param_model, args, do_blinding=any(blinded_fits)) diff --git a/rabbit/param_models/helpers.py b/rabbit/param_models/helpers.py index db004c37..ff499bef 100644 --- a/rabbit/param_models/helpers.py +++ b/rabbit/param_models/helpers.py @@ -14,3 +14,21 @@ def load_model(class_name, indata, *args, **kwargs): class_name, baseline_models, base_dir="rabbit.param_models" ) return model.parse_args(indata, *args, **kwargs) + + +def load_models(model_specs, indata, **kwargs): + """Load one or more param models and return a single model (or composite). + + Args: + model_specs: list of lists, e.g. [["Mu"], ["ABCD", "nonprompt", "ch_A", ...]] + indata: FitInputData instance + **kwargs: passed to each model's parse_args (e.g. from vars(args)) + """ + from rabbit.param_models.param_model import CompositeParamModel + + models = [load_model(spec[0], indata, *spec[1:], **kwargs) for spec in model_specs] + if len(models) == 1: + return models[0] + return CompositeParamModel( + models, allowNegativeParam=kwargs.get("allowNegativeParam", False) + ) diff --git a/rabbit/parsing.py b/rabbit/parsing.py index 96c63fab..0b5699af 100644 --- a/rabbit/parsing.py +++ b/rabbit/parsing.py @@ -333,9 +333,12 @@ def common_parser(): ) parser.add_argument( "--paramModel", - default=["Mu"], + default=None, nargs="+", - help="Specify param model to be used to introduce non standard parameterization", + action="append", + help="Specify param model to be used to introduce non standard parameterization. " + "Can be specified multiple times to combine models via CompositeParamModel, " + "e.g. '--paramModel Mu --paramModel ABCD nonprompt ch_A ch_B ch_C ch_D'.", ) parser.add_argument( From 94d39fbfd5124dd3d556e2b84bc4b8627816632d Mon Sep 17 00:00:00 2001 From: davidwalter2 Date: Sat, 11 Apr 2026 13:53:36 -0400 Subject: [PATCH 04/10] Add SmoothABCD param model with exponential polynomial parameterisation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit SmoothABCD replaces the per-bin free parameters of one nominated axis with an exponential polynomial val(x)=exp(p_0+p_1*x̃+...), reducing 3*n_bins parameters to 3*n_outer*(order+1). Default order=1 (log-linear). Outer bins (remaining axes after the smooth axis) have independent polynomial coefficients. Reuses _get_global_indices from abcd_model.py. Co-Authored-By: Claude Sonnet 4.6 --- CLAUDE.md | 1 + rabbit/param_models/helpers.py | 1 + rabbit/param_models/smooth_abcd_model.py | 337 +++++++++++++++++++++++ 3 files changed, 339 insertions(+) create mode 100644 rabbit/param_models/smooth_abcd_model.py diff --git a/CLAUDE.md b/CLAUDE.md index 30604d5f..c68b6663 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -81,6 +81,7 @@ Built-in models: - `Ones`: no parameters (all yields fixed to MC) - `Mixture`: mixing POIs between pairs of processes - `ABCD`: data-driven background estimation using four regions; `npoi=0`, `nnui=3*n_bins`. CLI: `--paramModel ABCD [ax:val ...] [ax:val ...] [ax:val ...] [ax:val ...]` where `ax:val` pairs optionally select a single bin along a named axis (e.g. `iso:0`). Regions A, B, C have free parameters; D is predicted as `a*c/b` times an MC correction factor. +- `SmoothABCD`: like ABCD but one axis is parameterised with an exponential polynomial `val(x)=exp(p_0+p_1·x̃+...)` instead of per-bin free parameters. CLI: `--paramModel SmoothABCD [order:N] ... `. Default order=1 (log-linear). Reduces parameters from `3·n_bins` to `3·n_outer·(order+1)`. Custom models are loaded by providing a dotted Python path (e.g. `--paramModel mymod.MyModel`); the module must be on `$PYTHONPATH` with an `__init__.py`. diff --git a/rabbit/param_models/helpers.py b/rabbit/param_models/helpers.py index ff499bef..c5ef36ee 100644 --- a/rabbit/param_models/helpers.py +++ b/rabbit/param_models/helpers.py @@ -6,6 +6,7 @@ "Mu": "param_model", "Mixture": "param_model", "ABCD": "abcd_model", + "SmoothABCD": "smooth_abcd_model", } diff --git a/rabbit/param_models/smooth_abcd_model.py b/rabbit/param_models/smooth_abcd_model.py new file mode 100644 index 00000000..78fe4801 --- /dev/null +++ b/rabbit/param_models/smooth_abcd_model.py @@ -0,0 +1,337 @@ +""" +SmoothABCD background estimation param model. + +Like ABCD, but the per-bin free parameters along one nominated smoothing axis +are replaced by an exponential polynomial: + + val_X(x) = exp(p_0 + p_1·x̃ + p_2·x̃² + ...) + +where x̃ are normalised bin centres of the smoothing axis in [0, 1]. + +This reduces parameters from 3·n_bins to 3·n_outer·(order+1), where +n_outer is the number of bins along the remaining (non-smooth) axes. +""" + +import itertools + +import numpy as np +import tensorflow as tf + +from rabbit.param_models.abcd_model import _get_global_indices +from rabbit.param_models.param_model import ParamModel + + +class SmoothABCD(ParamModel): + """ + Smooth ABCD background estimation model. + + Regions A, B, C have free parameters that follow an exponential polynomial + along the smoothing axis; region D is derived as d = a·c/b · mc_factor_D. + + Parameters are pure model nuisances (npoi=0, nnui=3·n_outer·(order+1)). + + CLI syntax: + --paramModel SmoothABCD [order:N] \\ + [ax:val ...] [ax:val ...] \\ + [ax:val ...] [ax:val ...] + + Python constructor: + SmoothABCD(indata, "pt", "nonprompt", + channel_A={"ch_fakes": {"iso": 0}}, + channel_B={"ch_fakes": {"iso": 1}}, + channel_C={"ch_C": {}}, + channel_D={"ch_D": {}}, + order=1) + """ + + def __init__( + self, + indata, + smoothing_axis, + abcd_process, + channel_A, + channel_B, + channel_C, + channel_D, + order=1, + **kwargs, + ): + """ + Args: + indata: FitInputData instance. + smoothing_axis: name of the axis to parameterise smoothly. + abcd_process: name of the background process (str). + channel_A/B/C/D: dicts {ch_name: {axis_name: bin_idx, ...}}. + order: polynomial degree in the exponent (default 1 = log-linear). + """ + self.indata = indata + self.order = order + + # Validate process + proc_name = ( + abcd_process.encode() if isinstance(abcd_process, str) else abcd_process + ) + if proc_name not in indata.procs: + raise ValueError( + f"Process '{abcd_process}' not found. Available: {list(indata.procs)}" + ) + proc_idx = int(np.where(np.array(indata.procs) == proc_name)[0][0]) + + # Build outer × smooth index structure for each region + # outer_global_indices[region] = list of length n_outer, + # each entry is an array of n_smooth global flat indices + outer_global_indices = {} + outer_axes_by_region = {} + smooth_size = None + + for region, ch_dict in zip( + "ABCD", [channel_A, channel_B, channel_C, channel_D] + ): + ch_name, axis_sel = next(iter(ch_dict.items())) + axes = indata.channel_info[ch_name]["axes"] + + # Free axes after the fixed axis selections + remaining = [a for a in axes if a.name not in axis_sel] + + smooth_ax = next((a for a in remaining if a.name == smoothing_axis), None) + if smooth_ax is None: + raise ValueError( + f"Smoothing axis '{smoothing_axis}' not found among free axes " + f"of channel '{ch_name}' for region {region}. " + f"Free axes: {[a.name for a in remaining]}" + ) + + n_smooth = smooth_ax.size + if smooth_size is None: + smooth_size = n_smooth + elif smooth_size != n_smooth: + raise ValueError( + f"Smoothing axis '{smoothing_axis}' has inconsistent size: " + f"expected {smooth_size}, got {n_smooth} in region {region}" + ) + + outer_axes = [a for a in remaining if a.name != smoothing_axis] + outer_axes_by_region[region] = outer_axes + + # For each outer-bin combination, gather the n_smooth global indices + # (smooth axis is left free → _get_global_indices returns n_smooth values) + region_outer_indices = [] + for outer_combo in itertools.product(*[range(a.size) for a in outer_axes]): + extended_sel = dict(axis_sel) + for ax, idx in zip(outer_axes, outer_combo): + extended_sel[ax.name] = idx + region_outer_indices.append( + _get_global_indices(indata, {ch_name: extended_sel}) + ) + + outer_global_indices[region] = region_outer_indices + + n_smooth = smooth_size + n_outer = len(outer_global_indices["A"]) + self.n_outer = n_outer + self.n_smooth = n_smooth + + for region in "BCD": + if len(outer_global_indices[region]) != n_outer: + raise ValueError( + f"Region {region} has {len(outer_global_indices[region])} outer bins, " + f"expected {n_outer}" + ) + + # Vandermonde matrix from normalised bin centres of the smooth axis + ch_name_A, axis_sel_A = next(iter(channel_A.items())) + axes_A = indata.channel_info[ch_name_A]["axes"] + smooth_ax_A = next(a for a in axes_A if a.name == smoothing_axis) + x_centers = np.array(smooth_ax_A.centers) + x_min, x_max = float(x_centers.min()), float(x_centers.max()) + x_norm = ( + (x_centers - x_min) / (x_max - x_min) + if x_max > x_min + else np.zeros_like(x_centers) + ) + # V[k, s] = x_norm[s]^k, shape [order+1, n_smooth] + vander = np.array( + [[float(x_norm[s]) ** k for s in range(n_smooth)] for k in range(order + 1)] + ) + self.vander = tf.constant(vander, dtype=indata.dtype) + + # MC correction factor for region D, shape [n_outer * n_smooth] + norm_dense = tf.sparse.to_dense(indata.norm) if indata.sparse else indata.norm + norm_proc = tf.cast(norm_dense[:, proc_idx], dtype=indata.dtype) + + def _gather_flat(region_outer_idx): + return tf.gather(norm_proc, np.concatenate(region_outer_idx)) + + mc_A = _gather_flat(outer_global_indices["A"]) + mc_B = _gather_flat(outer_global_indices["B"]) + mc_C = _gather_flat(outer_global_indices["C"]) + mc_D = _gather_flat(outer_global_indices["D"]) + + mc_B_safe = tf.where(mc_B == 0.0, tf.ones_like(mc_B), mc_B) + mc_D_safe = tf.where(mc_D == 0.0, tf.ones_like(mc_D), mc_D) + self.mc_factor_D = mc_A * mc_C / (mc_B_safe * mc_D_safe) + + # Scatter indices in [bin, proc] format, ordered outer × smooth + def _build_scatter_idx(region_outer_idx): + return [ + [int(region_outer_idx[m][s]), proc_idx] + for m in range(n_outer) + for s in range(n_smooth) + ] + + self._idx = {r: _build_scatter_idx(outer_global_indices[r]) for r in "ABCD"} + + # Per-region activity flags for non-full mode + self._active_nonfull = { + r: bool( + all( + int(outer_global_indices[r][m][s]) < indata.nbins + for m in range(n_outer) + for s in range(n_smooth) + ) + ) + for r in "ABCD" + } + + # Parameter names: region × outer_bin × poly_coeff + names = [] + for region, ch_dict in zip("ABC", [channel_A, channel_B, channel_C]): + ch_name, _ = next(iter(ch_dict.items())) + outer_axes = outer_axes_by_region[region] + combos = list(itertools.product(*[range(a.size) for a in outer_axes])) or [ + () + ] + for outer_combo in combos: + if outer_axes: + outer_label = "_".join( + f"{a.name}{i}" for a, i in zip(outer_axes, outer_combo) + ) + else: + outer_label = "outer0" + for k in range(order + 1): + names.append( + f"{abcd_process}_{ch_name}_{region}_{outer_label}_poly{k}".encode() + ) + + # Model attributes + self.nparams = 3 * n_outer * (order + 1) + self.npoi = 0 + self.nnui = self.nparams + self.params = np.array(names) + self.is_linear = False + self.allowNegativeParam = False + # Default: all coefficients 0 → exp(0) = 1 (MC template as starting point) + self.xparamdefault = tf.zeros([self.nparams], dtype=indata.dtype) + + @classmethod + def parse_args(cls, indata, *args, **kwargs): + """Parse CLI arguments for SmoothABCD. + + Syntax: + --paramModel SmoothABCD [order:N] \\ + [ax:val ...] [ax:val ...] \\ + [ax:val ...] [ax:val ...] + """ + if len(args) < 6: + raise ValueError( + "SmoothABCD expects: axis [order:N] process " + "ch_A [ax:val ...] ch_B [ax:val ...] ch_C [ax:val ...] ch_D [ax:val ...]" + ) + tokens = list(args) + smoothing_axis = tokens.pop(0) + + order = 1 + if tokens and tokens[0].startswith("order:"): + order = int(tokens.pop(0).split(":", 1)[1]) + + if not tokens: + raise ValueError("SmoothABCD: expected process name after axis/order") + process = tokens.pop(0) + + def _parse_one_region(tokens, region_name): + if not tokens or ":" in tokens[0]: + raise ValueError( + f"Expected channel name for region {region_name}, " + f"got '{tokens[0] if tokens else ''}'" + ) + ch_name = tokens.pop(0) + axis_sel = {} + while tokens and ":" in tokens[0]: + ax, val = tokens.pop(0).split(":", 1) + axis_sel[ax] = int(val) + return {ch_name: axis_sel} + + channel_A = _parse_one_region(tokens, "A") + channel_B = _parse_one_region(tokens, "B") + channel_C = _parse_one_region(tokens, "C") + channel_D = _parse_one_region(tokens, "D") + + if tokens: + raise ValueError(f"Unexpected extra arguments after channel_D: {tokens}") + + return cls( + indata, + smoothing_axis, + process, + channel_A, + channel_B, + channel_C, + channel_D, + order=order, + **kwargs, + ) + + def compute(self, param, full=False): + """Compute per-bin, per-process yield scale factors. + + Args: + param: 1D tensor of length 3 * n_outer * (order+1). + full: if True return [nbinsfull, nproc]; else [nbins, nproc]. + """ + n_coeffs = self.n_outer * (self.order + 1) + + # Reshape to [n_outer, order+1] for each region + params_A = tf.reshape(param[:n_coeffs], [self.n_outer, self.order + 1]) + params_B = tf.reshape( + param[n_coeffs : 2 * n_coeffs], [self.n_outer, self.order + 1] + ) + params_C = tf.reshape(param[2 * n_coeffs :], [self.n_outer, self.order + 1]) + + # Evaluate exp(polynomial): [n_outer, n_smooth] + # self.vander is a constant of shape [order+1, n_smooth] + a = tf.exp(tf.matmul(params_A, self.vander)) + b = tf.exp(tf.matmul(params_B, self.vander)) + c = tf.exp(tf.matmul(params_C, self.vander)) + + # Flatten to [n_outer * n_smooth] = [n_total_bins per region] + a_flat = tf.reshape(a, [-1]) + b_flat = tf.reshape(b, [-1]) + c_flat = tf.reshape(c, [-1]) + + b_safe = tf.where(b_flat == 0.0, tf.ones_like(b_flat) * 1e-10, b_flat) + d_flat = a_flat * c_flat / b_safe * self.mc_factor_D + + nbins = self.indata.nbinsfull if full else self.indata.nbins + rnorm = tf.ones([nbins, self.indata.nproc], dtype=self.indata.dtype) + + # Python-level control flow resolved at @tf.function trace time + indices = [] + updates = [] + for region, vals in [ + ("A", a_flat), + ("B", b_flat), + ("C", c_flat), + ("D", d_flat), + ]: + active = True if full else self._active_nonfull[region] + if active: + indices.extend(self._idx[region]) + updates.append(vals) + + if indices: + rnorm = tf.tensor_scatter_nd_update( + rnorm, + tf.constant(indices, dtype=tf.int64), + tf.concat(updates, axis=0), + ) + return rnorm From 65d193b7a210eb80b1dd3c32c854a0a516831921 Mon Sep 17 00:00:00 2001 From: davidwalter2 Date: Sat, 11 Apr 2026 14:24:03 -0400 Subject: [PATCH 05/10] Add ExtendedABCD param model and ABCD CI tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ExtendedABCD uses 6 regions (Ax, Bx extra sideband + standard A, B, C, D) with log-linear fake rate extrapolation: D = C * Ax * B² / (Bx * A²). Free parameters: a, b, c, ax, bx per bin (npoi=0, nnui=5*n_bins). CI: new make-abcd-tensors and abcd-fits jobs covering ABCD, SmoothABCD, and ExtendedABCD with Mu+ABCD composite fits and --saveHists. Co-Authored-By: Claude Sonnet 4.6 --- .github/workflows/main.yml | 66 ++++- CLAUDE.md | 1 + rabbit/param_models/extended_abcd_model.py | 271 +++++++++++++++++++++ rabbit/param_models/helpers.py | 1 + tests/make_extended_abcd_tensor.py | 120 +++++++++ 5 files changed, 458 insertions(+), 1 deletion(-) create mode 100644 rabbit/param_models/extended_abcd_model.py create mode 100644 tests/make_extended_abcd_tensor.py diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 63e6416f..b71bbd7f 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -543,9 +543,73 @@ jobs: rabbit_plot_likelihood_scan2D.py $RABBIT_OUTDIR/fitresults_scans.hdf5 -o $WEB_DIR/$PLOT_DIR --title Rabbit --subtitle Preliminary --params sig slope_signal --titlePos 0 --legPos 'lower right' + make-abcd-tensors: + runs-on: [self-hosted, linux, x64] + needs: setenv + + steps: + - env: + RABBIT_OUTDIR: ${{ needs.setenv.outputs.RABBIT_OUTDIR }} + PYTHONPATH: ${{ needs.setenv.outputs.PYTHONPATH }} + PATH: ${{ needs.setenv.outputs.PATH }} + run: | + echo "RABBIT_OUTDIR=${RABBIT_OUTDIR}" >> $GITHUB_ENV + echo "PYTHONPATH=${PYTHONPATH}" >> $GITHUB_ENV + echo "PATH=${PATH}" >> $GITHUB_ENV + + - uses: actions/checkout@v5 + + - name: make ABCD tensor + run: >- + python tests/make_abcd_tensor.py -o $RABBIT_OUTDIR/ + + - name: make extended ABCD tensor + run: >- + python tests/make_extended_abcd_tensor.py -o $RABBIT_OUTDIR/ + + abcd-fits: + runs-on: [self-hosted, linux, x64] + needs: [setenv, make-abcd-tensors] + + steps: + - env: + RABBIT_OUTDIR: ${{ needs.setenv.outputs.RABBIT_OUTDIR }} + PYTHONPATH: ${{ needs.setenv.outputs.PYTHONPATH }} + PATH: ${{ needs.setenv.outputs.PATH }} + run: | + echo "RABBIT_OUTDIR=${RABBIT_OUTDIR}" >> $GITHUB_ENV + echo "PYTHONPATH=${PYTHONPATH}" >> $GITHUB_ENV + echo "PATH=${PATH}" >> $GITHUB_ENV + + - uses: actions/checkout@v5 + + - name: ABCD fit + run: >- + rabbit_fit.py $RABBIT_OUTDIR/test_abcd_tensor/rabbit_input.hdf5 + -o $RABBIT_OUTDIR/ --postfix abcd + -t 0 --paramModel Mu + --paramModel ABCD nonprompt ch_A ch_B ch_C ch_D + --saveHists + + - name: SmoothABCD fit + run: >- + rabbit_fit.py $RABBIT_OUTDIR/test_abcd_tensor/rabbit_input.hdf5 + -o $RABBIT_OUTDIR/ --postfix smooth_abcd + -t 0 --paramModel Mu + --paramModel SmoothABCD pt nonprompt ch_A ch_B ch_C ch_D + --saveHists + + - name: ExtendedABCD fit + run: >- + rabbit_fit.py $RABBIT_OUTDIR/test_extended_abcd_tensor/rabbit_input.hdf5 + -o $RABBIT_OUTDIR/ --postfix extended_abcd + -t 0 --paramModel Mu + --paramModel ExtendedABCD nonprompt ch_Ax ch_Bx ch_A ch_B ch_C ch_D + --saveHists + copy-clean: runs-on: [self-hosted, linux, x64] - needs: [setenv, symmerizations, sparse-fits, alternative-fits, regularization, bsm, plotting, likelihoodscans] + needs: [setenv, symmerizations, sparse-fits, alternative-fits, regularization, bsm, plotting, likelihoodscans, abcd-fits] if: always() steps: - env: diff --git a/CLAUDE.md b/CLAUDE.md index c68b6663..c696de5c 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -82,6 +82,7 @@ Built-in models: - `Mixture`: mixing POIs between pairs of processes - `ABCD`: data-driven background estimation using four regions; `npoi=0`, `nnui=3*n_bins`. CLI: `--paramModel ABCD [ax:val ...] [ax:val ...] [ax:val ...] [ax:val ...]` where `ax:val` pairs optionally select a single bin along a named axis (e.g. `iso:0`). Regions A, B, C have free parameters; D is predicted as `a*c/b` times an MC correction factor. - `SmoothABCD`: like ABCD but one axis is parameterised with an exponential polynomial `val(x)=exp(p_0+p_1·x̃+...)` instead of per-bin free parameters. CLI: `--paramModel SmoothABCD [order:N] ... `. Default order=1 (log-linear). Reduces parameters from `3·n_bins` to `3·n_outer·(order+1)`. +- `ExtendedABCD`: 6-region ABCD using two sideband bins in the x direction (Ax, Bx further from signal, A/B in the middle). Fake rate is log-linearly extrapolated: `D = C·Ax·B² / (Bx·A²)`. `npoi=0`, `nnui=5·n_bins`. CLI: `--paramModel ExtendedABCD [ax:val ...] [ax:val ...] [ax:val ...] [ax:val ...] [ax:val ...] [ax:val ...]`. Custom models are loaded by providing a dotted Python path (e.g. `--paramModel mymod.MyModel`); the module must be on `$PYTHONPATH` with an `__init__.py`. diff --git a/rabbit/param_models/extended_abcd_model.py b/rabbit/param_models/extended_abcd_model.py new file mode 100644 index 00000000..b580c439 --- /dev/null +++ b/rabbit/param_models/extended_abcd_model.py @@ -0,0 +1,271 @@ +""" +ExtendedABCD background estimation param model. + +Uses 6 regions by adding a second sideband bin (Ax, Bx) in the x direction. +The fake rate is extrapolated log-linearly from the two sideband measurements +to the signal region: + + y=signal y=sideband +x=extra-sideband: Ax Bx +x=sideband: A B +x=signal: C D ← predicted + +D = C * Ax * B² / (Bx * A²) + +This follows from assuming the fake rate f(x) = B_x / A_x is log-linear in x, +so f(x_signal) = f(x_sideband)² / f(x_extra-sideband). +""" + +import numpy as np +import tensorflow as tf + +from rabbit.param_models.abcd_model import _build_param_names, _get_global_indices +from rabbit.param_models.param_model import ParamModel + + +class ExtendedABCD(ParamModel): + """ + Extended ABCD background estimation model with 6 regions. + + Free parameters a_i, b_i, c_i, ax_i, bx_i for regions A, B, C, Ax, Bx. + Region D is derived: D = C * Ax * B² / (Bx * A²) * mc_factor_D. + + Parameters are pure model nuisances (npoi=0, nnui=5*n_bins). Positivity + is enforced inside compute() via tf.square() on the raw fit variables. + + CLI syntax: + --paramModel ExtendedABCD \\ + [ax:val ...] [ax:val ...] \\ + [ax:val ...] [ax:val ...] \\ + [ax:val ...] [ax:val ...] + + Python constructor: + ExtendedABCD(indata, "nonprompt", + channel_A={"ch_fakes": {"iso": 1}}, + channel_B={"ch_fakes": {"iso": 2}}, + channel_C={"ch_C": {}}, + channel_D={"ch_D": {}}, + channel_Ax={"ch_fakes": {"iso": 0}}, + channel_Bx={"ch_fakes": {"iso": 3}}) + """ + + def __init__( + self, + indata, + abcd_process, + channel_A, + channel_B, + channel_C, + channel_D, + channel_Ax, + channel_Bx, + **kwargs, + ): + """ + Args: + indata: FitInputData instance. + abcd_process: name of the background process (str). + channel_A/B/C/D: dicts {ch_name: {axis_name: bin_idx, ...}}. + channel_Ax/Bx: extra sideband regions (further from signal than A/B). + """ + self.indata = indata + + # Validate process + proc_name = ( + abcd_process.encode() if isinstance(abcd_process, str) else abcd_process + ) + if proc_name not in indata.procs: + raise ValueError( + f"Process '{abcd_process}' not found. Available: {list(indata.procs)}" + ) + proc_idx = int(np.where(np.array(indata.procs) == proc_name)[0][0]) + + # Get global flat indices for all six regions + idx_A = _get_global_indices(indata, channel_A) + idx_B = _get_global_indices(indata, channel_B) + idx_C = _get_global_indices(indata, channel_C) + idx_D = _get_global_indices(indata, channel_D) + idx_Ax = _get_global_indices(indata, channel_Ax) + idx_Bx = _get_global_indices(indata, channel_Bx) + + n = len(idx_A) + for label, idx in [ + ("B", idx_B), + ("C", idx_C), + ("D", idx_D), + ("Ax", idx_Ax), + ("Bx", idx_Bx), + ]: + if len(idx) != n: + raise ValueError( + f"All regions must have the same number of bins; " + f"A={n} but {label}={len(idx)}" + ) + self.n_bins = n + + # Extract MC templates + norm_dense = tf.sparse.to_dense(indata.norm) if indata.sparse else indata.norm + norm_proc = tf.cast(norm_dense[:, proc_idx], dtype=indata.dtype) + + mc_A = tf.gather(norm_proc, idx_A) + mc_B = tf.gather(norm_proc, idx_B) + mc_C = tf.gather(norm_proc, idx_C) + mc_D = tf.gather(norm_proc, idx_D) + mc_Ax = tf.gather(norm_proc, idx_Ax) + mc_Bx = tf.gather(norm_proc, idx_Bx) + + # MC correction factor: D = C * Ax * B² / (Bx * A² * D_MC) * (numerator MC) + # rnorm_D = rnorm_C * rnorm_Ax * rnorm_B² / (rnorm_Bx * rnorm_A²) * mc_factor_D + # mc_factor_D = N_C * N_Ax * N_B² / (N_Bx * N_A² * N_D) + mc_A_safe = tf.where(mc_A == 0.0, tf.ones_like(mc_A), mc_A) + mc_Bx_safe = tf.where(mc_Bx == 0.0, tf.ones_like(mc_Bx), mc_Bx) + mc_D_safe = tf.where(mc_D == 0.0, tf.ones_like(mc_D), mc_D) + self.mc_factor_D = ( + mc_C * mc_Ax * mc_B**2 / (mc_Bx_safe * mc_A_safe**2 * mc_D_safe) + ) + + # Scatter indices for compute() + self._idx = { + "A": [[int(i), proc_idx] for i in idx_A], + "B": [[int(i), proc_idx] for i in idx_B], + "C": [[int(i), proc_idx] for i in idx_C], + "D": [[int(i), proc_idx] for i in idx_D], + "Ax": [[int(i), proc_idx] for i in idx_Ax], + "Bx": [[int(i), proc_idx] for i in idx_Bx], + } + + # Per-region activity flags for non-full mode + self._active_nonfull = { + r: bool(all(int(i) < indata.nbins for i in idxs)) + for r, idxs in zip( + ["A", "B", "C", "D", "Ax", "Bx"], + [idx_A, idx_B, idx_C, idx_D, idx_Ax, idx_Bx], + ) + } + + # Parameter names: A, B, C, Ax, Bx (5 regions × n_bins) + names = [] + for ch_dict, region_label in [ + (channel_A, "A"), + (channel_B, "B"), + (channel_C, "C"), + (channel_Ax, "Ax"), + (channel_Bx, "Bx"), + ]: + ch_name, axis_sel = next(iter(ch_dict.items())) + ch_axes = indata.channel_info[ch_name]["axes"] + ch_shape = [a.size for a in ch_axes] + names.extend( + _build_param_names(abcd_process, ch_name, ch_axes, axis_sel, ch_shape) + ) + + # Model attributes + self.nparams = 5 * n + self.npoi = 0 + self.nnui = 5 * n + self.params = np.array(names) + self.is_linear = False + self.allowNegativeParam = False + self.xparamdefault = tf.ones([5 * n], dtype=indata.dtype) + + @classmethod + def parse_args(cls, indata, *args, **kwargs): + """Parse CLI arguments for ExtendedABCD. + + Syntax: + --paramModel ExtendedABCD \\ + [ax:val ...] [ax:val ...] \\ + [ax:val ...] [ax:val ...] \\ + [ax:val ...] [ax:val ...] + """ + if len(args) < 7: + raise ValueError( + "ExtendedABCD expects: process " + "ch_Ax [ax:val ...] ch_Bx [ax:val ...] " + "ch_A [ax:val ...] ch_B [ax:val ...] " + "ch_C [ax:val ...] ch_D [ax:val ...]" + ) + tokens = list(args) + process = tokens.pop(0) + + def _parse_one_region(tokens, region_name): + if not tokens or ":" in tokens[0]: + raise ValueError( + f"Expected channel name for region {region_name}, " + f"got '{tokens[0] if tokens else ''}'" + ) + ch_name = tokens.pop(0) + axis_sel = {} + while tokens and ":" in tokens[0]: + ax, val = tokens.pop(0).split(":", 1) + axis_sel[ax] = int(val) + return {ch_name: axis_sel} + + channel_Ax = _parse_one_region(tokens, "Ax") + channel_Bx = _parse_one_region(tokens, "Bx") + channel_A = _parse_one_region(tokens, "A") + channel_B = _parse_one_region(tokens, "B") + channel_C = _parse_one_region(tokens, "C") + channel_D = _parse_one_region(tokens, "D") + + if tokens: + raise ValueError(f"Unexpected extra arguments after channel_D: {tokens}") + + return cls( + indata, + process, + channel_A, + channel_B, + channel_C, + channel_D, + channel_Ax, + channel_Bx, + **kwargs, + ) + + def compute(self, param, full=False): + """Compute per-bin, per-process yield scale factors. + + Args: + param: 1D tensor of length 5*n_bins (raw fit vars for a, b, c, ax, bx). + full: if True return [nbinsfull, nproc]; else [nbins, nproc]. + """ + n = self.n_bins + # Enforce positivity via squaring; raw param=1 → physical value=1 + a = tf.square(param[0 * n : 1 * n]) + b = tf.square(param[1 * n : 2 * n]) + c = tf.square(param[2 * n : 3 * n]) + ax = tf.square(param[3 * n : 4 * n]) + bx = tf.square(param[4 * n :]) + + # D = C * Ax * B² / (Bx * A²) * mc_factor_D + a_safe = tf.where(a == 0.0, tf.ones_like(a) * 1e-10, a) + bx_safe = tf.where(bx == 0.0, tf.ones_like(bx) * 1e-10, bx) + d = c * ax * b**2 / (bx_safe * a_safe**2) * self.mc_factor_D + + nbins = self.indata.nbinsfull if full else self.indata.nbins + rnorm = tf.ones([nbins, self.indata.nproc], dtype=self.indata.dtype) + + # Python-level control flow resolved at @tf.function trace time + indices = [] + updates = [] + for region, vals in [ + ("A", a), + ("B", b), + ("C", c), + ("D", d), + ("Ax", ax), + ("Bx", bx), + ]: + active = True if full else self._active_nonfull[region] + if active: + indices.extend(self._idx[region]) + updates.append(vals) + + if indices: + rnorm = tf.tensor_scatter_nd_update( + rnorm, + tf.constant(indices, dtype=tf.int64), + tf.concat(updates, axis=0), + ) + return rnorm diff --git a/rabbit/param_models/helpers.py b/rabbit/param_models/helpers.py index c5ef36ee..2acb4bd4 100644 --- a/rabbit/param_models/helpers.py +++ b/rabbit/param_models/helpers.py @@ -6,6 +6,7 @@ "Mu": "param_model", "Mixture": "param_model", "ABCD": "abcd_model", + "ExtendedABCD": "extended_abcd_model", "SmoothABCD": "smooth_abcd_model", } diff --git a/tests/make_extended_abcd_tensor.py b/tests/make_extended_abcd_tensor.py new file mode 100644 index 00000000..f66c5083 --- /dev/null +++ b/tests/make_extended_abcd_tensor.py @@ -0,0 +1,120 @@ +""" +Create a test HDF5 tensor with six ExtendedABCD regions. + +Layout (x=isolation bins, y=mt bins): + ch_Ax: extra-sideband isolation, low mt (x=extra-sideband, y=signal) + ch_Bx: extra-sideband isolation, high mt (x=extra-sideband, y=sideband) + ch_A: sideband isolation, low mt (x=sideband, y=signal) + ch_B: sideband isolation, high mt (x=sideband, y=sideband) + ch_C: signal isolation, low mt (x=signal, y=signal) + ch_D: signal isolation, high mt (x=signal, y=sideband) ← predicted + +The ExtendedABCD formula: D = C * Ax * B² / (Bx * A²). + +A signal process lives in all six channels (with small yield in control regions) +to allow a Mu POI alongside the background model. A nonprompt background lives +in all six channels. +""" + +import argparse +import os + +import hist +import numpy as np + +from rabbit import tensorwriter + +parser = argparse.ArgumentParser() +parser.add_argument("-o", "--output", default="./", help="output directory") +parser.add_argument( + "--outname", default="test_extended_abcd_tensor", help="output file name" +) +args = parser.parse_args() + +rng = np.random.default_rng(42) + +n_pt = 5 # number of pt bins +ax_pt = hist.axis.Regular(n_pt, 0, 50, name="pt") + + +def make_hist(n_events, mean_weight): + """Histogram filled with uniform pt values and given mean weight.""" + h = hist.Hist(ax_pt, storage=hist.storage.Weight()) + pt_vals = rng.uniform(0, 50, n_events) + weights = rng.normal(mean_weight, 0.1 * mean_weight, n_events) + h.fill(pt=pt_vals, weight=weights) + return h + + +def make_data(sig_h, bkg_h): + """Poisson-fluctuated data from signal + background templates.""" + h = hist.Hist(ax_pt, storage=hist.storage.Double()) + sig_vals = sig_h.values() if sig_h is not None else 0.0 + expected = sig_vals + bkg_h.values() + h.view()[:] = rng.poisson(np.maximum(expected, 0)) + return h + + +# Signal process: present in all channels (small in control regions) +h_sig_Ax = make_hist(200, 0.05) +h_sig_Bx = make_hist(200, 0.05) +h_sig_A = make_hist(500, 0.1) +h_sig_B = make_hist(500, 0.1) +h_sig_C = make_hist(5000, 1.0) +h_sig_D = make_hist(5000, 1.0) + +# Nonprompt background: set up so ExtendedABCD relation holds approximately +# f(x) = B/A is the "fake rate" per bin. +# f_extra = Bx/Ax (extra sideband) +# f_sideband = B/A +# f_signal = f_sideband² / f_extra = B²*Ax / (A²*Bx) +# D ≈ C * f_signal = C * B² * Ax / (A² * Bx) +h_np_Ax = make_hist(10000, 3.0) # extra sideband +h_np_Bx = make_hist(10000, 9.0) # Bx ~ 3x Ax → f_extra = 3 +h_np_A = make_hist(8000, 2.0) +h_np_B = make_hist(8000, 5.0) # B ~ 2.5x A → f_sb = 2.5 +h_np_C = make_hist(6000, 1.5) +# D ~ C * f_signal = C * 2.5² / 3 * correction ≈ C * 2.08 +h_np_D = make_hist(6000, 3.1) + +# Poisson-fluctuated data +h_data_Ax = make_data(h_sig_Ax, h_np_Ax) +h_data_Bx = make_data(h_sig_Bx, h_np_Bx) +h_data_A = make_data(h_sig_A, h_np_A) +h_data_B = make_data(h_sig_B, h_np_B) +h_data_C = make_data(h_sig_C, h_np_C) +h_data_D = make_data(h_sig_D, h_np_D) + +# --- Build tensor --- +writer = tensorwriter.TensorWriter(sparse=False) + +for ch in ["ch_Ax", "ch_Bx", "ch_A", "ch_B", "ch_C", "ch_D"]: + writer.add_channel(h_sig_C.axes, ch) + +# Data +writer.add_data(h_data_Ax, "ch_Ax") +writer.add_data(h_data_Bx, "ch_Bx") +writer.add_data(h_data_A, "ch_A") +writer.add_data(h_data_B, "ch_B") +writer.add_data(h_data_C, "ch_C") +writer.add_data(h_data_D, "ch_D") + +# Signal process (all channels, dominant in C and D) +writer.add_process(h_sig_Ax, "signal", "ch_Ax", signal=True) +writer.add_process(h_sig_Bx, "signal", "ch_Bx", signal=True) +writer.add_process(h_sig_A, "signal", "ch_A", signal=True) +writer.add_process(h_sig_B, "signal", "ch_B", signal=True) +writer.add_process(h_sig_C, "signal", "ch_C", signal=True) +writer.add_process(h_sig_D, "signal", "ch_D", signal=True) + +# Nonprompt background (all six channels) +writer.add_process(h_np_Ax, "nonprompt", "ch_Ax", signal=False) +writer.add_process(h_np_Bx, "nonprompt", "ch_Bx", signal=False) +writer.add_process(h_np_A, "nonprompt", "ch_A", signal=False) +writer.add_process(h_np_B, "nonprompt", "ch_B", signal=False) +writer.add_process(h_np_C, "nonprompt", "ch_C", signal=False) +writer.add_process(h_np_D, "nonprompt", "ch_D", signal=False) + +outdir = os.path.join(args.output, args.outname) +writer.write(outdir) +print(f"Written to {outdir}/rabbit_input.hdf5") From 00c46f1d2c5bc0094b1c11ec97fa51cd5e69344a Mon Sep 17 00:00:00 2001 From: davidwalter2 Date: Sat, 11 Apr 2026 14:34:23 -0400 Subject: [PATCH 06/10] Add SmoothExtendedABCD param model MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Like ExtendedABCD but all five free-parameter regions (A, B, C, Ax, Bx) are parameterised with an exponential polynomial along one smoothing axis: val_X(x) = exp(p_0 + p_1·x̃ + ...) Reduces parameters from 5·n_bins to 5·n_outer·(order+1). D is still derived from D = C·Ax·B² / (Bx·A²). CLI: SmoothExtendedABCD [order:N] . Added to CI abcd-fits job. Co-Authored-By: Claude Sonnet 4.6 --- .github/workflows/main.yml | 8 + CLAUDE.md | 1 + rabbit/param_models/helpers.py | 1 + .../smooth_extended_abcd_model.py | 386 ++++++++++++++++++ 4 files changed, 396 insertions(+) create mode 100644 rabbit/param_models/smooth_extended_abcd_model.py diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index b71bbd7f..b25b3323 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -607,6 +607,14 @@ jobs: --paramModel ExtendedABCD nonprompt ch_Ax ch_Bx ch_A ch_B ch_C ch_D --saveHists + - name: SmoothExtendedABCD fit + run: >- + rabbit_fit.py $RABBIT_OUTDIR/test_extended_abcd_tensor/rabbit_input.hdf5 + -o $RABBIT_OUTDIR/ --postfix smooth_extended_abcd + -t 0 --paramModel Mu + --paramModel SmoothExtendedABCD pt nonprompt ch_Ax ch_Bx ch_A ch_B ch_C ch_D + --saveHists + copy-clean: runs-on: [self-hosted, linux, x64] needs: [setenv, symmerizations, sparse-fits, alternative-fits, regularization, bsm, plotting, likelihoodscans, abcd-fits] diff --git a/CLAUDE.md b/CLAUDE.md index c696de5c..83960d2e 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -83,6 +83,7 @@ Built-in models: - `ABCD`: data-driven background estimation using four regions; `npoi=0`, `nnui=3*n_bins`. CLI: `--paramModel ABCD [ax:val ...] [ax:val ...] [ax:val ...] [ax:val ...]` where `ax:val` pairs optionally select a single bin along a named axis (e.g. `iso:0`). Regions A, B, C have free parameters; D is predicted as `a*c/b` times an MC correction factor. - `SmoothABCD`: like ABCD but one axis is parameterised with an exponential polynomial `val(x)=exp(p_0+p_1·x̃+...)` instead of per-bin free parameters. CLI: `--paramModel SmoothABCD [order:N] ... `. Default order=1 (log-linear). Reduces parameters from `3·n_bins` to `3·n_outer·(order+1)`. - `ExtendedABCD`: 6-region ABCD using two sideband bins in the x direction (Ax, Bx further from signal, A/B in the middle). Fake rate is log-linearly extrapolated: `D = C·Ax·B² / (Bx·A²)`. `npoi=0`, `nnui=5·n_bins`. CLI: `--paramModel ExtendedABCD [ax:val ...] [ax:val ...] [ax:val ...] [ax:val ...] [ax:val ...] [ax:val ...]`. +- `SmoothExtendedABCD`: like `ExtendedABCD` but all five free-parameter regions (A, B, C, Ax, Bx) are parameterised with an exponential polynomial along one smoothing axis. `npoi=0`, `nnui=5·n_outer·(order+1)`. CLI: `--paramModel SmoothExtendedABCD [order:N] [ax:val ...] [ax:val ...] [ax:val ...] [ax:val ...] [ax:val ...] [ax:val ...]`. Custom models are loaded by providing a dotted Python path (e.g. `--paramModel mymod.MyModel`); the module must be on `$PYTHONPATH` with an `__init__.py`. diff --git a/rabbit/param_models/helpers.py b/rabbit/param_models/helpers.py index 2acb4bd4..0f843028 100644 --- a/rabbit/param_models/helpers.py +++ b/rabbit/param_models/helpers.py @@ -8,6 +8,7 @@ "ABCD": "abcd_model", "ExtendedABCD": "extended_abcd_model", "SmoothABCD": "smooth_abcd_model", + "SmoothExtendedABCD": "smooth_extended_abcd_model", } diff --git a/rabbit/param_models/smooth_extended_abcd_model.py b/rabbit/param_models/smooth_extended_abcd_model.py new file mode 100644 index 00000000..ff406098 --- /dev/null +++ b/rabbit/param_models/smooth_extended_abcd_model.py @@ -0,0 +1,386 @@ +""" +SmoothExtendedABCD background estimation param model. + +Like ExtendedABCD, but the per-bin free parameters along one nominated smoothing +axis are replaced by an exponential polynomial: + + val_X(x) = exp(p_0 + p_1·x̃ + p_2·x̃² + ...) + +where x̃ are normalised bin centres of the smoothing axis in [0, 1]. + +This applies to all five free-parameter regions (A, B, C, Ax, Bx), reducing +parameters from 5·n_bins to 5·n_outer·(order+1). + + y=signal y=sideband +x=extra-sideband: Ax Bx +x=sideband: A B +x=signal: C D ← predicted + +D = C * Ax * B² / (Bx * A²) +""" + +import itertools + +import numpy as np +import tensorflow as tf + +from rabbit.param_models.abcd_model import _get_global_indices +from rabbit.param_models.param_model import ParamModel + + +class SmoothExtendedABCD(ParamModel): + """ + Smooth ExtendedABCD background estimation model with 6 regions. + + Regions A, B, C, Ax, Bx have free parameters that follow an exponential + polynomial along the smoothing axis; region D is derived: + D = C * Ax * B² / (Bx * A²) * mc_factor_D. + + Parameters are pure model nuisances (npoi=0, nnui=5·n_outer·(order+1)). + + CLI syntax: + --paramModel SmoothExtendedABCD [order:N] \\ + [ax:val ...] [ax:val ...] \\ + [ax:val ...] [ax:val ...] \\ + [ax:val ...] [ax:val ...] + + Python constructor: + SmoothExtendedABCD(indata, "pt", "nonprompt", + channel_A={"ch_fakes": {"iso": 1}}, + channel_B={"ch_fakes": {"iso": 2}}, + channel_C={"ch_C": {}}, + channel_D={"ch_D": {}}, + channel_Ax={"ch_fakes": {"iso": 0}}, + channel_Bx={"ch_fakes": {"iso": 3}}, + order=1) + """ + + def __init__( + self, + indata, + smoothing_axis, + abcd_process, + channel_A, + channel_B, + channel_C, + channel_D, + channel_Ax, + channel_Bx, + order=1, + **kwargs, + ): + """ + Args: + indata: FitInputData instance. + smoothing_axis: name of the axis to parameterise smoothly. + abcd_process: name of the background process (str). + channel_A/B/C/D: dicts {ch_name: {axis_name: bin_idx, ...}}. + channel_Ax/Bx: extra sideband regions (further from signal than A/B). + order: polynomial degree in the exponent (default 1 = log-linear). + """ + self.indata = indata + self.order = order + + # Validate process + proc_name = ( + abcd_process.encode() if isinstance(abcd_process, str) else abcd_process + ) + if proc_name not in indata.procs: + raise ValueError( + f"Process '{abcd_process}' not found. Available: {list(indata.procs)}" + ) + proc_idx = int(np.where(np.array(indata.procs) == proc_name)[0][0]) + + # Build outer × smooth index structure for all six regions + all_regions = ["A", "B", "C", "D", "Ax", "Bx"] + all_channels = [ + channel_A, + channel_B, + channel_C, + channel_D, + channel_Ax, + channel_Bx, + ] + + outer_global_indices = {} + outer_axes_by_region = {} + smooth_size = None + + for region, ch_dict in zip(all_regions, all_channels): + ch_name, axis_sel = next(iter(ch_dict.items())) + axes = indata.channel_info[ch_name]["axes"] + + remaining = [a for a in axes if a.name not in axis_sel] + + smooth_ax = next((a for a in remaining if a.name == smoothing_axis), None) + if smooth_ax is None: + raise ValueError( + f"Smoothing axis '{smoothing_axis}' not found among free axes " + f"of channel '{ch_name}' for region {region}. " + f"Free axes: {[a.name for a in remaining]}" + ) + + n_smooth = smooth_ax.size + if smooth_size is None: + smooth_size = n_smooth + elif smooth_size != n_smooth: + raise ValueError( + f"Smoothing axis '{smoothing_axis}' has inconsistent size: " + f"expected {smooth_size}, got {n_smooth} in region {region}" + ) + + outer_axes = [a for a in remaining if a.name != smoothing_axis] + outer_axes_by_region[region] = outer_axes + + region_outer_indices = [] + for outer_combo in itertools.product(*[range(a.size) for a in outer_axes]): + extended_sel = dict(axis_sel) + for ax, idx in zip(outer_axes, outer_combo): + extended_sel[ax.name] = idx + region_outer_indices.append( + _get_global_indices(indata, {ch_name: extended_sel}) + ) + + outer_global_indices[region] = region_outer_indices + + n_smooth = smooth_size + n_outer = len(outer_global_indices["A"]) + self.n_outer = n_outer + self.n_smooth = n_smooth + + for region in ["B", "C", "D", "Ax", "Bx"]: + if len(outer_global_indices[region]) != n_outer: + raise ValueError( + f"Region {region} has {len(outer_global_indices[region])} outer bins, " + f"expected {n_outer}" + ) + + # Vandermonde matrix from normalised bin centres of the smooth axis + ch_name_A, axis_sel_A = next(iter(channel_A.items())) + axes_A = indata.channel_info[ch_name_A]["axes"] + smooth_ax_A = next(a for a in axes_A if a.name == smoothing_axis) + x_centers = np.array(smooth_ax_A.centers) + x_min, x_max = float(x_centers.min()), float(x_centers.max()) + x_norm = ( + (x_centers - x_min) / (x_max - x_min) + if x_max > x_min + else np.zeros_like(x_centers) + ) + # V[k, s] = x_norm[s]^k, shape [order+1, n_smooth] + vander = np.array( + [[float(x_norm[s]) ** k for s in range(n_smooth)] for k in range(order + 1)] + ) + self.vander = tf.constant(vander, dtype=indata.dtype) + + # MC correction factor for region D, shape [n_outer * n_smooth] + norm_dense = tf.sparse.to_dense(indata.norm) if indata.sparse else indata.norm + norm_proc = tf.cast(norm_dense[:, proc_idx], dtype=indata.dtype) + + def _gather_flat(region_outer_idx): + return tf.gather(norm_proc, np.concatenate(region_outer_idx)) + + mc_A = _gather_flat(outer_global_indices["A"]) + mc_B = _gather_flat(outer_global_indices["B"]) + mc_C = _gather_flat(outer_global_indices["C"]) + mc_D = _gather_flat(outer_global_indices["D"]) + mc_Ax = _gather_flat(outer_global_indices["Ax"]) + mc_Bx = _gather_flat(outer_global_indices["Bx"]) + + mc_A_safe = tf.where(mc_A == 0.0, tf.ones_like(mc_A), mc_A) + mc_Bx_safe = tf.where(mc_Bx == 0.0, tf.ones_like(mc_Bx), mc_Bx) + mc_D_safe = tf.where(mc_D == 0.0, tf.ones_like(mc_D), mc_D) + self.mc_factor_D = ( + mc_C * mc_Ax * mc_B**2 / (mc_Bx_safe * mc_A_safe**2 * mc_D_safe) + ) + + # Scatter indices in [bin, proc] format, ordered outer × smooth + def _build_scatter_idx(region_outer_idx): + return [ + [int(region_outer_idx[m][s]), proc_idx] + for m in range(n_outer) + for s in range(n_smooth) + ] + + self._idx = { + r: _build_scatter_idx(outer_global_indices[r]) for r in all_regions + } + + # Per-region activity flags for non-full mode + self._active_nonfull = { + r: bool( + all( + int(outer_global_indices[r][m][s]) < indata.nbins + for m in range(n_outer) + for s in range(n_smooth) + ) + ) + for r in all_regions + } + + # Parameter names: region × outer_bin × poly_coeff (for the 5 free regions) + free_regions = ["A", "B", "C", "Ax", "Bx"] + free_channels = [channel_A, channel_B, channel_C, channel_Ax, channel_Bx] + names = [] + for region, ch_dict in zip(free_regions, free_channels): + ch_name, _ = next(iter(ch_dict.items())) + outer_axes = outer_axes_by_region[region] + combos = list(itertools.product(*[range(a.size) for a in outer_axes])) or [ + () + ] + for outer_combo in combos: + if outer_axes: + outer_label = "_".join( + f"{a.name}{i}" for a, i in zip(outer_axes, outer_combo) + ) + else: + outer_label = "outer0" + for k in range(order + 1): + names.append( + f"{abcd_process}_{ch_name}_{region}_{outer_label}_poly{k}".encode() + ) + + # Model attributes + self.nparams = 5 * n_outer * (order + 1) + self.npoi = 0 + self.nnui = self.nparams + self.params = np.array(names) + self.is_linear = False + self.allowNegativeParam = False + # Default: all coefficients 0 → exp(0) = 1 (MC template as starting point) + self.xparamdefault = tf.zeros([self.nparams], dtype=indata.dtype) + + @classmethod + def parse_args(cls, indata, *args, **kwargs): + """Parse CLI arguments for SmoothExtendedABCD. + + Syntax: + --paramModel SmoothExtendedABCD [order:N] \\ + [ax:val ...] [ax:val ...] \\ + [ax:val ...] [ax:val ...] \\ + [ax:val ...] [ax:val ...] + """ + if len(args) < 8: + raise ValueError( + "SmoothExtendedABCD expects: axis [order:N] process " + "ch_Ax [ax:val ...] ch_Bx [ax:val ...] " + "ch_A [ax:val ...] ch_B [ax:val ...] " + "ch_C [ax:val ...] ch_D [ax:val ...]" + ) + tokens = list(args) + smoothing_axis = tokens.pop(0) + + order = 1 + if tokens and tokens[0].startswith("order:"): + order = int(tokens.pop(0).split(":", 1)[1]) + + if not tokens: + raise ValueError( + "SmoothExtendedABCD: expected process name after axis/order" + ) + process = tokens.pop(0) + + def _parse_one_region(tokens, region_name): + if not tokens or ":" in tokens[0]: + raise ValueError( + f"Expected channel name for region {region_name}, " + f"got '{tokens[0] if tokens else ''}'" + ) + ch_name = tokens.pop(0) + axis_sel = {} + while tokens and ":" in tokens[0]: + ax, val = tokens.pop(0).split(":", 1) + axis_sel[ax] = int(val) + return {ch_name: axis_sel} + + channel_Ax = _parse_one_region(tokens, "Ax") + channel_Bx = _parse_one_region(tokens, "Bx") + channel_A = _parse_one_region(tokens, "A") + channel_B = _parse_one_region(tokens, "B") + channel_C = _parse_one_region(tokens, "C") + channel_D = _parse_one_region(tokens, "D") + + if tokens: + raise ValueError(f"Unexpected extra arguments after channel_D: {tokens}") + + return cls( + indata, + smoothing_axis, + process, + channel_A, + channel_B, + channel_C, + channel_D, + channel_Ax, + channel_Bx, + order=order, + **kwargs, + ) + + def compute(self, param, full=False): + """Compute per-bin, per-process yield scale factors. + + Args: + param: 1D tensor of length 5 * n_outer * (order+1). + full: if True return [nbinsfull, nproc]; else [nbins, nproc]. + """ + n_coeffs = self.n_outer * (self.order + 1) + + # Reshape to [n_outer, order+1] for each region + params_A = tf.reshape( + param[0 * n_coeffs : 1 * n_coeffs], [self.n_outer, self.order + 1] + ) + params_B = tf.reshape( + param[1 * n_coeffs : 2 * n_coeffs], [self.n_outer, self.order + 1] + ) + params_C = tf.reshape( + param[2 * n_coeffs : 3 * n_coeffs], [self.n_outer, self.order + 1] + ) + params_Ax = tf.reshape( + param[3 * n_coeffs : 4 * n_coeffs], [self.n_outer, self.order + 1] + ) + params_Bx = tf.reshape(param[4 * n_coeffs :], [self.n_outer, self.order + 1]) + + # Evaluate exp(polynomial): [n_outer, n_smooth] + a = tf.exp(tf.matmul(params_A, self.vander)) + b = tf.exp(tf.matmul(params_B, self.vander)) + c = tf.exp(tf.matmul(params_C, self.vander)) + ax = tf.exp(tf.matmul(params_Ax, self.vander)) + bx = tf.exp(tf.matmul(params_Bx, self.vander)) + + # Flatten to [n_outer * n_smooth] = [n_total_bins per region] + a_flat = tf.reshape(a, [-1]) + b_flat = tf.reshape(b, [-1]) + c_flat = tf.reshape(c, [-1]) + ax_flat = tf.reshape(ax, [-1]) + bx_flat = tf.reshape(bx, [-1]) + + a_safe = tf.where(a_flat == 0.0, tf.ones_like(a_flat) * 1e-10, a_flat) + bx_safe = tf.where(bx_flat == 0.0, tf.ones_like(bx_flat) * 1e-10, bx_flat) + d_flat = c_flat * ax_flat * b_flat**2 / (bx_safe * a_safe**2) * self.mc_factor_D + + nbins = self.indata.nbinsfull if full else self.indata.nbins + rnorm = tf.ones([nbins, self.indata.nproc], dtype=self.indata.dtype) + + # Python-level control flow resolved at @tf.function trace time + indices = [] + updates = [] + for region, vals in [ + ("A", a_flat), + ("B", b_flat), + ("C", c_flat), + ("D", d_flat), + ("Ax", ax_flat), + ("Bx", bx_flat), + ]: + active = True if full else self._active_nonfull[region] + if active: + indices.extend(self._idx[region]) + updates.append(vals) + + if indices: + rnorm = tf.tensor_scatter_nd_update( + rnorm, + tf.constant(indices, dtype=tf.int64), + tf.concat(updates, axis=0), + ) + return rnorm From 825e7ae2dbd55483b40345b6870db2ffa427586d Mon Sep 17 00:00:00 2001 From: davidwalter2 Date: Sat, 11 Apr 2026 15:07:02 -0400 Subject: [PATCH 07/10] Fix saturated chi2 ndf to subtract all free param model parameters Previously only npoi was subtracted; nnui (unconstrained model nuisances such as per-bin fake rates from ABCD) was not. For ABCD with npoi=0 the ndf was not reduced at all despite having nnui=3*n_bins free parameters. Use nparams (= npoi + nnui) instead. For the traditional Mu model nnui=0 so behaviour is unchanged. Co-Authored-By: Claude Sonnet 4.6 --- bin/rabbit_fit.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/bin/rabbit_fit.py b/bin/rabbit_fit.py index a7dd740c..d637c0ac 100755 --- a/bin/rabbit_fit.py +++ b/bin/rabbit_fit.py @@ -449,7 +449,9 @@ def fit(args, fitter, ws, dofit=True): nllvalreduced = fitter.reduced_nll().numpy() ndfsat = ( - tf.size(fitter.nobs) - fitter.param_model.npoi - fitter.indata.nsystnoconstraint + tf.size(fitter.nobs) + - fitter.param_model.nparams + - fitter.indata.nsystnoconstraint ).numpy() chi2_val = 2.0 * nllvalreduced From a7e89191322369d5a4f592b6a8c3e6f3eceb8392 Mon Sep 17 00:00:00 2001 From: davidwalter2 Date: Sat, 11 Apr 2026 15:15:47 -0400 Subject: [PATCH 08/10] =?UTF-8?q?Rename=20nnui=20=E2=86=92=20npou=20(param?= =?UTF-8?q?eters=20of=20uninterest)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit nnui stood for "number of nuisances"; npou stands for "parameters of uninterest", a clearer contrast to npoi (parameters of interest). Co-Authored-By: Claude Sonnet 4.6 --- CLAUDE.md | 8 ++++---- rabbit/fitter.py | 6 +++--- rabbit/impacts/nonprofiled_impacts.py | 2 +- rabbit/param_models/abcd_model.py | 4 ++-- rabbit/param_models/extended_abcd_model.py | 4 ++-- rabbit/param_models/param_model.py | 16 ++++++++-------- rabbit/param_models/smooth_abcd_model.py | 4 ++-- .../param_models/smooth_extended_abcd_model.py | 4 ++-- 8 files changed, 24 insertions(+), 24 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index 83960d2e..7a4ca401 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -74,16 +74,16 @@ Supports dense and sparse tensor representations, symmetric/asymmetric systemati Base class `ParamModel` defines `compute(param)` which returns a `[1, nproc]` tensor scaling process yields. Each model declares: - `nparams`: total parameters (POIs + model nuisances) - `npoi`: true parameters of interest (first `npoi` entries; reported as POIs in outputs) -- `nnui`: model nuisance parameters (`nparams - npoi`; reported as nuisances in outputs) +- `npou`: model nuisance parameters (`nparams - npoi`; reported as nuisances in outputs) Built-in models: - `Mu` (default): one POI per signal process - `Ones`: no parameters (all yields fixed to MC) - `Mixture`: mixing POIs between pairs of processes -- `ABCD`: data-driven background estimation using four regions; `npoi=0`, `nnui=3*n_bins`. CLI: `--paramModel ABCD [ax:val ...] [ax:val ...] [ax:val ...] [ax:val ...]` where `ax:val` pairs optionally select a single bin along a named axis (e.g. `iso:0`). Regions A, B, C have free parameters; D is predicted as `a*c/b` times an MC correction factor. +- `ABCD`: data-driven background estimation using four regions; `npoi=0`, `npou=3*n_bins`. CLI: `--paramModel ABCD [ax:val ...] [ax:val ...] [ax:val ...] [ax:val ...]` where `ax:val` pairs optionally select a single bin along a named axis (e.g. `iso:0`). Regions A, B, C have free parameters; D is predicted as `a*c/b` times an MC correction factor. - `SmoothABCD`: like ABCD but one axis is parameterised with an exponential polynomial `val(x)=exp(p_0+p_1·x̃+...)` instead of per-bin free parameters. CLI: `--paramModel SmoothABCD [order:N] ... `. Default order=1 (log-linear). Reduces parameters from `3·n_bins` to `3·n_outer·(order+1)`. -- `ExtendedABCD`: 6-region ABCD using two sideband bins in the x direction (Ax, Bx further from signal, A/B in the middle). Fake rate is log-linearly extrapolated: `D = C·Ax·B² / (Bx·A²)`. `npoi=0`, `nnui=5·n_bins`. CLI: `--paramModel ExtendedABCD [ax:val ...] [ax:val ...] [ax:val ...] [ax:val ...] [ax:val ...] [ax:val ...]`. -- `SmoothExtendedABCD`: like `ExtendedABCD` but all five free-parameter regions (A, B, C, Ax, Bx) are parameterised with an exponential polynomial along one smoothing axis. `npoi=0`, `nnui=5·n_outer·(order+1)`. CLI: `--paramModel SmoothExtendedABCD [order:N] [ax:val ...] [ax:val ...] [ax:val ...] [ax:val ...] [ax:val ...] [ax:val ...]`. +- `ExtendedABCD`: 6-region ABCD using two sideband bins in the x direction (Ax, Bx further from signal, A/B in the middle). Fake rate is log-linearly extrapolated: `D = C·Ax·B² / (Bx·A²)`. `npoi=0`, `npou=5·n_bins`. CLI: `--paramModel ExtendedABCD [ax:val ...] [ax:val ...] [ax:val ...] [ax:val ...] [ax:val ...] [ax:val ...]`. +- `SmoothExtendedABCD`: like `ExtendedABCD` but all five free-parameter regions (A, B, C, Ax, Bx) are parameterised with an exponential polynomial along one smoothing axis. `npoi=0`, `npou=5·n_outer·(order+1)`. CLI: `--paramModel SmoothExtendedABCD [order:N] [ax:val ...] [ax:val ...] [ax:val ...] [ax:val ...] [ax:val ...] [ax:val ...]`. Custom models are loaded by providing a dotted Python path (e.g. `--paramModel mymod.MyModel`); the module must be on `$PYTHONPATH` with an `__init__.py`. diff --git a/rabbit/fitter.py b/rabbit/fitter.py index 4930c499..d1a055fb 100644 --- a/rabbit/fitter.py +++ b/rabbit/fitter.py @@ -522,8 +522,8 @@ def get_theta(self): def get_model_nui(self): npoi = self.param_model.npoi - nnui = self.param_model.nnui - return self.x[npoi : npoi + nnui] + npou = self.param_model.npou + return self.x[npoi : npoi + npou] def get_poi(self): xpoi = self.x[: self.param_model.npoi] @@ -845,7 +845,7 @@ def impacts_parms(self, hess): nstat = ( self.param_model.npoi - + self.param_model.nnui + + self.param_model.npou + self.indata.nsystnoconstraint ) hess_stat = hess[:nstat, :nstat] diff --git a/rabbit/impacts/nonprofiled_impacts.py b/rabbit/impacts/nonprofiled_impacts.py index 89fc379e..650df202 100644 --- a/rabbit/impacts/nonprofiled_impacts.py +++ b/rabbit/impacts/nonprofiled_impacts.py @@ -43,7 +43,7 @@ def nonprofiled_impacts_parms( constraintweights: constraint weights for each nuisance parameter. systgroups: systematic group names. systgroupidxs: per-group lists of nuisance parameter indices. - nparams: total number of model parameters (nparams + nnui); offset from x index to theta0 index. + nparams: total number of model parameters (nparams + npou); offset from x index to theta0 index. minimize_fn: callable that runs the fit (no arguments). diagnostics: if True, log EDM after each minimization (requires loss_val_grad_hess_fn). loss_val_grad_hess_fn: callable returning (val, grad, hess); used only when diagnostics=True. diff --git a/rabbit/param_models/abcd_model.py b/rabbit/param_models/abcd_model.py index 13bfd7ad..b5f4329d 100644 --- a/rabbit/param_models/abcd_model.py +++ b/rabbit/param_models/abcd_model.py @@ -91,7 +91,7 @@ class ABCD(ParamModel): where mc_factor_D_i = N_A^MC * N_C^MC / (N_B^MC * N_D^MC) accounts for non-unity MC templates. - Parameters are pure model nuisances (npoi=0, nnui=3*n_bins). Positivity + Parameters are pure model nuisances (npoi=0, npou=3*n_bins). Positivity is enforced inside compute() via tf.square() on the raw fit variables. CLI syntax: @@ -196,7 +196,7 @@ def __init__( # Model attributes self.nparams = 3 * n self.npoi = 0 - self.nnui = 3 * n + self.npou = 3 * n self.params = np.array(names) self.is_linear = False self.allowNegativeParam = False diff --git a/rabbit/param_models/extended_abcd_model.py b/rabbit/param_models/extended_abcd_model.py index b580c439..aa10d85c 100644 --- a/rabbit/param_models/extended_abcd_model.py +++ b/rabbit/param_models/extended_abcd_model.py @@ -30,7 +30,7 @@ class ExtendedABCD(ParamModel): Free parameters a_i, b_i, c_i, ax_i, bx_i for regions A, B, C, Ax, Bx. Region D is derived: D = C * Ax * B² / (Bx * A²) * mc_factor_D. - Parameters are pure model nuisances (npoi=0, nnui=5*n_bins). Positivity + Parameters are pure model nuisances (npoi=0, npou=5*n_bins). Positivity is enforced inside compute() via tf.square() on the raw fit variables. CLI syntax: @@ -162,7 +162,7 @@ def __init__( # Model attributes self.nparams = 5 * n self.npoi = 0 - self.nnui = 5 * n + self.npou = 5 * n self.params = np.array(names) self.is_linear = False self.allowNegativeParam = False diff --git a/rabbit/param_models/param_model.py b/rabbit/param_models/param_model.py index d6dd59d2..f4596540 100644 --- a/rabbit/param_models/param_model.py +++ b/rabbit/param_models/param_model.py @@ -11,9 +11,9 @@ def __init__(self, indata, *args, **kwargs): self.indata = indata # # a param model must set these attributes - # self.nparams = # total number of parameters (npoi + nnui) + # self.nparams = # total number of parameters (npoi + npou) # self.npoi = # number of true parameters of interest (POIs), reported as POIs in outputs - # self.nnui = # number of model nuisance parameters (= nparams - npoi) + # self.npou = # number of model nuisance parameters (= nparams - npoi) # self.params = # list of names for all parameters (POIs first, then model nuisances) # self.xparamdefault = # default values for all parameters (length nparams) # self.is_linear = # define if the model is linear in the parameters @@ -35,7 +35,7 @@ def set_param_default(self, expectSignal, allowNegativeParam=False): """ Set default parameter values, used by different param models. Only the first npoi entries (true POIs) support the squaring transform; - model nuisance parameters (nnui entries) are always stored directly. + model nuisance parameters (npou entries) are always stored directly. """ paramdefault = tf.ones([self.nparams], dtype=self.indata.dtype) if expectSignal is not None: @@ -80,7 +80,7 @@ def __init__( self.nparams = sum([m.nparams for m in param_models]) self.npoi = sum([m.npoi for m in param_models]) - self.nnui = sum([m.nnui for m in param_models]) + self.npou = sum([m.npou for m in param_models]) self.params = np.concatenate([m.params for m in param_models]) @@ -110,7 +110,7 @@ def __init__(self, indata, **kwargs): self.indata = indata self.nparams = 0 self.npoi = 0 - self.nnui = 0 + self.npou = 0 self.params = np.array([]) self.xparamdefault = tf.zeros([0], dtype=self.indata.dtype) @@ -133,7 +133,7 @@ def __init__(self, indata, expectSignal=None, allowNegativeParam=False, **kwargs self.nparams = self.indata.nsignals self.npoi = self.nparams - self.nnui = 0 + self.npou = 0 self.params = np.array([s for s in self.indata.signals]) @@ -206,7 +206,7 @@ def __init__( self.nparams = len(primary_processes) self.npoi = self.nparams - self.nnui = 0 + self.npou = 0 self.params = np.array( [ f"{p}_{c}_mixing".encode() @@ -279,7 +279,7 @@ def __init__( ] ) self.npoi = self.nparams - self.nnui = 0 + self.npou = 0 names = [] for k, v in self.channel_info_mapping.items(): diff --git a/rabbit/param_models/smooth_abcd_model.py b/rabbit/param_models/smooth_abcd_model.py index 78fe4801..cb7b3515 100644 --- a/rabbit/param_models/smooth_abcd_model.py +++ b/rabbit/param_models/smooth_abcd_model.py @@ -28,7 +28,7 @@ class SmoothABCD(ParamModel): Regions A, B, C have free parameters that follow an exponential polynomial along the smoothing axis; region D is derived as d = a·c/b · mc_factor_D. - Parameters are pure model nuisances (npoi=0, nnui=3·n_outer·(order+1)). + Parameters are pure model nuisances (npoi=0, npou=3·n_outer·(order+1)). CLI syntax: --paramModel SmoothABCD [order:N] \\ @@ -216,7 +216,7 @@ def _build_scatter_idx(region_outer_idx): # Model attributes self.nparams = 3 * n_outer * (order + 1) self.npoi = 0 - self.nnui = self.nparams + self.npou = self.nparams self.params = np.array(names) self.is_linear = False self.allowNegativeParam = False diff --git a/rabbit/param_models/smooth_extended_abcd_model.py b/rabbit/param_models/smooth_extended_abcd_model.py index ff406098..5b99c87a 100644 --- a/rabbit/param_models/smooth_extended_abcd_model.py +++ b/rabbit/param_models/smooth_extended_abcd_model.py @@ -36,7 +36,7 @@ class SmoothExtendedABCD(ParamModel): polynomial along the smoothing axis; region D is derived: D = C * Ax * B² / (Bx * A²) * mc_factor_D. - Parameters are pure model nuisances (npoi=0, nnui=5·n_outer·(order+1)). + Parameters are pure model nuisances (npoi=0, npou=5·n_outer·(order+1)). CLI syntax: --paramModel SmoothExtendedABCD [order:N] \\ @@ -242,7 +242,7 @@ def _build_scatter_idx(region_outer_idx): # Model attributes self.nparams = 5 * n_outer * (order + 1) self.npoi = 0 - self.nnui = self.nparams + self.npou = self.nparams self.params = np.array(names) self.is_linear = False self.allowNegativeParam = False From 89fb7aff05cc8f38f6fc0a7f336d93647f653b18 Mon Sep 17 00:00:00 2001 From: davidwalter2 Date: Sat, 11 Apr 2026 15:24:44 -0400 Subject: [PATCH 09/10] Make nparams a @property (npoi + npou) on ParamModel base class Enforces the invariant nparams = npoi + npou without redundant manual assignment. Removes all self.nparams = ... lines from subclasses. Co-Authored-By: Claude Sonnet 4.6 --- rabbit/param_models/abcd_model.py | 1 - rabbit/param_models/extended_abcd_model.py | 1 - rabbit/param_models/param_model.py | 29 ++++++++++--------- rabbit/param_models/smooth_abcd_model.py | 3 +- .../smooth_extended_abcd_model.py | 3 +- 5 files changed, 17 insertions(+), 20 deletions(-) diff --git a/rabbit/param_models/abcd_model.py b/rabbit/param_models/abcd_model.py index b5f4329d..aec1c4e5 100644 --- a/rabbit/param_models/abcd_model.py +++ b/rabbit/param_models/abcd_model.py @@ -194,7 +194,6 @@ def __init__( ) # Model attributes - self.nparams = 3 * n self.npoi = 0 self.npou = 3 * n self.params = np.array(names) diff --git a/rabbit/param_models/extended_abcd_model.py b/rabbit/param_models/extended_abcd_model.py index aa10d85c..4afcce39 100644 --- a/rabbit/param_models/extended_abcd_model.py +++ b/rabbit/param_models/extended_abcd_model.py @@ -160,7 +160,6 @@ def __init__( ) # Model attributes - self.nparams = 5 * n self.npoi = 0 self.npou = 5 * n self.params = np.array(names) diff --git a/rabbit/param_models/param_model.py b/rabbit/param_models/param_model.py index f4596540..abd0c70c 100644 --- a/rabbit/param_models/param_model.py +++ b/rabbit/param_models/param_model.py @@ -11,14 +11,18 @@ def __init__(self, indata, *args, **kwargs): self.indata = indata # # a param model must set these attributes - # self.nparams = # total number of parameters (npoi + npou) # self.npoi = # number of true parameters of interest (POIs), reported as POIs in outputs - # self.npou = # number of model nuisance parameters (= nparams - npoi) + # self.npou = # number of model nuisance parameters (parameters of uninterest) # self.params = # list of names for all parameters (POIs first, then model nuisances) # self.xparamdefault = # default values for all parameters (length nparams) # self.is_linear = # define if the model is linear in the parameters # self.allowNegativeParam = # define if the POI parameters can be negative or not + @property + def nparams(self): + """Total number of parameters: npoi + npou.""" + return self.npoi + self.npou + # class function to parse strings as given by the argparse input e.g. --paramModel ... @classmethod def parse_args(cls, indata, *args, **kwargs): @@ -78,7 +82,6 @@ def __init__( self.param_models = param_models - self.nparams = sum([m.nparams for m in param_models]) self.npoi = sum([m.npoi for m in param_models]) self.npou = sum([m.npou for m in param_models]) @@ -108,7 +111,6 @@ class Ones(ParamModel): def __init__(self, indata, **kwargs): self.indata = indata - self.nparams = 0 self.npoi = 0 self.npou = 0 self.params = np.array([]) @@ -131,8 +133,7 @@ class Mu(ParamModel): def __init__(self, indata, expectSignal=None, allowNegativeParam=False, **kwargs): self.indata = indata - self.nparams = self.indata.nsignals - self.npoi = self.nparams + self.npoi = self.indata.nsignals self.npou = 0 self.params = np.array([s for s in self.indata.signals]) @@ -204,8 +205,7 @@ def __init__( )[0] self.all_idx = np.concatenate([self.primary_idxs, self.complementary_idxs]) - self.nparams = len(primary_processes) - self.npoi = self.nparams + self.npoi = len(primary_processes) self.npou = 0 self.params = np.array( [ @@ -272,13 +272,14 @@ def __init__( self.indata = indata self.channel_info_mapping = channel_info - self.nparams = np.sum( - [ - np.prod([a.size for a in v["axes"]]) if len(v["axes"]) else 1 - for v in channel_info.values() - ] + self.npoi = int( + np.sum( + [ + np.prod([a.size for a in v["axes"]]) if len(v["axes"]) else 1 + for v in channel_info.values() + ] + ) ) - self.npoi = self.nparams self.npou = 0 names = [] diff --git a/rabbit/param_models/smooth_abcd_model.py b/rabbit/param_models/smooth_abcd_model.py index cb7b3515..3d3c4644 100644 --- a/rabbit/param_models/smooth_abcd_model.py +++ b/rabbit/param_models/smooth_abcd_model.py @@ -214,9 +214,8 @@ def _build_scatter_idx(region_outer_idx): ) # Model attributes - self.nparams = 3 * n_outer * (order + 1) self.npoi = 0 - self.npou = self.nparams + self.npou = 3 * n_outer * (order + 1) self.params = np.array(names) self.is_linear = False self.allowNegativeParam = False diff --git a/rabbit/param_models/smooth_extended_abcd_model.py b/rabbit/param_models/smooth_extended_abcd_model.py index 5b99c87a..0f6ade4f 100644 --- a/rabbit/param_models/smooth_extended_abcd_model.py +++ b/rabbit/param_models/smooth_extended_abcd_model.py @@ -240,9 +240,8 @@ def _build_scatter_idx(region_outer_idx): ) # Model attributes - self.nparams = 5 * n_outer * (order + 1) self.npoi = 0 - self.npou = self.nparams + self.npou = 5 * n_outer * (order + 1) self.params = np.array(names) self.is_linear = False self.allowNegativeParam = False From 87feeabd11f32d321eb51c71c21a1bbab5c79527 Mon Sep 17 00:00:00 2001 From: davidwalter2 Date: Sat, 11 Apr 2026 17:03:45 -0400 Subject: [PATCH 10/10] Add IsoMT convenience wrappers for all four ABCD param models ABCDIsoMT, ExtendedABCDIsoMT, SmoothABCDIsoMT, SmoothExtendedABCDIsoMT hardcode the WRemnants nonprompt convention (y=relIso, x=mt, smoothing=pt) so the model can be specified with just a process and single channel name: --paramModel SmoothExtendedABCDIsoMT nonprompt ch_fakes Co-Authored-By: Claude Sonnet 4.6 --- rabbit/param_models/abcd_isomtmt_model.py | 209 ++++++++++++++++++++++ rabbit/param_models/helpers.py | 4 + 2 files changed, 213 insertions(+) create mode 100644 rabbit/param_models/abcd_isomtmt_model.py diff --git a/rabbit/param_models/abcd_isomtmt_model.py b/rabbit/param_models/abcd_isomtmt_model.py new file mode 100644 index 00000000..58eb5617 --- /dev/null +++ b/rabbit/param_models/abcd_isomtmt_model.py @@ -0,0 +1,209 @@ +""" +IsoMT convenience wrappers for the four ABCD param models. + +Hardcode the axis conventions used in the WRemnants nonprompt background +estimation: + - y-axis: "relIso" (bin 0 = tight/signal, bin 1 = loose/sideband) + - x-axis: "mt" (last bin = signal, earlier bins = sidebands) + - smoothing axis: "pt" (for the smooth variants only) + +All four wrappers accept a single channel name and derive the per-region +channel dicts automatically. The extended variants require at least 3 mt bins; +the plain variants require at least 2. + +CLI syntax: + --paramModel ABCDIsoMT + --paramModel ExtendedABCDIsoMT + --paramModel SmoothABCDIsoMT [order:N] + --paramModel SmoothExtendedABCDIsoMT [order:N] + +Region layout (shared by all four): + + relIso=signal(0) relIso=sideband(1) +mt=extra-sb: Ax Bx +mt=sideband: A B +mt=signal: C D ← predicted +""" + +from rabbit.param_models.abcd_model import ABCD +from rabbit.param_models.extended_abcd_model import ExtendedABCD +from rabbit.param_models.smooth_abcd_model import SmoothABCD +from rabbit.param_models.smooth_extended_abcd_model import SmoothExtendedABCD + + +def _isomtmt_regions(indata, channel_name): + """Return (mt_signal_idx, mt_sideband_idx) from the channel's mt axis size.""" + axes = indata.channel_info[channel_name]["axes"] + mt_ax = next((a for a in axes if a.name == "mt"), None) + if mt_ax is None: + raise ValueError( + f"Channel '{channel_name}' has no axis named 'mt'. " + f"Available axes: {[a.name for a in axes]}" + ) + if next((a for a in axes if a.name == "relIso"), None) is None: + raise ValueError( + f"Channel '{channel_name}' has no axis named 'relIso'. " + f"Available axes: {[a.name for a in axes]}" + ) + return mt_ax.size - 1, mt_ax.size - 2, mt_ax.size - 3 + + +def _parse_isomtmt_args(tokens): + """Parse [order:N] from a token list.""" + order = 1 + if tokens and tokens[0].startswith("order:"): + order = int(tokens.pop(0).split(":", 1)[1]) + if len(tokens) < 2: + raise ValueError("Expected ") + process = tokens.pop(0) + channel = tokens.pop(0) + if tokens: + raise ValueError(f"Unexpected extra arguments: {tokens}") + return order, process, channel + + +class ABCDIsoMT(ABCD): + """ + ABCD in the (mt × relIso) plane for a single channel. + + Uses the last two mt bins as signal / sideband and relIso bins 0 / 1 as + signal / sideband. All other axes in the channel become free per-bin + parameters (outer axes of the ABCD model). + """ + + def __init__(self, indata, abcd_process, channel_name, **kwargs): + mt_s, mt_sb, _ = _isomtmt_regions(indata, channel_name) + if mt_sb < 0: + raise ValueError(f"Channel '{channel_name}' mt axis has fewer than 2 bins") + channel_A = {channel_name: {"relIso": 1, "mt": mt_sb}} + channel_B = {channel_name: {"relIso": 0, "mt": mt_sb}} + channel_C = {channel_name: {"relIso": 1, "mt": mt_s}} + channel_D = {channel_name: {"relIso": 0, "mt": mt_s}} + super().__init__( + indata, abcd_process, channel_A, channel_B, channel_C, channel_D, **kwargs + ) + + @classmethod + def parse_args(cls, indata, *args, **kwargs): + tokens = list(args) + if len(tokens) < 2: + raise ValueError("ABCDIsoMT expects: ") + process, channel = tokens[0], tokens[1] + if len(tokens) > 2: + raise ValueError(f"Unexpected extra arguments: {tokens[2:]}") + return cls(indata, process, channel, **kwargs) + + +class ExtendedABCDIsoMT(ExtendedABCD): + """ + ExtendedABCD in the (mt × relIso) plane for a single channel. + + Uses the last three mt bins as signal / sideband / extra-sideband and + relIso bins 0 / 1 as signal / sideband. + """ + + def __init__(self, indata, abcd_process, channel_name, **kwargs): + mt_s, mt_sb, mt_xsb = _isomtmt_regions(indata, channel_name) + if mt_xsb < 0: + raise ValueError(f"Channel '{channel_name}' mt axis has fewer than 3 bins") + channel_Ax = {channel_name: {"relIso": 0, "mt": mt_xsb}} + channel_Bx = {channel_name: {"relIso": 1, "mt": mt_xsb}} + channel_A = {channel_name: {"relIso": 0, "mt": mt_sb}} + channel_B = {channel_name: {"relIso": 1, "mt": mt_sb}} + channel_C = {channel_name: {"relIso": 0, "mt": mt_s}} + channel_D = {channel_name: {"relIso": 1, "mt": mt_s}} + super().__init__( + indata, + abcd_process, + channel_A, + channel_B, + channel_C, + channel_D, + channel_Ax, + channel_Bx, + **kwargs, + ) + + @classmethod + def parse_args(cls, indata, *args, **kwargs): + tokens = list(args) + if len(tokens) < 2: + raise ValueError("ExtendedABCDIsoMT expects: ") + process, channel = tokens[0], tokens[1] + if len(tokens) > 2: + raise ValueError(f"Unexpected extra arguments: {tokens[2:]}") + return cls(indata, process, channel, **kwargs) + + +class SmoothABCDIsoMT(SmoothABCD): + """ + SmoothABCD in the (mt × relIso) plane, smoothed along the pt axis. + + Uses the last two mt bins and relIso bins 0 / 1. The pt axis is the + smoothing axis; all remaining axes become outer axes. + """ + + def __init__(self, indata, abcd_process, channel_name, order=1, **kwargs): + mt_s, mt_sb, _ = _isomtmt_regions(indata, channel_name) + if mt_sb < 0: + raise ValueError(f"Channel '{channel_name}' mt axis has fewer than 2 bins") + channel_A = {channel_name: {"relIso": 1, "mt": mt_sb}} + channel_B = {channel_name: {"relIso": 0, "mt": mt_sb}} + channel_C = {channel_name: {"relIso": 1, "mt": mt_s}} + channel_D = {channel_name: {"relIso": 0, "mt": mt_s}} + super().__init__( + indata, + "pt", + abcd_process, + channel_A, + channel_B, + channel_C, + channel_D, + order=order, + **kwargs, + ) + + @classmethod + def parse_args(cls, indata, *args, **kwargs): + tokens = list(args) + order, process, channel = _parse_isomtmt_args(tokens) + return cls(indata, process, channel, order=order, **kwargs) + + +class SmoothExtendedABCDIsoMT(SmoothExtendedABCD): + """ + SmoothExtendedABCD in the (mt × relIso) plane, smoothed along the pt axis. + + Uses the last three mt bins and relIso bins 0 / 1. The pt axis is the + smoothing axis; all remaining axes become outer axes. + """ + + def __init__(self, indata, abcd_process, channel_name, order=1, **kwargs): + mt_s, mt_sb, mt_xsb = _isomtmt_regions(indata, channel_name) + if mt_xsb < 0: + raise ValueError(f"Channel '{channel_name}' mt axis has fewer than 3 bins") + channel_Ax = {channel_name: {"relIso": 0, "mt": mt_xsb}} + channel_Bx = {channel_name: {"relIso": 1, "mt": mt_xsb}} + channel_A = {channel_name: {"relIso": 0, "mt": mt_sb}} + channel_B = {channel_name: {"relIso": 1, "mt": mt_sb}} + channel_C = {channel_name: {"relIso": 0, "mt": mt_s}} + channel_D = {channel_name: {"relIso": 1, "mt": mt_s}} + super().__init__( + indata, + "pt", + abcd_process, + channel_A, + channel_B, + channel_C, + channel_D, + channel_Ax, + channel_Bx, + order=order, + **kwargs, + ) + + @classmethod + def parse_args(cls, indata, *args, **kwargs): + tokens = list(args) + order, process, channel = _parse_isomtmt_args(tokens) + return cls(indata, process, channel, order=order, **kwargs) diff --git a/rabbit/param_models/helpers.py b/rabbit/param_models/helpers.py index 0f843028..77b7f682 100644 --- a/rabbit/param_models/helpers.py +++ b/rabbit/param_models/helpers.py @@ -9,6 +9,10 @@ "ExtendedABCD": "extended_abcd_model", "SmoothABCD": "smooth_abcd_model", "SmoothExtendedABCD": "smooth_extended_abcd_model", + "ABCDIsoMT": "abcd_isomtmt_model", + "ExtendedABCDIsoMT": "abcd_isomtmt_model", + "SmoothABCDIsoMT": "abcd_isomtmt_model", + "SmoothExtendedABCDIsoMT": "abcd_isomtmt_model", }