From 32fef325955cc6b7a1d461e9c9313c2f8ff7d3f4 Mon Sep 17 00:00:00 2001 From: Luca Lavezzo Date: Thu, 11 Jun 2026 17:48:27 -0400 Subject: [PATCH 01/16] Blinding groups and opt-in Gaussian priors on ParamModel parameters --blindingGroup: regex-defined groups of parameters that share a single deterministic blinding offset (seeded from the group regex string), so relative pulls and differences between matched parameters stay meaningful while absolute values remain blinded. Parameters matching both --unblind and --blindingGroup are a configuration error and abort the fit. match_regexp_params now returns the union of exact and regex matches (deduplicated, order-preserving) instead of short-circuiting on the first exact match, so mixed exact/regex expression lists behave consistently. --paramModelPriors: opt-in Gaussian priors on ParamModel parameters. The Fitter reads optional prior_sigmas / prior_means attributes from the ParamModel and adds the corresponding 0.5*(x-mu)^2/sigma^2 penalty to the NLL constraint term; the applied priors are stored in the output metadata. Off by default; without the flag all ParamModel params float free as before. Co-Authored-By: Claude Fable 5 --- bin/rabbit_fit.py | 11 ++ rabbit/fitter.py | 189 +++++++++++++++++++++++++---- rabbit/param_models/param_model.py | 10 ++ rabbit/parsing.py | 25 ++++ 4 files changed, 210 insertions(+), 25 deletions(-) diff --git a/bin/rabbit_fit.py b/bin/rabbit_fit.py index 7c2fb30..66eec18 100755 --- a/bin/rabbit_fit.py +++ b/bin/rabbit_fit.py @@ -823,6 +823,17 @@ def main(): "nois": ifitter.parms[ifitter.param_model.nparams :][indata.noiidxs], } + # ParamModel Gaussian priors (if --paramModelPriors was used and the + # model declared sigmas). Stored as a small dict so downstream tooling + # can see what was applied without parsing the rabbit log. + if getattr(ifitter, "param_prior_active", False): + meta["param_priors"] = { + "params": ifitter.param_model.params, # all nparams names + "mask": ifitter.param_prior_mask_np, # bool array + "sigmas": ifitter.param_prior_sigmas_np, # NaN where mask False + "means": ifitter.param_prior_means_np, # NaN where mask False + } + with workspace.Workspace( args.outpath, args.outname, diff --git a/rabbit/fitter.py b/rabbit/fitter.py index d9258ae..bc53cc1 100644 --- a/rabbit/fitter.py +++ b/rabbit/fitter.py @@ -31,23 +31,28 @@ def match_regexp_params(regular_expressions, parameter_names): + # Match parameter names against a list of expressions where each entry may + # be either an exact parameter name or a regex (re.match, anchored to the + # start). Returns the union of exact and regex matches, preserving the + # parameter_names order and de-duplicating. Mixing exact and regex entries + # in the same call is supported. if isinstance(regular_expressions, str): regular_expressions = [regular_expressions] - # Check for exact matches first - exact_matches = [ - s for expr in regular_expressions for s in parameter_names if s.decode() == expr - ] - if exact_matches: - return exact_matches - - # Fall back to regex matching + exact_lookup = set(regular_expressions) compiled_expressions = [re.compile(expr) for expr in regular_expressions] - return [ - s - for s in parameter_names - if any(regex.match(s.decode()) for regex in compiled_expressions) - ] + + matched = [] + seen = set() + for s in parameter_names: + decoded = s.decode() if hasattr(s, "decode") else s + if decoded in exact_lookup or any( + r.match(decoded) for r in compiled_expressions + ): + if decoded not in seen: + seen.add(decoded) + matched.append(s) + return matched class FitterCallback: @@ -207,7 +212,9 @@ def __init__( param_model, options.setConstraintMinimum, unblind=options.unblind, + blinding_group=getattr(options, "blindingGroup", []), freeze_parameters=options.freezeParameters, + options=options, ) self.nexpnom = tf.Variable( @@ -219,7 +226,9 @@ def init_fit_parms( param_model, set_constraint_minimum=[], unblind=False, + blinding_group=[], freeze_parameters=[], + options=None, ): self.param_model = param_model @@ -246,7 +255,7 @@ def init_fit_parms( trainable=False, name="offset_theta", ) - self.init_blinding_values(unblind) + self.init_blinding_values(unblind, blinding_group) self.parms = np.concatenate([self.param_model.params, self.indata.systs]) @@ -311,6 +320,79 @@ def init_fit_parms( self.indata.dtype, ) + # ParamModel Gaussian priors (opt-in via --paramModelPriors). + # When enabled, the Fitter reads two optional attributes from the + # ParamModel: + # - prior_sigmas: np.ndarray shape (nparams,). Entries that are + # finite and > 0 are Gaussian-constrained at that width; NaN / + # non-finite / 0 entries leave the corresponding parameter free. + # - prior_means : np.ndarray shape (nparams,), optional. + # Defaults to param_model.xparamdefault when not provided. + # The penalty 0.5 * (x - μ)^2 / σ^2 is added to _compute_lc. + self.param_prior_active = False + self.param_prior_mask_np = None + self.param_prior_sigmas_np = None + self.param_prior_means_np = None + if getattr(options, "paramModelPriors", False): + pm = self.param_model + sigmas = getattr(pm, "prior_sigmas", None) + if sigmas is not None: + sigmas = np.asarray(sigmas, dtype=np.float64) + if sigmas.shape != (pm.nparams,): + raise ValueError( + f"param_model.prior_sigmas must have shape ({pm.nparams},); " + f"got {sigmas.shape}" + ) + means = getattr(pm, "prior_means", None) + if means is None: + means = pm.xparamdefault.numpy().astype(np.float64) + else: + means = np.asarray(means, dtype=np.float64) + if means.shape != (pm.nparams,): + raise ValueError( + f"param_model.prior_means must have shape ({pm.nparams},); " + f"got {means.shape}" + ) + mask = np.isfinite(sigmas) & (sigmas > 0) + n_priored = int(mask.sum()) + if n_priored > 0: + self.param_prior_active = True + self.param_prior_mask_np = mask.copy() + self.param_prior_sigmas_np = np.where(mask, sigmas, np.nan) + self.param_prior_means_np = np.where(mask, means, np.nan) + inv2_np = np.where(mask, 1.0 / sigmas**2, 0.0) + self.param_prior_mask = tf.constant(mask, dtype=tf.bool) + self.param_prior_inv2 = tf.constant( + inv2_np, dtype=self.indata.dtype + ) + self.param_prior_means = tf.constant( + np.where(mask, means, 0.0), dtype=self.indata.dtype + ) + self.param_prior_log_sigma2 = tf.constant( + np.where(mask, np.log(np.where(mask, sigmas, 1.0) ** 2), 0.0), + dtype=self.indata.dtype, + ) + logger.info( + f"[paramPriors] applying Gaussian priors to " + f"{n_priored}/{pm.nparams} ParamModel params:" + ) + for i, p in enumerate(pm.params): + if mask[i]: + name = p.decode() if isinstance(p, bytes) else str(p) + logger.info( + f" {name}: μ={means[i]:.4g} σ={sigmas[i]:.4g}" + ) + else: + logger.info( + "[paramPriors] --paramModelPriors set but prior_sigmas " + "has no finite entries; no priors applied." + ) + else: + logger.info( + "[paramPriors] --paramModelPriors set but ParamModel does " + "not declare prior_sigmas; no priors applied." + ) + # constraint minima for nuisance parameters self.theta0 = tf.Variable( self.theta0default, @@ -466,14 +548,16 @@ def defreeze_params(self, unfrozen_parmeter_expressions): ] self.update_frozen_params() - def init_blinding_values(self, unblind_parameter_expressions=[]): + def init_blinding_values( + self, unblind_parameter_expressions=[], blinding_group_expressions=[] + ): logger.debug(f"Unblind parameters with {unblind_parameter_expressions}") + all_param_names = [ + *self.param_model.params[: self.param_model.npoi], + *[self.indata.systs[i] for i in self.indata.noiidxs], + ] unblind_parameters = match_regexp_params( - unblind_parameter_expressions, - [ - *self.param_model.params[: self.param_model.npoi], - *[self.indata.systs[i] for i in self.indata.noiidxs], - ], + unblind_parameter_expressions, all_param_names ) # check if dataset is an integer (i.e. if it is real data or not) and use this to choose the random seed @@ -498,14 +582,47 @@ def deterministic_random_from_string(s, mean=0.0, std=5.0): value = rng.normal(loc=mean, scale=std) return value + # Build a map: param name -> seed string. Default is the param's own name; for + # parameters matching a blinding group the seed is the group regex string, so all + # members of the group share an identical deterministic offset (preserving relative + # differences while blinding absolute values). + if isinstance(blinding_group_expressions, str): + blinding_group_expressions = [blinding_group_expressions] + param_to_seed = {} + param_to_group = {} + for expr in blinding_group_expressions: + matched = match_regexp_params(expr, all_param_names) + for p in matched: + if p in param_to_seed: + continue # first matching group wins + param_to_seed[p] = expr + param_to_group[p] = expr + + # Refuse to run if --unblind and --blindingGroup match the same parameter: + # the user's intent is ambiguous (unblind it, or share a group offset?), and + # silently picking one risks accidentally unblinding a sensitive parameter. + overlap = [p for p in unblind_parameters if p in param_to_group] + if overlap: + details = ", ".join( + f"{p.decode() if isinstance(p, bytes) else p}" + f" (group '{param_to_group[p]}')" + for p in overlap + ) + raise RuntimeError( + "The following parameters match both --unblind and --blindingGroup; " + "refusing to proceed to avoid an ambiguous (un)blinding. " + f"Tighten the regexes to make the intent explicit: {details}" + ) + # multiply offset to nois self._blinding_values_theta = np.zeros(self.indata.nsyst, dtype=np.float64) for i in self.indata.noiidxs: param = self.indata.systs[i] if param in unblind_parameters: continue - logger.debug(f"Blind parameter {param}") - value = deterministic_random_from_string(param) + seed = param_to_seed.get(param, param) + logger.debug(f"Blind parameter {param} (seed='{seed}')") + value = deterministic_random_from_string(seed) self._blinding_values_theta[i] = value # add offset to pois @@ -514,8 +631,9 @@ def deterministic_random_from_string(s, mean=0.0, std=5.0): param = self.param_model.params[i] if param in unblind_parameters: continue - logger.debug(f"Blind signal strength modifier for {param}") - value = deterministic_random_from_string(param) + seed = param_to_seed.get(param, param) + logger.debug(f"Blind signal strength modifier for {param} (seed='{seed}')") + value = deterministic_random_from_string(seed) self._blinding_values_poi[i] = np.exp(value) def set_blinding_offsets(self, blind=True): @@ -1834,7 +1952,28 @@ def _compute_lc(self, full_nll=False): # normalization factor for normal distribution: log(1/sqrt(2*pi)) = -0.9189385332046727 lc = lc + 0.9189385332046727 * self.indata.constraintweights - return tf.reduce_sum(lc) + total = tf.reduce_sum(lc) + + # ParamModel Gaussian priors (opt-in via --paramModelPriors). + # Per-entry mask is baked into param_prior_inv2 (zero where mask is + # False), so this is a single elementwise op with no branching. + if self.param_prior_active: + pm_params = tf.concat([self.get_poi(), self.get_model_nui()], axis=0) + lc_pm = ( + 0.5 + * self.param_prior_inv2 + * tf.square(pm_params - self.param_prior_means) + ) + if full_nll: + # 0.5*log(2π σ²) on entries with a prior. The log(σ²) term is + # precomputed (zero where masked off). + mask_dtype = tf.cast(self.param_prior_mask, lc_pm.dtype) + lc_pm = lc_pm + mask_dtype * ( + 0.9189385332046727 + 0.5 * self.param_prior_log_sigma2 + ) + total = total + tf.reduce_sum(lc_pm) + + return total def _compute_lbeta(self, beta, full_nll=False): return self.bbstat.lbeta(beta, full_nll=full_nll) diff --git a/rabbit/param_models/param_model.py b/rabbit/param_models/param_model.py index 1f33596..d8653a2 100644 --- a/rabbit/param_models/param_model.py +++ b/rabbit/param_models/param_model.py @@ -17,6 +17,16 @@ def __init__(self, indata, *args, **kwargs): # 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 + # + # # optional: Gaussian priors on the model's parameters + # # consumed by Fitter when --paramModelPriors is set (otherwise ignored). + # self.prior_sigmas = # np.ndarray, shape (nparams,). Entries that are + # # finite and > 0 are Gaussian-constrained at that + # # width; NaN / non-finite / 0 entries leave the + # # corresponding parameter free. Convention is up + # # to the model (e.g. POIs free, POUs constrained). + # self.prior_means = # np.ndarray, shape (nparams,). Optional; defaults + # # to self.xparamdefault when not provided. @property def nparams(self): diff --git a/rabbit/parsing.py b/rabbit/parsing.py index 085ba1e..9503917 100644 --- a/rabbit/parsing.py +++ b/rabbit/parsing.py @@ -246,6 +246,17 @@ def common_parser(): action="store_true", help="Don't compute the estimated distance to minimum as fit quality evaluation", ) + parser.add_argument( + "--paramModelPriors", + default=False, + action="store_true", + help="Opt in to Gaussian priors on ParamModel parameters. When set, " + "the Fitter reads param_model.prior_sigmas (np.ndarray of shape " + "(nparams,), entries finite > 0 → prior of that width, NaN → free) " + "and param_model.prior_means (defaults to param_model.xparamdefault) " + "and adds the corresponding penalty to the NLL constraint term. By " + "default the feature is off and ParamModel params float free.", + ) parser.add_argument( "--forceLinear", default=False, @@ -273,6 +284,20 @@ def common_parser(): E.g. use '--unblind ^signal$' to unblind a parameter named signal or '--unblind' to unblind all. """, ) + parser.add_argument( + "--blindingGroup", + type=str, + default=[], + nargs="*", + help=""" + Specify list of regex defining groups of parameters that share a single deterministic + blinding offset (seeded from the regex string itself). Useful to keep relative pulls / + differences between matched parameters meaningful while still blinding their absolute + values. E.g. '--blindingGroup ^alphaS_y\\d+$' applies the same offset to all alphaS + rapidity-bin parameters. Parameters not matching any group keep per-name blinding. + Overlap with --unblind is treated as a configuration error and aborts the fit. + """, + ) parser.add_argument( "--setConstraintMinimum", default=[], From 81f73c706d74d8b3192bff607bf4fdecd3fd9220 Mon Sep 17 00:00:00 2001 From: Luca Lavezzo Date: Fri, 12 Jun 2026 07:40:08 -0400 Subject: [PATCH 02/16] Address review: drop --paramModelPriors, the ParamModel decides Per review on WMass/rabbit#133 (parsing.py:252): there is no rabbit-side CLA anymore. If a ParamModel declares prior_sigmas, the priors are applied; whether and how to enable them (e.g. a token in the --paramModel spec) is the model's own decision. Co-Authored-By: Claude Fable 5 --- bin/rabbit_fit.py | 6 +- rabbit/fitter.py | 108 ++++++++++++----------------- rabbit/param_models/param_model.py | 6 +- rabbit/parsing.py | 11 --- 4 files changed, 53 insertions(+), 78 deletions(-) diff --git a/bin/rabbit_fit.py b/bin/rabbit_fit.py index 66eec18..743d385 100755 --- a/bin/rabbit_fit.py +++ b/bin/rabbit_fit.py @@ -823,9 +823,9 @@ def main(): "nois": ifitter.parms[ifitter.param_model.nparams :][indata.noiidxs], } - # ParamModel Gaussian priors (if --paramModelPriors was used and the - # model declared sigmas). Stored as a small dict so downstream tooling - # can see what was applied without parsing the rabbit log. + # ParamModel Gaussian priors (if the model declared sigmas). Stored as + # a small dict so downstream tooling can see what was applied without + # parsing the rabbit log. if getattr(ifitter, "param_prior_active", False): meta["param_priors"] = { "params": ifitter.param_model.params, # all nparams names diff --git a/rabbit/fitter.py b/rabbit/fitter.py index bc53cc1..86b8403 100644 --- a/rabbit/fitter.py +++ b/rabbit/fitter.py @@ -214,7 +214,6 @@ def __init__( unblind=options.unblind, blinding_group=getattr(options, "blindingGroup", []), freeze_parameters=options.freezeParameters, - options=options, ) self.nexpnom = tf.Variable( @@ -228,7 +227,6 @@ def init_fit_parms( unblind=False, blinding_group=[], freeze_parameters=[], - options=None, ): self.param_model = param_model @@ -320,78 +318,64 @@ def init_fit_parms( self.indata.dtype, ) - # ParamModel Gaussian priors (opt-in via --paramModelPriors). - # When enabled, the Fitter reads two optional attributes from the - # ParamModel: + # ParamModel Gaussian priors. The ParamModel itself decides whether + # its parameters carry priors, by declaring two optional attributes: # - prior_sigmas: np.ndarray shape (nparams,). Entries that are # finite and > 0 are Gaussian-constrained at that width; NaN / # non-finite / 0 entries leave the corresponding parameter free. # - prior_means : np.ndarray shape (nparams,), optional. # Defaults to param_model.xparamdefault when not provided. - # The penalty 0.5 * (x - μ)^2 / σ^2 is added to _compute_lc. + # If the model declares priors they are applied; how to enable or + # disable them (e.g. a token in the --paramModel spec) is up to the + # model. The penalty 0.5 * (x - μ)^2 / σ^2 is added to _compute_lc. self.param_prior_active = False self.param_prior_mask_np = None self.param_prior_sigmas_np = None self.param_prior_means_np = None - if getattr(options, "paramModelPriors", False): - pm = self.param_model - sigmas = getattr(pm, "prior_sigmas", None) - if sigmas is not None: - sigmas = np.asarray(sigmas, dtype=np.float64) - if sigmas.shape != (pm.nparams,): + pm = self.param_model + sigmas = getattr(pm, "prior_sigmas", None) + if sigmas is not None: + sigmas = np.asarray(sigmas, dtype=np.float64) + if sigmas.shape != (pm.nparams,): + raise ValueError( + f"param_model.prior_sigmas must have shape ({pm.nparams},); " + f"got {sigmas.shape}" + ) + means = getattr(pm, "prior_means", None) + if means is None: + means = pm.xparamdefault.numpy().astype(np.float64) + else: + means = np.asarray(means, dtype=np.float64) + if means.shape != (pm.nparams,): raise ValueError( - f"param_model.prior_sigmas must have shape ({pm.nparams},); " - f"got {sigmas.shape}" - ) - means = getattr(pm, "prior_means", None) - if means is None: - means = pm.xparamdefault.numpy().astype(np.float64) - else: - means = np.asarray(means, dtype=np.float64) - if means.shape != (pm.nparams,): - raise ValueError( - f"param_model.prior_means must have shape ({pm.nparams},); " - f"got {means.shape}" - ) - mask = np.isfinite(sigmas) & (sigmas > 0) - n_priored = int(mask.sum()) - if n_priored > 0: - self.param_prior_active = True - self.param_prior_mask_np = mask.copy() - self.param_prior_sigmas_np = np.where(mask, sigmas, np.nan) - self.param_prior_means_np = np.where(mask, means, np.nan) - inv2_np = np.where(mask, 1.0 / sigmas**2, 0.0) - self.param_prior_mask = tf.constant(mask, dtype=tf.bool) - self.param_prior_inv2 = tf.constant( - inv2_np, dtype=self.indata.dtype - ) - self.param_prior_means = tf.constant( - np.where(mask, means, 0.0), dtype=self.indata.dtype - ) - self.param_prior_log_sigma2 = tf.constant( - np.where(mask, np.log(np.where(mask, sigmas, 1.0) ** 2), 0.0), - dtype=self.indata.dtype, + f"param_model.prior_means must have shape ({pm.nparams},); " + f"got {means.shape}" ) - logger.info( - f"[paramPriors] applying Gaussian priors to " - f"{n_priored}/{pm.nparams} ParamModel params:" - ) - for i, p in enumerate(pm.params): - if mask[i]: - name = p.decode() if isinstance(p, bytes) else str(p) - logger.info( - f" {name}: μ={means[i]:.4g} σ={sigmas[i]:.4g}" - ) - else: - logger.info( - "[paramPriors] --paramModelPriors set but prior_sigmas " - "has no finite entries; no priors applied." - ) - else: + mask = np.isfinite(sigmas) & (sigmas > 0) + n_priored = int(mask.sum()) + if n_priored > 0: + self.param_prior_active = True + self.param_prior_mask_np = mask.copy() + self.param_prior_sigmas_np = np.where(mask, sigmas, np.nan) + self.param_prior_means_np = np.where(mask, means, np.nan) + inv2_np = np.where(mask, 1.0 / sigmas**2, 0.0) + self.param_prior_mask = tf.constant(mask, dtype=tf.bool) + self.param_prior_inv2 = tf.constant(inv2_np, dtype=self.indata.dtype) + self.param_prior_means = tf.constant( + np.where(mask, means, 0.0), dtype=self.indata.dtype + ) + self.param_prior_log_sigma2 = tf.constant( + np.where(mask, np.log(np.where(mask, sigmas, 1.0) ** 2), 0.0), + dtype=self.indata.dtype, + ) logger.info( - "[paramPriors] --paramModelPriors set but ParamModel does " - "not declare prior_sigmas; no priors applied." + f"[paramPriors] applying Gaussian priors to " + f"{n_priored}/{pm.nparams} ParamModel params:" ) + for i, p in enumerate(pm.params): + if mask[i]: + name = p.decode() if isinstance(p, bytes) else str(p) + logger.info(f" {name}: μ={means[i]:.4g} σ={sigmas[i]:.4g}") # constraint minima for nuisance parameters self.theta0 = tf.Variable( @@ -1954,7 +1938,7 @@ def _compute_lc(self, full_nll=False): total = tf.reduce_sum(lc) - # ParamModel Gaussian priors (opt-in via --paramModelPriors). + # ParamModel Gaussian priors (applied when the model declares them). # Per-entry mask is baked into param_prior_inv2 (zero where mask is # False), so this is a single elementwise op with no branching. if self.param_prior_active: diff --git a/rabbit/param_models/param_model.py b/rabbit/param_models/param_model.py index d8653a2..c47c16c 100644 --- a/rabbit/param_models/param_model.py +++ b/rabbit/param_models/param_model.py @@ -18,8 +18,10 @@ def __init__(self, indata, *args, **kwargs): # self.is_linear = # define if the model is linear in the parameters # self.allowNegativeParam = # define if the POI parameters can be negative or not # - # # optional: Gaussian priors on the model's parameters - # # consumed by Fitter when --paramModelPriors is set (otherwise ignored). + # # optional: Gaussian priors on the model's parameters. + # # If declared, the Fitter applies them automatically; the model + # # itself decides whether (and for which parameters) to declare + # # priors, e.g. via its own --paramModel spec tokens. # self.prior_sigmas = # np.ndarray, shape (nparams,). Entries that are # # finite and > 0 are Gaussian-constrained at that # # width; NaN / non-finite / 0 entries leave the diff --git a/rabbit/parsing.py b/rabbit/parsing.py index 9503917..ad894a8 100644 --- a/rabbit/parsing.py +++ b/rabbit/parsing.py @@ -246,17 +246,6 @@ def common_parser(): action="store_true", help="Don't compute the estimated distance to minimum as fit quality evaluation", ) - parser.add_argument( - "--paramModelPriors", - default=False, - action="store_true", - help="Opt in to Gaussian priors on ParamModel parameters. When set, " - "the Fitter reads param_model.prior_sigmas (np.ndarray of shape " - "(nparams,), entries finite > 0 → prior of that width, NaN → free) " - "and param_model.prior_means (defaults to param_model.xparamdefault) " - "and adds the corresponding penalty to the NLL constraint term. By " - "default the feature is off and ParamModel params float free.", - ) parser.add_argument( "--forceLinear", default=False, From 81e18029c91321ac6a46677508b949f93695c13b Mon Sep 17 00:00:00 2001 From: Luca Lavezzo Date: Fri, 12 Jun 2026 08:47:03 -0400 Subject: [PATCH 03/16] Address review: drop _np prior attributes, use .numpy() on TF constants Per review on WMass/rabbit#133: param_prior_sigmas / param_prior_means are stored once as TF constants holding the priors as declared (NaN where no prior); the compute-safe masked forms are derived inside _compute_lc, and the output metadata reads the arrays via .numpy(). Co-Authored-By: Claude Fable 5 --- bin/rabbit_fit.py | 6 +++--- rabbit/fitter.py | 48 +++++++++++++++++++++++------------------------ 2 files changed, 26 insertions(+), 28 deletions(-) diff --git a/bin/rabbit_fit.py b/bin/rabbit_fit.py index 743d385..2af5caa 100755 --- a/bin/rabbit_fit.py +++ b/bin/rabbit_fit.py @@ -829,9 +829,9 @@ def main(): if getattr(ifitter, "param_prior_active", False): meta["param_priors"] = { "params": ifitter.param_model.params, # all nparams names - "mask": ifitter.param_prior_mask_np, # bool array - "sigmas": ifitter.param_prior_sigmas_np, # NaN where mask False - "means": ifitter.param_prior_means_np, # NaN where mask False + "mask": ifitter.param_prior_mask.numpy(), # bool array + "sigmas": ifitter.param_prior_sigmas.numpy(), # NaN where mask False + "means": ifitter.param_prior_means.numpy(), # NaN where mask False } with workspace.Workspace( diff --git a/rabbit/fitter.py b/rabbit/fitter.py index 86b8403..6fdff96 100644 --- a/rabbit/fitter.py +++ b/rabbit/fitter.py @@ -328,10 +328,10 @@ def init_fit_parms( # If the model declares priors they are applied; how to enable or # disable them (e.g. a token in the --paramModel spec) is up to the # model. The penalty 0.5 * (x - μ)^2 / σ^2 is added to _compute_lc. + # param_prior_sigmas / param_prior_means hold the priors as declared + # (NaN where no prior); the compute-safe forms are derived in + # _compute_lc. self.param_prior_active = False - self.param_prior_mask_np = None - self.param_prior_sigmas_np = None - self.param_prior_means_np = None pm = self.param_model sigmas = getattr(pm, "prior_sigmas", None) if sigmas is not None: @@ -355,18 +355,12 @@ def init_fit_parms( n_priored = int(mask.sum()) if n_priored > 0: self.param_prior_active = True - self.param_prior_mask_np = mask.copy() - self.param_prior_sigmas_np = np.where(mask, sigmas, np.nan) - self.param_prior_means_np = np.where(mask, means, np.nan) - inv2_np = np.where(mask, 1.0 / sigmas**2, 0.0) self.param_prior_mask = tf.constant(mask, dtype=tf.bool) - self.param_prior_inv2 = tf.constant(inv2_np, dtype=self.indata.dtype) - self.param_prior_means = tf.constant( - np.where(mask, means, 0.0), dtype=self.indata.dtype + self.param_prior_sigmas = tf.constant( + np.where(mask, sigmas, np.nan), dtype=self.indata.dtype ) - self.param_prior_log_sigma2 = tf.constant( - np.where(mask, np.log(np.where(mask, sigmas, 1.0) ** 2), 0.0), - dtype=self.indata.dtype, + self.param_prior_means = tf.constant( + np.where(mask, means, np.nan), dtype=self.indata.dtype ) logger.info( f"[paramPriors] applying Gaussian priors to " @@ -1939,21 +1933,25 @@ def _compute_lc(self, full_nll=False): total = tf.reduce_sum(lc) # ParamModel Gaussian priors (applied when the model declares them). - # Per-entry mask is baked into param_prior_inv2 (zero where mask is - # False), so this is a single elementwise op with no branching. + # param_prior_sigmas / param_prior_means are NaN where no prior is + # declared; derive compute-safe forms via the mask (tf.where, so the + # NaN entries never enter the sum). if self.param_prior_active: - pm_params = tf.concat([self.get_poi(), self.get_model_nui()], axis=0) - lc_pm = ( - 0.5 - * self.param_prior_inv2 - * tf.square(pm_params - self.param_prior_means) + mask = self.param_prior_mask + zeros = tf.zeros_like(self.param_prior_sigmas) + inv2 = tf.where( + mask, tf.math.reciprocal(tf.square(self.param_prior_sigmas)), zeros ) + means = tf.where(mask, self.param_prior_means, zeros) + pm_params = tf.concat([self.get_poi(), self.get_model_nui()], axis=0) + lc_pm = 0.5 * inv2 * tf.square(pm_params - means) if full_nll: - # 0.5*log(2π σ²) on entries with a prior. The log(σ²) term is - # precomputed (zero where masked off). - mask_dtype = tf.cast(self.param_prior_mask, lc_pm.dtype) - lc_pm = lc_pm + mask_dtype * ( - 0.9189385332046727 + 0.5 * self.param_prior_log_sigma2 + # 0.5*log(2π σ²) on entries with a prior. + lc_pm = lc_pm + tf.where( + mask, + 0.9189385332046727 + + tf.math.log(self.param_prior_sigmas), + zeros, ) total = total + tf.reduce_sum(lc_pm) From b3fcb6019946b12483dd9584ee16fb98a4e9e3ae Mon Sep 17 00:00:00 2001 From: Luca Lavezzo Date: Fri, 12 Jun 2026 09:08:42 -0400 Subject: [PATCH 04/16] Apply pre-commit hook formatting Co-Authored-By: Claude Fable 5 --- rabbit/fitter.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/rabbit/fitter.py b/rabbit/fitter.py index 6fdff96..cff7f7f 100644 --- a/rabbit/fitter.py +++ b/rabbit/fitter.py @@ -1949,8 +1949,7 @@ def _compute_lc(self, full_nll=False): # 0.5*log(2π σ²) on entries with a prior. lc_pm = lc_pm + tf.where( mask, - 0.9189385332046727 - + tf.math.log(self.param_prior_sigmas), + 0.9189385332046727 + tf.math.log(self.param_prior_sigmas), zeros, ) total = total + tf.reduce_sum(lc_pm) From 212a835513a9058c0f629ac9af958d3ef8ffe737 Mon Sep 17 00:00:00 2001 From: Luca Lavezzo Date: Fri, 12 Jun 2026 13:14:20 -0400 Subject: [PATCH 05/16] CompositeParamModel: preserve the [POIs | POUs] parameter layout CompositeParamModel summed npoi but concatenated params in model order, so a POU-carrying model placed before a POI-carrying one leaked its POUs into the composite's POI slice. Every npoi-sliced consumer then misbehaved: get_poi() squared the POUs (allowNegativeParam=False default) and, on real data, blinded them as "signal strength modifiers" -- visible with --computeSaturatedProjectionTests + a custom npou>0 param model, where the saturated refit ran with squared+blinded theory parameters (frozen ones unable to compensate), invalidating the saturated p-value. Fix: assemble the composite per-block, [poi(m1), poi(m2), ... | pou(m1), pou(m2), ...], and reassemble each submodel's native [poi | pou] vector in compute(). Legacy-valid orderings produce identical layouts. Additionally: - allowNegativeParam is derived from the POI-carrying submodels only (mixed flags raise: the fitter applies a single squaring transform to the POI block). load_models' any() over all models had the same trap. - submodel prior_sigmas/prior_means propagate through the permutation (NaN = no prior), so a composite refit keeps the same prior penalty. - param_impact_groups are merged (name-based, permutation-safe). - is_linear no longer claims linearity for products of >1 parameter-dependent factor or sqrt-stored POIs. - blinding log message: "signal strength modifier" -> "parameter" (review request). Tested in tests/test_composite_param_model.py (layout, compute slicing vs manual evaluation, gradient flow, legacy-ordering invariance, flag derivation/conflicts, prior+group propagation). Co-Authored-By: Claude Fable 5 --- rabbit/fitter.py | 2 +- rabbit/param_models/helpers.py | 6 +- rabbit/param_models/param_model.py | 123 ++++++++++++++++++++++++++-- tests/test_composite_param_model.py | 122 +++++++++++++++++++++++++++ 4 files changed, 241 insertions(+), 12 deletions(-) create mode 100644 tests/test_composite_param_model.py diff --git a/rabbit/fitter.py b/rabbit/fitter.py index cff7f7f..010ff4b 100644 --- a/rabbit/fitter.py +++ b/rabbit/fitter.py @@ -610,7 +610,7 @@ def deterministic_random_from_string(s, mean=0.0, std=5.0): if param in unblind_parameters: continue seed = param_to_seed.get(param, param) - logger.debug(f"Blind signal strength modifier for {param} (seed='{seed}')") + logger.debug(f"Blind parameter {param} (seed='{seed}')") value = deterministic_random_from_string(seed) self._blinding_values_poi[i] = np.exp(value) diff --git a/rabbit/param_models/helpers.py b/rabbit/param_models/helpers.py index 53a1a9e..b1fce72 100644 --- a/rabbit/param_models/helpers.py +++ b/rabbit/param_models/helpers.py @@ -39,5 +39,7 @@ def load_models(model_specs, indata, **kwargs): models = [load_model(spec[0], indata, *spec[1:], **kwargs) for spec in model_specs] if len(models) == 1: return models[0] - allow_neg = any(getattr(m, "allowNegativeParam", False) for m in models) - return CompositeParamModel(models, allowNegativeParam=allow_neg) + # allowNegativeParam is derived inside CompositeParamModel from the + # POI-carrying submodels only (an `any` over all models would let a + # POU-only model's flag corrupt another model's sqrt-stored POIs). + return CompositeParamModel(models) diff --git a/rabbit/param_models/param_model.py b/rabbit/param_models/param_model.py index c47c16c..ab1fb26 100644 --- a/rabbit/param_models/param_model.py +++ b/rabbit/param_models/param_model.py @@ -83,13 +83,33 @@ def set_param_default(self, expectSignal, allowNegativeParam=False): class CompositeParamModel(ParamModel): """ - multiply different param models together + Multiply different param models together. + + The fitter-facing parameter layout MUST be [all POIs | all POUs] — every + npoi-sliced consumer assumes it (get_poi squaring & blinding, impacts, + output reporting). Submodels are therefore concatenated per-block: + + params = [poi(m1), poi(m2), ..., pou(m1), pou(m2), ...] + + and compute() reassembles each submodel's native [poi | pou] vector from + the two blocks. (A naive per-model concatenation would place an earlier + model's POUs inside the composite's POI slice whenever a POU-carrying + model precedes a POI-carrying one — the fitter would then silently + square and blind those POUs.) + + ``allowNegativeParam`` governs the fitter's single squaring transform of + the composite POI block, so it is derived from the POI-carrying + submodels (which must agree); POU-only submodels don't vote. Passing it + explicitly raises if it conflicts with the submodels. + + Submodel ``prior_sigmas`` / ``prior_means`` declarations and + ``param_impact_groups`` are propagated through the same permutation. """ def __init__( self, param_models, - allowNegativeParam=False, + allowNegativeParam=None, ): self.param_models = param_models @@ -97,20 +117,105 @@ def __init__( self.npoi = sum([m.npoi 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]) + self.params = np.concatenate( + [np.asarray(m.params[: m.npoi]) for m in param_models] + + [np.asarray(m.params[m.npoi :]) for m in param_models] + ) - self.allowNegativeParam = allowNegativeParam + self.xparamdefault = tf.concat( + [m.xparamdefault[: m.npoi] for m in param_models] + + [m.xparamdefault[m.npoi :] for m in param_models], + axis=0, + ) - self.is_linear = self.nparams == 0 or self.allowNegativeParam + poi_flags = {bool(m.allowNegativeParam) for m in param_models if m.npoi > 0} + if len(poi_flags) > 1: + raise ValueError( + "CompositeParamModel: submodels with POIs disagree on " + "allowNegativeParam; the fitter applies a single squaring " + "transform to the composite POI block, so a mix cannot be " + "represented." + ) + derived = next(iter(poi_flags)) if poi_flags else None + if allowNegativeParam is None: + self.allowNegativeParam = derived if derived is not None else False + else: + if derived is not None and bool(allowNegativeParam) != derived: + raise ValueError( + f"CompositeParamModel: explicit allowNegativeParam=" + f"{allowNegativeParam} conflicts with the POI-carrying " + f"submodels (allowNegativeParam={derived})." + ) + self.allowNegativeParam = bool(allowNegativeParam) + + # The product of >1 parameter-dependent factors is nonlinear, and the + # squaring transform (allowNegativeParam=False) makes even a single + # linear factor quadratic in the stored parameters. + n_active = sum(1 for m in param_models if m.nparams > 0) + self.is_linear = self.nparams == 0 or ( + self.allowNegativeParam + and n_active <= 1 + and all(m.is_linear for m in param_models if m.nparams > 0) + ) - self.xparamdefault = tf.concat([m.xparamdefault for m in param_models], axis=0) + # Gaussian priors: propagate submodel declarations (NaN = no prior), + # permuted like params. Means default to each submodel's own + # xparamdefault, matching the Fitter's fallback. + if any(getattr(m, "prior_sigmas", None) is not None for m in param_models): + sigmas = [] + means = [] + for m in param_models: + s = getattr(m, "prior_sigmas", None) + s = ( + np.asarray(s, dtype=np.float64) + if s is not None + else np.full(m.nparams, np.nan) + ) + mu = getattr(m, "prior_means", None) + mu = ( + np.asarray(mu, dtype=np.float64) + if mu is not None + else np.asarray(m.xparamdefault).astype(np.float64) + ) + sigmas.append(s) + means.append(mu) + self.prior_sigmas = np.concatenate( + [s[: m.npoi] for s, m in zip(sigmas, param_models)] + + [s[m.npoi :] for s, m in zip(sigmas, param_models)] + ) + self.prior_means = np.concatenate( + [v[: m.npoi] for v, m in zip(means, param_models)] + + [v[m.npoi :] for v, m in zip(means, param_models)] + ) + + # impact groups are name-based, so a plain merge survives the + # permutation + groups = {} + for m in param_models: + for k, v in getattr(m, "param_impact_groups", {}).items(): + groups.setdefault(k, []).extend(v) + if groups: + self.param_impact_groups = groups def compute(self, param, full=False): - start = 0 + poi_offset = 0 + pou_offset = self.npoi results = [] for m in self.param_models: - results.append(m.compute(param[start : start + m.nparams], full)) - start += m.nparams + parts = [] + if m.npoi > 0: + parts.append(param[poi_offset : poi_offset + m.npoi]) + if m.npou > 0: + parts.append(param[pou_offset : pou_offset + m.npou]) + if len(parts) > 1: + mparam = tf.concat(parts, axis=0) + elif parts: + mparam = parts[0] + else: + mparam = param[0:0] + results.append(m.compute(mparam, full)) + poi_offset += m.npoi + pou_offset += m.npou rnorm = functools.reduce(lambda a, b: a * b, results) return rnorm diff --git a/tests/test_composite_param_model.py b/tests/test_composite_param_model.py new file mode 100644 index 0000000..ea98913 --- /dev/null +++ b/tests/test_composite_param_model.py @@ -0,0 +1,122 @@ +"""CompositeParamModel must preserve the fitter-facing [POIs | POUs] layout. + +Regression test for the bug where a POU-carrying model placed before a +POI-carrying one (e.g. a custom npou>0 param model composited with +SaturatedProjectModel by --computeSaturatedProjectionTests) leaked its POUs +into the composite's POI slice — the fitter then silently squared and +blinded them. +""" + +import numpy as np +import pytest +import tensorflow as tf + +from rabbit.param_models.param_model import CompositeParamModel + + +class PouOnly: + """Mimics a POU-only model (npoi=0), e.g. a theory-parameter model.""" + + npoi = 0 + npou = 2 + nparams = 2 + params = np.array([b"a1", b"a2"]) + xparamdefault = tf.constant([0.4, 2.0], dtype=tf.float64) + allowNegativeParam = True + is_linear = False + prior_sigmas = np.array([0.5, np.nan]) + param_impact_groups = {"groupA": [b"a1", b"a2"]} + w = tf.constant([1.0, 10.0], dtype=tf.float64) + + def compute(self, param, full=False): + return 1.0 + tf.reduce_sum(param * self.w) + + +class PoiAndPou: + """Mimics a POI-carrying model (e.g. SaturatedProjectModel + one POU).""" + + npoi = 2 + npou = 1 + nparams = 3 + params = np.array([b"b1", b"b2", b"bp"]) + xparamdefault = tf.constant([1.0, 1.0, 3.0], dtype=tf.float64) + allowNegativeParam = False + is_linear = False + w = tf.constant([100.0, 1000.0, 10000.0], dtype=tf.float64) + + def compute(self, param, full=False): + return 1.0 + tf.reduce_sum(param * self.w) + + +def test_pou_model_first_layout(): + # the configuration that was broken: POU-carrying model first + c = CompositeParamModel([PouOnly(), PoiAndPou()]) + + assert c.npoi == 2 and c.npou == 3 + assert list(c.params) == [b"b1", b"b2", b"a1", b"a2", b"bp"] + # the first npoi names must be exactly the POIs (this is the slice that + # get_poi squaring, blinding, and output reporting act on) + assert list(c.params[: c.npoi]) == [b"b1", b"b2"] + np.testing.assert_allclose(c.xparamdefault.numpy(), [1.0, 1.0, 0.4, 2.0, 3.0]) + + +def test_compute_reassembles_native_order(): + A, B = PouOnly(), PoiAndPou() + c = CompositeParamModel([A, B]) + param = tf.constant([1.5, 2.5, -0.3, 1.1, 4.0], dtype=tf.float64) + expected = A.compute(tf.constant([-0.3, 1.1], dtype=tf.float64)) * B.compute( + tf.constant([1.5, 2.5, 4.0], dtype=tf.float64) + ) + np.testing.assert_allclose(c.compute(param).numpy(), expected.numpy(), rtol=1e-14) + + # gradient flows through the permutation + with tf.GradientTape() as t: + t.watch(param) + val = c.compute(param) + g = t.gradient(val, param).numpy() + assert np.all(np.isfinite(g)) and np.all(g != 0) + + +def test_legacy_valid_ordering_unchanged(): + A, B = PouOnly(), PoiAndPou() + c = CompositeParamModel([B, A]) + assert list(c.params) == [b"b1", b"b2", b"bp", b"a1", b"a2"] + param = tf.constant([1.5, 2.5, 4.0, -0.3, 1.1], dtype=tf.float64) + expected = A.compute(tf.constant([-0.3, 1.1], dtype=tf.float64)) * B.compute( + tf.constant([1.5, 2.5, 4.0], dtype=tf.float64) + ) + np.testing.assert_allclose(c.compute(param).numpy(), expected.numpy(), rtol=1e-14) + + +def test_allow_negative_derived_from_poi_models_only(): + # the POU-only model has allowNegativeParam=True, but only the + # POI-carrying model's flag matters for the squaring transform + c = CompositeParamModel([PouOnly(), PoiAndPou()]) + assert c.allowNegativeParam is False + + +def test_priors_and_impact_groups_propagated(): + c = CompositeParamModel([PouOnly(), PoiAndPou()]) + np.testing.assert_allclose(c.prior_sigmas, [np.nan, np.nan, 0.5, np.nan, np.nan]) + np.testing.assert_allclose(c.prior_means, [1.0, 1.0, 0.4, 2.0, 3.0]) + assert c.param_impact_groups == {"groupA": [b"a1", b"a2"]} + + +def test_conflicting_flags_raise(): + class PoiNeg(PoiAndPou): + allowNegativeParam = True + + with pytest.raises(ValueError): + CompositeParamModel([PoiAndPou(), PoiNeg()]) + + with pytest.raises(ValueError): + CompositeParamModel([PouOnly(), PoiAndPou()], allowNegativeParam=True) + + # explicit flag consistent with submodels is accepted + c = CompositeParamModel([PouOnly(), PoiAndPou()], allowNegativeParam=False) + assert c.allowNegativeParam is False + # POU-only composite: explicit flag respected (no POI carrier to conflict) + assert ( + CompositeParamModel([PouOnly()], allowNegativeParam=True).allowNegativeParam + is True + ) From 2bbbb90a066eb8e8fae0c4e616e616a528388efa Mon Sep 17 00:00:00 2001 From: Luca Lavezzo Date: Fri, 12 Jun 2026 16:13:28 -0400 Subject: [PATCH 06/16] Unify ParamModel priors with the nuisance constraint structure MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ParamModel priors now use the same constraint structure as the nuisance constraints: param_cw holds the per-parameter constraint weight (1/sigma^2; 0 = free) and param_x0 the constraint center (the prior mean) — the ParamModel-block analogues of indata.constraintweights and theta0. Consequences: - _compute_lc penalizes cw * 0.5 * (param - x0)^2, same form as the nuisance term. - prefit variances are 1/cw uniformly, so priored params get proper prefit uncertainties (pulls/constraints work without special cases). - Toys treat the prior like any auxiliary constraint: frequentist toys fluctuate param_x0 around its default with width sqrt(1/cw) exactly like theta0, and bayesian toys sample the priored param values from their priors. Previously priors were static in toys, silently underestimating the toy spread of anything correlated with a priored parameter. This also fixes the long-standing FIXMEs in bayesassign / frequentistassign: constraint centers now fluctuate around their defaults (not zero) with width sqrt(1/cw), and unconstrained entries are no longer randomized with unit width. - Gaussian priors on POIs require allowNegativeParam=True; with the squared storage the penalty would silently apply to sqrt(poi), so the Fitter raises instead. Verified on the test tensor: identical postfit optimum as the previous implementation, prior width visible in the prefit variances, per-toy fluctuation of the prior centers, and unchanged behavior for models without priors. Co-Authored-By: Claude Fable 5 --- rabbit/fitter.py | 206 ++++++++++++++++++----------- rabbit/param_models/param_model.py | 4 +- 2 files changed, 129 insertions(+), 81 deletions(-) diff --git a/rabbit/fitter.py b/rabbit/fitter.py index 010ff4b..4695a98 100644 --- a/rabbit/fitter.py +++ b/rabbit/fitter.py @@ -281,43 +281,6 @@ def init_fit_parms( self.x = tf.Variable(xdefault, trainable=True, name="x") - # Per-parameter prefit variance vector. Always allocated; the - # prefit covariance is intrinsically diagonal so this is the - # only form needed for prefit uncertainties. - self.var_prefit = tf.Variable( - self.prefit_variance( - unconstrained_err=self.prefit_unconstrained_nuisance_uncertainty - ), - trainable=False, - name="var_prefit", - ) - - # Full parameter covariance matrix. Allocated only when the - # postfit Hessian will actually be computed; otherwise None to - # avoid the O(npar^2) allocation (94 GB for 108k parameters). - if self.compute_cov: - self.cov = tf.Variable( - tf.linalg.diag(self.var_prefit), - trainable=False, - name="cov", - ) - else: - self.cov = None - - # regularization - self.regularizers = [] - # one common regularization strength parameter - self.tau = tf.Variable(1.0, trainable=True, name="tau", dtype=tf.float64) - - # External likelihood terms (additive g^T x + 0.5 x^T H x - # contributions to the NLL). See rabbit.external_likelihood for - # the construction helper and the matching scalar evaluator. - self.external_terms = external_likelihood.build_tf_external_terms( - self.indata.external_terms, - self.parms, - self.indata.dtype, - ) - # ParamModel Gaussian priors. The ParamModel itself decides whether # its parameters carry priors, by declaring two optional attributes: # - prior_sigmas: np.ndarray shape (nparams,). Entries that are @@ -327,12 +290,19 @@ def init_fit_parms( # Defaults to param_model.xparamdefault when not provided. # If the model declares priors they are applied; how to enable or # disable them (e.g. a token in the --paramModel spec) is up to the - # model. The penalty 0.5 * (x - μ)^2 / σ^2 is added to _compute_lc. - # param_prior_sigmas / param_prior_means hold the priors as declared - # (NaN where no prior); the compute-safe forms are derived in - # _compute_lc. + # model. + # + # Priors use the same constraint structure as the nuisance + # constraints: param_cw is the per-parameter constraint weight + # (1/sigma^2; 0 = free) and param_x0 the constraint center (the + # prior mean) — the ParamModel-block analogues of + # indata.constraintweights and theta0. _compute_lc adds + # 0.5 * cw * (param - x0)^2, prefit variances are 1/cw, and the toy + # randomization treats the centers exactly like theta0. self.param_prior_active = False pm = self.param_model + param_cw_np = np.zeros(pm.nparams, dtype=np.float64) + param_x0_np = np.zeros(pm.nparams, dtype=np.float64) sigmas = getattr(pm, "prior_sigmas", None) if sigmas is not None: sigmas = np.asarray(sigmas, dtype=np.float64) @@ -354,6 +324,14 @@ def init_fit_parms( mask = np.isfinite(sigmas) & (sigmas > 0) n_priored = int(mask.sum()) if n_priored > 0: + if mask[: pm.npoi].any() and not pm.allowNegativeParam: + raise ValueError( + "Gaussian priors on POIs require allowNegativeParam=" + "True: with allowNegativeParam=False the stored " + "parameter is sqrt(poi), so the Gaussian penalty " + "would apply to sqrt(poi) rather than to the POI " + "itself." + ) self.param_prior_active = True self.param_prior_mask = tf.constant(mask, dtype=tf.bool) self.param_prior_sigmas = tf.constant( @@ -362,6 +340,8 @@ def init_fit_parms( self.param_prior_means = tf.constant( np.where(mask, means, np.nan), dtype=self.indata.dtype ) + param_cw_np = np.where(mask, 1.0 / sigmas**2, 0.0) + param_x0_np = np.where(mask, means, 0.0) logger.info( f"[paramPriors] applying Gaussian priors to " f"{n_priored}/{pm.nparams} ParamModel params:" @@ -370,6 +350,50 @@ def init_fit_parms( if mask[i]: name = p.decode() if isinstance(p, bytes) else str(p) logger.info(f" {name}: μ={means[i]:.4g} σ={sigmas[i]:.4g}") + self.param_cw = tf.constant(param_cw_np, dtype=self.indata.dtype) + self.param_x0default = tf.constant(param_x0_np, dtype=self.indata.dtype) + # constraint centers for the ParamModel block; fluctuated in toys + # together with theta0 + self.param_x0 = tf.Variable( + self.param_x0default, trainable=False, name="param_x0" + ) + + # Per-parameter prefit variance vector. Always allocated; the + # prefit covariance is intrinsically diagonal so this is the + # only form needed for prefit uncertainties. + self.var_prefit = tf.Variable( + self.prefit_variance( + unconstrained_err=self.prefit_unconstrained_nuisance_uncertainty + ), + trainable=False, + name="var_prefit", + ) + + # Full parameter covariance matrix. Allocated only when the + # postfit Hessian will actually be computed; otherwise None to + # avoid the O(npar^2) allocation (94 GB for 108k parameters). + if self.compute_cov: + self.cov = tf.Variable( + tf.linalg.diag(self.var_prefit), + trainable=False, + name="cov", + ) + else: + self.cov = None + + # regularization + self.regularizers = [] + # one common regularization strength parameter + self.tau = tf.Variable(1.0, trainable=True, name="tau", dtype=tf.float64) + + # External likelihood terms (additive g^T x + 0.5 x^T H x + # contributions to the NLL). See rabbit.external_likelihood for + # the construction helper and the matching scalar evaluator. + self.external_terms = external_likelihood.build_tf_external_terms( + self.indata.external_terms, + self.parms, + self.indata.dtype, + ) # constraint minima for nuisance parameters self.theta0 = tf.Variable( @@ -678,21 +702,25 @@ def get_x(self): def prefit_variance(self, unconstrained_err=0.0): """Per-parameter prefit variance vector of length npar. - Free parameters (POIs and unconstrained nuisances) are assigned a - placeholder variance of unconstrained_err**2 (zero by default). - Constrained nuisances take their variance from the constraint - term (1 / constraintweight). + Free parameters (ParamModel params without a prior and unconstrained + nuisances) are assigned a placeholder variance of unconstrained_err**2 + (zero by default). Constrained entries take their variance from the + constraint term (1 / constraint weight), uniformly for priored + ParamModel params (param_cw) and constrained nuisances + (indata.constraintweights). """ - var_poi = ( - tf.ones([self.param_model.nparams], dtype=self.indata.dtype) - * unconstrained_err**2 + var_param = tf.where( + self.param_cw == 0.0, + unconstrained_err**2 + * tf.ones([self.param_model.nparams], dtype=self.indata.dtype), + tf.math.reciprocal(self.param_cw), ) var_theta = tf.where( self.indata.constraintweights == 0.0, unconstrained_err**2, tf.math.reciprocal(self.indata.constraintweights), ) - return tf.concat([var_poi, var_theta], axis=0) + return tf.concat([var_param, var_theta], axis=0) def prefit_covariance(self, unconstrained_err=0.0): """Full prefit covariance as a tf.linalg.LinearOperatorDiag. @@ -732,7 +760,10 @@ def set_nobs(self, values, variances=None): self.lognobs.assign(tf.math.log(nobssafe)) def theta0defaultassign(self): + # reset all constraint centers: the nuisance block and the + # ParamModel-prior block self.theta0.assign(self.theta0default) + self.param_x0.assign(self.param_x0default) def xdefaultassign(self): if self.param_model.nparams == 0: @@ -763,33 +794,55 @@ def defaultassign(self): reg.set_expectations(xinit, nexp0) def bayesassign(self): - # FIXME use theta0 as the mean and constraintweight to scale the width + # Sample the parameter values from their priors: width sqrt(1/cw) + # around the constraint centers wherever cw > 0 (constrained + # nuisances and priored ParamModel params); free entries stay at + # their defaults. + theta_part = self.theta0 + tf.sqrt(self.var_theta0) * tf.random.normal( + shape=self.theta0.shape, dtype=self.theta0.dtype + ) if self.param_model.nparams == 0: - self.x.assign( - self.theta0 - + tf.random.normal(shape=self.theta0.shape, dtype=self.theta0.dtype) - ) + self.x.assign(theta_part) else: - self.x.assign( - tf.concat( - [ - self.param_model.xparamdefault, - self.theta0 - + tf.random.normal( - shape=self.theta0.shape, dtype=self.theta0.dtype - ), - ], - axis=0, - ) + param_sigma = tf.where( + self.param_cw > 0, + tf.math.rsqrt(self.param_cw), + tf.zeros_like(self.param_cw), ) + param_part = tf.where( + self.param_cw > 0, + self.param_x0 + + param_sigma + * tf.random.normal( + shape=self.param_x0.shape, dtype=self.param_x0.dtype + ), + self.param_model.xparamdefault, + ) + self.x.assign(tf.concat([param_part, theta_part], axis=0)) self.bbstat.randomize_bayes() def frequentistassign(self): - # FIXME use theta as the mean and constraintweight to scale the width + # Fluctuate the constraint centers around their defaults with the + # prefit constraint widths, uniformly for the nuisance block + # (theta0) and the ParamModel-prior block (param_x0). Entries with + # cw = 0 (unconstrained) are not randomized. self.theta0.assign( - tf.random.normal(shape=self.theta0.shape, dtype=self.theta0.dtype) + self.theta0default + + tf.sqrt(self.var_theta0) + * tf.random.normal(shape=self.theta0.shape, dtype=self.theta0.dtype) ) + if self.param_prior_active: + param_sigma = tf.where( + self.param_cw > 0, + tf.math.rsqrt(self.param_cw), + tf.zeros_like(self.param_cw), + ) + self.param_x0.assign( + self.param_x0default + + param_sigma + * tf.random.normal(shape=self.param_x0.shape, dtype=self.param_x0.dtype) + ) self.bbstat.randomize_frequentist() def toyassign( @@ -1932,25 +1985,18 @@ def _compute_lc(self, full_nll=False): total = tf.reduce_sum(lc) - # ParamModel Gaussian priors (applied when the model declares them). - # param_prior_sigmas / param_prior_means are NaN where no prior is - # declared; derive compute-safe forms via the mask (tf.where, so the - # NaN entries never enter the sum). + # ParamModel Gaussian priors (applied when the model declares them): + # same structure as the nuisance constraint above, with param_cw and + # param_x0 in the roles of constraintweights and theta0. if self.param_prior_active: - mask = self.param_prior_mask - zeros = tf.zeros_like(self.param_prior_sigmas) - inv2 = tf.where( - mask, tf.math.reciprocal(tf.square(self.param_prior_sigmas)), zeros - ) - means = tf.where(mask, self.param_prior_means, zeros) pm_params = tf.concat([self.get_poi(), self.get_model_nui()], axis=0) - lc_pm = 0.5 * inv2 * tf.square(pm_params - means) + lc_pm = self.param_cw * 0.5 * tf.square(pm_params - self.param_x0) if full_nll: # 0.5*log(2π σ²) on entries with a prior. lc_pm = lc_pm + tf.where( - mask, + self.param_prior_mask, 0.9189385332046727 + tf.math.log(self.param_prior_sigmas), - zeros, + tf.zeros_like(lc_pm), ) total = total + tf.reduce_sum(lc_pm) diff --git a/rabbit/param_models/param_model.py b/rabbit/param_models/param_model.py index ab1fb26..70b68c8 100644 --- a/rabbit/param_models/param_model.py +++ b/rabbit/param_models/param_model.py @@ -21,7 +21,9 @@ def __init__(self, indata, *args, **kwargs): # # optional: Gaussian priors on the model's parameters. # # If declared, the Fitter applies them automatically; the model # # itself decides whether (and for which parameters) to declare - # # priors, e.g. via its own --paramModel spec tokens. + # # priors, e.g. via its own --paramModel spec tokens. Priors on + # # POIs require allowNegativeParam=True (with the squared storage + # # the penalty would apply to sqrt(poi), so the Fitter raises). # self.prior_sigmas = # np.ndarray, shape (nparams,). Entries that are # # finite and > 0 are Gaussian-constrained at that # # width; NaN / non-finite / 0 entries leave the From 65164bb7c89e2b91804d68777d290f3ab4d65863 Mon Sep 17 00:00:00 2001 From: Luca Lavezzo Date: Fri, 12 Jun 2026 16:24:53 -0400 Subject: [PATCH 07/16] Propagate ParamModel priors through the global impacts The prior centers (param_x0) are auxiliary observables like theta0, so they now appear as global-impact sources, one column per priored param (named _prior, appended after the systs / syst groups in the output axes): - likelihood-based global impacts: _compute_global_impacts_x0 already differentiates the full constraint term; the ParamModel-block rows (nonzero where priors are declared) are now kept instead of discarded, in per-1-prefit-sigma units automatically via sqrt(d2lc/dx2) = 1/sigma. Their variance is subtracted from the residual data-stat term, which previously absorbed it. - gaussian global impacts: _dxdvars additionally differentiates with respect to param_x0 and the impacts are dx/dparam_x0 * sigma. - observable (histogram) impacts: both variants extended through _dndvars with the chain-rule term pdndx @ dxdparam_x0. - profiled chi2: the residual covariance in _residuals_profiled gains the prior-center variance contribution. Verified on the test tensor: the likelihood and Gaussian paths agree exactly on the prior impact, and the variance closure sum(sources^2) + stat^2 + binByBinStat^2 = sigma_tot^2 holds to all printed digits. Output schemas are unchanged for fits without priors. Co-Authored-By: Claude Fable 5 --- rabbit/fitter.py | 74 ++++++++++++++++++++++++----- rabbit/impacts/global_impacts.py | 80 +++++++++++++++++++++++++++++--- rabbit/workspace.py | 20 ++++++-- 3 files changed, 153 insertions(+), 21 deletions(-) diff --git a/rabbit/fitter.py b/rabbit/fitter.py index 4695a98..942ba99 100644 --- a/rabbit/fitter.py +++ b/rabbit/fitter.py @@ -352,6 +352,11 @@ def init_fit_parms( logger.info(f" {name}: μ={means[i]:.4g} σ={sigmas[i]:.4g}") self.param_cw = tf.constant(param_cw_np, dtype=self.indata.dtype) self.param_x0default = tf.constant(param_x0_np, dtype=self.indata.dtype) + # indices (into the ParamModel block) of the priored params; impact + # sources in the global-impacts calculations + self.param_prior_idxs = tf.constant( + np.where(param_cw_np > 0)[0], dtype=tf.int64 + ) # constraint centers for the ParamModel block; fluctuated in toys # together with theta0 self.param_x0 = tf.Variable( @@ -1118,11 +1123,21 @@ def global_impacts_parms(self): self.bbstat.binByBinStatMode, self.globalImpactsFromJVP, self.cov, + param_prior_idxs=( + self.param_prior_idxs if self.param_prior_active else None + ), ) @tf.function def gaussian_global_impacts_parms(self): - dxdtheta0, dxdnobs, dxdbeta0 = self._dxdvars() + dxdtheta0, dxdnobs, dxdbeta0, dxdparam_x0 = self._dxdvars() + + if self.param_prior_active: + dxdparam0 = tf.gather(dxdparam_x0, self.param_prior_idxs, axis=1) + varparam0 = tf.gather(self._var_param_x0(), self.param_prior_idxs) + else: + dxdparam0 = None + varparam0 = None impacts, impacts_grouped = global_impacts.gaussian_global_impacts_parms( dxdtheta0, @@ -1144,6 +1159,8 @@ def gaussian_global_impacts_parms(self): self.bbstat.beta_shape, self.indata.systgroupidxs, self.data_cov_inv, + dxdparam0=dxdparam0, + varparam0=varparam0, ) return impacts, impacts_grouped @@ -1340,14 +1357,14 @@ def _pd2ldbeta2(self, profile=False): def _dxdvars(self): with tf.GradientTape() as t2: - t2.watch([self.theta0, self.nobs, self.bbstat.beta0]) + t2.watch([self.theta0, self.nobs, self.bbstat.beta0, self.param_x0]) with tf.GradientTape() as t1: - t1.watch([self.theta0, self.nobs, self.bbstat.beta0]) + t1.watch([self.theta0, self.nobs, self.bbstat.beta0, self.param_x0]) val = self._compute_loss() grad = t1.gradient(val, self.x) - pd2ldxdtheta0, pd2ldxdnobs, pd2ldxdbeta0 = t2.jacobian( + pd2ldxdtheta0, pd2ldxdnobs, pd2ldxdbeta0, pd2ldxdparam_x0 = t2.jacobian( grad, - [self.theta0, self.nobs, self.bbstat.beta0], + [self.theta0, self.nobs, self.bbstat.beta0, self.param_x0], unconnected_gradients="zero", ) @@ -1355,8 +1372,9 @@ def _dxdvars(self): dxdtheta0 = -self.cov @ pd2ldxdtheta0 dxdnobs = -self.cov @ pd2ldxdnobs dxdbeta0 = -self.cov @ tf.reshape(pd2ldxdbeta0, [pd2ldxdbeta0.shape[0], -1]) + dxdparam_x0 = -self.cov @ pd2ldxdparam_x0 - return dxdtheta0, dxdnobs, dxdbeta0 + return dxdtheta0, dxdnobs, dxdbeta0, dxdparam_x0 def _dndvars(self, fun): with tf.GradientTape() as t: @@ -1371,13 +1389,27 @@ def _dndvars(self, fun): ) # apply chain rule to take into account correlations with the fit parameters - dxdtheta0, dxdnobs, dxdbeta0 = self._dxdvars() + dxdtheta0, dxdnobs, dxdbeta0, dxdparam_x0 = self._dxdvars() dndtheta0 = pdndtheta0 + pdndx @ dxdtheta0 dndnobs = pdndnobs + pdndx @ dxdnobs dndbeta0 = tf.reshape(pdndbeta0, [pdndbeta0.shape[0], -1]) + pdndx @ dxdbeta0 - - return n, dndtheta0, dndnobs, dndbeta0 + # the yields depend on param_x0 only through the fit parameters, so + # there is no partial term + dndparam_x0 = pdndx @ dxdparam_x0 + + return n, dndtheta0, dndnobs, dndbeta0, dndparam_x0 + + def _var_param_x0(self): + """Prefit variance of the ParamModel constraint centers (1/cw for + priored params, 0 for free ones).""" + return tf.where( + self.param_cw > 0, + tf.math.reciprocal( + tf.where(self.param_cw > 0, self.param_cw, tf.ones_like(self.param_cw)) + ), + tf.zeros_like(self.param_cw), + ) def _compute_expected( self, fun_exp, inclusive=True, profile=False, full=True, need_observables=True @@ -1494,6 +1526,9 @@ def compute_derivatives(dvars): pdexpdbeta, pd2ldbeta2_pdexpdbeta if pdexpdbeta is not None else None, self.prefit_unconstrained_nuisance_uncertainty, + param_prior_idxs=( + self.param_prior_idxs if self.param_prior_active else None + ), ) else: impacts = None @@ -1510,7 +1545,13 @@ def fun_n(): need_observables=need_observables, ) - _, dndtheta0, dndnobs, dndbeta0 = self._dndvars(fun_n) + _, dndtheta0, dndnobs, dndbeta0, dndparam_x0 = self._dndvars(fun_n) + if self.param_prior_active: + dndparam0 = tf.gather(dndparam_x0, self.param_prior_idxs, axis=1) + varparam0 = tf.gather(self._var_param_x0(), self.param_prior_idxs) + else: + dndparam0 = None + varparam0 = None impacts_gaussian, impacts_gaussian_grouped = ( global_impacts.gaussian_global_impacts_obs( dndtheta0, @@ -1529,6 +1570,8 @@ def fun_n(): self.bbstat.beta_shape, self.indata.systgroupidxs, self.data_cov_inv, + dxdparam0=dndparam0, + varparam0=varparam0, ) ) else: @@ -1809,10 +1852,19 @@ def fun_res(): observed = fun(None, self.nobs) return expected - observed - residuals, dresdtheta0, dresdnobs, dresdbeta0 = self._dndvars(fun_res) + residuals, dresdtheta0, dresdnobs, dresdbeta0, dresdparam_x0 = self._dndvars( + fun_res + ) res_cov = dresdtheta0 @ (self.var_theta0[:, None] * tf.transpose(dresdtheta0)) + if self.param_prior_active: + # ParamModel prior centers contribute like additional auxiliary + # observables with variance 1/cw + res_cov += dresdparam_x0 @ ( + self._var_param_x0()[:, None] * tf.transpose(dresdparam_x0) + ) + if self.covarianceFit: res_cov_stat = dresdnobs @ tf.linalg.solve( self.data_cov_inv, tf.transpose(dresdnobs) diff --git a/rabbit/impacts/global_impacts.py b/rabbit/impacts/global_impacts.py index 47c48a3..4f3402a 100644 --- a/rabbit/impacts/global_impacts.py +++ b/rabbit/impacts/global_impacts.py @@ -201,8 +201,14 @@ def _compute_grouped_impacts( impacts_nobs, impacts_beta0_total, impacts_beta0_process, + impacts_param0_sq=None, ): - """Assemble the grouped impacts tensor from all contributions.""" + """Assemble the grouped impacts tensor from all contributions. + + impacts_param0_sq holds the squared impacts of the ParamModel prior + centers (one column per priored param); they are appended at the END of + the grouped axis, matching the column order of the output axes. + """ if bin_by_bin_stat: impacts_grouped = tf.stack([impacts_nobs, impacts_beta0_total], axis=-1) if bin_by_bin_stat_mode == "full": @@ -223,6 +229,11 @@ def _compute_grouped_impacts( impacts_grouped_syst = tf.transpose(impacts_grouped_syst) impacts_grouped = tf.concat([impacts_grouped_syst, impacts_grouped], axis=-1) + if impacts_param0_sq is not None: + impacts_grouped = tf.concat( + [impacts_grouped, tf.sqrt(impacts_param0_sq)], axis=-1 + ) + return impacts_grouped @@ -241,6 +252,7 @@ def global_impacts_parms( bin_by_bin_stat_mode, global_impacts_from_jvp, cov, + param_prior_idxs=None, ): idxs_poi = tf.range(nsignal_params, dtype=tf.int64) idxs_noi = tf.constant(nmodel_params + noiidxs, dtype=tf.int64) @@ -267,12 +279,25 @@ def global_impacts_parms( pd2ldbeta2_pdexpdbeta=None, ) + # impacts_x0 covers the full constraint term: the nuisance block AND any + # ParamModel prior penalties (rows in the ParamModel block, nonzero where + # the model declares priors). Per-1-sigma units come out automatically + # since sc = sqrt(d2lc/dx2) = sqrt(cw) = 1/sigma. impacts_x0 = _compute_global_impacts_x0(x, compute_lc_fn, cov_dexpdx) impacts_theta0 = tf.transpose(impacts_x0[nmodel_params:]) impacts_theta0_sq = tf.square(impacts_theta0) - var_theta0 = tf.reduce_sum(impacts_theta0_sq, axis=-1) - var_nobs = var_total - var_theta0 + var_x0 = tf.reduce_sum(impacts_theta0_sq, axis=-1) + + impacts_param0_sq = None + if param_prior_idxs is not None: + impacts_param0 = tf.transpose( + tf.gather(impacts_x0[:nmodel_params], param_prior_idxs, axis=0) + ) + impacts_param0_sq = tf.square(impacts_param0) + var_x0 += tf.reduce_sum(impacts_param0_sq, axis=-1) + + var_nobs = var_total - var_x0 if bin_by_bin_stat: var_nobs -= var_beta0 @@ -284,8 +309,12 @@ def global_impacts_parms( tf.sqrt(var_nobs), impacts_beta0_total, impacts_beta0_process, + impacts_param0_sq=impacts_param0_sq, ) + if param_prior_idxs is not None: + impacts_theta0 = tf.concat([impacts_theta0, impacts_param0], axis=-1) + return impacts_theta0, impacts_grouped @@ -309,6 +338,7 @@ def global_impacts_obs( pdexpdbeta=None, pd2ldbeta2_pdexpdbeta=None, prefit_unconstrained_nuisance_uncertainty=0.0, + param_prior_idxs=None, ): """ Global impacts on observable bins, used inside _expected_with_variance. @@ -360,8 +390,17 @@ def global_impacts_obs( impacts_theta0 = tf.transpose(impacts_x0[nmodel_params:]) impacts_theta0_sq = tf.square(impacts_theta0) - var_theta0 = tf.reduce_sum(impacts_theta0_sq, axis=-1) - var_nobs = expvar_flat - var_theta0 + var_x0 = tf.reduce_sum(impacts_theta0_sq, axis=-1) + + impacts_param0_sq = None + if param_prior_idxs is not None: + impacts_param0 = tf.transpose( + tf.gather(impacts_x0[:nmodel_params], param_prior_idxs, axis=0) + ) + impacts_param0_sq = tf.square(impacts_param0) + var_x0 += tf.reduce_sum(impacts_param0_sq, axis=-1) + + var_nobs = expvar_flat - var_x0 if bin_by_bin_stat: var_nobs -= var_beta0 @@ -373,8 +412,12 @@ def global_impacts_obs( tf.sqrt(var_nobs), impacts_beta0_total, impacts_beta0_process, + impacts_param0_sq=impacts_param0_sq, ) + if param_prior_idxs is not None: + impacts_theta0 = tf.concat([impacts_theta0, impacts_param0], axis=-1) + impacts = tf.reshape(impacts_theta0, [*expvar_shape, impacts_theta0.shape[-1]]) impacts_grouped = tf.reshape( impacts_grouped, [*expvar_shape, impacts_grouped.shape[-1]] @@ -395,6 +438,8 @@ def _gaussian_global_impacts( beta_shape, systgroupidxs, data_cov_inv=None, + dxdparam0=None, + varparam0=None, ): if data_cov_inv is not None: data_cov = tf.linalg.inv(data_cov_inv) @@ -438,7 +483,18 @@ def _gaussian_global_impacts( impacts_grouped_syst = tf.transpose(impacts_grouped_syst) impacts_grouped = tf.concat([impacts_grouped_syst, impacts_grouped], axis=1) - return dxdtheta0, impacts_grouped + impacts = dxdtheta0 + if dxdparam0 is not None: + # impacts of the ParamModel prior centers, per 1 prefit sigma; one + # column per priored param, appended at the END of the source axis + # and of the grouped axis (matching the output axes). + impacts_param0 = dxdparam0 * tf.sqrt(varparam0) + impacts = tf.concat([impacts, impacts_param0], axis=-1) + if impacts_grouped.shape.rank == 1: + impacts_grouped = impacts_grouped[..., None] + impacts_grouped = tf.concat([impacts_grouped, tf.abs(impacts_param0)], axis=-1) + + return impacts, impacts_grouped def gaussian_global_impacts_parms( @@ -456,6 +512,8 @@ def gaussian_global_impacts_parms( beta_shape, systgroupidxs, data_cov_inv=None, + dxdparam0=None, + varparam0=None, ): # compute impacts for pois and nois dxdtheta0 = _gather_poi_noi_vector( @@ -463,6 +521,10 @@ def gaussian_global_impacts_parms( ) dxdnobs = _gather_poi_noi_vector(dxdnobs, noiidxs, nsignal_params, nmodel_params) dxdbeta0 = _gather_poi_noi_vector(dxdbeta0, noiidxs, nsignal_params, nmodel_params) + if dxdparam0 is not None: + dxdparam0 = _gather_poi_noi_vector( + dxdparam0, noiidxs, nsignal_params, nmodel_params + ) return _gaussian_global_impacts( dxdtheta0, @@ -476,6 +538,8 @@ def gaussian_global_impacts_parms( beta_shape, systgroupidxs, data_cov_inv, + dxdparam0=dxdparam0, + varparam0=varparam0, ) @@ -491,6 +555,8 @@ def gaussian_global_impacts_obs( beta_shape, systgroupidxs, data_cov_inv=None, + dxdparam0=None, + varparam0=None, ): return _gaussian_global_impacts( dndtheta0, @@ -504,4 +570,6 @@ def gaussian_global_impacts_obs( beta_shape, systgroupidxs, data_cov_inv, + dxdparam0=dxdparam0, + varparam0=varparam0, ) diff --git a/rabbit/workspace.py b/rabbit/workspace.py index e069cfa..9995774 100644 --- a/rabbit/workspace.py +++ b/rabbit/workspace.py @@ -22,8 +22,9 @@ def getGroupedImpactsAxes( impact_names.append("binByBinStat") if per_process: impact_names.extend([f"binByBinStat{p}" for p in indata.procs.astype(str)]) - # ParamModel parameter groups are appended at the end, matching the column - # order produced by traditional_impacts.impacts_parms. + # ParamModel-related columns (parameter groups for the traditional + # impacts, prior sources for the global impacts) are appended at the + # end, matching the column order produced by the impact calculations. if extra_groups: impact_names.extend(extra_groups) return hist.axis.StrCategory(impact_names, name="impacts") @@ -63,10 +64,20 @@ def __init__(self, outdir, outname, fitter, postfix=None): # some information for the impact histograms systs = list(fitter.indata.systs.astype(str)) parms = list(fitter.parms.astype(str)) + # ParamModel prior centers act as additional global-impact sources; + # their columns are appended after the systs / syst groups. + if fitter.param_prior_active: + prior_idxs = fitter.param_prior_idxs.numpy() + prior_names = [f"{parms[int(i)]}_prior" for i in prior_idxs] + else: + prior_names = [] self.impact_axis = hist.axis.StrCategory(parms, name="impacts") - self.global_impact_axis = hist.axis.StrCategory(systs, name="impacts") + self.global_impact_axis = hist.axis.StrCategory( + systs + prior_names, name="impacts" + ) # ParamModel impact groups (e.g. SCETlib NP gamma_nu / F_eff) are only - # computed by the traditional impacts, so extend that axis only. + # computed by the traditional impacts, so they extend that axis only; + # the prior sources extend the global axes only. param_impact_group_names = [ name for name, _ in fitter._resolved_param_impact_groups() ] @@ -80,6 +91,7 @@ def __init__(self, outdir, outname, fitter, postfix=None): fitter.indata, bin_by_bin_stat=fitter.bbstat.enabled, per_process=fitter.bbstat.binByBinStatMode == "full", + extra_groups=prior_names, ) self.extension = "hdf5" From c674aaea9fdcfdd4b3af4207d4883402e9624dac Mon Sep 17 00:00:00 2001 From: Luca Lavezzo Date: Fri, 12 Jun 2026 16:33:15 -0400 Subject: [PATCH 08/16] Single constraint term over the full parameter vector MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit _compute_lc is now one expression: cw * 0.5 * (x - x0)^2 summed over the full effective parameter vector [poi, model_nui, theta], with cw = concat(param_cw, constraintweights) and x0 = concat(param_x0, theta0). No special-casing of the ParamModel block — a model without priors simply contributes cw = 0 entries. The full-NLL normalization is generalized to 0.5*log(2 pi / cw) for constrained entries (identical to the previous 0.9189*cw for the standard cw in {0, 1}). var_param_x0 is promoted to an attribute (the analogue of var_theta0) and the toy randomization of param_x0 is unconditional, mirroring theta0. Co-Authored-By: Claude Fable 5 --- rabbit/fitter.py | 88 +++++++++++++++++++----------------------------- 1 file changed, 34 insertions(+), 54 deletions(-) diff --git a/rabbit/fitter.py b/rabbit/fitter.py index 942ba99..cfd3d1f 100644 --- a/rabbit/fitter.py +++ b/rabbit/fitter.py @@ -362,6 +362,13 @@ def init_fit_parms( self.param_x0 = tf.Variable( self.param_x0default, trainable=False, name="param_x0" ) + # prefit variance of the ParamModel constraint centers (1/cw for + # priored params, 0 for free ones); the analogue of var_theta0 + self.var_param_x0 = tf.where( + self.param_cw == 0.0, + tf.zeros_like(self.param_cw), + tf.math.reciprocal(self.param_cw), + ) # Per-parameter prefit variance vector. Always allocated; the # prefit covariance is intrinsically diagonal so this is the @@ -809,15 +816,10 @@ def bayesassign(self): if self.param_model.nparams == 0: self.x.assign(theta_part) else: - param_sigma = tf.where( - self.param_cw > 0, - tf.math.rsqrt(self.param_cw), - tf.zeros_like(self.param_cw), - ) param_part = tf.where( self.param_cw > 0, self.param_x0 - + param_sigma + + tf.sqrt(self.var_param_x0) * tf.random.normal( shape=self.param_x0.shape, dtype=self.param_x0.dtype ), @@ -837,17 +839,11 @@ def frequentistassign(self): + tf.sqrt(self.var_theta0) * tf.random.normal(shape=self.theta0.shape, dtype=self.theta0.dtype) ) - if self.param_prior_active: - param_sigma = tf.where( - self.param_cw > 0, - tf.math.rsqrt(self.param_cw), - tf.zeros_like(self.param_cw), - ) - self.param_x0.assign( - self.param_x0default - + param_sigma - * tf.random.normal(shape=self.param_x0.shape, dtype=self.param_x0.dtype) - ) + self.param_x0.assign( + self.param_x0default + + tf.sqrt(self.var_param_x0) + * tf.random.normal(shape=self.param_x0.shape, dtype=self.param_x0.dtype) + ) self.bbstat.randomize_frequentist() def toyassign( @@ -1134,7 +1130,7 @@ def gaussian_global_impacts_parms(self): if self.param_prior_active: dxdparam0 = tf.gather(dxdparam_x0, self.param_prior_idxs, axis=1) - varparam0 = tf.gather(self._var_param_x0(), self.param_prior_idxs) + varparam0 = tf.gather(self.var_param_x0, self.param_prior_idxs) else: dxdparam0 = None varparam0 = None @@ -1400,17 +1396,6 @@ def _dndvars(self, fun): return n, dndtheta0, dndnobs, dndbeta0, dndparam_x0 - def _var_param_x0(self): - """Prefit variance of the ParamModel constraint centers (1/cw for - priored params, 0 for free ones).""" - return tf.where( - self.param_cw > 0, - tf.math.reciprocal( - tf.where(self.param_cw > 0, self.param_cw, tf.ones_like(self.param_cw)) - ), - tf.zeros_like(self.param_cw), - ) - def _compute_expected( self, fun_exp, inclusive=True, profile=False, full=True, need_observables=True ): @@ -1548,7 +1533,7 @@ def fun_n(): _, dndtheta0, dndnobs, dndbeta0, dndparam_x0 = self._dndvars(fun_n) if self.param_prior_active: dndparam0 = tf.gather(dndparam_x0, self.param_prior_idxs, axis=1) - varparam0 = tf.gather(self._var_param_x0(), self.param_prior_idxs) + varparam0 = tf.gather(self.var_param_x0, self.param_prior_idxs) else: dndparam0 = None varparam0 = None @@ -1862,7 +1847,7 @@ def fun_res(): # ParamModel prior centers contribute like additional auxiliary # observables with variance 1/cw res_cov += dresdparam_x0 @ ( - self._var_param_x0()[:, None] * tf.transpose(dresdparam_x0) + self.var_param_x0[:, None] * tf.transpose(dresdparam_x0) ) if self.covarianceFit: @@ -2028,31 +2013,26 @@ def reduced_nll(self): return self._compute_nll(full_nll=False) def _compute_lc(self, full_nll=False): - # constraints - theta = self.get_theta() - lc = self.indata.constraintweights * 0.5 * tf.square(theta - self.theta0) + # One constraint term over the full effective parameter vector + # [poi, model_nui, theta]: the ParamModel block is constrained by + # param_cw / param_x0 (declared priors; cw = 0 -> free) and the + # nuisance block by indata.constraintweights / theta0. + xeff = self.get_x() + cw = tf.concat([self.param_cw, self.indata.constraintweights], axis=0) + x0 = tf.concat([self.param_x0, self.theta0], axis=0) + lc = cw * 0.5 * tf.square(xeff - x0) if full_nll: - # normalization factor for normal distribution: log(1/sqrt(2*pi)) = -0.9189385332046727 - lc = lc + 0.9189385332046727 * self.indata.constraintweights - - total = tf.reduce_sum(lc) - - # ParamModel Gaussian priors (applied when the model declares them): - # same structure as the nuisance constraint above, with param_cw and - # param_x0 in the roles of constraintweights and theta0. - if self.param_prior_active: - pm_params = tf.concat([self.get_poi(), self.get_model_nui()], axis=0) - lc_pm = self.param_cw * 0.5 * tf.square(pm_params - self.param_x0) - if full_nll: - # 0.5*log(2π σ²) on entries with a prior. - lc_pm = lc_pm + tf.where( - self.param_prior_mask, - 0.9189385332046727 + tf.math.log(self.param_prior_sigmas), - tf.zeros_like(lc_pm), - ) - total = total + tf.reduce_sum(lc_pm) + # normalization factor 0.5*log(2*pi*sigma^2) for constrained + # entries, with sigma^2 = 1/cw and + # log(1/sqrt(2*pi)) = -0.9189385332046727 + lc = lc + tf.where( + cw > 0, + 0.9189385332046727 + - 0.5 * tf.math.log(tf.where(cw > 0, cw, tf.ones_like(cw))), + tf.zeros_like(lc), + ) - return total + return tf.reduce_sum(lc) def _compute_lbeta(self, beta, full_nll=False): return self.bbstat.lbeta(beta, full_nll=full_nll) From 6433ad7d8c60817edefdb1ab3436c311ca2b0747 Mon Sep 17 00:00:00 2001 From: Luca Lavezzo Date: Fri, 12 Jun 2026 17:07:39 -0400 Subject: [PATCH 09/16] Unified cw / x0 over the full parameter vector, replacing theta0 The constraint structure is now held in two full-length objects, index-aligned with x: cw constraint weights (1/sigma^2 for priored ParamModel params, indata.constraintweights for systs; 0 = unconstrained) x0 constraint centers (prior means / theta0), a tf.Variable var_x0 prefit center variances (1/cw; 0 where free) fitter.theta0, param_x0, param_cw, var_theta0 and var_param_x0 are gone. Every consumer is migrated, in most cases with less code since the +- nparams index offsets cancel: - _compute_lc is one term: cw * 0.5 * (x - x0)^2. - prefit_variance is a single tf.where over cw. - frequentistassign / bayesassign are single full-vector expressions. - _dxdvars / _dndvars differentiate w.r.t. x0 once; the dxdtheta0 / dxdparam0 pair and its plumbing through the gaussian global impacts collapse into a dxdx0 split helper. - _residuals_profiled has a single x0-variance term. - nonprofiled_impacts and global_asym_impacts operate on x0 directly (x0[idx] instead of theta0[idx - nparams]). - the saturated-projection re-init preserves the toy-randomized centers via the x0 syst-block slice. No behavior change: postfit results, impacts, closures and the global-asym test suite are identical to the previous two-block implementation. External code that accessed fitter.theta0 must use fitter.x0[fitter.param_model.nparams:] instead. Co-Authored-By: Claude Fable 5 --- bin/rabbit_fit.py | 10 +- rabbit/fitter.py | 242 +++++++++++--------------- rabbit/impacts/global_asym_impacts.py | 26 +-- rabbit/impacts/global_impacts.py | 47 +++-- rabbit/impacts/nonprofiled_impacts.py | 24 ++- 5 files changed, 165 insertions(+), 184 deletions(-) diff --git a/bin/rabbit_fit.py b/bin/rabbit_fit.py index 2af5caa..b6097da 100755 --- a/bin/rabbit_fit.py +++ b/bin/rabbit_fit.py @@ -410,7 +410,13 @@ def save_hists(args, mappings, fitter, ws, prefit=True, profile=False): fitter_saturated = copy.deepcopy(fitter) - toy_theta0 = tf.identity(fitter_saturated.theta0) + # preserve the (possibly toy-randomized) nuisance constraint + # centers across the re-init; the ParamModel block of x0 is + # rebuilt by init_fit_parms since the composite model has its + # own params + toy_theta0 = tf.identity( + fitter_saturated.x0[fitter_saturated.param_model.nparams :] + ) saved_regularizers = fitter_saturated.regularizers saved_tau = float(fitter_saturated.tau.numpy()) fitter_saturated.init_fit_parms( @@ -419,7 +425,7 @@ def save_hists(args, mappings, fitter, ws, prefit=True, profile=False): unblind=args.unblind, freeze_parameters=args.freezeParameters, ) - fitter_saturated.theta0.assign(toy_theta0) + fitter_saturated.x0[composite_model.nparams :].assign(toy_theta0) fitter_saturated.regularizers = saved_regularizers fitter_saturated.tau.assign(saved_tau) diff --git a/rabbit/fitter.py b/rabbit/fitter.py index cfd3d1f..84e1d0a 100644 --- a/rabbit/fitter.py +++ b/rabbit/fitter.py @@ -293,12 +293,12 @@ def init_fit_parms( # model. # # Priors use the same constraint structure as the nuisance - # constraints: param_cw is the per-parameter constraint weight - # (1/sigma^2; 0 = free) and param_x0 the constraint center (the - # prior mean) — the ParamModel-block analogues of - # indata.constraintweights and theta0. _compute_lc adds - # 0.5 * cw * (param - x0)^2, prefit variances are 1/cw, and the toy - # randomization treats the centers exactly like theta0. + # constraints: declared priors set the constraint weight + # (1/sigma^2; 0 = free) and the constraint center (the prior mean) + # of the ParamModel block in the full-length cw / x0 vectors built + # below. _compute_lc adds 0.5 * cw * (x - x0)^2 over the full + # parameter vector, prefit variances are 1/cw, and the toy + # randomization fluctuates the centers wherever cw > 0. self.param_prior_active = False pm = self.param_model param_cw_np = np.zeros(pm.nparams, dtype=np.float64) @@ -350,24 +350,37 @@ def init_fit_parms( if mask[i]: name = p.decode() if isinstance(p, bytes) else str(p) logger.info(f" {name}: μ={means[i]:.4g} σ={sigmas[i]:.4g}") - self.param_cw = tf.constant(param_cw_np, dtype=self.indata.dtype) - self.param_x0default = tf.constant(param_x0_np, dtype=self.indata.dtype) # indices (into the ParamModel block) of the priored params; impact # sources in the global-impacts calculations self.param_prior_idxs = tf.constant( np.where(param_cw_np > 0)[0], dtype=tf.int64 ) - # constraint centers for the ParamModel block; fluctuated in toys - # together with theta0 - self.param_x0 = tf.Variable( - self.param_x0default, trainable=False, name="param_x0" + + # Unified constraint structure over the full parameter vector, + # index-aligned with x = [params | systs]: + # cw constraint weights (1/sigma^2; 0 = unconstrained) + # x0 constraint centers (prior means / theta0), fluctuated in + # toys wherever cw > 0 + # var_x0 prefit variance of the centers (1/cw; 0 where free) + self.cw = tf.concat( + [ + tf.constant(param_cw_np, dtype=self.indata.dtype), + self.indata.constraintweights, + ], + axis=0, + ) + self.x0default = tf.concat( + [ + tf.constant(param_x0_np, dtype=self.indata.dtype), + self.theta0default, + ], + axis=0, ) - # prefit variance of the ParamModel constraint centers (1/cw for - # priored params, 0 for free ones); the analogue of var_theta0 - self.var_param_x0 = tf.where( - self.param_cw == 0.0, - tf.zeros_like(self.param_cw), - tf.math.reciprocal(self.param_cw), + self.x0 = tf.Variable(self.x0default, trainable=False, name="x0") + self.var_x0 = tf.where( + self.cw == 0.0, + tf.zeros_like(self.cw), + tf.math.reciprocal(self.cw), ) # Per-parameter prefit variance vector. Always allocated; the @@ -407,18 +420,6 @@ def init_fit_parms( self.indata.dtype, ) - # constraint minima for nuisance parameters - self.theta0 = tf.Variable( - self.theta0default, - trainable=False, - name="theta0", - ) - self.var_theta0 = tf.where( - self.indata.constraintweights == 0.0, - tf.zeros_like(self.indata.constraintweights), - tf.math.reciprocal(self.indata.constraintweights), - ) - # for freezing parameters self.frozen_params = [] self.frozen_params_mask = tf.Variable( @@ -714,25 +715,16 @@ def get_x(self): def prefit_variance(self, unconstrained_err=0.0): """Per-parameter prefit variance vector of length npar. - Free parameters (ParamModel params without a prior and unconstrained - nuisances) are assigned a placeholder variance of unconstrained_err**2 - (zero by default). Constrained entries take their variance from the - constraint term (1 / constraint weight), uniformly for priored - ParamModel params (param_cw) and constrained nuisances - (indata.constraintweights). + Unconstrained entries (cw = 0: ParamModel params without a prior and + unconstrained nuisances) are assigned a placeholder variance of + unconstrained_err**2 (zero by default); constrained entries take + their variance from the constraint term (1 / cw). """ - var_param = tf.where( - self.param_cw == 0.0, - unconstrained_err**2 - * tf.ones([self.param_model.nparams], dtype=self.indata.dtype), - tf.math.reciprocal(self.param_cw), + return tf.where( + self.cw == 0.0, + unconstrained_err**2 * tf.ones_like(self.cw), + tf.math.reciprocal(self.cw), ) - var_theta = tf.where( - self.indata.constraintweights == 0.0, - unconstrained_err**2, - tf.math.reciprocal(self.indata.constraintweights), - ) - return tf.concat([var_param, var_theta], axis=0) def prefit_covariance(self, unconstrained_err=0.0): """Full prefit covariance as a tf.linalg.LinearOperatorDiag. @@ -771,18 +763,22 @@ def set_nobs(self, values, variances=None): nobssafe = tf.where(values == 0.0, tf.constant(1.0, dtype=values.dtype), values) self.lognobs.assign(tf.math.log(nobssafe)) - def theta0defaultassign(self): - # reset all constraint centers: the nuisance block and the - # ParamModel-prior block - self.theta0.assign(self.theta0default) - self.param_x0.assign(self.param_x0default) + def x0defaultassign(self): + # reset all constraint centers + self.x0.assign(self.x0default) def xdefaultassign(self): if self.param_model.nparams == 0: - self.x.assign(self.theta0) + self.x.assign(self.x0) else: self.x.assign( - tf.concat([self.param_model.xparamdefault, self.theta0], axis=0) + tf.concat( + [ + self.param_model.xparamdefault, + self.x0[self.param_model.nparams :], + ], + axis=0, + ) ) def defaultassign(self): @@ -792,7 +788,7 @@ def defaultassign(self): self.var_prefit.assign(var_pre) if self.cov is not None: self.cov.assign(tf.linalg.diag(var_pre)) - self.theta0defaultassign() + self.x0defaultassign() if self.bbstat.enabled: self.bbstat.beta0_default_assign() self.bbstat.beta_default_assign() @@ -807,42 +803,34 @@ def defaultassign(self): def bayesassign(self): # Sample the parameter values from their priors: width sqrt(1/cw) - # around the constraint centers wherever cw > 0 (constrained - # nuisances and priored ParamModel params); free entries stay at - # their defaults. - theta_part = self.theta0 + tf.sqrt(self.var_theta0) * tf.random.normal( - shape=self.theta0.shape, dtype=self.theta0.dtype - ) + # around the constraint centers wherever cw > 0; free entries stay + # at their defaults (for cw = 0 nuisances the default is the x0 + # entry itself). if self.param_model.nparams == 0: - self.x.assign(theta_part) + defaults = self.x0 else: - param_part = tf.where( - self.param_cw > 0, - self.param_x0 - + tf.sqrt(self.var_param_x0) - * tf.random.normal( - shape=self.param_x0.shape, dtype=self.param_x0.dtype - ), - self.param_model.xparamdefault, + defaults = tf.concat( + [ + self.param_model.xparamdefault, + self.x0[self.param_model.nparams :], + ], + axis=0, ) - self.x.assign(tf.concat([param_part, theta_part], axis=0)) + sampled = self.x0 + tf.sqrt(self.var_x0) * tf.random.normal( + shape=self.x0.shape, dtype=self.x0.dtype + ) + self.x.assign(tf.where(self.cw > 0, sampled, defaults)) self.bbstat.randomize_bayes() def frequentistassign(self): # Fluctuate the constraint centers around their defaults with the - # prefit constraint widths, uniformly for the nuisance block - # (theta0) and the ParamModel-prior block (param_x0). Entries with - # cw = 0 (unconstrained) are not randomized. - self.theta0.assign( - self.theta0default - + tf.sqrt(self.var_theta0) - * tf.random.normal(shape=self.theta0.shape, dtype=self.theta0.dtype) - ) - self.param_x0.assign( - self.param_x0default - + tf.sqrt(self.var_param_x0) - * tf.random.normal(shape=self.param_x0.shape, dtype=self.param_x0.dtype) + # prefit constraint widths; entries with cw = 0 (unconstrained) are + # not randomized. + self.x0.assign( + self.x0default + + tf.sqrt(self.var_x0) + * tf.random.normal(shape=self.x0.shape, dtype=self.x0.dtype) ) self.bbstat.randomize_frequentist() @@ -1126,20 +1114,13 @@ def global_impacts_parms(self): @tf.function def gaussian_global_impacts_parms(self): - dxdtheta0, dxdnobs, dxdbeta0, dxdparam_x0 = self._dxdvars() - - if self.param_prior_active: - dxdparam0 = tf.gather(dxdparam_x0, self.param_prior_idxs, axis=1) - varparam0 = tf.gather(self.var_param_x0, self.param_prior_idxs) - else: - dxdparam0 = None - varparam0 = None + dxdx0, dxdnobs, dxdbeta0 = self._dxdvars() impacts, impacts_grouped = global_impacts.gaussian_global_impacts_parms( - dxdtheta0, + dxdx0, dxdnobs, dxdbeta0, - self.var_theta0, + self.var_x0, self.nobs if self.varnobs is None else self.varnobs, ( 1.0 @@ -1155,8 +1136,9 @@ def gaussian_global_impacts_parms(self): self.bbstat.beta_shape, self.indata.systgroupidxs, self.data_cov_inv, - dxdparam0=dxdparam0, - varparam0=varparam0, + param_prior_idxs=( + self.param_prior_idxs if self.param_prior_active else None + ), ) return impacts, impacts_grouped @@ -1309,13 +1291,12 @@ def global_asym_impacts_parms( def nonprofiled_impacts_parms(self, unconstrained_err=1.0): return nonprofiled_impacts.nonprofiled_impacts_parms( self.x, - self.theta0, + self.x0, self.frozen_indices, self.frozen_params, - self.indata.constraintweights, + self.cw, self.indata.systgroups, self.indata.systgroupidxs, - self.param_model.nparams, self.minimize, self.diagnostics, self.loss_val_grad_hess, @@ -1353,48 +1334,44 @@ def _pd2ldbeta2(self, profile=False): def _dxdvars(self): with tf.GradientTape() as t2: - t2.watch([self.theta0, self.nobs, self.bbstat.beta0, self.param_x0]) + t2.watch([self.x0, self.nobs, self.bbstat.beta0]) with tf.GradientTape() as t1: - t1.watch([self.theta0, self.nobs, self.bbstat.beta0, self.param_x0]) + t1.watch([self.x0, self.nobs, self.bbstat.beta0]) val = self._compute_loss() grad = t1.gradient(val, self.x) - pd2ldxdtheta0, pd2ldxdnobs, pd2ldxdbeta0, pd2ldxdparam_x0 = t2.jacobian( + pd2ldxdx0, pd2ldxdnobs, pd2ldxdbeta0 = t2.jacobian( grad, - [self.theta0, self.nobs, self.bbstat.beta0, self.param_x0], + [self.x0, self.nobs, self.bbstat.beta0], unconnected_gradients="zero", ) # cov is inverse hesse, thus cov ~ d2xd2l - dxdtheta0 = -self.cov @ pd2ldxdtheta0 + dxdx0 = -self.cov @ pd2ldxdx0 dxdnobs = -self.cov @ pd2ldxdnobs dxdbeta0 = -self.cov @ tf.reshape(pd2ldxdbeta0, [pd2ldxdbeta0.shape[0], -1]) - dxdparam_x0 = -self.cov @ pd2ldxdparam_x0 - return dxdtheta0, dxdnobs, dxdbeta0, dxdparam_x0 + return dxdx0, dxdnobs, dxdbeta0 def _dndvars(self, fun): with tf.GradientTape() as t: - t.watch([self.theta0, self.nobs, self.bbstat.beta0]) + t.watch([self.x0, self.nobs, self.bbstat.beta0]) n = fun() n_flat = tf.reshape(n, (-1,)) - pdndx, pdndtheta0, pdndnobs, pdndbeta0 = t.jacobian( + pdndx, pdndx0, pdndnobs, pdndbeta0 = t.jacobian( n_flat, - [self.x, self.theta0, self.nobs, self.bbstat.beta0], + [self.x, self.x0, self.nobs, self.bbstat.beta0], unconnected_gradients="zero", ) # apply chain rule to take into account correlations with the fit parameters - dxdtheta0, dxdnobs, dxdbeta0, dxdparam_x0 = self._dxdvars() + dxdx0, dxdnobs, dxdbeta0 = self._dxdvars() - dndtheta0 = pdndtheta0 + pdndx @ dxdtheta0 + dndx0 = pdndx0 + pdndx @ dxdx0 dndnobs = pdndnobs + pdndx @ dxdnobs dndbeta0 = tf.reshape(pdndbeta0, [pdndbeta0.shape[0], -1]) + pdndx @ dxdbeta0 - # the yields depend on param_x0 only through the fit parameters, so - # there is no partial term - dndparam_x0 = pdndx @ dxdparam_x0 - return n, dndtheta0, dndnobs, dndbeta0, dndparam_x0 + return n, dndx0, dndnobs, dndbeta0 def _compute_expected( self, fun_exp, inclusive=True, profile=False, full=True, need_observables=True @@ -1530,19 +1507,13 @@ def fun_n(): need_observables=need_observables, ) - _, dndtheta0, dndnobs, dndbeta0, dndparam_x0 = self._dndvars(fun_n) - if self.param_prior_active: - dndparam0 = tf.gather(dndparam_x0, self.param_prior_idxs, axis=1) - varparam0 = tf.gather(self.var_param_x0, self.param_prior_idxs) - else: - dndparam0 = None - varparam0 = None + _, dndx0, dndnobs, dndbeta0 = self._dndvars(fun_n) impacts_gaussian, impacts_gaussian_grouped = ( global_impacts.gaussian_global_impacts_obs( - dndtheta0, + dndx0, dndnobs, dndbeta0, - self.var_theta0, + self.var_x0, self.nobs if self.varnobs is None else self.varnobs, ( 1.0 @@ -1550,13 +1521,15 @@ def fun_n(): or not self.bbstat.enabled else 1.0 / self.bbstat.kstat ), + self.param_model.nparams, self.bbstat.enabled, self.bbstat.binByBinStatMode, self.bbstat.beta_shape, self.indata.systgroupidxs, self.data_cov_inv, - dxdparam0=dndparam0, - varparam0=varparam0, + param_prior_idxs=( + self.param_prior_idxs if self.param_prior_active else None + ), ) ) else: @@ -1837,18 +1810,9 @@ def fun_res(): observed = fun(None, self.nobs) return expected - observed - residuals, dresdtheta0, dresdnobs, dresdbeta0, dresdparam_x0 = self._dndvars( - fun_res - ) + residuals, dresdx0, dresdnobs, dresdbeta0 = self._dndvars(fun_res) - res_cov = dresdtheta0 @ (self.var_theta0[:, None] * tf.transpose(dresdtheta0)) - - if self.param_prior_active: - # ParamModel prior centers contribute like additional auxiliary - # observables with variance 1/cw - res_cov += dresdparam_x0 @ ( - self.var_param_x0[:, None] * tf.transpose(dresdparam_x0) - ) + res_cov = dresdx0 @ (self.var_x0[:, None] * tf.transpose(dresdx0)) if self.covarianceFit: res_cov_stat = dresdnobs @ tf.linalg.solve( @@ -2017,10 +1981,8 @@ def _compute_lc(self, full_nll=False): # [poi, model_nui, theta]: the ParamModel block is constrained by # param_cw / param_x0 (declared priors; cw = 0 -> free) and the # nuisance block by indata.constraintweights / theta0. - xeff = self.get_x() - cw = tf.concat([self.param_cw, self.indata.constraintweights], axis=0) - x0 = tf.concat([self.param_x0, self.theta0], axis=0) - lc = cw * 0.5 * tf.square(xeff - x0) + cw = self.cw + lc = cw * 0.5 * tf.square(self.get_x() - self.x0) if full_nll: # normalization factor 0.5*log(2*pi*sigma^2) for constrained # entries, with sigma^2 = 1/cw and diff --git a/rabbit/impacts/global_asym_impacts.py b/rabbit/impacts/global_asym_impacts.py index a169cc0..07478e6 100644 --- a/rabbit/impacts/global_asym_impacts.py +++ b/rabbit/impacts/global_asym_impacts.py @@ -89,8 +89,8 @@ def global_asym_impacts_parms( # Snapshot postfit nominal state to restore between iterations. x_nom = tf.identity(fitter.x.value()) - theta0_nom = tf.identity(fitter.theta0.value()) - theta0_nom_np = theta0_nom.numpy() + x0_nom = tf.identity(fitter.x0.value()) + x0_nom_np = x0_nom.numpy() x_nom_np = x_nom.numpy() logger.info( @@ -100,11 +100,11 @@ def global_asym_impacts_parms( ) # Optional Gaussian-approximation warm-start. - # dxdtheta0 has shape [npar, nsyst]; column i gives the linearised + # dxdx0 has shape [npar, npar]; column nparams + i gives the linearised # response of the postfit minimum to a unit shift of theta0[i]. Computing # it once before the loop is the same cost as one --gaussianGlobalImpacts # call. - dxdtheta0_np = None + dxdx0_np = None if linear_warmstart: if fitter.cov is None: raise RuntimeError( @@ -112,10 +112,10 @@ def global_asym_impacts_parms( "(incompatible with --noHessian)." ) t_lws = time.perf_counter() - dxdtheta0_tf, _, _ = fitter._dxdvars() - dxdtheta0_np = dxdtheta0_tf.numpy() + dxdx0_tf, _, _ = fitter._dxdvars() + dxdx0_np = dxdx0_tf.numpy() logger.info( - f"global_asym_impacts: dxdtheta0 prepared in " + f"global_asym_impacts: dxdx0 prepared in " f"{time.perf_counter() - t_lws:.2f}s" ) @@ -133,21 +133,21 @@ def global_asym_impacts_parms( shift = float(sign) * float(sigma) # Always shift the constraint center for nuisance i by `shift`. - theta0_shifted = theta0_nom_np.copy() - theta0_shifted[i] += shift + x0_shifted = x0_nom_np.copy() + x0_shifted[nparams + i] += shift # Warm-start x. With linear_warmstart, use the Gaussian-approx new - # minimum x_nom + dxdtheta0[:, i] * shift -- on near-Gaussian + # minimum x_nom + dxdx0[:, nparams+i] * shift -- on near-Gaussian # nuisances this lands at the new minimum to within roundoff. # Without it, just shift x[nparams+i] by `shift` so the nuisance # itself starts at the new constraint center. if linear_warmstart: - x_shifted = x_nom_np + dxdtheta0_np[:, i] * shift + x_shifted = x_nom_np + dxdx0_np[:, nparams + i] * shift else: x_shifted = x_nom_np.copy() x_shifted[nparams + i] += shift - fitter.theta0.assign(theta0_shifted) + fitter.x0.assign(x0_shifted) fitter.x.assign(x_shifted) try: @@ -165,7 +165,7 @@ def global_asym_impacts_parms( logger.info(f" took {t_per[k]:.2f}s") # Restore the fit state so downstream postfit computations see the nominal. - fitter.theta0.assign(theta0_nom) + fitter.x0.assign(x0_nom) fitter.x.assign(x_nom) if fitter.bbstat.enabled: fitter._profile_beta() diff --git a/rabbit/impacts/global_impacts.py b/rabbit/impacts/global_impacts.py index 4f3402a..ab20431 100644 --- a/rabbit/impacts/global_impacts.py +++ b/rabbit/impacts/global_impacts.py @@ -497,11 +497,25 @@ def _gaussian_global_impacts( return impacts, impacts_grouped +def _split_x0_sources(dxdx0, varx0, nmodel_params, param_prior_idxs): + """Split full-x0 derivatives into the nuisance block and the optional + ParamModel-prior source columns.""" + dxdtheta0 = dxdx0[:, nmodel_params:] + vartheta0 = varx0[nmodel_params:] + if param_prior_idxs is not None: + dxdparam0 = tf.gather(dxdx0, param_prior_idxs, axis=1) + varparam0 = tf.gather(varx0, param_prior_idxs) + else: + dxdparam0 = None + varparam0 = None + return dxdtheta0, vartheta0, dxdparam0, varparam0 + + def gaussian_global_impacts_parms( - dxdtheta0, + dxdx0, dxdnobs, dxdbeta0, - vartheta0, + varx0, varnobs, varbeta0, nsignal_params, @@ -512,19 +526,16 @@ def gaussian_global_impacts_parms( beta_shape, systgroupidxs, data_cov_inv=None, - dxdparam0=None, - varparam0=None, + param_prior_idxs=None, ): # compute impacts for pois and nois - dxdtheta0 = _gather_poi_noi_vector( - dxdtheta0, noiidxs, nsignal_params, nmodel_params - ) + dxdx0 = _gather_poi_noi_vector(dxdx0, noiidxs, nsignal_params, nmodel_params) dxdnobs = _gather_poi_noi_vector(dxdnobs, noiidxs, nsignal_params, nmodel_params) dxdbeta0 = _gather_poi_noi_vector(dxdbeta0, noiidxs, nsignal_params, nmodel_params) - if dxdparam0 is not None: - dxdparam0 = _gather_poi_noi_vector( - dxdparam0, noiidxs, nsignal_params, nmodel_params - ) + + dxdtheta0, vartheta0, dxdparam0, varparam0 = _split_x0_sources( + dxdx0, varx0, nmodel_params, param_prior_idxs + ) return _gaussian_global_impacts( dxdtheta0, @@ -544,20 +555,24 @@ def gaussian_global_impacts_parms( def gaussian_global_impacts_obs( - dndtheta0, + dndx0, dndnobs, dndbeta0, - vartheta0, + varx0, varnobs, varbeta0, + nmodel_params, bin_by_bin_stat, bin_by_bin_stat_mode, beta_shape, systgroupidxs, data_cov_inv=None, - dxdparam0=None, - varparam0=None, + param_prior_idxs=None, ): + dndtheta0, vartheta0, dndparam0, varparam0 = _split_x0_sources( + dndx0, varx0, nmodel_params, param_prior_idxs + ) + return _gaussian_global_impacts( dndtheta0, dndnobs, @@ -570,6 +585,6 @@ def gaussian_global_impacts_obs( beta_shape, systgroupidxs, data_cov_inv, - dxdparam0=dxdparam0, + dxdparam0=dndparam0, varparam0=varparam0, ) diff --git a/rabbit/impacts/nonprofiled_impacts.py b/rabbit/impacts/nonprofiled_impacts.py index 650df20..7286b04 100644 --- a/rabbit/impacts/nonprofiled_impacts.py +++ b/rabbit/impacts/nonprofiled_impacts.py @@ -22,13 +22,12 @@ def _envelope(values): def nonprofiled_impacts_parms( x, - theta0, + x0, frozen_indices, frozen_params, - constraintweights, + cw, systgroups, systgroupidxs, - nparams, minimize_fn, diagnostics=False, loss_val_grad_hess_fn=None, @@ -37,13 +36,12 @@ def nonprofiled_impacts_parms( """ Args: x: TF Variable holding all fit parameters (POIs + nuisances). - theta0: TF Variable list of nuisance parameter central values. + x0: TF Variable of constraint centers, index-aligned with x. frozen_indices: indices (into x) of the frozen parameters. frozen_params: names of the frozen parameters. - constraintweights: constraint weights for each nuisance parameter. + cw: constraint weights, index-aligned with x. systgroups: systematic group names. systgroupidxs: per-group lists of nuisance parameter indices. - 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. @@ -53,21 +51,21 @@ def nonprofiled_impacts_parms( x_tmp_tiled = tf.tile(tf.reshape(x_tmp, [1, 1, -1]), [len(frozen_indices), 2, 1]) nonprofiled_impacts = tf.Variable(x_tmp_tiled) - theta0_tmp = tf.identity(theta0.value()) + x0_tmp = tf.identity(x0.value()) - err_theta = tf.where( - constraintweights == 0.0, + err_x0 = tf.where( + cw == 0.0, unconstrained_err, - tf.math.reciprocal(constraintweights), + tf.math.reciprocal(cw), ) for i, idx in enumerate(frozen_indices): logger.info(f"Now at parameter {frozen_params[i]}") for j, sign in enumerate((1, -1)): - variation = sign * err_theta[idx - nparams] + theta0_tmp[idx - nparams] + variation = sign * err_x0[idx] + x0_tmp[idx] # vary the non-profiled parameter - theta0[idx - nparams].assign(variation) + x0[idx].assign(variation) x[idx].assign( variation ) # this should not be needed but should accelerate the minimization @@ -83,7 +81,7 @@ def nonprofiled_impacts_parms( x.assign(x_tmp) # back to original value - theta0[idx - nparams].assign(theta0_tmp[idx - nparams]) + x0[idx].assign(x0_tmp[idx]) impact_group_names = [] impact_groups = [] From 3f2341b0a676423ef92930e7fc8fca7805f8dbfe Mon Sep 17 00:00:00 2001 From: Luca Lavezzo Date: Fri, 12 Jun 2026 17:37:08 -0400 Subject: [PATCH 10/16] Parameter matching via re.fullmatch, log resolved unblind parameters An expression that names one parameter exactly can no longer also match parameters whose names merely extend it (prefix matches could silently unblind more than intended); families are matched with an explicit pattern, e.g. alphaS.*. The parameters an --unblind expression resolves to are now reported at INFO level. Co-Authored-By: Claude Fable 5 --- rabbit/fitter.py | 22 +++++++++++++++++----- rabbit/parsing.py | 12 ++++++++---- 2 files changed, 25 insertions(+), 9 deletions(-) diff --git a/rabbit/fitter.py b/rabbit/fitter.py index 84e1d0a..bd8791f 100644 --- a/rabbit/fitter.py +++ b/rabbit/fitter.py @@ -32,10 +32,14 @@ def match_regexp_params(regular_expressions, parameter_names): # Match parameter names against a list of expressions where each entry may - # be either an exact parameter name or a regex (re.match, anchored to the - # start). Returns the union of exact and regex matches, preserving the - # parameter_names order and de-duplicating. Mixing exact and regex entries - # in the same call is supported. + # be either an exact parameter name or a regex matched against the FULL + # parameter name (re.fullmatch). Full anchoring means an expression that + # names one parameter exactly can never also match parameters whose names + # merely extend it (important for --unblind, where a prefix match would + # silently unblind more than intended); match a family of parameters with + # an explicit pattern, e.g. 'alphaS.*'. Returns the union of exact and + # regex matches, preserving the parameter_names order and de-duplicating. + # Mixing exact and regex entries in the same call is supported. if isinstance(regular_expressions, str): regular_expressions = [regular_expressions] @@ -47,7 +51,7 @@ def match_regexp_params(regular_expressions, parameter_names): for s in parameter_names: decoded = s.decode() if hasattr(s, "decode") else s if decoded in exact_lookup or any( - r.match(decoded) for r in compiled_expressions + r.fullmatch(decoded) for r in compiled_expressions ): if decoded not in seen: seen.add(decoded) @@ -574,6 +578,14 @@ def init_blinding_values( unblind_parameters = match_regexp_params( unblind_parameter_expressions, all_param_names ) + # unblinding is sensitive: always report exactly which parameters the + # expressions resolved to, so an over-broad pattern is visible + if unblind_parameters: + unblind_names = [ + p.decode() if isinstance(p, bytes) else str(p) + for p in unblind_parameters + ] + logger.info(f"Unblinding {len(unblind_names)} parameters: {unblind_names}") # check if dataset is an integer (i.e. if it is real data or not) and use this to choose the random seed is_dataobs_int = np.sum( diff --git a/rabbit/parsing.py b/rabbit/parsing.py index ad894a8..7bf37f2 100644 --- a/rabbit/parsing.py +++ b/rabbit/parsing.py @@ -269,8 +269,10 @@ def common_parser(): nargs="*", action=OptionalListAction, help=""" - Specify list of regex to unblind matching parameters of interest. - E.g. use '--unblind ^signal$' to unblind a parameter named signal or '--unblind' to unblind all. + Specify list of expressions to unblind matching parameters of interest. + Each entry is an exact parameter name or a regex matched against the full parameter + name (use e.g. 'alphaS.*' to match a family of parameters). + E.g. use '--unblind signal' to unblind a parameter named signal or '--unblind' to unblind all. """, ) parser.add_argument( @@ -279,7 +281,8 @@ def common_parser(): default=[], nargs="*", help=""" - Specify list of regex defining groups of parameters that share a single deterministic + Specify list of regex (matched against the full parameter name) defining groups of + parameters that share a single deterministic blinding offset (seeded from the regex string itself). Useful to keep relative pulls / differences between matched parameters meaningful while still blinding their absolute values. E.g. '--blindingGroup ^alphaS_y\\d+$' applies the same offset to all alphaS @@ -300,7 +303,8 @@ def common_parser(): default=[], nargs="+", help=""" - Specify list of regex to freeze matching parameters of interest. + Specify list of expressions to freeze matching parameters of interest + (exact names or regex matched against the full parameter name). """, ) parser.add_argument( From 32b8d3cc2c288f406e7213e9f610c2642bf0c2e7 Mon Sep 17 00:00:00 2001 From: Luca Lavezzo Date: Fri, 12 Jun 2026 17:37:23 -0400 Subject: [PATCH 11/16] Fix nonprofiled variation width for non-unit constraint weights The variation magnitude used 1/cw (the variance) instead of 1/sqrt(cw) (the sigma); identical for the 0/1 tensor constraint weights, but wrong for ParamModel priors where cw = 1/sigma^2 is genuinely non-unit. Co-Authored-By: Claude Fable 5 --- rabbit/impacts/nonprofiled_impacts.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/rabbit/impacts/nonprofiled_impacts.py b/rabbit/impacts/nonprofiled_impacts.py index 7286b04..870c4fc 100644 --- a/rabbit/impacts/nonprofiled_impacts.py +++ b/rabbit/impacts/nonprofiled_impacts.py @@ -53,10 +53,12 @@ def nonprofiled_impacts_parms( x0_tmp = tf.identity(x0.value()) + # prefit sigma = 1/sqrt(cw); the distinction from the variance 1/cw + # matters since ParamModel priors introduce genuinely non-unit cw err_x0 = tf.where( cw == 0.0, unconstrained_err, - tf.math.reciprocal(cw), + tf.math.rsqrt(cw), ) for i, idx in enumerate(frozen_indices): From 7bd0d81b5afc9324e3795439c3387ee7c14ab221 Mon Sep 17 00:00:00 2001 From: Luca Lavezzo Date: Fri, 12 Jun 2026 17:37:29 -0400 Subject: [PATCH 12/16] Preserve blinding groups and prior centers across the saturated re-init The saturated-projection re-init dropped --blindingGroup (falling back to per-name blinding seeds) and reset the ParamModel block of x0 to its defaults, losing toy-fluctuated prior centers. Both are now carried through the composite [POIs | POUs] permutation. Co-Authored-By: Claude Fable 5 --- bin/rabbit_fit.py | 28 ++++++++++++++++++++-------- 1 file changed, 20 insertions(+), 8 deletions(-) diff --git a/bin/rabbit_fit.py b/bin/rabbit_fit.py index b6097da..847b6f2 100755 --- a/bin/rabbit_fit.py +++ b/bin/rabbit_fit.py @@ -410,22 +410,34 @@ def save_hists(args, mappings, fitter, ws, prefit=True, profile=False): fitter_saturated = copy.deepcopy(fitter) - # preserve the (possibly toy-randomized) nuisance constraint - # centers across the re-init; the ParamModel block of x0 is - # rebuilt by init_fit_parms since the composite model has its - # own params - toy_theta0 = tf.identity( - fitter_saturated.x0[fitter_saturated.param_model.nparams :] - ) + # preserve the (possibly toy-randomized) constraint centers + # across the re-init: the theta block maps 1:1, and the + # original model's prior centers land at [0:npoi] and + # [composite.npoi : composite.npoi+npou] in the composite + # [POIs | POUs] layout; the saturated model's own params keep + # the freshly initialized centers + orig_model = fitter_saturated.param_model + toy_x0 = tf.identity(fitter_saturated.x0.value()) saved_regularizers = fitter_saturated.regularizers saved_tau = float(fitter_saturated.tau.numpy()) fitter_saturated.init_fit_parms( composite_model, args.setConstraintMinimum, unblind=args.unblind, + blinding_group=args.blindingGroup, freeze_parameters=args.freezeParameters, ) - fitter_saturated.x0[composite_model.nparams :].assign(toy_theta0) + fitter_saturated.x0[composite_model.nparams :].assign( + toy_x0[orig_model.nparams :] + ) + if orig_model.npoi > 0: + fitter_saturated.x0[: orig_model.npoi].assign( + toy_x0[: orig_model.npoi] + ) + if orig_model.npou > 0: + fitter_saturated.x0[ + composite_model.npoi : composite_model.npoi + orig_model.npou + ].assign(toy_x0[orig_model.npoi : orig_model.nparams]) fitter_saturated.regularizers = saved_regularizers fitter_saturated.tau.assign(saved_tau) From 1c8e4f63db1939e355345bbeb33abfddcfa48f78 Mon Sep 17 00:00:00 2001 From: Luca Lavezzo Date: Mon, 15 Jun 2026 10:18:27 -0400 Subject: [PATCH 13/16] Address review: build prior arrays in indata dtype, center free x0 at default - prior_sigmas/means and the cw/x0 numpy intermediates are built in self.indata.dtype (via as_numpy_dtype) rather than a hardcoded np.float64, matching how the rest of init_fit_parms respects the input dtype (davidwalter2 review on #139). - free (unpriored) entries of the ParamModel x0 block now default to the parameter own default instead of 0. cw = 0 makes this irrelevant to the likelihood, but nonprofiled_impacts reads x0 as a parameter natural center, so a frozen free param is now varied around its default rather than around 0. With no priors declared x0default now equals xdefault. - fix a stale theta0 reference in the global_asym_impacts docstring. Co-Authored-By: Claude Fable 5 --- rabbit/fitter.py | 20 ++++++++++++++------ rabbit/impacts/global_asym_impacts.py | 4 ++-- 2 files changed, 16 insertions(+), 8 deletions(-) diff --git a/rabbit/fitter.py b/rabbit/fitter.py index bd8791f..565c641 100644 --- a/rabbit/fitter.py +++ b/rabbit/fitter.py @@ -305,11 +305,18 @@ def init_fit_parms( # randomization fluctuates the centers wherever cw > 0. self.param_prior_active = False pm = self.param_model - param_cw_np = np.zeros(pm.nparams, dtype=np.float64) - param_x0_np = np.zeros(pm.nparams, dtype=np.float64) + np_dtype = self.indata.dtype.as_numpy_dtype + param_cw_np = np.zeros(pm.nparams, dtype=np_dtype) + # Constraint centers default to each parameter's own default; declared + # priors override the masked entries with their prior mean below. + # Keeping the default (rather than 0) for the free entries matters for + # consumers that read x0 as a parameter's natural center (e.g. + # nonprofiled impacts), even though cw = 0 makes it irrelevant to the + # likelihood itself. + param_x0_np = pm.xparamdefault.numpy().astype(np_dtype) sigmas = getattr(pm, "prior_sigmas", None) if sigmas is not None: - sigmas = np.asarray(sigmas, dtype=np.float64) + sigmas = np.asarray(sigmas, dtype=np_dtype) if sigmas.shape != (pm.nparams,): raise ValueError( f"param_model.prior_sigmas must have shape ({pm.nparams},); " @@ -317,9 +324,9 @@ def init_fit_parms( ) means = getattr(pm, "prior_means", None) if means is None: - means = pm.xparamdefault.numpy().astype(np.float64) + means = pm.xparamdefault.numpy().astype(np_dtype) else: - means = np.asarray(means, dtype=np.float64) + means = np.asarray(means, dtype=np_dtype) if means.shape != (pm.nparams,): raise ValueError( f"param_model.prior_means must have shape ({pm.nparams},); " @@ -345,7 +352,8 @@ def init_fit_parms( np.where(mask, means, np.nan), dtype=self.indata.dtype ) param_cw_np = np.where(mask, 1.0 / sigmas**2, 0.0) - param_x0_np = np.where(mask, means, 0.0) + # priored entries -> prior mean; free entries keep their default + param_x0_np = np.where(mask, means, param_x0_np) logger.info( f"[paramPriors] applying Gaussian priors to " f"{n_priored}/{pm.nparams} ParamModel params:" diff --git a/rabbit/impacts/global_asym_impacts.py b/rabbit/impacts/global_asym_impacts.py index 07478e6..465f0b3 100644 --- a/rabbit/impacts/global_asym_impacts.py +++ b/rabbit/impacts/global_asym_impacts.py @@ -54,11 +54,11 @@ def global_asym_impacts_parms( signs=(-1, 1), linear_warmstart=False, ): - """Run a per-nuisance theta0-shift + re-fit and assemble the asymmetric + """Run a per-nuisance x0-shift + re-fit and assemble the asymmetric global impact tensor. Args: - fitter: the Fitter instance (used for x, theta0 and minimize). + fitter: the Fitter instance (used for x, x0 and minimize). selected_idxs: indices into the syst axis (0..nsyst-1) of nuisances to scan. selected_names: names of those nuisances (bytes), used as impact-axis From 68c6f0efa5dfb5f402c02ef7dfcf68ad14b236f1 Mon Sep 17 00:00:00 2001 From: Luca Lavezzo Date: Mon, 15 Jun 2026 10:49:47 -0400 Subject: [PATCH 14/16] Differentiate global impacts only w.r.t. constraint sources; scan priored params in asym impacts Item 5 (exact + cheaper): x0 enters the NLL only through cw*(x-x0)^2, so the loss derivative w.r.t. any center with cw = 0 is identically zero. _dxdvars now perturbs only the constraint sources (the nuisance block plus priored params) via a scattered tangent and returns the response split into (dxdtheta0, dxdparam0), shrinking the Jacobian from [npar, npar] to [npar, nsyst (+ n_prior)]; the gaussian global-impacts helpers take the pre-split pieces (dropping _split_x0_sources). _dndvars likewise drops its x0 watch since the forward model has no x0 dependence (verified d(yields)/dx0 = 0 exactly). Validated bit-identical to the previous code on the test tensor: parms gaussian/likelihood impacts, observable impacts, and the profiled-chi2 residual covariance all unchanged. Item 6b: global_asym_impacts now scans priored ParamModel params as sources too (full-x indexed, labelled _prior), matching the source set of gaussian_global_impacts_parms, with ParamModel impact groups folded into the grouped envelopes. The per-source center shift is scaled by the source prefit width sqrt(var_x0) (1 for unit-constrained nuisances, the prior sigma for priored params), so a unit-sigma asym impact stays in the same units as the gaussian one. On the test tensor the sig_prior asym impact (-1.10e-2, +1.11e-2) reproduces the gaussian value (1.106e-2) to <1%, the small spread being genuine non-Gaussianity; nuisance impacts are unchanged. Co-Authored-By: Claude Fable 5 --- rabbit/fitter.py | 195 +++++++++++++++++++------- rabbit/impacts/global_asym_impacts.py | 121 ++++++++++------ rabbit/impacts/global_impacts.py | 50 +++---- 3 files changed, 242 insertions(+), 124 deletions(-) diff --git a/rabbit/fitter.py b/rabbit/fitter.py index 565c641..05e43f7 100644 --- a/rabbit/fitter.py +++ b/rabbit/fitter.py @@ -1134,13 +1134,16 @@ def global_impacts_parms(self): @tf.function def gaussian_global_impacts_parms(self): - dxdx0, dxdnobs, dxdbeta0 = self._dxdvars() + dxdtheta0, dxdparam0, dxdnobs, dxdbeta0 = self._dxdvars() + vartheta0, varparam0 = self._x0_source_vars() impacts, impacts_grouped = global_impacts.gaussian_global_impacts_parms( - dxdx0, + dxdtheta0, + dxdparam0, dxdnobs, dxdbeta0, - self.var_x0, + vartheta0, + varparam0, self.nobs if self.varnobs is None else self.varnobs, ( 1.0 @@ -1156,9 +1159,6 @@ def gaussian_global_impacts_parms(self): self.bbstat.beta_shape, self.indata.systgroupidxs, self.data_cov_inv, - param_prior_idxs=( - self.param_prior_idxs if self.param_prior_active else None - ), ) return impacts, impacts_grouped @@ -1260,9 +1260,12 @@ def global_asym_impacts_parms( ): """Fully likelihood-based asymmetric global impacts. - For each selected nuisance i, shift theta0[i] by +/- sigma (in units of - the prefit constraint width) and re-run the full fit. POI shifts at - each sign are the asymmetric global impacts. + For each selected source, shift its constraint center x0[idx] by +/- + sigma (in units of the prefit constraint width) and re-run the full + fit. POI/NOI shifts at each sign are the asymmetric global impacts. + Sources are the constrained nuisances and any priored ParamModel + params (exposed as _prior, matching the source set of + gaussian_global_impacts_parms so the two agree in the Gaussian limit). Unconstrained nuisances (constraintweight = 0) are always skipped: they have no prefit sigma, and their theta0 does not enter the NLL, @@ -1270,13 +1273,14 @@ def global_asym_impacts_parms( (zero impact at the cost of two full fits). Args: - include: optional regex(es) restricting which nuisances to scan. - exclude: optional regex(es) excluding nuisances from the scan. + include: optional regex(es) restricting which sources to scan + (matched against nuisance names and bare priored-param names). + exclude: optional regex(es) excluding sources from the scan. sigma: shift magnitude in prefit-sigma units. linear_warmstart: experimental, see global_asym_impacts.global_asym_impacts_parms. """ - nsyst = self.indata.nsyst + nparams = self.param_model.nparams cw = self.indata.constraintweights.numpy() syst_names = np.array(self.indata.systs).astype(bytes) @@ -1284,26 +1288,59 @@ def global_asym_impacts_parms( # finite prefit sigma to shift by, and the refit would be a no-op. selected = cw > 0 if include is not None: - keep = match_regexp_params(include, syst_names) - keep_set = set(keep) + keep_set = set(match_regexp_params(include, syst_names)) selected &= np.array([n in keep_set for n in syst_names]) if exclude is not None: - drop = match_regexp_params(exclude, syst_names) - drop_set = set(drop) + drop_set = set(match_regexp_params(exclude, syst_names)) selected &= np.array([n not in drop_set for n in syst_names]) - selected_idxs = np.where(selected)[0] - selected_names = syst_names[selected_idxs] + syst_sel = np.where(selected)[0] + # nuisance sources, in full-x coordinates (nparams + syst index) + src_x_idxs = [nparams + int(i) for i in syst_sel] + src_names = [syst_names[int(i)] for i in syst_sel] + n_nuis = len(src_x_idxs) + + # priored ParamModel params are sources too, labelled _prior to + # match the gaussian/likelihood global impacts; same include/exclude, + # matched against the bare param name. param_prior_idxs are already + # full-x indices (the param block leads x). + if self.param_prior_active: + parms = np.array(self.parms).astype(bytes) + prior_idxs = self.param_prior_idxs.numpy().astype(int) + prior_names = parms[prior_idxs] + psel = np.ones(len(prior_idxs), dtype=bool) + if include is not None: + keep_set = set(match_regexp_params(include, prior_names)) + psel &= np.array([n in keep_set for n in prior_names]) + if exclude is not None: + drop_set = set(match_regexp_params(exclude, prior_names)) + psel &= np.array([n not in drop_set for n in prior_names]) + for p, nm, ok in zip(prior_idxs, prior_names, psel): + if ok: + src_x_idxs.append(int(p)) + src_names.append(nm + b"_prior") + + # group membership in full-x indices: syst groups + ParamModel impact + # groups (the module function keeps only the scanned members). + group_members = {} + for gname, gidxs in zip(self.indata.systgroups, self.indata.systgroupidxs): + group_members[gname] = [nparams + int(i) for i in np.asarray(gidxs)] + for label, idxs in self._resolved_param_impact_groups(): + key = label.encode() if isinstance(label, str) else label + group_members.setdefault(key, []) + group_members[key].extend(int(i) for i in idxs) logger.info( - f"global_asym_impacts_parms: selected {len(selected_idxs)}/{nsyst} " - f"nuisances (unconstrained nuisances always excluded)" + f"global_asym_impacts_parms: selected {len(src_x_idxs)} sources " + f"({n_nuis} nuisances + {len(src_x_idxs) - n_nuis} priored params; " + f"unconstrained nuisances always excluded)" ) return global_asym_impacts.global_asym_impacts_parms( self, - selected_idxs, - selected_names, + src_x_idxs, + src_names, + group_members=group_members, sigma=sigma, linear_warmstart=linear_warmstart, ) @@ -1352,46 +1389,91 @@ def _pd2ldbeta2(self, profile=False): pd2ldbeta2 = t2.gradient(pdldbeta, self.bbstat.ubeta) return pd2ldbeta2 + def _x0_constraint_source_idxs(self): + # Full-x indices of the constraint sources reported by the global + # impacts, in column order [nuisance/theta block | priored params]. + # x0 enters the NLL only through cw * (x - x0)^2, so the derivative + # w.r.t. any center with cw = 0 is identically zero; dropping those + # columns is exact, not an approximation, and keeps the Jacobian at + # [npar, nsyst (+ n_prior)] instead of [npar, npar]. The param block + # sits first in x, so param_prior_idxs are already full-x indices. + nparams = self.param_model.nparams + theta_idxs = tf.range(nparams, nparams + self.indata.nsyst, dtype=tf.int64) + if self.param_prior_active: + return tf.concat([theta_idxs, self.param_prior_idxs], axis=0) + return theta_idxs + + def _x0_source_vars(self): + # Prefit variances of the constraint sources, in the same column + # order / split as _dxdvars / _dndvars: (vartheta0, varparam0). + vartheta0 = self.var_x0[self.param_model.nparams :] + if self.param_prior_active: + varparam0 = tf.gather(self.var_x0, self.param_prior_idxs) + else: + varparam0 = None + return vartheta0, varparam0 + def _dxdvars(self): + # Differentiate the loss only w.r.t. the constraint sources (the + # nuisance block plus any priored params) by perturbing those x0 + # entries with `ds` and scattering back into the full vector; see + # _x0_constraint_source_idxs. The resulting columns are ordered + # [theta block | priored params] and split below. + src_idxs = self._x0_constraint_source_idxs() + n_src = self.indata.nsyst + ( + int(self.param_prior_idxs.shape[0]) if self.param_prior_active else 0 + ) + scatter_idx = src_idxs[:, None] with tf.GradientTape() as t2: - t2.watch([self.x0, self.nobs, self.bbstat.beta0]) + ds = tf.zeros([n_src], dtype=self.indata.dtype) + t2.watch([ds, self.nobs, self.bbstat.beta0]) with tf.GradientTape() as t1: - t1.watch([self.x0, self.nobs, self.bbstat.beta0]) - val = self._compute_loss() + t1.watch([ds, self.nobs, self.bbstat.beta0]) + x0 = self.x0 + tf.scatter_nd(scatter_idx, ds, self.x0.shape) + val = self._compute_loss(x0=x0) grad = t1.gradient(val, self.x) - pd2ldxdx0, pd2ldxdnobs, pd2ldxdbeta0 = t2.jacobian( + pd2ldxds, pd2ldxdnobs, pd2ldxdbeta0 = t2.jacobian( grad, - [self.x0, self.nobs, self.bbstat.beta0], + [ds, self.nobs, self.bbstat.beta0], unconnected_gradients="zero", ) # cov is inverse hesse, thus cov ~ d2xd2l - dxdx0 = -self.cov @ pd2ldxdx0 + dxds = -self.cov @ pd2ldxds dxdnobs = -self.cov @ pd2ldxdnobs dxdbeta0 = -self.cov @ tf.reshape(pd2ldxdbeta0, [pd2ldxdbeta0.shape[0], -1]) - return dxdx0, dxdnobs, dxdbeta0 + nsyst = self.indata.nsyst + dxdtheta0 = dxds[:, :nsyst] + dxdparam0 = dxds[:, nsyst:] if self.param_prior_active else None + + return dxdtheta0, dxdparam0, dxdnobs, dxdbeta0 def _dndvars(self, fun): + # The forward model n depends on the parameters x but not on the + # constraint centers x0, so dn/dx0 = (dn/dx)(dx/dx0): differentiate n + # only w.r.t. x / nobs / beta0 here and pick up the x0 dependence + # through _dxdvars below. with tf.GradientTape() as t: - t.watch([self.x0, self.nobs, self.bbstat.beta0]) + t.watch([self.nobs, self.bbstat.beta0]) n = fun() n_flat = tf.reshape(n, (-1,)) - pdndx, pdndx0, pdndnobs, pdndbeta0 = t.jacobian( + pdndx, pdndnobs, pdndbeta0 = t.jacobian( n_flat, - [self.x, self.x0, self.nobs, self.bbstat.beta0], + [self.x, self.nobs, self.bbstat.beta0], unconnected_gradients="zero", ) # apply chain rule to take into account correlations with the fit parameters - dxdx0, dxdnobs, dxdbeta0 = self._dxdvars() + dxdtheta0, dxdparam0, dxdnobs, dxdbeta0 = self._dxdvars() - dndx0 = pdndx0 + pdndx @ dxdx0 + dndtheta0 = pdndx @ dxdtheta0 + dndparam0 = pdndx @ dxdparam0 if dxdparam0 is not None else None dndnobs = pdndnobs + pdndx @ dxdnobs dndbeta0 = tf.reshape(pdndbeta0, [pdndbeta0.shape[0], -1]) + pdndx @ dxdbeta0 - return n, dndx0, dndnobs, dndbeta0 + return n, dndtheta0, dndparam0, dndnobs, dndbeta0 def _compute_expected( self, fun_exp, inclusive=True, profile=False, full=True, need_observables=True @@ -1527,13 +1609,16 @@ def fun_n(): need_observables=need_observables, ) - _, dndx0, dndnobs, dndbeta0 = self._dndvars(fun_n) + _, dndtheta0, dndparam0, dndnobs, dndbeta0 = self._dndvars(fun_n) + vartheta0, varparam0 = self._x0_source_vars() impacts_gaussian, impacts_gaussian_grouped = ( global_impacts.gaussian_global_impacts_obs( - dndx0, + dndtheta0, + dndparam0, dndnobs, dndbeta0, - self.var_x0, + vartheta0, + varparam0, self.nobs if self.varnobs is None else self.varnobs, ( 1.0 @@ -1541,15 +1626,11 @@ def fun_n(): or not self.bbstat.enabled else 1.0 / self.bbstat.kstat ), - self.param_model.nparams, self.bbstat.enabled, self.bbstat.binByBinStatMode, self.bbstat.beta_shape, self.indata.systgroupidxs, self.data_cov_inv, - param_prior_idxs=( - self.param_prior_idxs if self.param_prior_active else None - ), ) ) else: @@ -1830,9 +1911,16 @@ def fun_res(): observed = fun(None, self.nobs) return expected - observed - residuals, dresdx0, dresdnobs, dresdbeta0 = self._dndvars(fun_res) + residuals, dresdtheta0, dresdparam0, dresdnobs, dresdbeta0 = self._dndvars( + fun_res + ) + vartheta0, varparam0 = self._x0_source_vars() - res_cov = dresdx0 @ (self.var_x0[:, None] * tf.transpose(dresdx0)) + # source variances enter only where cw > 0; the dropped (free) columns + # contributed exactly zero to the old full-x0 quadratic form. + res_cov = dresdtheta0 @ (vartheta0[:, None] * tf.transpose(dresdtheta0)) + if dresdparam0 is not None: + res_cov += dresdparam0 @ (varparam0[:, None] * tf.transpose(dresdparam0)) if self.covarianceFit: res_cov_stat = dresdnobs @ tf.linalg.solve( @@ -1996,13 +2084,16 @@ def full_nll(self): def reduced_nll(self): return self._compute_nll(full_nll=False) - def _compute_lc(self, full_nll=False): + def _compute_lc(self, full_nll=False, x0=None): # One constraint term over the full effective parameter vector # [poi, model_nui, theta]: the ParamModel block is constrained by # param_cw / param_x0 (declared priors; cw = 0 -> free) and the # nuisance block by indata.constraintweights / theta0. + # x0 defaults to self.x0; _dxdvars passes a perturbed copy to + # differentiate the loss w.r.t. a subset of the constraint centers. + x0 = self.x0 if x0 is None else x0 cw = self.cw - lc = cw * 0.5 * tf.square(self.get_x() - self.x0) + lc = cw * 0.5 * tf.square(self.get_x() - x0) if full_nll: # normalization factor 0.5*log(2*pi*sigma^2) for constrained # entries, with sigma^2 = 1/cw and @@ -2049,7 +2140,7 @@ def _compute_ln(self, nexp, full_nll=False): ) return ln - def _compute_nll_components(self, profile=True, full_nll=False): + def _compute_nll_components(self, profile=True, full_nll=False, x0=None): nexpfullcentral, _, beta = self._compute_yields_with_beta( profile=profile, compute_norm=False, @@ -2060,7 +2151,7 @@ def _compute_nll_components(self, profile=True, full_nll=False): ln = self._compute_ln(nexp, full_nll) - lc = self._compute_lc(full_nll) + lc = self._compute_lc(full_nll, x0=x0) lbeta = self._compute_lbeta(beta, full_nll) @@ -2082,9 +2173,9 @@ def _compute_external_nll(self): self.external_terms, self.x, self.indata.dtype ) - def _compute_nll(self, profile=True, full_nll=False): + def _compute_nll(self, profile=True, full_nll=False, x0=None): ln, lc, lbeta, lpenalty, beta = self._compute_nll_components( - profile=profile, full_nll=full_nll + profile=profile, full_nll=full_nll, x0=x0 ) l = ln + lc @@ -2099,8 +2190,8 @@ def _compute_nll(self, profile=True, full_nll=False): l = l + lext return l - def _compute_loss(self, profile=True): - return self._compute_nll(profile=profile) + def _compute_loss(self, profile=True, x0=None): + return self._compute_nll(profile=profile, x0=x0) def _make_tf_functions(self): # Build tf.function wrappers at instance construction time so that diff --git a/rabbit/impacts/global_asym_impacts.py b/rabbit/impacts/global_asym_impacts.py index 465f0b3..4cc8a78 100644 --- a/rabbit/impacts/global_asym_impacts.py +++ b/rabbit/impacts/global_asym_impacts.py @@ -48,28 +48,40 @@ def _envelope(values): def global_asym_impacts_parms( fitter, - selected_idxs, + selected_x_idxs, selected_names, + group_members=None, sigma=1.0, signs=(-1, 1), linear_warmstart=False, ): - """Run a per-nuisance x0-shift + re-fit and assemble the asymmetric - global impact tensor. + """Run a per-source x0-shift + re-fit and assemble the asymmetric global + impact tensor. + + A "source" is any constraint center: a constrained nuisance or a priored + ParamModel parameter. For each one its center x0[idx] is shifted by +/- + sigma and the full fit is re-run; the resulting POI/NOI shifts are the + asymmetric global impacts. This mirrors the sources of + gaussian_global_impacts_parms (which exposes priored params as + _prior columns), so the two agree in the Gaussian limit. Args: fitter: the Fitter instance (used for x, x0 and minimize). - selected_idxs: indices into the syst axis (0..nsyst-1) of nuisances to - scan. - selected_names: names of those nuisances (bytes), used as impact-axis - labels. + selected_x_idxs: full-x indices of the sources to scan (a nuisance i + is at fitter.param_model.nparams + i; a priored param is at its + own position in the leading param block). + selected_names: labels for those sources (bytes), used as impact-axis + labels (priored params are labelled _prior by the caller). + group_members: optional dict {group_name(bytes): [full-x idxs]} used + to build the grouped quadrature envelopes. Covers both syst groups + and ParamModel impact groups. sigma: shift magnitude in units of the prefit constraint width (constraints are unit-sigma in rabbit, so 1.0 = 1 prefit sigma). signs: sequence (down, up). Bin 0 of axis_downUpVar -> first sign. linear_warmstart: experimental. If True, warm-start each refit at - x_nom + dxdtheta0[:, i] * shift, the Gaussian-approximation new - minimum for the shifted theta0. Should drastically reduce the - number of optimizer iterations on near-Gaussian nuisances. + x_nom + dxds[:, source] * shift, the Gaussian-approximation new + minimum for the shifted center. Should drastically reduce the + number of optimizer iterations on near-Gaussian sources. Requires fitter.cov to exist (same prerequisite as --gaussianGlobalImpacts). @@ -78,15 +90,25 @@ def global_asym_impacts_parms( impacts: np.ndarray of shape (n_scanned, 2, n_total_params). Axis 1 is [down, up] matching axis_downUpVar. group_names: np.ndarray of bytes for groups containing scanned - nuisances (plus a trailing "Total"). + sources (plus a trailing "Total"). impacts_grouped: np.ndarray of shape (n_groups, 2, n_total_params). """ - n_scanned = len(selected_idxs) + if group_members is None: + group_members = {} + selected_x_idxs = [int(idx) for idx in selected_x_idxs] + n_scanned = len(selected_x_idxs) n_total = len(fitter.parms) impacts = np.zeros((n_scanned, 2, n_total)) nparams = fitter.param_model.nparams + # Prefit width of each constraint center: sqrt(var_x0) = 1/sqrt(cw). For + # nuisances cw = 1 so this is 1 (the historical "1.0 = 1 sigma"); for + # priored params it is the prior sigma, so a unit-sigma shift moves x0 by + # sigma, not by 1. Scaling by it keeps the asym impact in the same per-1- + # prefit-sigma units as gaussian_global_impacts_parms. + src_sigma_np = np.sqrt(fitter.var_x0.numpy()) + # Snapshot postfit nominal state to restore between iterations. x_nom = tf.identity(fitter.x.value()) x0_nom = tf.identity(fitter.x0.value()) @@ -94,17 +116,18 @@ def global_asym_impacts_parms( x_nom_np = x_nom.numpy() logger.info( - f"global_asym_impacts: shifting theta0 by +/- {sigma} sigma and " - f"re-fitting for {n_scanned} nuisances" + f"global_asym_impacts: shifting constraint centers by +/- {sigma} sigma " + f"and re-fitting for {n_scanned} sources" + (" (linear warm-start enabled)" if linear_warmstart else "") ) - # Optional Gaussian-approximation warm-start. - # dxdx0 has shape [npar, npar]; column nparams + i gives the linearised - # response of the postfit minimum to a unit shift of theta0[i]. Computing - # it once before the loop is the same cost as one --gaussianGlobalImpacts - # call. - dxdx0_np = None + # Optional Gaussian-approximation warm-start. _dxdvars returns the postfit + # response to a unit shift of each constraint center, pre-split into the + # nuisance block (dxdtheta0[:, i]) and the priored-param columns + # (dxdparam0[:, j]); we map each scanned source's full-x index to its + # column. Computing it once is the same cost as one + # --gaussianGlobalImpacts call. + warmstart_col = None if linear_warmstart: if fitter.cov is None: raise RuntimeError( @@ -112,40 +135,57 @@ def global_asym_impacts_parms( "(incompatible with --noHessian)." ) t_lws = time.perf_counter() - dxdx0_tf, _, _ = fitter._dxdvars() - dxdx0_np = dxdx0_tf.numpy() + dxdtheta0_tf, dxdparam0_tf, _, _ = fitter._dxdvars() + dxdtheta0_np = dxdtheta0_tf.numpy() + dxdparam0_np = dxdparam0_tf.numpy() if dxdparam0_tf is not None else None + prior_col = ( + {int(p): j for j, p in enumerate(fitter.param_prior_idxs.numpy())} + if fitter.param_prior_active + else {} + ) + + def _warmstart_col(idx): + if idx >= nparams: + return dxdtheta0_np[:, idx - nparams] + return dxdparam0_np[:, prior_col[idx]] + + warmstart_col = _warmstart_col + logger.info( - f"global_asym_impacts: dxdx0 prepared in " + f"global_asym_impacts: dxds prepared in " f"{time.perf_counter() - t_lws:.2f}s" ) t_per = np.zeros(n_scanned, dtype=np.float64) t_total0 = time.perf_counter() - for k, i in enumerate(selected_idxs): - i = int(i) + for k, idx in enumerate(selected_x_idxs): name = selected_names[k] name_str = name.decode() if isinstance(name, bytes) else name - logger.info(f" [{k + 1}/{n_scanned}] theta0-shift refit for {name_str}") + logger.info(f" [{k + 1}/{n_scanned}] x0-shift refit for {name_str}") + + # shift the center by `sigma` prefit sigmas of this source (1 for + # unit-constrained nuisances, the prior sigma for priored params). + width = src_sigma_np[idx] t0 = time.perf_counter() for j, sign in enumerate(signs): - shift = float(sign) * float(sigma) + shift = float(sign) * float(sigma) * float(width) - # Always shift the constraint center for nuisance i by `shift`. + # Always shift the constraint center for source idx by `shift`. x0_shifted = x0_nom_np.copy() - x0_shifted[nparams + i] += shift + x0_shifted[idx] += shift # Warm-start x. With linear_warmstart, use the Gaussian-approx new - # minimum x_nom + dxdx0[:, nparams+i] * shift -- on near-Gaussian - # nuisances this lands at the new minimum to within roundoff. - # Without it, just shift x[nparams+i] by `shift` so the nuisance - # itself starts at the new constraint center. - if linear_warmstart: - x_shifted = x_nom_np + dxdx0_np[:, nparams + i] * shift + # minimum x_nom + dxds[:, source] * shift -- on near-Gaussian + # sources this lands at the new minimum to within roundoff. + # Without it, just shift x[idx] by `shift` so the parameter itself + # starts at the new constraint center. + if warmstart_col is not None: + x_shifted = x_nom_np + warmstart_col(idx) * shift else: x_shifted = x_nom_np.copy() - x_shifted[nparams + i] += shift + x_shifted[idx] += shift fitter.x0.assign(x0_shifted) fitter.x.assign(x_shifted) @@ -175,17 +215,16 @@ def global_asym_impacts_parms( logger.info( f"global_asym_impacts: total {t_total:.1f}s " f"(mean {t_per.mean():.2f}s, min {t_per.min():.2f}s, " - f"max {t_per.max():.2f}s per nuisance)" + f"max {t_per.max():.2f}s per source)" ) # Grouped impacts via quadrature envelope, separately for down/up. - selected_set = set(int(idx) for idx in selected_idxs) - pos_in_scanned = {int(idx): k for k, idx in enumerate(selected_idxs)} + selected_set = set(selected_x_idxs) + pos_in_scanned = {idx: k for k, idx in enumerate(selected_x_idxs)} group_names = [] group_impacts = [] - for gname, gidxs in zip(fitter.indata.systgroups, fitter.indata.systgroupidxs): - gidxs = np.asarray(gidxs).astype(int) + for gname, gidxs in group_members.items(): in_scanned = [pos_in_scanned[i] for i in gidxs if int(i) in selected_set] if not in_scanned: continue diff --git a/rabbit/impacts/global_impacts.py b/rabbit/impacts/global_impacts.py index ab20431..b682891 100644 --- a/rabbit/impacts/global_impacts.py +++ b/rabbit/impacts/global_impacts.py @@ -497,25 +497,13 @@ def _gaussian_global_impacts( return impacts, impacts_grouped -def _split_x0_sources(dxdx0, varx0, nmodel_params, param_prior_idxs): - """Split full-x0 derivatives into the nuisance block and the optional - ParamModel-prior source columns.""" - dxdtheta0 = dxdx0[:, nmodel_params:] - vartheta0 = varx0[nmodel_params:] - if param_prior_idxs is not None: - dxdparam0 = tf.gather(dxdx0, param_prior_idxs, axis=1) - varparam0 = tf.gather(varx0, param_prior_idxs) - else: - dxdparam0 = None - varparam0 = None - return dxdtheta0, vartheta0, dxdparam0, varparam0 - - def gaussian_global_impacts_parms( - dxdx0, + dxdtheta0, + dxdparam0, dxdnobs, dxdbeta0, - varx0, + vartheta0, + varparam0, varnobs, varbeta0, nsignal_params, @@ -526,16 +514,20 @@ def gaussian_global_impacts_parms( beta_shape, systgroupidxs, data_cov_inv=None, - param_prior_idxs=None, ): - # compute impacts for pois and nois - dxdx0 = _gather_poi_noi_vector(dxdx0, noiidxs, nsignal_params, nmodel_params) + # The derivative columns come pre-split from the Fitter as + # [nuisance/theta block] and [priored params]; here we only gather the + # poi/noi rows whose impacts are reported (dxdparam0 is None when no + # priors are declared). + dxdtheta0 = _gather_poi_noi_vector( + dxdtheta0, noiidxs, nsignal_params, nmodel_params + ) dxdnobs = _gather_poi_noi_vector(dxdnobs, noiidxs, nsignal_params, nmodel_params) dxdbeta0 = _gather_poi_noi_vector(dxdbeta0, noiidxs, nsignal_params, nmodel_params) - - dxdtheta0, vartheta0, dxdparam0, varparam0 = _split_x0_sources( - dxdx0, varx0, nmodel_params, param_prior_idxs - ) + if dxdparam0 is not None: + dxdparam0 = _gather_poi_noi_vector( + dxdparam0, noiidxs, nsignal_params, nmodel_params + ) return _gaussian_global_impacts( dxdtheta0, @@ -555,24 +547,20 @@ def gaussian_global_impacts_parms( def gaussian_global_impacts_obs( - dndx0, + dndtheta0, + dndparam0, dndnobs, dndbeta0, - varx0, + vartheta0, + varparam0, varnobs, varbeta0, - nmodel_params, bin_by_bin_stat, bin_by_bin_stat_mode, beta_shape, systgroupidxs, data_cov_inv=None, - param_prior_idxs=None, ): - dndtheta0, vartheta0, dndparam0, varparam0 = _split_x0_sources( - dndx0, varx0, nmodel_params, param_prior_idxs - ) - return _gaussian_global_impacts( dndtheta0, dndnobs, From 002f8dabe8dee1a7082f846f2b14191f19c24a13 Mon Sep 17 00:00:00 2001 From: Luca Lavezzo Date: Mon, 22 Jun 2026 14:50:04 -0400 Subject: [PATCH 15/16] Carry global-impact sources as one collection (address review) Per davidwalter2 review: treat the ParamModel params and the systs on equal footing instead of carrying / splitting two collections. - Fitter builds a single x0_source_idxs (the constrained centers, ordered [nuisance block | priored params]) and var_x0_sources, derived from cw. - _dxdvars / _dndvars return ONE dxdx0 / dndx0 over those sources instead of the (dxdtheta0, dxdparam0) pair; the helper methods _x0_constraint_source_idxs and _x0_source_vars are gone. - the likelihood global impacts gather impacts_x0 over source_idxs in one step (no theta-block slice + separate prior gather); the gaussian impacts take one dxdx0 / varx0 and compute the per-source impact uniformly as dxdx0 * sqrt(var) (systs have var = 1, priored params carry their sigma). - the only remaining nuisance-vs-param boundary is a local slice when assembling the grouped axis in its historical order [syst groups | stat | bbb | priors], driven by n_param_sources. - global_asym_impacts warm-start maps each source to its single dxdx0 column. - drop the duplicated param_prior_sigmas / _means / _mask Fitter attributes; priors live only on the model, and the output metadata reads them there. param_prior_active is now derived from cw. Validated bit-identical to the previous code on the test tensor: likelihood and gaussian global impacts (parms + observables), and the asym sig_prior impact still matches the gaussian one to <1%. Source axis unchanged. Co-Authored-By: Claude Fable 5 --- bin/rabbit_fit.py | 27 +++-- rabbit/fitter.py | 154 +++++++++++--------------- rabbit/impacts/global_asym_impacts.py | 26 ++--- rabbit/impacts/global_impacts.py | 142 +++++++++++------------- 4 files changed, 156 insertions(+), 193 deletions(-) diff --git a/bin/rabbit_fit.py b/bin/rabbit_fit.py index 847b6f2..2213950 100755 --- a/bin/rabbit_fit.py +++ b/bin/rabbit_fit.py @@ -305,9 +305,9 @@ def make_parser(): default=False, action="store_true", help="EXPERIMENTAL: warm-start each --globalAsymImpacts refit at the " - "Gaussian-approximation new minimum x_nom + dxdtheta0[:, i] * shift " + "Gaussian-approximation new minimum x_nom + dxdx0[:, source] * shift " "(same Jacobian as --gaussianGlobalImpacts). On near-Gaussian " - "nuisances this should reduce per-nuisance refit cost by 10-50x. " + "sources this should reduce per-source refit cost by 10-50x. " "Adds one --gaussianGlobalImpacts-equivalent precompute up front. " "Off by default until validated on real tensors.", ) @@ -841,15 +841,24 @@ def main(): "nois": ifitter.parms[ifitter.param_model.nparams :][indata.noiidxs], } - # ParamModel Gaussian priors (if the model declared sigmas). Stored as - # a small dict so downstream tooling can see what was applied without - # parsing the rabbit log. + # ParamModel Gaussian priors (if the model declared sigmas). Read straight + # from the model (the single source of truth) so downstream tooling can see + # what was applied without parsing the rabbit log. if getattr(ifitter, "param_prior_active", False): + pm = ifitter.param_model + sigmas = np.asarray(pm.prior_sigmas, dtype=np.float64) + mask = np.isfinite(sigmas) & (sigmas > 0) + means = getattr(pm, "prior_means", None) + means = ( + np.asarray(pm.xparamdefault).astype(np.float64) + if means is None + else np.asarray(means, dtype=np.float64) + ) meta["param_priors"] = { - "params": ifitter.param_model.params, # all nparams names - "mask": ifitter.param_prior_mask.numpy(), # bool array - "sigmas": ifitter.param_prior_sigmas.numpy(), # NaN where mask False - "means": ifitter.param_prior_means.numpy(), # NaN where mask False + "params": pm.params, # all nparams names + "mask": mask, # bool array + "sigmas": np.where(mask, sigmas, np.nan), # NaN where no prior + "means": np.where(mask, means, np.nan), # NaN where no prior } with workspace.Workspace( diff --git a/rabbit/fitter.py b/rabbit/fitter.py index 05e43f7..cd8b4d5 100644 --- a/rabbit/fitter.py +++ b/rabbit/fitter.py @@ -303,7 +303,6 @@ def init_fit_parms( # below. _compute_lc adds 0.5 * cw * (x - x0)^2 over the full # parameter vector, prefit variances are 1/cw, and the toy # randomization fluctuates the centers wherever cw > 0. - self.param_prior_active = False pm = self.param_model np_dtype = self.indata.dtype.as_numpy_dtype param_cw_np = np.zeros(pm.nparams, dtype=np_dtype) @@ -314,6 +313,9 @@ def init_fit_parms( # nonprofiled impacts), even though cw = 0 makes it irrelevant to the # likelihood itself. param_x0_np = pm.xparamdefault.numpy().astype(np_dtype) + # The priors live ONLY on the model (prior_sigmas / prior_means); here + # they are folded into the constraint weights / centers below, no copy + # of the prior arrays is kept on the Fitter. sigmas = getattr(pm, "prior_sigmas", None) if sigmas is not None: sigmas = np.asarray(sigmas, dtype=np_dtype) @@ -343,14 +345,6 @@ def init_fit_parms( "would apply to sqrt(poi) rather than to the POI " "itself." ) - self.param_prior_active = True - self.param_prior_mask = tf.constant(mask, dtype=tf.bool) - self.param_prior_sigmas = tf.constant( - np.where(mask, sigmas, np.nan), dtype=self.indata.dtype - ) - self.param_prior_means = tf.constant( - np.where(mask, means, np.nan), dtype=self.indata.dtype - ) param_cw_np = np.where(mask, 1.0 / sigmas**2, 0.0) # priored entries -> prior mean; free entries keep their default param_x0_np = np.where(mask, means, param_x0_np) @@ -362,11 +356,6 @@ def init_fit_parms( if mask[i]: name = p.decode() if isinstance(p, bytes) else str(p) logger.info(f" {name}: μ={means[i]:.4g} σ={sigmas[i]:.4g}") - # indices (into the ParamModel block) of the priored params; impact - # sources in the global-impacts calculations - self.param_prior_idxs = tf.constant( - np.where(param_cw_np > 0)[0], dtype=tf.int64 - ) # Unified constraint structure over the full parameter vector, # index-aligned with x = [params | systs]: @@ -395,6 +384,26 @@ def init_fit_parms( tf.math.reciprocal(self.cw), ) + # Global-impact "sources" are the constraint centers x0, carried as a + # SINGLE collection so params and systs are treated identically by the + # impact code. They are ordered [full nuisance block | priored params] + # to preserve the historical source axis (all systs, then any priored + # params). param_prior_idxs (priored params within the leading param + # block) is derived from cw, so the priors stay defined only on the + # model. + self.param_prior_idxs = tf.constant( + np.where(param_cw_np > 0)[0].astype(np.int64), dtype=tf.int64 + ) + self.param_prior_active = bool(int(self.param_prior_idxs.shape[0]) > 0) + self.x0_source_idxs = tf.concat( + [ + tf.range(pm.nparams, len(self.parms), dtype=tf.int64), + self.param_prior_idxs, + ], + axis=0, + ) + self.var_x0_sources = tf.gather(self.var_x0, self.x0_source_idxs) + # Per-parameter prefit variance vector. Always allocated; the # prefit covariance is intrinsically diagonal so this is the # only form needed for prefit uncertainties. @@ -1127,23 +1136,19 @@ def global_impacts_parms(self): self.bbstat.binByBinStatMode, self.globalImpactsFromJVP, self.cov, - param_prior_idxs=( - self.param_prior_idxs if self.param_prior_active else None - ), + self.x0_source_idxs, + n_param_sources=int(self.param_prior_idxs.shape[0]), ) @tf.function def gaussian_global_impacts_parms(self): - dxdtheta0, dxdparam0, dxdnobs, dxdbeta0 = self._dxdvars() - vartheta0, varparam0 = self._x0_source_vars() + dxdx0, dxdnobs, dxdbeta0 = self._dxdvars() impacts, impacts_grouped = global_impacts.gaussian_global_impacts_parms( - dxdtheta0, - dxdparam0, + dxdx0, dxdnobs, dxdbeta0, - vartheta0, - varparam0, + self.var_x0_sources, self.nobs if self.varnobs is None else self.varnobs, ( 1.0 @@ -1158,7 +1163,8 @@ def gaussian_global_impacts_parms(self): self.bbstat.binByBinStatMode, self.bbstat.beta_shape, self.indata.systgroupidxs, - self.data_cov_inv, + n_param_sources=int(self.param_prior_idxs.shape[0]), + data_cov_inv=self.data_cov_inv, ) return impacts, impacts_grouped @@ -1389,40 +1395,16 @@ def _pd2ldbeta2(self, profile=False): pd2ldbeta2 = t2.gradient(pdldbeta, self.bbstat.ubeta) return pd2ldbeta2 - def _x0_constraint_source_idxs(self): - # Full-x indices of the constraint sources reported by the global - # impacts, in column order [nuisance/theta block | priored params]. - # x0 enters the NLL only through cw * (x - x0)^2, so the derivative - # w.r.t. any center with cw = 0 is identically zero; dropping those - # columns is exact, not an approximation, and keeps the Jacobian at - # [npar, nsyst (+ n_prior)] instead of [npar, npar]. The param block - # sits first in x, so param_prior_idxs are already full-x indices. - nparams = self.param_model.nparams - theta_idxs = tf.range(nparams, nparams + self.indata.nsyst, dtype=tf.int64) - if self.param_prior_active: - return tf.concat([theta_idxs, self.param_prior_idxs], axis=0) - return theta_idxs - - def _x0_source_vars(self): - # Prefit variances of the constraint sources, in the same column - # order / split as _dxdvars / _dndvars: (vartheta0, varparam0). - vartheta0 = self.var_x0[self.param_model.nparams :] - if self.param_prior_active: - varparam0 = tf.gather(self.var_x0, self.param_prior_idxs) - else: - varparam0 = None - return vartheta0, varparam0 - def _dxdvars(self): - # Differentiate the loss only w.r.t. the constraint sources (the - # nuisance block plus any priored params) by perturbing those x0 - # entries with `ds` and scattering back into the full vector; see - # _x0_constraint_source_idxs. The resulting columns are ordered - # [theta block | priored params] and split below. - src_idxs = self._x0_constraint_source_idxs() - n_src = self.indata.nsyst + ( - int(self.param_prior_idxs.shape[0]) if self.param_prior_active else 0 - ) + # Response of the postfit minimum to a unit shift of each constraint + # center, as a SINGLE collection of columns aligned with + # x0_source_idxs (the nuisance block plus any priored params). x0 + # enters the NLL only through cw * (x - x0)^2, so the derivative + # w.r.t. any center with cw = 0 is identically zero; restricting the + # differentiation to these sources is exact, not an approximation, + # and keeps the Jacobian at [npar, n_sources] instead of [npar, npar]. + src_idxs = self.x0_source_idxs + n_src = int(src_idxs.shape[0]) scatter_idx = src_idxs[:, None] with tf.GradientTape() as t2: ds = tf.zeros([n_src], dtype=self.indata.dtype) @@ -1432,28 +1414,24 @@ def _dxdvars(self): x0 = self.x0 + tf.scatter_nd(scatter_idx, ds, self.x0.shape) val = self._compute_loss(x0=x0) grad = t1.gradient(val, self.x) - pd2ldxds, pd2ldxdnobs, pd2ldxdbeta0 = t2.jacobian( + pd2ldxdx0, pd2ldxdnobs, pd2ldxdbeta0 = t2.jacobian( grad, [ds, self.nobs, self.bbstat.beta0], unconnected_gradients="zero", ) # cov is inverse hesse, thus cov ~ d2xd2l - dxds = -self.cov @ pd2ldxds + dxdx0 = -self.cov @ pd2ldxdx0 dxdnobs = -self.cov @ pd2ldxdnobs dxdbeta0 = -self.cov @ tf.reshape(pd2ldxdbeta0, [pd2ldxdbeta0.shape[0], -1]) - nsyst = self.indata.nsyst - dxdtheta0 = dxds[:, :nsyst] - dxdparam0 = dxds[:, nsyst:] if self.param_prior_active else None - - return dxdtheta0, dxdparam0, dxdnobs, dxdbeta0 + return dxdx0, dxdnobs, dxdbeta0 def _dndvars(self, fun): # The forward model n depends on the parameters x but not on the # constraint centers x0, so dn/dx0 = (dn/dx)(dx/dx0): differentiate n # only w.r.t. x / nobs / beta0 here and pick up the x0 dependence - # through _dxdvars below. + # through _dxdvars below. dndx0 is the single source collection. with tf.GradientTape() as t: t.watch([self.nobs, self.bbstat.beta0]) n = fun() @@ -1466,14 +1444,13 @@ def _dndvars(self, fun): ) # apply chain rule to take into account correlations with the fit parameters - dxdtheta0, dxdparam0, dxdnobs, dxdbeta0 = self._dxdvars() + dxdx0, dxdnobs, dxdbeta0 = self._dxdvars() - dndtheta0 = pdndx @ dxdtheta0 - dndparam0 = pdndx @ dxdparam0 if dxdparam0 is not None else None + dndx0 = pdndx @ dxdx0 dndnobs = pdndnobs + pdndx @ dxdnobs dndbeta0 = tf.reshape(pdndbeta0, [pdndbeta0.shape[0], -1]) + pdndx @ dxdbeta0 - return n, dndtheta0, dndparam0, dndnobs, dndbeta0 + return n, dndx0, dndnobs, dndbeta0 def _compute_expected( self, fun_exp, inclusive=True, profile=False, full=True, need_observables=True @@ -1587,11 +1564,14 @@ def compute_derivatives(dvars): expvar_flat, expvar.shape, profile, - pdexpdbeta, - pd2ldbeta2_pdexpdbeta if pdexpdbeta is not None else None, - self.prefit_unconstrained_nuisance_uncertainty, - param_prior_idxs=( - self.param_prior_idxs if self.param_prior_active else None + self.x0_source_idxs, + n_param_sources=int(self.param_prior_idxs.shape[0]), + pdexpdbeta=pdexpdbeta, + pd2ldbeta2_pdexpdbeta=( + pd2ldbeta2_pdexpdbeta if pdexpdbeta is not None else None + ), + prefit_unconstrained_nuisance_uncertainty=( + self.prefit_unconstrained_nuisance_uncertainty ), ) else: @@ -1609,16 +1589,13 @@ def fun_n(): need_observables=need_observables, ) - _, dndtheta0, dndparam0, dndnobs, dndbeta0 = self._dndvars(fun_n) - vartheta0, varparam0 = self._x0_source_vars() + _, dndx0, dndnobs, dndbeta0 = self._dndvars(fun_n) impacts_gaussian, impacts_gaussian_grouped = ( global_impacts.gaussian_global_impacts_obs( - dndtheta0, - dndparam0, + dndx0, dndnobs, dndbeta0, - vartheta0, - varparam0, + self.var_x0_sources, self.nobs if self.varnobs is None else self.varnobs, ( 1.0 @@ -1630,7 +1607,8 @@ def fun_n(): self.bbstat.binByBinStatMode, self.bbstat.beta_shape, self.indata.systgroupidxs, - self.data_cov_inv, + n_param_sources=int(self.param_prior_idxs.shape[0]), + data_cov_inv=self.data_cov_inv, ) ) else: @@ -1911,16 +1889,12 @@ def fun_res(): observed = fun(None, self.nobs) return expected - observed - residuals, dresdtheta0, dresdparam0, dresdnobs, dresdbeta0 = self._dndvars( - fun_res - ) - vartheta0, varparam0 = self._x0_source_vars() + residuals, dresdx0, dresdnobs, dresdbeta0 = self._dndvars(fun_res) - # source variances enter only where cw > 0; the dropped (free) columns - # contributed exactly zero to the old full-x0 quadratic form. - res_cov = dresdtheta0 @ (vartheta0[:, None] * tf.transpose(dresdtheta0)) - if dresdparam0 is not None: - res_cov += dresdparam0 @ (varparam0[:, None] * tf.transpose(dresdparam0)) + # one source collection; var_x0_sources weights each source by its + # prefit variance (the dropped cw = 0 columns contributed exactly zero + # to the old full-x0 quadratic form). + res_cov = dresdx0 @ (self.var_x0_sources[:, None] * tf.transpose(dresdx0)) if self.covarianceFit: res_cov_stat = dresdnobs @ tf.linalg.solve( diff --git a/rabbit/impacts/global_asym_impacts.py b/rabbit/impacts/global_asym_impacts.py index 4cc8a78..394af56 100644 --- a/rabbit/impacts/global_asym_impacts.py +++ b/rabbit/impacts/global_asym_impacts.py @@ -100,8 +100,6 @@ def global_asym_impacts_parms( n_total = len(fitter.parms) impacts = np.zeros((n_scanned, 2, n_total)) - nparams = fitter.param_model.nparams - # Prefit width of each constraint center: sqrt(var_x0) = 1/sqrt(cw). For # nuisances cw = 1 so this is 1 (the historical "1.0 = 1 sigma"); for # priored params it is the prior sigma, so a unit-sigma shift moves x0 by @@ -122,11 +120,10 @@ def global_asym_impacts_parms( ) # Optional Gaussian-approximation warm-start. _dxdvars returns the postfit - # response to a unit shift of each constraint center, pre-split into the - # nuisance block (dxdtheta0[:, i]) and the priored-param columns - # (dxdparam0[:, j]); we map each scanned source's full-x index to its - # column. Computing it once is the same cost as one - # --gaussianGlobalImpacts call. + # response to a unit shift of each constraint center as a single collection + # of columns aligned with fitter.x0_source_idxs; we map each scanned + # source's full-x index to its column. Computing it once is the same cost + # as one --gaussianGlobalImpacts call. warmstart_col = None if linear_warmstart: if fitter.cov is None: @@ -135,19 +132,12 @@ def global_asym_impacts_parms( "(incompatible with --noHessian)." ) t_lws = time.perf_counter() - dxdtheta0_tf, dxdparam0_tf, _, _ = fitter._dxdvars() - dxdtheta0_np = dxdtheta0_tf.numpy() - dxdparam0_np = dxdparam0_tf.numpy() if dxdparam0_tf is not None else None - prior_col = ( - {int(p): j for j, p in enumerate(fitter.param_prior_idxs.numpy())} - if fitter.param_prior_active - else {} - ) + dxdx0_tf, _, _ = fitter._dxdvars() + dxdx0_np = dxdx0_tf.numpy() + src_col = {int(s): k for k, s in enumerate(fitter.x0_source_idxs.numpy())} def _warmstart_col(idx): - if idx >= nparams: - return dxdtheta0_np[:, idx - nparams] - return dxdparam0_np[:, prior_col[idx]] + return dxdx0_np[:, src_col[idx]] warmstart_col = _warmstart_col diff --git a/rabbit/impacts/global_impacts.py b/rabbit/impacts/global_impacts.py index b682891..fdead04 100644 --- a/rabbit/impacts/global_impacts.py +++ b/rabbit/impacts/global_impacts.py @@ -197,18 +197,24 @@ def _compute_grouped_impacts( bin_by_bin_stat, bin_by_bin_stat_mode, systgroupidxs, - impacts_theta0_sq, + impacts_sources_sq, impacts_nobs, impacts_beta0_total, impacts_beta0_process, - impacts_param0_sq=None, + n_param_sources=0, ): """Assemble the grouped impacts tensor from all contributions. - impacts_param0_sq holds the squared impacts of the ParamModel prior - centers (one column per priored param); they are appended at the END of - the grouped axis, matching the column order of the output axes. + impacts_sources_sq holds the squared per-source impacts as a single + collection ordered [nuisance block | priored params]. The trailing + n_param_sources columns (the priored params) are appended at the END of + the grouped axis as individual sources, matching the output axes; the + leading nuisance columns are combined into the systematic groups. The + split here is only to assemble the grouped axis in its historical order. """ + nsyst = int(impacts_sources_sq.shape[-1]) - n_param_sources + impacts_theta0_sq = impacts_sources_sq[:, :nsyst] + impacts_param0_sq = impacts_sources_sq[:, nsyst:] if n_param_sources else None if bin_by_bin_stat: impacts_grouped = tf.stack([impacts_nobs, impacts_beta0_total], axis=-1) if bin_by_bin_stat_mode == "full": @@ -252,7 +258,8 @@ def global_impacts_parms( bin_by_bin_stat_mode, global_impacts_from_jvp, cov, - param_prior_idxs=None, + source_idxs, + n_param_sources=0, ): idxs_poi = tf.range(nsignal_params, dtype=tf.int64) idxs_noi = tf.constant(nmodel_params + noiidxs, dtype=tf.int64) @@ -279,23 +286,16 @@ def global_impacts_parms( pd2ldbeta2_pdexpdbeta=None, ) - # impacts_x0 covers the full constraint term: the nuisance block AND any - # ParamModel prior penalties (rows in the ParamModel block, nonzero where - # the model declares priors). Per-1-sigma units come out automatically - # since sc = sqrt(d2lc/dx2) = sqrt(cw) = 1/sigma. + # impacts_x0 covers the full constraint term over the whole parameter + # vector. The reported sources (constrained nuisances + priored params) + # are a single collection picked out by source_idxs; params and systs are + # treated identically here. Per-1-sigma units come out automatically since + # sc = sqrt(d2lc/dx2) = sqrt(cw) = 1/sigma. impacts_x0 = _compute_global_impacts_x0(x, compute_lc_fn, cov_dexpdx) - impacts_theta0 = tf.transpose(impacts_x0[nmodel_params:]) + impacts_sources = tf.transpose(tf.gather(impacts_x0, source_idxs, axis=0)) - impacts_theta0_sq = tf.square(impacts_theta0) - var_x0 = tf.reduce_sum(impacts_theta0_sq, axis=-1) - - impacts_param0_sq = None - if param_prior_idxs is not None: - impacts_param0 = tf.transpose( - tf.gather(impacts_x0[:nmodel_params], param_prior_idxs, axis=0) - ) - impacts_param0_sq = tf.square(impacts_param0) - var_x0 += tf.reduce_sum(impacts_param0_sq, axis=-1) + impacts_sources_sq = tf.square(impacts_sources) + var_x0 = tf.reduce_sum(impacts_sources_sq, axis=-1) var_nobs = var_total - var_x0 if bin_by_bin_stat: @@ -305,17 +305,14 @@ def global_impacts_parms( bin_by_bin_stat, bin_by_bin_stat_mode, systgroupidxs, - impacts_theta0_sq, + impacts_sources_sq, tf.sqrt(var_nobs), impacts_beta0_total, impacts_beta0_process, - impacts_param0_sq=impacts_param0_sq, + n_param_sources=n_param_sources, ) - if param_prior_idxs is not None: - impacts_theta0 = tf.concat([impacts_theta0, impacts_param0], axis=-1) - - return impacts_theta0, impacts_grouped + return impacts_sources, impacts_grouped def global_impacts_obs( @@ -335,10 +332,11 @@ def global_impacts_obs( expvar_flat, expvar_shape, profile, + source_idxs, + n_param_sources=0, pdexpdbeta=None, pd2ldbeta2_pdexpdbeta=None, prefit_unconstrained_nuisance_uncertainty=0.0, - param_prior_idxs=None, ): """ Global impacts on observable bins, used inside _expected_with_variance. @@ -387,18 +385,11 @@ def global_impacts_obs( ) impacts_x0 = _compute_global_impacts_x0(x, compute_lc_fn, cov_dexpdx) - impacts_theta0 = tf.transpose(impacts_x0[nmodel_params:]) + # single source collection (constrained nuisances + priored params) + impacts_sources = tf.transpose(tf.gather(impacts_x0, source_idxs, axis=0)) - impacts_theta0_sq = tf.square(impacts_theta0) - var_x0 = tf.reduce_sum(impacts_theta0_sq, axis=-1) - - impacts_param0_sq = None - if param_prior_idxs is not None: - impacts_param0 = tf.transpose( - tf.gather(impacts_x0[:nmodel_params], param_prior_idxs, axis=0) - ) - impacts_param0_sq = tf.square(impacts_param0) - var_x0 += tf.reduce_sum(impacts_param0_sq, axis=-1) + impacts_sources_sq = tf.square(impacts_sources) + var_x0 = tf.reduce_sum(impacts_sources_sq, axis=-1) var_nobs = expvar_flat - var_x0 if bin_by_bin_stat: @@ -408,17 +399,14 @@ def global_impacts_obs( bin_by_bin_stat, bin_by_bin_stat_mode, systgroupidxs, - impacts_theta0_sq, + impacts_sources_sq, tf.sqrt(var_nobs), impacts_beta0_total, impacts_beta0_process, - impacts_param0_sq=impacts_param0_sq, + n_param_sources=n_param_sources, ) - if param_prior_idxs is not None: - impacts_theta0 = tf.concat([impacts_theta0, impacts_param0], axis=-1) - - impacts = tf.reshape(impacts_theta0, [*expvar_shape, impacts_theta0.shape[-1]]) + impacts = tf.reshape(impacts_sources, [*expvar_shape, impacts_sources.shape[-1]]) impacts_grouped = tf.reshape( impacts_grouped, [*expvar_shape, impacts_grouped.shape[-1]] ) @@ -427,20 +415,34 @@ def global_impacts_obs( def _gaussian_global_impacts( - dxdtheta0, + dxdx0, dxdnobs, dxdbeta0, - vartheta0, + varx0, varnobs, varbeta0, bin_by_bin_stat, bin_by_bin_stat_mode, beta_shape, systgroupidxs, + n_param_sources=0, data_cov_inv=None, - dxdparam0=None, - varparam0=None, ): + # dxdx0 / varx0 are the single source collection ordered + # [nuisance block | priored params]. The leading nuisance columns and the + # trailing priored-param columns are split out only to assemble the + # grouped axis in its historical order [syst groups | stat | bbb | priors]; + # the per-source impacts themselves are computed uniformly. + nsyst = int(dxdx0.shape[-1]) - n_param_sources + dxdtheta0 = dxdx0[:, :nsyst] + vartheta0 = varx0[:nsyst] + if n_param_sources: + dxdparam0 = dxdx0[:, nsyst:] + varparam0 = varx0[nsyst:] + else: + dxdparam0 = None + varparam0 = None + if data_cov_inv is not None: data_cov = tf.linalg.inv(data_cov_inv) # equivalent to tf.linalg.diag_part(dxdnobs @ data_cov @ tf.transpose(dxdnobs)) but avoiding computing full matrix @@ -498,12 +500,10 @@ def _gaussian_global_impacts( def gaussian_global_impacts_parms( - dxdtheta0, - dxdparam0, + dxdx0, dxdnobs, dxdbeta0, - vartheta0, - varparam0, + varx0, varnobs, varbeta0, nsignal_params, @@ -513,66 +513,56 @@ def gaussian_global_impacts_parms( bin_by_bin_stat_mode, beta_shape, systgroupidxs, + n_param_sources=0, data_cov_inv=None, ): - # The derivative columns come pre-split from the Fitter as - # [nuisance/theta block] and [priored params]; here we only gather the - # poi/noi rows whose impacts are reported (dxdparam0 is None when no - # priors are declared). - dxdtheta0 = _gather_poi_noi_vector( - dxdtheta0, noiidxs, nsignal_params, nmodel_params - ) + # dxdx0 is the single source-derivative collection; gather the poi/noi + # rows whose impacts are reported (columns / sources are untouched). + dxdx0 = _gather_poi_noi_vector(dxdx0, noiidxs, nsignal_params, nmodel_params) dxdnobs = _gather_poi_noi_vector(dxdnobs, noiidxs, nsignal_params, nmodel_params) dxdbeta0 = _gather_poi_noi_vector(dxdbeta0, noiidxs, nsignal_params, nmodel_params) - if dxdparam0 is not None: - dxdparam0 = _gather_poi_noi_vector( - dxdparam0, noiidxs, nsignal_params, nmodel_params - ) return _gaussian_global_impacts( - dxdtheta0, + dxdx0, dxdnobs, dxdbeta0, - vartheta0, + varx0, varnobs, varbeta0, bin_by_bin_stat, bin_by_bin_stat_mode, beta_shape, systgroupidxs, + n_param_sources, data_cov_inv, - dxdparam0=dxdparam0, - varparam0=varparam0, ) def gaussian_global_impacts_obs( - dndtheta0, - dndparam0, + dndx0, dndnobs, dndbeta0, - vartheta0, - varparam0, + varx0, varnobs, varbeta0, bin_by_bin_stat, bin_by_bin_stat_mode, beta_shape, systgroupidxs, + n_param_sources=0, data_cov_inv=None, ): return _gaussian_global_impacts( - dndtheta0, + dndx0, dndnobs, dndbeta0, - vartheta0, + varx0, varnobs, varbeta0, bin_by_bin_stat, bin_by_bin_stat_mode, beta_shape, systgroupidxs, + n_param_sources, data_cov_inv, - dxdparam0=dndparam0, - varparam0=varparam0, ) From 32902cd6f655fd22807ab3d0e8972ca81fe88367 Mon Sep 17 00:00:00 2001 From: Luca Lavezzo Date: Tue, 23 Jun 2026 15:07:41 -0400 Subject: [PATCH 16/16] Address review: work with the full x0 everywhere, drop the source subset Per davidwalter2's latest comments, simplify to a single full-length x0 instead of carrying a reduced 'source' set: - _dxdvars / _dndvars differentiate w.r.t. the whole self.x0 (no scatter / perturbation, no x0= loss argument); columns for unconstrained centers (cw = 0) are exactly zero and carried along. - the global impacts work with the full impacts_x0 directly: impacts_x0_sq = square(impacts_x0), var_x0 = reduce_sum, no gather; the per-source axis is now the whole parameter list (matching the traditional impacts). gaussian per-source = dxdx0 * sqrt(var_x0) (one expression). - grouping is one operation over full-x column-index lists (syst groups shifted by nmodel_params, ParamModel impact groups in the param block) so param and syst sources combine uniformly; the grouped global axis now uses the param-impact groups (like the traditional grouped impacts). - drop x0_source_idxs / var_x0_sources; param_prior_idxs kept only for the asym scan. xdefaultassign / bayesassign use self.x0default. - nits: drop the param-model-specific prior comment; meta priors use indata.dtype; rename dxds -> dxdx0 in global_asym_impacts. Validated: lint clean, pytest (15) + the five CI script-tests pass, and the ParamModel-prior impacts are unchanged (likelihood = gaussian source[sig] = 1.105769e-02, asym (-1.097e-2, +1.114e-2)). NOTE for review: the global-impact source axis is now the full parameter list (was systs + _prior); priored params appear under their own name with unconstrained params as exact-zero columns. Co-Authored-By: Claude Fable 5 --- bin/rabbit_fit.py | 7 +- rabbit/fitter.py | 157 +++++++-------------- rabbit/impacts/global_asym_impacts.py | 34 ++--- rabbit/impacts/global_impacts.py | 188 ++++++++++++-------------- rabbit/workspace.py | 22 +-- 5 files changed, 162 insertions(+), 246 deletions(-) diff --git a/bin/rabbit_fit.py b/bin/rabbit_fit.py index 92071b2..772d6c6 100755 --- a/bin/rabbit_fit.py +++ b/bin/rabbit_fit.py @@ -882,13 +882,14 @@ def main(): # what was applied without parsing the rabbit log. if getattr(ifitter, "param_prior_active", False): pm = ifitter.param_model - sigmas = np.asarray(pm.prior_sigmas, dtype=np.float64) + np_dtype = ifitter.indata.dtype.as_numpy_dtype + sigmas = np.asarray(pm.prior_sigmas, dtype=np_dtype) mask = np.isfinite(sigmas) & (sigmas > 0) means = getattr(pm, "prior_means", None) means = ( - np.asarray(pm.xparamdefault).astype(np.float64) + np.asarray(pm.xparamdefault).astype(np_dtype) if means is None - else np.asarray(means, dtype=np.float64) + else np.asarray(means, dtype=np_dtype) ) meta["param_priors"] = { "params": pm.params, # all nparams names diff --git a/rabbit/fitter.py b/rabbit/fitter.py index 6c4f440..a5d8be2 100644 --- a/rabbit/fitter.py +++ b/rabbit/fitter.py @@ -298,24 +298,11 @@ def init_fit_parms( self.x = tf.Variable(xdefault, trainable=True, name="x") - # ParamModel Gaussian priors. The ParamModel itself decides whether - # its parameters carry priors, by declaring two optional attributes: - # - prior_sigmas: np.ndarray shape (nparams,). Entries that are - # finite and > 0 are Gaussian-constrained at that width; NaN / - # non-finite / 0 entries leave the corresponding parameter free. - # - prior_means : np.ndarray shape (nparams,), optional. - # Defaults to param_model.xparamdefault when not provided. - # If the model declares priors they are applied; how to enable or - # disable them (e.g. a token in the --paramModel spec) is up to the - # model. - # - # Priors use the same constraint structure as the nuisance - # constraints: declared priors set the constraint weight - # (1/sigma^2; 0 = free) and the constraint center (the prior mean) - # of the ParamModel block in the full-length cw / x0 vectors built - # below. _compute_lc adds 0.5 * cw * (x - x0)^2 over the full - # parameter vector, prefit variances are 1/cw, and the toy - # randomization fluctuates the centers wherever cw > 0. + # ParamModel Gaussian priors (declared by the model via prior_sigmas / + # prior_means; see ParamModel). They are folded into the same + # constraint structure as the nuisances: a declared prior sets the + # constraint weight (1/sigma^2) and center (prior mean) of the + # ParamModel block in the full-length cw / x0 vectors built below. pm = self.param_model np_dtype = self.indata.dtype.as_numpy_dtype param_cw_np = np.zeros(pm.nparams, dtype=np_dtype) @@ -397,25 +384,13 @@ def init_fit_parms( tf.math.reciprocal(self.cw), ) - # Global-impact "sources" are the constraint centers x0, carried as a - # SINGLE collection so params and systs are treated identically by the - # impact code. They are ordered [full nuisance block | priored params] - # to preserve the historical source axis (all systs, then any priored - # params). param_prior_idxs (priored params within the leading param - # block) is derived from cw, so the priors stay defined only on the - # model. + # Indices (within the leading param block) of the priored params, + # derived from cw so the priors stay defined only on the model. Used by + # the asymmetric global impacts to scan the prior centers as sources. self.param_prior_idxs = tf.constant( np.where(param_cw_np > 0)[0].astype(np.int64), dtype=tf.int64 ) self.param_prior_active = bool(int(self.param_prior_idxs.shape[0]) > 0) - self.x0_source_idxs = tf.concat( - [ - tf.range(pm.nparams, len(self.parms), dtype=tf.int64), - self.param_prior_idxs, - ], - axis=0, - ) - self.var_x0_sources = tf.gather(self.var_x0, self.x0_source_idxs) # Per-parameter prefit variance vector. Always allocated; the # prefit covariance is intrinsically diagonal so this is the @@ -815,18 +790,9 @@ def x0defaultassign(self): self.x0.assign(self.x0default) def xdefaultassign(self): - if self.param_model.nparams == 0: - self.x.assign(self.x0) - else: - self.x.assign( - tf.concat( - [ - self.param_model.xparamdefault, - self.x0[self.param_model.nparams :], - ], - axis=0, - ) - ) + # start every parameter at its constraint center (prior mean / theta0 + # default, and the model default for unpriored params) + self.x.assign(self.x0default) def defaultassign(self): var_pre = self.prefit_variance( @@ -850,23 +816,12 @@ def defaultassign(self): def bayesassign(self): # Sample the parameter values from their priors: width sqrt(1/cw) - # around the constraint centers wherever cw > 0; free entries stay - # at their defaults (for cw = 0 nuisances the default is the x0 - # entry itself). - if self.param_model.nparams == 0: - defaults = self.x0 - else: - defaults = tf.concat( - [ - self.param_model.xparamdefault, - self.x0[self.param_model.nparams :], - ], - axis=0, - ) + # around the constraint centers wherever cw > 0; free entries (cw = 0) + # stay at their default constraint center. sampled = self.x0 + tf.sqrt(self.var_x0) * tf.random.normal( shape=self.x0.shape, dtype=self.x0.dtype ) - self.x.assign(tf.where(self.cw > 0, sampled, defaults)) + self.x.assign(tf.where(self.cw > 0, sampled, self.x0default)) self.bbstat.randomize_bayes() @@ -1139,6 +1094,7 @@ def impacts_parms(self, hess): @tf.function def global_impacts_parms(self): + param_groupidxs = [idxs for _, idxs in self._resolved_param_impact_groups()] return global_impacts.global_impacts_parms( self.x, self.bbstat.ubeta, @@ -1154,19 +1110,19 @@ def global_impacts_parms(self): self.bbstat.binByBinStatMode, self.globalImpactsFromJVP, self.cov, - self.x0_source_idxs, - n_param_sources=int(self.param_prior_idxs.shape[0]), + param_groupidxs=param_groupidxs, ) @tf.function def gaussian_global_impacts_parms(self): dxdx0, dxdnobs, dxdbeta0 = self._dxdvars() + param_groupidxs = [idxs for _, idxs in self._resolved_param_impact_groups()] impacts, impacts_grouped = global_impacts.gaussian_global_impacts_parms( dxdx0, dxdnobs, dxdbeta0, - self.var_x0_sources, + self.var_x0, self.nobs if self.varnobs is None else self.varnobs, ( 1.0 @@ -1181,7 +1137,7 @@ def gaussian_global_impacts_parms(self): self.bbstat.binByBinStatMode, self.bbstat.beta_shape, self.indata.systgroupidxs, - n_param_sources=int(self.param_prior_idxs.shape[0]), + param_groupidxs=param_groupidxs, data_cov_inv=self.data_cov_inv, ) @@ -1415,26 +1371,18 @@ def _pd2ldbeta2(self, profile=False): def _dxdvars(self): # Response of the postfit minimum to a unit shift of each constraint - # center, as a SINGLE collection of columns aligned with - # x0_source_idxs (the nuisance block plus any priored params). x0 - # enters the NLL only through cw * (x - x0)^2, so the derivative - # w.r.t. any center with cw = 0 is identically zero; restricting the - # differentiation to these sources is exact, not an approximation, - # and keeps the Jacobian at [npar, n_sources] instead of [npar, npar]. - src_idxs = self.x0_source_idxs - n_src = int(src_idxs.shape[0]) - scatter_idx = src_idxs[:, None] + # center x0, over the whole parameter vector. x0 enters the NLL only + # through cw * (x - x0)^2, so columns for unconstrained centers (cw = 0) + # are exactly zero and carried along harmlessly. with tf.GradientTape() as t2: - ds = tf.zeros([n_src], dtype=self.indata.dtype) - t2.watch([ds, self.nobs, self.bbstat.beta0]) + t2.watch([self.x0, self.nobs, self.bbstat.beta0]) with tf.GradientTape() as t1: - t1.watch([ds, self.nobs, self.bbstat.beta0]) - x0 = self.x0 + tf.scatter_nd(scatter_idx, ds, self.x0.shape) - val = self._compute_loss(x0=x0) + t1.watch([self.x0, self.nobs, self.bbstat.beta0]) + val = self._compute_loss() grad = t1.gradient(val, self.x) pd2ldxdx0, pd2ldxdnobs, pd2ldxdbeta0 = t2.jacobian( grad, - [ds, self.nobs, self.bbstat.beta0], + [self.x0, self.nobs, self.bbstat.beta0], unconnected_gradients="zero", ) @@ -1446,25 +1394,21 @@ def _dxdvars(self): return dxdx0, dxdnobs, dxdbeta0 def _dndvars(self, fun): - # The forward model n depends on the parameters x but not on the - # constraint centers x0, so dn/dx0 = (dn/dx)(dx/dx0): differentiate n - # only w.r.t. x / nobs / beta0 here and pick up the x0 dependence - # through _dxdvars below. dndx0 is the single source collection. with tf.GradientTape() as t: - t.watch([self.nobs, self.bbstat.beta0]) + t.watch([self.x0, self.nobs, self.bbstat.beta0]) n = fun() n_flat = tf.reshape(n, (-1,)) - pdndx, pdndnobs, pdndbeta0 = t.jacobian( + pdndx, pdndx0, pdndnobs, pdndbeta0 = t.jacobian( n_flat, - [self.x, self.nobs, self.bbstat.beta0], + [self.x, self.x0, self.nobs, self.bbstat.beta0], unconnected_gradients="zero", ) # apply chain rule to take into account correlations with the fit parameters dxdx0, dxdnobs, dxdbeta0 = self._dxdvars() - dndx0 = pdndx @ dxdx0 + dndx0 = pdndx0 + pdndx @ dxdx0 dndnobs = pdndnobs + pdndx @ dxdnobs dndbeta0 = tf.reshape(pdndbeta0, [pdndbeta0.shape[0], -1]) + pdndx @ dxdbeta0 @@ -1564,6 +1508,7 @@ def compute_derivatives(dvars): expvar = tf.reshape(expvar_flat, tf.shape(expected)) + param_groupidxs = [idxs for _, idxs in self._resolved_param_impact_groups()] if compute_global_impacts: impacts, impacts_grouped = global_impacts.global_impacts_obs( self.x, @@ -1582,8 +1527,7 @@ def compute_derivatives(dvars): expvar_flat, expvar.shape, profile, - self.x0_source_idxs, - n_param_sources=int(self.param_prior_idxs.shape[0]), + param_groupidxs=param_groupidxs, pdexpdbeta=pdexpdbeta, pd2ldbeta2_pdexpdbeta=( pd2ldbeta2_pdexpdbeta if pdexpdbeta is not None else None @@ -1613,7 +1557,7 @@ def fun_n(): dndx0, dndnobs, dndbeta0, - self.var_x0_sources, + self.var_x0, self.nobs if self.varnobs is None else self.varnobs, ( 1.0 @@ -1625,7 +1569,8 @@ def fun_n(): self.bbstat.binByBinStatMode, self.bbstat.beta_shape, self.indata.systgroupidxs, - n_param_sources=int(self.param_prior_idxs.shape[0]), + self.param_model.nparams, + param_groupidxs=param_groupidxs, data_cov_inv=self.data_cov_inv, ) ) @@ -1909,10 +1854,9 @@ def fun_res(): residuals, dresdx0, dresdnobs, dresdbeta0 = self._dndvars(fun_res) - # one source collection; var_x0_sources weights each source by its - # prefit variance (the dropped cw = 0 columns contributed exactly zero - # to the old full-x0 quadratic form). - res_cov = dresdx0 @ (self.var_x0_sources[:, None] * tf.transpose(dresdx0)) + # dresdx0 spans the full parameter vector; var_x0 weights each center + # by its prefit variance (cw = 0 entries contribute exactly zero). + res_cov = dresdx0 @ (self.var_x0[:, None] * tf.transpose(dresdx0)) if self.covarianceFit: res_cov_stat = dresdnobs @ tf.linalg.solve( @@ -2076,16 +2020,13 @@ def full_nll(self): def reduced_nll(self): return self._compute_nll(full_nll=False) - def _compute_lc(self, full_nll=False, x0=None): + def _compute_lc(self, full_nll=False): # One constraint term over the full effective parameter vector - # [poi, model_nui, theta]: the ParamModel block is constrained by - # param_cw / param_x0 (declared priors; cw = 0 -> free) and the - # nuisance block by indata.constraintweights / theta0. - # x0 defaults to self.x0; _dxdvars passes a perturbed copy to - # differentiate the loss w.r.t. a subset of the constraint centers. - x0 = self.x0 if x0 is None else x0 + # [poi, model_nui, theta]: the ParamModel block is constrained by the + # declared priors (cw = 0 -> free) and the nuisance block by + # indata.constraintweights, all folded into cw / x0. cw = self.cw - lc = cw * 0.5 * tf.square(self.get_x() - x0) + lc = cw * 0.5 * tf.square(self.get_x() - self.x0) if full_nll: # normalization factor 0.5*log(2*pi*sigma^2) for constrained # entries, with sigma^2 = 1/cw and @@ -2132,7 +2073,7 @@ def _compute_ln(self, nexp, full_nll=False): ) return ln - def _compute_nll_components(self, profile=True, full_nll=False, x0=None): + def _compute_nll_components(self, profile=True, full_nll=False): nexpfullcentral, _, beta = self._compute_yields_with_beta( profile=profile, compute_norm=False, @@ -2143,7 +2084,7 @@ def _compute_nll_components(self, profile=True, full_nll=False, x0=None): ln = self._compute_ln(nexp, full_nll) - lc = self._compute_lc(full_nll, x0=x0) + lc = self._compute_lc(full_nll) lbeta = self._compute_lbeta(beta, full_nll) @@ -2165,9 +2106,9 @@ def _compute_external_nll(self): self.external_terms, self.x, self.indata.dtype ) - def _compute_nll(self, profile=True, full_nll=False, x0=None): + def _compute_nll(self, profile=True, full_nll=False): ln, lc, lbeta, lpenalty, beta = self._compute_nll_components( - profile=profile, full_nll=full_nll, x0=x0 + profile=profile, full_nll=full_nll ) l = ln + lc @@ -2182,8 +2123,8 @@ def _compute_nll(self, profile=True, full_nll=False, x0=None): l = l + lext return l - def _compute_loss(self, profile=True, x0=None): - return self._compute_nll(profile=profile, x0=x0) + def _compute_loss(self, profile=True): + return self._compute_nll(profile=profile) def _make_tf_functions(self): # Build tf.function wrappers at instance construction time so that diff --git a/rabbit/impacts/global_asym_impacts.py b/rabbit/impacts/global_asym_impacts.py index 394af56..e799f35 100644 --- a/rabbit/impacts/global_asym_impacts.py +++ b/rabbit/impacts/global_asym_impacts.py @@ -79,7 +79,7 @@ def global_asym_impacts_parms( (constraints are unit-sigma in rabbit, so 1.0 = 1 prefit sigma). signs: sequence (down, up). Bin 0 of axis_downUpVar -> first sign. linear_warmstart: experimental. If True, warm-start each refit at - x_nom + dxds[:, source] * shift, the Gaussian-approximation new + x_nom + dxdx0[:, idx] * shift, the Gaussian-approximation new minimum for the shifted center. Should drastically reduce the number of optimizer iterations on near-Gaussian sources. Requires fitter.cov to exist (same prerequisite as @@ -119,12 +119,11 @@ def global_asym_impacts_parms( + (" (linear warm-start enabled)" if linear_warmstart else "") ) - # Optional Gaussian-approximation warm-start. _dxdvars returns the postfit - # response to a unit shift of each constraint center as a single collection - # of columns aligned with fitter.x0_source_idxs; we map each scanned - # source's full-x index to its column. Computing it once is the same cost - # as one --gaussianGlobalImpacts call. - warmstart_col = None + # Optional Gaussian-approximation warm-start. dxdx0[:, j] is the postfit + # response to a unit shift of constraint center x0[j], so a scanned + # source's full-x index directly selects its column. Computing it once is + # the same cost as one --gaussianGlobalImpacts call. + dxdx0_np = None if linear_warmstart: if fitter.cov is None: raise RuntimeError( @@ -134,15 +133,8 @@ def global_asym_impacts_parms( t_lws = time.perf_counter() dxdx0_tf, _, _ = fitter._dxdvars() dxdx0_np = dxdx0_tf.numpy() - src_col = {int(s): k for k, s in enumerate(fitter.x0_source_idxs.numpy())} - - def _warmstart_col(idx): - return dxdx0_np[:, src_col[idx]] - - warmstart_col = _warmstart_col - logger.info( - f"global_asym_impacts: dxds prepared in " + f"global_asym_impacts: dxdx0 prepared in " f"{time.perf_counter() - t_lws:.2f}s" ) @@ -167,12 +159,12 @@ def _warmstart_col(idx): x0_shifted[idx] += shift # Warm-start x. With linear_warmstart, use the Gaussian-approx new - # minimum x_nom + dxds[:, source] * shift -- on near-Gaussian - # sources this lands at the new minimum to within roundoff. - # Without it, just shift x[idx] by `shift` so the parameter itself - # starts at the new constraint center. - if warmstart_col is not None: - x_shifted = x_nom_np + warmstart_col(idx) * shift + # minimum x_nom + dxdx0[:, idx] * shift -- on near-Gaussian sources + # this lands at the new minimum to within roundoff. Without it, + # just shift x[idx] by `shift` so the parameter itself starts at + # the new constraint center. + if dxdx0_np is not None: + x_shifted = x_nom_np + dxdx0_np[:, idx] * shift else: x_shifted = x_nom_np.copy() x_shifted[idx] += shift diff --git a/rabbit/impacts/global_impacts.py b/rabbit/impacts/global_impacts.py index fdead04..36d36be 100644 --- a/rabbit/impacts/global_impacts.py +++ b/rabbit/impacts/global_impacts.py @@ -193,28 +193,58 @@ def _compute_global_impacts_x0(x, compute_lc_fn, cov_dexpdx): return sc @ cov_dexpdx +def _grouped_columns(impacts_x0_sq, group_idxs): + """Quadrature-sum impacts_x0_sq over each group's columns -> [n_rows, + n_groups]. group_idxs is a (ragged) list of full-x column-index lists, so + syst and ParamModel sources are handled identically and a group may mix + both.""" + return tf.transpose( + tf.map_fn( + lambda idxs: _compute_global_impact_group(impacts_x0_sq, idxs), + tf.ragged.constant(group_idxs, dtype=tf.int64), + fn_output_signature=tf.TensorSpec( + shape=(impacts_x0_sq.shape[0],), dtype=impacts_x0_sq.dtype + ), + ) + ) + + +def _prepend_append_groups( + impacts_grouped, impacts_x0_sq, systgroupidxs, nmodel_params, param_groupidxs +): + """Prepend the syst groups and append the ParamModel impact groups to the + stat/bbb block, in the getGroupedImpactsAxes order + [syst groups | stat | bbb | param groups]. Groups are full-x column-index + lists into impacts_x0_sq -- syst groups shifted into the nuisance block by + nmodel_params, param groups indexing the leading param block -- so both are + combined the same way and a group could mix the two.""" + if len(systgroupidxs): + syst_cols = [[nmodel_params + int(i) for i in g] for g in systgroupidxs] + impacts_grouped = tf.concat( + [_grouped_columns(impacts_x0_sq, syst_cols), impacts_grouped], axis=-1 + ) + if param_groupidxs is not None and len(param_groupidxs): + param_cols = [[int(i) for i in g] for g in param_groupidxs] + impacts_grouped = tf.concat( + [impacts_grouped, _grouped_columns(impacts_x0_sq, param_cols)], axis=-1 + ) + return impacts_grouped + + def _compute_grouped_impacts( bin_by_bin_stat, bin_by_bin_stat_mode, systgroupidxs, - impacts_sources_sq, + nmodel_params, + param_groupidxs, + impacts_x0_sq, impacts_nobs, impacts_beta0_total, impacts_beta0_process, - n_param_sources=0, ): - """Assemble the grouped impacts tensor from all contributions. - - impacts_sources_sq holds the squared per-source impacts as a single - collection ordered [nuisance block | priored params]. The trailing - n_param_sources columns (the priored params) are appended at the END of - the grouped axis as individual sources, matching the output axes; the - leading nuisance columns are combined into the systematic groups. The - split here is only to assemble the grouped axis in its historical order. - """ - nsyst = int(impacts_sources_sq.shape[-1]) - n_param_sources - impacts_theta0_sq = impacts_sources_sq[:, :nsyst] - impacts_param0_sq = impacts_sources_sq[:, nsyst:] if n_param_sources else None + """Assemble the grouped impacts from the full per-source squared impacts + (likelihood path). impacts_x0_sq is [n_rows, npar] over the whole + parameter vector.""" if bin_by_bin_stat: impacts_grouped = tf.stack([impacts_nobs, impacts_beta0_total], axis=-1) if bin_by_bin_stat_mode == "full": @@ -224,23 +254,9 @@ def _compute_grouped_impacts( else: impacts_grouped = impacts_nobs[..., None] - if len(systgroupidxs): - impacts_grouped_syst = tf.map_fn( - lambda idxs: _compute_global_impact_group(impacts_theta0_sq, idxs), - tf.ragged.constant(systgroupidxs, dtype=tf.int64), - fn_output_signature=tf.TensorSpec( - shape=(impacts_theta0_sq.shape[0],), dtype=impacts_theta0_sq.dtype - ), - ) - impacts_grouped_syst = tf.transpose(impacts_grouped_syst) - impacts_grouped = tf.concat([impacts_grouped_syst, impacts_grouped], axis=-1) - - if impacts_param0_sq is not None: - impacts_grouped = tf.concat( - [impacts_grouped, tf.sqrt(impacts_param0_sq)], axis=-1 - ) - - return impacts_grouped + return _prepend_append_groups( + impacts_grouped, impacts_x0_sq, systgroupidxs, nmodel_params, param_groupidxs + ) def global_impacts_parms( @@ -258,8 +274,7 @@ def global_impacts_parms( bin_by_bin_stat_mode, global_impacts_from_jvp, cov, - source_idxs, - n_param_sources=0, + param_groupidxs=None, ): idxs_poi = tf.range(nsignal_params, dtype=tf.int64) idxs_noi = tf.constant(nmodel_params + noiidxs, dtype=tf.int64) @@ -286,16 +301,15 @@ def global_impacts_parms( pd2ldbeta2_pdexpdbeta=None, ) - # impacts_x0 covers the full constraint term over the whole parameter - # vector. The reported sources (constrained nuisances + priored params) - # are a single collection picked out by source_idxs; params and systs are - # treated identically here. Per-1-sigma units come out automatically since - # sc = sqrt(d2lc/dx2) = sqrt(cw) = 1/sigma. - impacts_x0 = _compute_global_impacts_x0(x, compute_lc_fn, cov_dexpdx) - impacts_sources = tf.transpose(tf.gather(impacts_x0, source_idxs, axis=0)) + # impacts_x0 is the per-source impact over the WHOLE parameter vector + # (one column per parameter); unconstrained entries (cw = 0) are exactly + # zero. Per-1-sigma units come out automatically since + # sc = sqrt(d2lc/dx2) = sqrt(cw) = 1/sigma. Params and systs are treated + # identically -- no source subset is carried. + impacts_x0 = tf.transpose(_compute_global_impacts_x0(x, compute_lc_fn, cov_dexpdx)) - impacts_sources_sq = tf.square(impacts_sources) - var_x0 = tf.reduce_sum(impacts_sources_sq, axis=-1) + impacts_x0_sq = tf.square(impacts_x0) + var_x0 = tf.reduce_sum(impacts_x0_sq, axis=-1) var_nobs = var_total - var_x0 if bin_by_bin_stat: @@ -305,14 +319,15 @@ def global_impacts_parms( bin_by_bin_stat, bin_by_bin_stat_mode, systgroupidxs, - impacts_sources_sq, + nmodel_params, + param_groupidxs, + impacts_x0_sq, tf.sqrt(var_nobs), impacts_beta0_total, impacts_beta0_process, - n_param_sources=n_param_sources, ) - return impacts_sources, impacts_grouped + return impacts_x0, impacts_grouped def global_impacts_obs( @@ -332,8 +347,7 @@ def global_impacts_obs( expvar_flat, expvar_shape, profile, - source_idxs, - n_param_sources=0, + param_groupidxs=None, pdexpdbeta=None, pd2ldbeta2_pdexpdbeta=None, prefit_unconstrained_nuisance_uncertainty=0.0, @@ -384,12 +398,11 @@ def global_impacts_obs( pd2ldbeta2_pdexpdbeta, ) - impacts_x0 = _compute_global_impacts_x0(x, compute_lc_fn, cov_dexpdx) - # single source collection (constrained nuisances + priored params) - impacts_sources = tf.transpose(tf.gather(impacts_x0, source_idxs, axis=0)) + # per-source impacts over the whole parameter vector (zero where cw = 0) + impacts_x0 = tf.transpose(_compute_global_impacts_x0(x, compute_lc_fn, cov_dexpdx)) - impacts_sources_sq = tf.square(impacts_sources) - var_x0 = tf.reduce_sum(impacts_sources_sq, axis=-1) + impacts_x0_sq = tf.square(impacts_x0) + var_x0 = tf.reduce_sum(impacts_x0_sq, axis=-1) var_nobs = expvar_flat - var_x0 if bin_by_bin_stat: @@ -399,14 +412,15 @@ def global_impacts_obs( bin_by_bin_stat, bin_by_bin_stat_mode, systgroupidxs, - impacts_sources_sq, + nmodel_params, + param_groupidxs, + impacts_x0_sq, tf.sqrt(var_nobs), impacts_beta0_total, impacts_beta0_process, - n_param_sources=n_param_sources, ) - impacts = tf.reshape(impacts_sources, [*expvar_shape, impacts_sources.shape[-1]]) + impacts = tf.reshape(impacts_x0, [*expvar_shape, impacts_x0.shape[-1]]) impacts_grouped = tf.reshape( impacts_grouped, [*expvar_shape, impacts_grouped.shape[-1]] ) @@ -425,23 +439,16 @@ def _gaussian_global_impacts( bin_by_bin_stat_mode, beta_shape, systgroupidxs, - n_param_sources=0, + nmodel_params, + param_groupidxs=None, data_cov_inv=None, ): - # dxdx0 / varx0 are the single source collection ordered - # [nuisance block | priored params]. The leading nuisance columns and the - # trailing priored-param columns are split out only to assemble the - # grouped axis in its historical order [syst groups | stat | bbb | priors]; - # the per-source impacts themselves are computed uniformly. - nsyst = int(dxdx0.shape[-1]) - n_param_sources - dxdtheta0 = dxdx0[:, :nsyst] - vartheta0 = varx0[:nsyst] - if n_param_sources: - dxdparam0 = dxdx0[:, nsyst:] - varparam0 = varx0[nsyst:] - else: - dxdparam0 = None - varparam0 = None + # Per-source impact over the whole parameter vector: a unit-sigma shift of + # source j moves the poi/noi by dxdx0[:, j] * sqrt(var_x0[j]). Unit- + # constrained nuisances have var = 1, priored params their prior sigma, + # unconstrained entries var = 0 (zero column). No param/syst split. + impacts = dxdx0 * tf.sqrt(varx0) + impacts_x0_sq = tf.square(impacts) if data_cov_inv is not None: data_cov = tf.linalg.inv(data_cov_inv) @@ -470,31 +477,11 @@ def _gaussian_global_impacts( [impacts_grouped, impacts_beta0_process], axis=-1 ) else: - impacts_grouped = impacts_data_stat + impacts_grouped = impacts_data_stat[..., None] - if len(systgroupidxs): - dxdtheta0_squared = tf.square(dxdtheta0) * vartheta0 - - impacts_grouped_syst = tf.map_fn( - lambda idxs: _compute_global_impact_group(dxdtheta0_squared, idxs), - tf.ragged.constant(systgroupidxs, dtype=tf.int64), - fn_output_signature=tf.TensorSpec( - shape=(dxdtheta0_squared.shape[0],), dtype=tf.float64 - ), - ) - impacts_grouped_syst = tf.transpose(impacts_grouped_syst) - impacts_grouped = tf.concat([impacts_grouped_syst, impacts_grouped], axis=1) - - impacts = dxdtheta0 - if dxdparam0 is not None: - # impacts of the ParamModel prior centers, per 1 prefit sigma; one - # column per priored param, appended at the END of the source axis - # and of the grouped axis (matching the output axes). - impacts_param0 = dxdparam0 * tf.sqrt(varparam0) - impacts = tf.concat([impacts, impacts_param0], axis=-1) - if impacts_grouped.shape.rank == 1: - impacts_grouped = impacts_grouped[..., None] - impacts_grouped = tf.concat([impacts_grouped, tf.abs(impacts_param0)], axis=-1) + impacts_grouped = _prepend_append_groups( + impacts_grouped, impacts_x0_sq, systgroupidxs, nmodel_params, param_groupidxs + ) return impacts, impacts_grouped @@ -513,10 +500,10 @@ def gaussian_global_impacts_parms( bin_by_bin_stat_mode, beta_shape, systgroupidxs, - n_param_sources=0, + param_groupidxs=None, data_cov_inv=None, ): - # dxdx0 is the single source-derivative collection; gather the poi/noi + # dxdx0 is the full per-source derivative collection; gather the poi/noi # rows whose impacts are reported (columns / sources are untouched). dxdx0 = _gather_poi_noi_vector(dxdx0, noiidxs, nsignal_params, nmodel_params) dxdnobs = _gather_poi_noi_vector(dxdnobs, noiidxs, nsignal_params, nmodel_params) @@ -533,7 +520,8 @@ def gaussian_global_impacts_parms( bin_by_bin_stat_mode, beta_shape, systgroupidxs, - n_param_sources, + nmodel_params, + param_groupidxs, data_cov_inv, ) @@ -549,7 +537,8 @@ def gaussian_global_impacts_obs( bin_by_bin_stat_mode, beta_shape, systgroupidxs, - n_param_sources=0, + nmodel_params, + param_groupidxs=None, data_cov_inv=None, ): return _gaussian_global_impacts( @@ -563,6 +552,7 @@ def gaussian_global_impacts_obs( bin_by_bin_stat_mode, beta_shape, systgroupidxs, - n_param_sources, + nmodel_params, + param_groupidxs, data_cov_inv, ) diff --git a/rabbit/workspace.py b/rabbit/workspace.py index 9995774..3d68902 100644 --- a/rabbit/workspace.py +++ b/rabbit/workspace.py @@ -62,22 +62,14 @@ def __init__(self, outdir, outname, fitter, postfix=None): self.noiidxs = fitter.indata.noiidxs # some information for the impact histograms - systs = list(fitter.indata.systs.astype(str)) parms = list(fitter.parms.astype(str)) - # ParamModel prior centers act as additional global-impact sources; - # their columns are appended after the systs / syst groups. - if fitter.param_prior_active: - prior_idxs = fitter.param_prior_idxs.numpy() - prior_names = [f"{parms[int(i)]}_prior" for i in prior_idxs] - else: - prior_names = [] + # The global impacts report a source per parameter over the whole + # vector (unconstrained params are exactly zero), same axis as the + # per-parameter (traditional) impacts. self.impact_axis = hist.axis.StrCategory(parms, name="impacts") - self.global_impact_axis = hist.axis.StrCategory( - systs + prior_names, name="impacts" - ) - # ParamModel impact groups (e.g. SCETlib NP gamma_nu / F_eff) are only - # computed by the traditional impacts, so they extend that axis only; - # the prior sources extend the global axes only. + self.global_impact_axis = hist.axis.StrCategory(parms, name="impacts") + # ParamModel impact groups (e.g. SCETlib NP gamma_nu / F_eff) extend + # both the traditional and the global grouped axes. param_impact_group_names = [ name for name, _ in fitter._resolved_param_impact_groups() ] @@ -91,7 +83,7 @@ def __init__(self, outdir, outname, fitter, postfix=None): fitter.indata, bin_by_bin_stat=fitter.bbstat.enabled, per_process=fitter.bbstat.binByBinStatMode == "full", - extra_groups=prior_names, + extra_groups=param_impact_group_names, ) self.extension = "hdf5"