diff --git a/README.md b/README.md index 7fcdfbe..32d527d 100644 --- a/README.md +++ b/README.md @@ -126,7 +126,7 @@ Bin-by-bin statistical uncertainties on the templates are added by default and c Perform mappings on the parameters and observables (the histogram bins in the (masked) channels). Baseline mappings are defined in `rabbit/mappings/` and can be called in `rabbit_fit` with the `--mapping` or `-m` option e.g. `-m Select ch0 -m Project ch1 b`. The first argument is the mapping name followed by arguments passed into the mapping. -Available physics models are +Available mappings are: * `BaseMapping`: Compute histograms in all bins and all channels. * `Select`: To select histograms of a channel, and perform a selection of processes and bins, supporting rebinning. * `Project`: To project histograms to lower dimensions, respecting the covariance matrix across bins. @@ -149,8 +149,18 @@ Custom mappings can be defined. They can be specified with the full path to the custom mapping e.g. `-m custom_mapping.MyCustomMapping`. The path must be accessable from your `$PYTHONPATH` variable and an `__ini__.py` file must be in the directory. -### Physics models -TBD +### POI models +POI models can be used to introduce paramter of interests (POIs) and modify the number of predicted events in the fit. +Baseline models are defined in `rabbit/poi_models/` and can be called in `rabbit_fit` with the `--poiModel` option, e.g. `--poiModel Mu`. +Only one POI model at a time can be used at a time. +Available POI models are: +* `Mu`: Scale the number of events for each signal process with an unconstrained parameter, and background proesses with 1. This is the default model +* `Ones`: Return ones, i.e. leave the number of predicted events the same. +* `Mixture`: Scale the `primary` processes by `x` and the `complementary` processes by `1-x` + +Custom POI models can be defined. +They can be specified with the full path to the custom mapping e.g. `--poiModel custom_model.MyCustomModel`. +The path must be accessable from your `$PYTHONPATH` variable and an `__ini__.py` file must be in the directory. ## Fit diagnostics diff --git a/bin/rabbit_fit.py b/bin/rabbit_fit.py index cf2c0b2..da3723f 100755 --- a/bin/rabbit_fit.py +++ b/bin/rabbit_fit.py @@ -14,6 +14,7 @@ from rabbit import fitter, inputdata, io_tools, workspace from rabbit.mappings import helpers as mh from rabbit.mappings import mapping as mp +from rabbit.poi_models import helpers as ph from rabbit.tfhelpers import edmval_cov from wums import output_tools, logging # isort: skip @@ -115,11 +116,14 @@ def make_parser(): ) parser.add_argument( "--expectSignal", - default=1.0, - type=float, - help="rate multiplier for signal expectation (used for fit starting values and for toys)", + default=None, + nargs=2, + action="append", + help=""" + Specify tuple with key and value to be passed to POI model (used for fit starting values and for toys). + E.g. '--expectSignal BSM 0.0 --expectSignal SM 1.0' + """, ) - parser.add_argument("--POIMode", default="mu", help="mode for POI's") parser.add_argument( "--allowNegativePOI", default=False, @@ -280,6 +284,12 @@ def make_parser(): nargs="*", help="run fit on pseudo data with the given name", ) + parser.add_argument( + "--poiModel", + default=["Mu"], + nargs="+", + help="Specify POI model to be used to introduce non standard parameterization", + ) parser.add_argument( "-m", "--mapping", @@ -529,7 +539,7 @@ def fit(args, fitter, ws, dofit=True): nllvalreduced = fitter.reduced_nll().numpy() ndfsat = ( - tf.size(fitter.nobs) - fitter.npoi - fitter.indata.nsystnoconstraint + tf.size(fitter.nobs) - fitter.poi_model.npoi - fitter.indata.nsystnoconstraint ).numpy() chi2 = 2.0 * nllvalreduced @@ -653,7 +663,11 @@ def main(): blinded_fits = [f == 0 or (f > 0 and args.toysDataMode == "observed") for f in fits] indata = inputdata.FitInputData(args.filename, args.pseudoData) - ifitter = fitter.Fitter(indata, args, do_blinding=any(blinded_fits)) + + margs = args.poiModel + poi_model = ph.load_model(margs[0], indata, *margs[1:], **vars(args)) + + ifitter = fitter.Fitter(indata, poi_model, args, do_blinding=any(blinded_fits)) # mappings for observables and parameters if len(args.mapping) == 0 and args.saveHists: @@ -661,7 +675,7 @@ def main(): args.mapping = [["BaseMapping"]] mappings = [] for margs in args.mapping: - mapping = mh.instance_from_class(margs[0], indata, *margs[1:]) + mapping = mh.load_mapping(margs[0], indata, *margs[1:]) mappings.append(mapping) if args.compositeMapping: @@ -676,9 +690,9 @@ def main(): meta = { "meta_info": output_tools.make_meta_info_dict(args=args), "meta_info_input": ifitter.indata.metadata, - "signals": ifitter.indata.signals, + "pois": ifitter.poi_model.pois, "procs": ifitter.indata.procs, - "nois": ifitter.parms[ifitter.npoi :][indata.noiidxs], + "nois": ifitter.parms[ifitter.poi_model.npoi :][indata.noiidxs], } with workspace.Workspace( diff --git a/rabbit/common.py b/rabbit/common.py index b07df42..448e99b 100644 --- a/rabbit/common.py +++ b/rabbit/common.py @@ -1,3 +1,4 @@ +import importlib import pathlib import re @@ -18,3 +19,29 @@ def natural_sort_dict(dictionary): sorted_keys = natural_sort(dictionary.keys()) sorted_dict = {key: dictionary[key] for key in sorted_keys} return sorted_dict + + +def load_class_from_module(class_name, class_module_dict, base_dir): + if "." in class_name: + # import from full relative or abslute path + parts = class_name.split(".") + module_name = ".".join(parts[:-1]) + class_name = parts[-1] + else: + # import one of the baseline classes + if class_name not in class_module_dict: + raise ValueError( + f"Class {class_name} not found, available classes are {class_module_dict.keys()}" + ) + module_name = f"{base_dir}.{class_module_dict[class_name]}" + + # Try to import the module + module = importlib.import_module(module_name) + + this_class = getattr(module, class_name, None) + if this_class is None: + raise AttributeError( + f"Class '{class_name}' not found in module '{module_name}'." + ) + + return this_class diff --git a/rabbit/fitter.py b/rabbit/fitter.py index 9a8d87d..ddb053f 100644 --- a/rabbit/fitter.py +++ b/rabbit/fitter.py @@ -46,7 +46,7 @@ class Fitter: valid_bin_by_bin_stat_types = ["gamma", "normal-additive", "normal-multiplicative"] valid_systematic_types = ["log_normal", "normal"] - def __init__(self, indata, options, do_blinding=False): + def __init__(self, indata, poi_model, options, do_blinding=False): self.indata = indata self.globalImpactsFromJVP = not options.globalImpactsDisableJVP @@ -105,25 +105,12 @@ def __init__(self, indata, options, do_blinding=False): options.prefitUnconstrainedNuisanceUncertainty ) - self.pois = [] - - if options.POIMode == "mu": - self.npoi = self.indata.nsignals - poidefault = options.expectSignal * tf.ones( - [self.npoi], dtype=self.indata.dtype - ) - for signal in self.indata.signals: - self.pois.append(signal) - elif options.POIMode == "none": - self.npoi = 0 - poidefault = tf.zeros([], dtype=self.indata.dtype) - else: - raise Exception("unsupported POIMode") + self.poi_model = poi_model self.do_blinding = do_blinding if self.do_blinding: self._blinding_offsets_poi = tf.Variable( - tf.ones([self.npoi], dtype=self.indata.dtype), + tf.ones([self.poi_model.npoi], dtype=self.indata.dtype), trainable=False, name="offset_poi", ) @@ -134,21 +121,14 @@ def __init__(self, indata, options, do_blinding=False): ) self.init_blinding_values(options.unblind) - self.parms = np.concatenate([self.pois, self.indata.systs]) + self.parms = np.concatenate([self.poi_model.pois, self.indata.systs]) self.init_frozen_params(options.freezeParameters) - self.allowNegativePOI = options.allowNegativePOI - - if self.allowNegativePOI: - self.xpoidefault = poidefault - else: - self.xpoidefault = tf.sqrt(poidefault) - # tf variable containing all fit parameters thetadefault = tf.zeros([self.indata.nsyst], dtype=self.indata.dtype) - if self.npoi > 0: - xdefault = tf.concat([self.xpoidefault, thetadefault], axis=0) + if self.poi_model.npoi > 0: + xdefault = tf.concat([self.poi_model.xpoidefault, thetadefault], axis=0) else: xdefault = thetadefault @@ -275,7 +255,7 @@ def __init__(self, indata, options, do_blinding=False): # determine if problem is linear (ie likelihood is purely quadratic) self.is_linear = ( (self.chisqFit or self.covarianceFit) - and (self.npoi == 0 or self.allowNegativePOI) + and self.poi_model.is_linear and self.indata.symmetric_tensor and self.indata.systematic_type == "normal" and ((not self.binByBinStat) or self.binByBinStatType == "normal-additive") @@ -331,8 +311,8 @@ def deterministic_random_from_string(s, mean=0.0, std=5.0): self._blinding_values_theta[i] = value # add offset to pois - self._blinding_values_poi = np.ones(self.npoi, dtype=np.float64) - for i in range(self.npoi): + self._blinding_values_poi = np.ones(self.poi_model.npoi, dtype=np.float64) + for i in range(self.poi_model.npoi): param = self.indata.signals[i] if param in unblind_parameters: continue @@ -347,21 +327,23 @@ def set_blinding_offsets(self, blind=True): self._blinding_offsets_poi.assign(self._blinding_values_poi) self._blinding_offsets_theta.assign(self._blinding_values_theta) else: - self._blinding_offsets_poi.assign(np.ones(self.npoi, dtype=np.float64)) + self._blinding_offsets_poi.assign( + np.ones(self.poi_model.npoi, dtype=np.float64) + ) self._blinding_offsets_theta.assign( np.zeros(self.indata.nsyst, dtype=np.float64) ) def get_blinded_theta(self): - theta = self.x[self.npoi :] + theta = self.x[self.poi_model.npoi :] if self.do_blinding: return theta + self._blinding_offsets_theta else: return theta def get_blinded_poi(self): - xpoi = self.x[: self.npoi] - if self.allowNegativePOI: + xpoi = self.x[: self.poi_model.npoi] + if self.poi_model.allowNegativePOI: poi = xpoi else: poi = tf.square(xpoi) @@ -378,7 +360,7 @@ def _default_beta0(self): def prefit_covariance(self, unconstrained_err=0.0): # free parameters are taken to have zero uncertainty for the purposes of prefit uncertainties - var_poi = tf.zeros([self.npoi], dtype=self.indata.dtype) + var_poi = tf.zeros([self.poi_model.npoi], dtype=self.indata.dtype) # nuisances have their uncertainty taken from the constraint term, but unconstrained nuisances # are set to a placeholder uncertainty (zero by default) for the purposes of prefit uncertainties @@ -427,10 +409,10 @@ def theta0defaultassign(self): self.theta0.assign(tf.zeros([self.indata.nsyst], dtype=self.theta0.dtype)) def xdefaultassign(self): - if self.npoi == 0: + if self.poi_model.npoi == 0: self.x.assign(self.theta0) else: - self.x.assign(tf.concat([self.xpoidefault, self.theta0], axis=0)) + self.x.assign(tf.concat([self.poi_model.xpoidefault, self.theta0], axis=0)) def beta0defaultassign(self): self.set_beta0(self._default_beta0()) @@ -454,7 +436,7 @@ def defaultassign(self): def bayesassign(self): # FIXME use theta0 as the mean and constraintweight to scale the width - if self.npoi == 0: + if self.poi_model.npoi == 0: self.x.assign( self.theta0 + tf.random.normal(shape=self.theta0.shape, dtype=self.theta0.dtype) @@ -463,7 +445,7 @@ def bayesassign(self): self.x.assign( tf.concat( [ - self.xpoidefault, + self.poi_model.xpoidefault, self.theta0 + tf.random.normal( shape=self.theta0.shape, dtype=self.theta0.dtype @@ -652,10 +634,11 @@ def nonprofiled_impacts_parms(self, unconstrained_err=1.0): for j, sign in enumerate((1, -1)): variation = ( - sign * err_theta[idx - self.npoi] + theta0_tmp[idx - self.npoi] + sign * err_theta[idx - self.poi_model.npoi] + + theta0_tmp[idx - self.poi_model.npoi] ) # vary the non-profile parameter - self.theta0[idx - self.npoi].assign(variation) + self.theta0[idx - self.poi_model.npoi].assign(variation) self.x[idx].assign( variation ) # this should not be needed but should accelerates the minimization @@ -671,7 +654,9 @@ def nonprofiled_impacts_parms(self, unconstrained_err=1.0): self.x.assign(x_tmp) # back to original value - self.theta0[idx - self.npoi].assign(theta0_tmp[idx - self.npoi]) + self.theta0[idx - self.poi_model.npoi].assign( + theta0_tmp[idx - self.poi_model.npoi] + ) # grouped nonprofiled impacts @tf.function @@ -712,7 +697,9 @@ def envelope(values): ) def _compute_impact_group(self, v, idxs): - cov_reduced = tf.gather(self.cov[self.npoi :, self.npoi :], idxs, axis=0) + cov_reduced = tf.gather( + self.cov[self.poi_model.npoi :, self.poi_model.npoi :], idxs, axis=0 + ) cov_reduced = tf.gather(cov_reduced, idxs, axis=1) v_reduced = tf.gather(v, idxs, axis=1) invC_v = tf.linalg.solve(cov_reduced, tf.transpose(v_reduced)) @@ -720,10 +707,10 @@ def _compute_impact_group(self, v, idxs): return tf.sqrt(v_invC_v) def _gather_poi_noi_vector(self, v): - v_poi = v[: self.npoi] + v_poi = v[: self.poi_model.npoi] # protection for constained NOIs, set them to 0 mask = (self.indata.noiidxs >= 0) & ( - self.indata.noiidxs < tf.shape(v[self.npoi :])[0] + self.indata.noiidxs < tf.shape(v[self.poi_model.npoi :])[0] ) safe_idxs = tf.where(mask, self.indata.noiidxs, 0) mask = tf.cast(mask, v.dtype) @@ -733,7 +720,7 @@ def _gather_poi_noi_vector(self, v): [tf.shape(mask), tf.ones(tf.rank(v) - 1, dtype=tf.int32)], axis=0 ), ) - v_noi = tf.gather(v[self.npoi :], safe_idxs) * mask + v_noi = tf.gather(v[self.poi_model.npoi :], safe_idxs) * mask v_gathered = tf.concat([v_poi, v_noi], axis=0) return v_gathered @@ -743,7 +730,7 @@ def impacts_parms(self, hess): v = self._gather_poi_noi_vector(self.cov) impacts = v / tf.reshape(tf.sqrt(tf.linalg.diag_part(self.cov)), [1, -1]) - nstat = self.npoi + self.indata.nsystnoconstraint + nstat = self.poi_model.npoi + self.indata.nsystnoconstraint hess_stat = hess[:nstat, :nstat] inv_hess_stat = tf.linalg.inv(hess_stat) @@ -772,7 +759,9 @@ def impacts_parms(self, hess): if len(self.indata.systgroupidxs): impacts_grouped_syst = tf.map_fn( - lambda idxs: self._compute_impact_group(v[:, self.npoi :], idxs), + lambda idxs: self._compute_impact_group( + v[:, self.poi_model.npoi :], idxs + ), tf.ragged.constant(self.indata.systgroupidxs, dtype=tf.int32), fn_output_signature=tf.TensorSpec( shape=(impacts.shape[0],), dtype=tf.float64 @@ -893,8 +882,10 @@ def _compute_global_impacts_x0(self, cov_dexpdx): def global_impacts_parms(self): # TODO migrate this to a mapping to avoid the below code which is largely duplicated - idxs_poi = tf.range(self.npoi, dtype=tf.int64) - idxs_noi = tf.constant(self.npoi + self.indata.noiidxs, dtype=tf.int64) + idxs_poi = tf.range(self.poi_model.npoi, dtype=tf.int64) + idxs_noi = tf.constant( + self.poi_model.npoi + self.indata.noiidxs, dtype=tf.int64 + ) idxsout = tf.concat([idxs_poi, idxs_noi], axis=0) dexpdx = tf.one_hot(idxsout, depth=self.cov.shape[0], dtype=self.cov.dtype) @@ -919,7 +910,7 @@ def global_impacts_parms(self): impacts_beta0_total = tf.sqrt(var_beta0) impacts_x0 = self._compute_global_impacts_x0(cov_dexpdx) - impacts_theta0 = impacts_x0[self.npoi :] + impacts_theta0 = impacts_x0[self.poi_model.npoi :] impacts_theta0 = tf.transpose(impacts_theta0) impacts = impacts_theta0 @@ -1144,7 +1135,7 @@ def compute_derivatives(dvars): ) impacts_x0 = self._compute_global_impacts_x0(cov_dexpdx) - impacts_theta0 = impacts_x0[self.npoi :] + impacts_theta0 = impacts_x0[self.poi_model.npoi :] impacts_theta0 = tf.transpose(impacts_theta0) impacts = impacts_theta0 @@ -1236,23 +1227,18 @@ def _expected_variations( return expvars - def _compute_yields_noBBB(self, compute_norm=False, full=True): - # compute_norm: compute yields for each process, otherwise inclusive + def _compute_yields_noBBB(self, full=True): # full: compute yields inclduing masked channels poi = self.get_blinded_poi() theta = self.get_blinded_theta() theta = tf.where( - self.frozen_params_mask[self.npoi :], tf.stop_gradient(theta), theta - ) - - rnorm = tf.concat( - [poi, tf.ones([self.indata.nproc - poi.shape[0]], dtype=self.indata.dtype)], - axis=0, + self.frozen_params_mask[self.poi_model.npoi :], + tf.stop_gradient(theta), + theta, ) - mrnorm = tf.expand_dims(rnorm, -1) - ernorm = tf.reshape(rnorm, [1, -1]) + rnorm = self.poi_model.compute(poi) normcentral = None if self.indata.symmetric_tensor: @@ -1281,7 +1267,7 @@ def _compute_yields_noBBB(self, compute_norm=False, full=True): snorm * self.indata.norm.values ) elif self.indata.systematic_type == "normal": - snormnorm_sparse = self.indata.norm * ernorm + snormnorm_sparse = self.indata.norm * rnorm snormnorm_sparse = snormnorm_sparse.with_values( snormnorm_sparse.values + logsnorm ) @@ -1292,15 +1278,12 @@ def _compute_yields_noBBB(self, compute_norm=False, full=True): ) if self.indata.systematic_type == "log_normal": - nexpcentral = tf.sparse.sparse_dense_matmul(snormnorm_sparse, mrnorm) - nexpcentral = tf.squeeze(nexpcentral, -1) - if compute_norm: - snormnorm = tf.sparse.to_dense(snormnorm_sparse) - normcentral = ernorm * snormnorm + snormnorm = tf.sparse.to_dense(snormnorm_sparse) + normcentral = rnorm * snormnorm elif self.indata.systematic_type == "normal": - if compute_norm: - normcentral = tf.sparse.to_dense(snormnorm_sparse) - nexpcentral = tf.sparse.reduce_sum(snormnorm_sparse, axis=-1) + normcentral = tf.sparse.to_dense(snormnorm_sparse) + + nexpcentral = tf.reduce_sum(normcentral, axis=-1) else: if full or self.indata.nbinsmasked == 0: nbins = self.indata.nbinsfull @@ -1328,20 +1311,16 @@ def _compute_yields_noBBB(self, compute_norm=False, full=True): if self.indata.systematic_type == "log_normal": snorm = tf.exp(logsnorm) snormnorm = snorm * norm - nexpcentral = tf.matmul(snormnorm, mrnorm) - nexpcentral = tf.squeeze(nexpcentral, -1) - if compute_norm: - normcentral = ernorm * snormnorm + normcentral = rnorm * snormnorm elif self.indata.systematic_type == "normal": - normcentral = norm * ernorm + logsnorm - nexpcentral = tf.reduce_sum(normcentral, axis=-1) + normcentral = norm * rnorm + logsnorm + + nexpcentral = tf.reduce_sum(normcentral, axis=-1) return nexpcentral, normcentral def _compute_yields_with_beta(self, profile=True, compute_norm=False, full=True): - nexp, norm = self._compute_yields_noBBB( - compute_norm or self.binByBinStatMode == "full", full=full - ) + nexp, norm = self._compute_yields_noBBB(full=full) if self.binByBinStat: if profile: diff --git a/rabbit/io_tools.py b/rabbit/io_tools.py index 57c7743..0c9c1c2 100644 --- a/rabbit/io_tools.py +++ b/rabbit/io_tools.py @@ -26,7 +26,9 @@ def get_fitresult(fitresult_filename, result=None, meta=False): def get_poi_names(meta): - return np.concatenate((meta["signals"], meta["nois"])).astype(str) + return np.concatenate((meta.get("pois", meta.get("signals")), meta["nois"])).astype( + str + ) def get_syst_labels(fitresult): diff --git a/rabbit/mappings/helpers.py b/rabbit/mappings/helpers.py index bd35e89..7bff423 100644 --- a/rabbit/mappings/helpers.py +++ b/rabbit/mappings/helpers.py @@ -1,4 +1,3 @@ -import importlib import re import hist @@ -6,7 +5,7 @@ import tensorflow as tf from wums import boostHistHelpers as hh -from rabbit import tfhelpers +from rabbit import common, tfhelpers # dictionary with class name and the corresponding filename where it is defined baseline_mappings = { @@ -22,29 +21,10 @@ } -def instance_from_class(class_name, *args, **kwargs): - if "." in class_name: - # import from full relative or abslute path - parts = class_name.split(".") - module_name = ".".join(parts[:-1]) - class_name = parts[-1] - else: - # import one of the baseline mappings - if class_name not in baseline_mappings: - raise ValueError( - f"Mapping {class_name} not found, available baseline mappings are {baseline_mappings.keys()}" - ) - module_name = f"rabbit.mappings.{baseline_mappings[class_name]}" - - # Try to import the module - module = importlib.import_module(module_name) - - mapping = getattr(module, class_name, None) - if mapping is None: - raise AttributeError( - f"Class '{class_name}' not found in module '{module_name}'." - ) - +def load_mapping(class_name, *args, **kwargs): + mapping = common.load_class_from_module( + class_name, baseline_mappings, base_dir="rabbit.mappings" + ) return mapping.parse_args(*args, **kwargs) diff --git a/rabbit/poi_models/__init__.py b/rabbit/poi_models/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/rabbit/poi_models/helpers.py b/rabbit/poi_models/helpers.py new file mode 100644 index 0000000..83836f3 --- /dev/null +++ b/rabbit/poi_models/helpers.py @@ -0,0 +1,15 @@ +from rabbit import common + +# dictionary with class name and the corresponding filename where it is defined +baseline_models = { + "Ones": "poi_model", + "Mu": "poi_model", + "Mixture": "poi_model", +} + + +def load_model(class_name, indata, *args, **kwargs): + model = common.load_class_from_module( + class_name, baseline_models, base_dir="rabbit.poi_models" + ) + return model.parse_args(indata, *args, **kwargs) diff --git a/rabbit/poi_models/poi_model.py b/rabbit/poi_models/poi_model.py new file mode 100644 index 0000000..2547c56 --- /dev/null +++ b/rabbit/poi_models/poi_model.py @@ -0,0 +1,196 @@ +import numpy as np +import tensorflow as tf + + +class POIModel: + + def __init__(self, indata, *args, **kwargs): + self.indata = indata + + # # a POI model must set these attribues + # self.npoi = # number of parameters of interest (POIs) + # self.pois = # list of names for the POIs + # self.xpoidefault = # default values for the POIs + # self.is_linear = # define if the model is linear in the POIs + # self.allowNegativePOI = # define if the POI can be negative or not + + # class function to parse strings as given by the argparse input e.g. --poiModel ... + @classmethod + def parse_args(cls, indata, *args, **kwargs): + return cls(indata, *args, **kwargs) + + def compute(self, poi): + """ + Compute an array for the rate per process + :param params: 1D tensor of explicit parameters in the fit + :return 2D tensor to be multiplied with [proc,bin] tensor + """ + + def set_poi_default(self, expectSignal, allowNegativePOI=False): + """ + Set default POI values, used by different POI models + """ + poidefault = tf.ones([self.npoi], dtype=self.indata.dtype) + if expectSignal is not None: + indices = [] + updates = [] + for signal, value in expectSignal: + if signal.encode() not in self.pois: + raise ValueError( + f"{signal.encode()} not in list of POIs: {self.pois}" + ) + idx = np.where(np.isin(self.pois, signal.encode()))[0][0] + + indices.append([idx]) + updates.append(float(value)) + + poidefault = tf.tensor_scatter_nd_update(poidefault, indices, updates) + + if allowNegativePOI: + self.xpoidefault = poidefault + else: + self.xpoidefault = tf.sqrt(poidefault) + + +class Ones(POIModel): + """ + multiply all processes with ones + """ + + def __init__(self, indata, **kwargs): + self.indata = indata + self.npoi = 0 + self.pois = np.array([]) + self.poidefault = tf.zeros([], dtype=self.indata.dtype) + + self.allowNegativePOI = False + + def compute(self, poi): + rnorm = tf.ones(self.indata.nproc, dtype=self.indata.dtype) + rnorm = tf.reshape(rnorm, [1, -1]) + return rnorm + + +class Mu(POIModel): + """ + multiply unconstrained parameter to signal processes, and ones otherwise + """ + + def __init__(self, indata, expectSignal=1, allowNegativePOI=False, **kwargs): + self.indata = indata + + self.npoi = self.indata.nsignals + + self.pois = np.array([s for s in self.indata.signals]) + + self.allowNegativePOI = allowNegativePOI + + self.is_linear = self.npoi == 0 or self.allowNegativePOI + + self.set_poi_default(expectSignal, allowNegativePOI) + + def compute(self, poi): + rnorm = tf.concat( + [poi, tf.ones([self.indata.nproc - poi.shape[0]], dtype=self.indata.dtype)], + axis=0, + ) + + rnorm = tf.reshape(rnorm, [1, -1]) + return rnorm + + +class Mixture(POIModel): + """ + Based on unconstrained parameters x_i + multiply `primary` process by x_i + multiply `complementary` process by 1-x_i + """ + + def __init__( + self, + indata, + primary_processes, + complementary_processes, + expectSignal=0, + allowNegativePOI=False, + **kwargs, + ): + self.indata = indata + + if type(primary_processes) == str: + primary_processes = [primary_processes] + + if type(complementary_processes) == str: + complementary_processes = [complementary_processes] + + primary_processes = np.array(primary_processes).astype("S") + complementary_processes = np.array(complementary_processes).astype("S") + + if len(primary_processes) != len(complementary_processes): + raise ValueError( + f"Length of pimary and complementary processes has to be the same, but got {len(primary_processes)} and {len(complementary_processes)}" + ) + + if any(n not in self.indata.procs for n in primary_processes): + not_found = [n for n in primary_processes if n not in self.indata.procs] + raise ValueError(f"{not_found} not found in processes {self.indata.procs}") + + if any(n not in self.indata.procs for n in complementary_processes): + not_found = [ + n for n in complementary_processes if n not in self.indata.procs + ] + raise ValueError(f"{not_found} not found in processes {self.indata.procs}") + + self.primary_idxs = np.where(np.isin(self.indata.procs, primary_processes))[0] + self.complementary_idxs = np.where( + np.isin(self.indata.procs, complementary_processes) + )[0] + self.all_idx = np.concatenate([self.primary_idxs, self.complementary_idxs]) + + self.npoi = len(primary_processes) + self.pois = np.array( + [ + f"{p}_{c}_mixing".encode() + for p, c in zip( + primary_processes.astype(str), complementary_processes.astype(str) + ) + ] + ) + + self.allowNegativePOI = allowNegativePOI + self.is_linear = False + + self.set_poi_default(expectSignal, allowNegativePOI) + + @classmethod + def parse_args(cls, indata, *args, **kwargs): + """ + parsing the input arguments into the constructor, is has to be called as + --poiModel Mixture ,,... ,,... + to introduce a mixing parameter for proc_0 with proc_a, and proc_1 with proc_b, etc. + """ + + if len(args) != 2: + raise ValueError( + f"Expected exactly 2 arguments for Mixture model but got {len(args)}" + ) + + primaries = args[0].split(",") + complementaries = args[1].split(",") + + return cls(indata, primaries, complementaries, **kwargs) + + def compute(self, poi): + + ones = tf.ones(self.npoi, dtype=self.indata.dtype) + updates = tf.concat([ones * poi, ones * (1 - poi)], axis=0) + + # Single scatter update + rnorm = tf.tensor_scatter_nd_update( + tf.ones(self.indata.nproc, dtype=self.indata.dtype), + self.all_idx[:, None], + updates, + ) + + rnorm = tf.reshape(rnorm, [1, -1]) + return rnorm diff --git a/rabbit/workspace.py b/rabbit/workspace.py index 7827f1d..e384d98 100644 --- a/rabbit/workspace.py +++ b/rabbit/workspace.py @@ -69,7 +69,7 @@ def __init__(self, outdir, outname, fitter, postfix=None): ) self.parms = fitter.parms - self.npoi = fitter.npoi + self.npoi = fitter.poi_model.npoi self.noiidxs = fitter.indata.noiidxs self.extension = "hdf5"