diff --git a/bin/rabbit_fit.py b/bin/rabbit_fit.py
index 7d52bd9..772d6c6 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.",
)
@@ -410,16 +410,34 @@ 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) 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.theta0.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)
@@ -859,6 +877,27 @@ def main():
"nois": ifitter.parms[ifitter.param_model.nparams :][indata.noiidxs],
}
+ # 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
+ 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_dtype)
+ if means is None
+ else np.asarray(means, dtype=np_dtype)
+ )
+ meta["param_priors"] = {
+ "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(
args.outpath,
args.outname,
diff --git a/rabbit/fitter.py b/rabbit/fitter.py
index fd0120d..a5d8be2 100644
--- a/rabbit/fitter.py
+++ b/rabbit/fitter.py
@@ -298,6 +298,100 @@ def init_fit_parms(
self.x = tf.Variable(xdefault, trainable=True, name="x")
+ # 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)
+ # 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)
+ # 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)
+ 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_dtype)
+ else:
+ means = np.asarray(means, dtype=np_dtype)
+ 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:
+ 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."
+ )
+ 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)
+ 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}")
+
+ # 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,
+ )
+ 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),
+ )
+
+ # 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)
+
# Per-parameter prefit variance vector. Always allocated; the
# prefit covariance is intrinsically diagonal so this is the
# only form needed for prefit uncertainties.
@@ -335,18 +429,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(
@@ -655,21 +737,16 @@ 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).
+ 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_poi = (
- tf.ones([self.param_model.nparams], dtype=self.indata.dtype)
- * unconstrained_err**2
- )
- var_theta = tf.where(
- self.indata.constraintweights == 0.0,
- unconstrained_err**2,
- tf.math.reciprocal(self.indata.constraintweights),
+ return tf.where(
+ self.cw == 0.0,
+ unconstrained_err**2 * tf.ones_like(self.cw),
+ tf.math.reciprocal(self.cw),
)
- return tf.concat([var_poi, var_theta], axis=0)
def prefit_covariance(self, unconstrained_err=0.0):
"""Full prefit covariance as a tf.linalg.LinearOperatorDiag.
@@ -708,16 +785,14 @@ 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):
- self.theta0.assign(self.theta0default)
+ 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)
- else:
- self.x.assign(
- tf.concat([self.param_model.xparamdefault, self.theta0], 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(
@@ -726,7 +801,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()
@@ -740,32 +815,24 @@ def defaultassign(self):
reg.set_expectations(xinit, nexp0)
def bayesassign(self):
- # FIXME use theta0 as the mean and constraintweight to scale the width
- if self.param_model.nparams == 0:
- self.x.assign(
- self.theta0
- + tf.random.normal(shape=self.theta0.shape, dtype=self.theta0.dtype)
- )
- 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,
- )
- )
+ # Sample the parameter values from their priors: width sqrt(1/cw)
+ # 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, self.x0default))
self.bbstat.randomize_bayes()
def frequentistassign(self):
- # FIXME use theta as the mean and constraintweight to scale the width
- self.theta0.assign(
- tf.random.normal(shape=self.theta0.shape, dtype=self.theta0.dtype)
+ # Fluctuate the constraint centers around their defaults with the
+ # 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()
@@ -1027,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,
@@ -1042,17 +1110,19 @@ def global_impacts_parms(self):
self.bbstat.binByBinStatMode,
self.globalImpactsFromJVP,
self.cov,
+ param_groupidxs=param_groupidxs,
)
@tf.function
def gaussian_global_impacts_parms(self):
- dxdtheta0, dxdnobs, dxdbeta0 = self._dxdvars()
+ 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(
- dxdtheta0,
+ dxdx0,
dxdnobs,
dxdbeta0,
- self.var_theta0,
+ self.var_x0,
self.nobs if self.varnobs is None else self.varnobs,
(
1.0
@@ -1067,7 +1137,8 @@ def gaussian_global_impacts_parms(self):
self.bbstat.binByBinStatMode,
self.bbstat.beta_shape,
self.indata.systgroupidxs,
- self.data_cov_inv,
+ param_groupidxs=param_groupidxs,
+ data_cov_inv=self.data_cov_inv,
)
return impacts, impacts_grouped
@@ -1169,9 +1240,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,
@@ -1179,13 +1253,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)
@@ -1193,26 +1268,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,
)
@@ -1220,13 +1328,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,
@@ -1263,45 +1370,49 @@ def _pd2ldbeta2(self, profile=False):
return pd2ldbeta2
def _dxdvars(self):
+ # Response of the postfit minimum to a unit shift of each constraint
+ # 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:
- t2.watch([self.theta0, self.nobs, self.bbstat.beta0])
+ t2.watch([self.x0, self.nobs, self.bbstat.beta0])
with tf.GradientTape() as t1:
- t1.watch([self.theta0, self.nobs, self.bbstat.beta0])
+ t1.watch([self.x0, self.nobs, self.bbstat.beta0])
val = self._compute_loss()
grad = t1.gradient(val, self.x)
- pd2ldxdtheta0, pd2ldxdnobs, pd2ldxdbeta0 = t2.jacobian(
+ pd2ldxdx0, pd2ldxdnobs, pd2ldxdbeta0 = t2.jacobian(
grad,
- [self.theta0, self.nobs, self.bbstat.beta0],
+ [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])
- return dxdtheta0, dxdnobs, dxdbeta0
+ 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 = 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
- return n, dndtheta0, dndnobs, dndbeta0
+ return n, dndx0, dndnobs, dndbeta0
def _compute_expected(
self, fun_exp, inclusive=True, profile=False, full=True, need_observables=True
@@ -1397,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,
@@ -1415,9 +1527,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_groupidxs=param_groupidxs,
+ pdexpdbeta=pdexpdbeta,
+ pd2ldbeta2_pdexpdbeta=(
+ pd2ldbeta2_pdexpdbeta if pdexpdbeta is not None else None
+ ),
+ prefit_unconstrained_nuisance_uncertainty=(
+ self.prefit_unconstrained_nuisance_uncertainty
+ ),
)
else:
impacts = None
@@ -1434,13 +1551,13 @@ def fun_n():
need_observables=need_observables,
)
- _, dndtheta0, dndnobs, dndbeta0 = self._dndvars(fun_n)
+ _, 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
@@ -1452,7 +1569,9 @@ def fun_n():
self.bbstat.binByBinStatMode,
self.bbstat.beta_shape,
self.indata.systgroupidxs,
- self.data_cov_inv,
+ self.param_model.nparams,
+ param_groupidxs=param_groupidxs,
+ data_cov_inv=self.data_cov_inv,
)
)
else:
@@ -1733,9 +1852,11 @@ def fun_res():
observed = fun(None, self.nobs)
return expected - observed
- residuals, dresdtheta0, dresdnobs, dresdbeta0 = self._dndvars(fun_res)
+ residuals, dresdx0, dresdnobs, dresdbeta0 = self._dndvars(fun_res)
- res_cov = dresdtheta0 @ (self.var_theta0[:, None] * tf.transpose(dresdtheta0))
+ # 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(
@@ -1900,12 +2021,22 @@ 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 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() - self.x0)
if full_nll:
- # normalization factor for normal distribution: log(1/sqrt(2*pi)) = -0.9189385332046727
- lc = lc + 0.9189385332046727 * self.indata.constraintweights
+ # 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 tf.reduce_sum(lc)
diff --git a/rabbit/impacts/global_asym_impacts.py b/rabbit/impacts/global_asym_impacts.py
index a169cc0..e799f35 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 theta0-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, theta0 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.
+ fitter: the Fitter instance (used for x, x0 and minimize).
+ 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 + 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
--gaussianGlobalImpacts).
@@ -78,33 +90,40 @@ 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())
- 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(
- 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.
- # dxdtheta0 has shape [npar, nsyst]; column 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
+ # 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(
@@ -112,42 +131,45 @@ 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"
)
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`.
- theta0_shifted = theta0_nom_np.copy()
- theta0_shifted[i] += shift
+ # Always shift the constraint center for source idx by `shift`.
+ x0_shifted = x0_nom_np.copy()
+ x0_shifted[idx] += shift
# Warm-start x. With linear_warmstart, use the Gaussian-approx new
- # minimum x_nom + dxdtheta0[:, 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
+ # 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[nparams + i] += shift
+ x_shifted[idx] += shift
- fitter.theta0.assign(theta0_shifted)
+ fitter.x0.assign(x0_shifted)
fitter.x.assign(x_shifted)
try:
@@ -165,7 +187,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()
@@ -175,17 +197,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 47c48a3..36d36be 100644
--- a/rabbit/impacts/global_impacts.py
+++ b/rabbit/impacts/global_impacts.py
@@ -193,16 +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_theta0_sq,
+ nmodel_params,
+ param_groupidxs,
+ impacts_x0_sq,
impacts_nobs,
impacts_beta0_total,
impacts_beta0_process,
):
- """Assemble the grouped impacts tensor from all contributions."""
+ """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":
@@ -212,18 +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)
-
- return impacts_grouped
+ return _prepend_append_groups(
+ impacts_grouped, impacts_x0_sq, systgroupidxs, nmodel_params, param_groupidxs
+ )
def global_impacts_parms(
@@ -241,6 +274,7 @@ def global_impacts_parms(
bin_by_bin_stat_mode,
global_impacts_from_jvp,
cov,
+ param_groupidxs=None,
):
idxs_poi = tf.range(nsignal_params, dtype=tf.int64)
idxs_noi = tf.constant(nmodel_params + noiidxs, dtype=tf.int64)
@@ -267,12 +301,17 @@ def global_impacts_parms(
pd2ldbeta2_pdexpdbeta=None,
)
- impacts_x0 = _compute_global_impacts_x0(x, compute_lc_fn, cov_dexpdx)
- impacts_theta0 = tf.transpose(impacts_x0[nmodel_params:])
+ # 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_x0_sq = tf.square(impacts_x0)
+ var_x0 = tf.reduce_sum(impacts_x0_sq, axis=-1)
- impacts_theta0_sq = tf.square(impacts_theta0)
- var_theta0 = tf.reduce_sum(impacts_theta0_sq, axis=-1)
- var_nobs = var_total - var_theta0
+ var_nobs = var_total - var_x0
if bin_by_bin_stat:
var_nobs -= var_beta0
@@ -280,13 +319,15 @@ def global_impacts_parms(
bin_by_bin_stat,
bin_by_bin_stat_mode,
systgroupidxs,
- impacts_theta0_sq,
+ nmodel_params,
+ param_groupidxs,
+ impacts_x0_sq,
tf.sqrt(var_nobs),
impacts_beta0_total,
impacts_beta0_process,
)
- return impacts_theta0, impacts_grouped
+ return impacts_x0, impacts_grouped
def global_impacts_obs(
@@ -306,6 +347,7 @@ def global_impacts_obs(
expvar_flat,
expvar_shape,
profile,
+ param_groupidxs=None,
pdexpdbeta=None,
pd2ldbeta2_pdexpdbeta=None,
prefit_unconstrained_nuisance_uncertainty=0.0,
@@ -356,12 +398,13 @@ def global_impacts_obs(
pd2ldbeta2_pdexpdbeta,
)
- impacts_x0 = _compute_global_impacts_x0(x, compute_lc_fn, cov_dexpdx)
- impacts_theta0 = tf.transpose(impacts_x0[nmodel_params:])
+ # 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_theta0_sq = tf.square(impacts_theta0)
- var_theta0 = tf.reduce_sum(impacts_theta0_sq, axis=-1)
- var_nobs = expvar_flat - var_theta0
+ 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:
var_nobs -= var_beta0
@@ -369,13 +412,15 @@ def global_impacts_obs(
bin_by_bin_stat,
bin_by_bin_stat_mode,
systgroupidxs,
- impacts_theta0_sq,
+ nmodel_params,
+ param_groupidxs,
+ impacts_x0_sq,
tf.sqrt(var_nobs),
impacts_beta0_total,
impacts_beta0_process,
)
- impacts = tf.reshape(impacts_theta0, [*expvar_shape, impacts_theta0.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]]
)
@@ -384,18 +429,27 @@ 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,
+ nmodel_params,
+ param_groupidxs=None,
data_cov_inv=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)
# equivalent to tf.linalg.diag_part(dxdnobs @ data_cov @ tf.transpose(dxdnobs)) but avoiding computing full matrix
@@ -423,29 +477,20 @@ 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_grouped = _prepend_append_groups(
+ impacts_grouped, impacts_x0_sq, systgroupidxs, nmodel_params, param_groupidxs
+ )
- return dxdtheta0, impacts_grouped
+ return impacts, impacts_grouped
def gaussian_global_impacts_parms(
- dxdtheta0,
+ dxdx0,
dxdnobs,
dxdbeta0,
- vartheta0,
+ varx0,
varnobs,
varbeta0,
nsignal_params,
@@ -455,53 +500,59 @@ def gaussian_global_impacts_parms(
bin_by_bin_stat_mode,
beta_shape,
systgroupidxs,
+ param_groupidxs=None,
data_cov_inv=None,
):
- # compute impacts for pois and nois
- dxdtheta0 = _gather_poi_noi_vector(
- dxdtheta0, noiidxs, nsignal_params, nmodel_params
- )
+ # 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)
dxdbeta0 = _gather_poi_noi_vector(dxdbeta0, 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,
+ nmodel_params,
+ param_groupidxs,
data_cov_inv,
)
def gaussian_global_impacts_obs(
- dndtheta0,
+ dndx0,
dndnobs,
dndbeta0,
- vartheta0,
+ varx0,
varnobs,
varbeta0,
bin_by_bin_stat,
bin_by_bin_stat_mode,
beta_shape,
systgroupidxs,
+ nmodel_params,
+ param_groupidxs=None,
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,
+ nmodel_params,
+ param_groupidxs,
data_cov_inv,
)
diff --git a/rabbit/impacts/nonprofiled_impacts.py b/rabbit/impacts/nonprofiled_impacts.py
index 650df20..870c4fc 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,23 @@ 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,
+ # 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(constraintweights),
+ tf.math.rsqrt(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 +83,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 = []
diff --git a/rabbit/param_models/param_model.py b/rabbit/param_models/param_model.py
index f8455d0..70b68c8 100644
--- a/rabbit/param_models/param_model.py
+++ b/rabbit/param_models/param_model.py
@@ -17,6 +17,20 @@ 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.
+ # # 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 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
+ # # 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/workspace.py b/rabbit/workspace.py
index e069cfa..3d68902 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")
@@ -61,12 +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))
+ # 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, 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.
+ 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()
]
@@ -80,6 +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=param_impact_group_names,
)
self.extension = "hdf5"