smoofit.model.Model

class smoofit.model.Model(do_bblite=False, seed=0)[source]

Main class defining the statistical model (likelihood function), fitting etc.

Methods

__init__([do_bblite, seed])

Constructor

add_constraint(fn)

Register additional constraint function

add_proc(proc)

Register a Process with the model

batch_nll(values, nObs, nGlob)

batch_pred(values)

Compute the summed yields of all processes using a batch of parameter values, including the effect of the BB-lite parameters (if activated).

channel_sub_idxs(channel_name)

Return the indices of a channel in the yield arrays

compile([compile_pred, compile_hess])

JIT-compile the NLL function and the analytical NLL gradient

concat_obs(nObs)

Concatenate observed yields across all channels in the order expected by the prepared model

correlate_nuisances(v1, v2, corr)

Specify the correlation between a pair of nuisance parameters

cov(values, nObs[, nGlob])

Compute the fit covariance matrix analytically

default_glob()

Return default vector of global observables

enable_bblite()

Enable the computation of bin-by-bin Monte-Carlo statistical uncertainties using the Barlow-Beeston lite method

fit(nObs[, nGlob, init, store_hess, ...])

Fit the model to the data

get_bblite_errs(values, sumw2)

Compute the contribution to the predicted yields due to the fluctuated bin-by-bin statistics

minos_bounds(var, best_fit, nObs[, nGlob, ...])

Find the confidence interval on a parameter using the Minos method

minos_direction(direction, best_fit, nObs[, ...])

Find the level crossing of the profile NLL in a specific direction

nll(values, nObs[, nGlob])

Evaluate the negative log-likelihood (NLL) of the model

nll_crossings(direction, best_fit, nObs[, ...])

Find the level crossing of the NLL in a specific direction

nll_fast(values, nObs, nGlob)

Evaluate the negative log-likelihood (NLL) of the model

nll_grad(values, nObs, nGlob)

nll_hess(values, nObs, nGlob)

pred(values)

Compute the summed yields of all processes given the parameter values, including the effect of the BB-lite parameters (if activated).

prepare()

Prepare the model for inference

sample_from_covariance(best_fit, nSamples[, ...])

Sample from the fit covariance matrix

split_obs(nObs)

Split prediction array into channels

toy_glob(values[, nToys])

Generate random toys for the global observables associated with the model's nuisance parameters

toy_pred(values[, nToys])

Generate random toys for the predicted yields

values_from_dict([values])

Build a global parameter vector

var_index(var)

Return the index of a Variable within the compiled model

var_sub_index(var)

Return the sub-indices of a Variable within the global parameter array of the compiled model
__init__(do_bblite=False, seed=0)[source]

Constructor

Parameters

  • do_bblite (bool) – Enable the computation of bin-by-bin Monte-Carlo statistical uncertainties using the Barlow-Beeston lite method. For more details see Model.enable_bblite()

  • seed (int) – Initial seed for random toy generation. See here for more information about random numbers in jax

add_constraint(fn)[source]

Register additional constraint function

The constraint function should return a scalar, and takes as argument three jax arrays (obviously, not all of them have to be used):

  • the parameters of the model (concatenated as a single 1D array)

  • the observed yields across all channels (concatenated as a single 1D array)

  • the vector of global observables

Sign convention: the constraint is added to the log-likelihood.

Any number of constraints can be added, but it is more effective if it is possible to express a single function in array form.

Parameters

fn (Callable[[DeviceArrayBase, DeviceArrayBase, DeviceArrayBase], DeviceArrayBase]) – the constraint function

add_proc(proc)[source]

Register a Process with the model

Further modifying the process (e.g. adding channel contributions etc.) is still possible until Model.prepare() is called.

batch_nll(values, nObs, nGlob)[source]
batch_pred(values)[source]

Compute the summed yields of all processes using a batch of parameter values, including the effect of the BB-lite parameters (if activated).

Note

This method is only available after the model has been prepared!

Parameters

values (DeviceArrayBase) – 2D array where the first axis is the batch dimension, i.e. every row specifies parameter values in the order expected by the prepared model

Return type

DeviceArrayBase

Returns

a 2D jnp.DeviceArray where the entries (i,:) contain the predicted yields, across all channels and bins, given the parameter values in row i of values

channel_sub_idxs(channel_name)[source]

Return the indices of a channel in the yield arrays

The yield arrays of different channels of a process are concatenated. This function returns the indices of a given channel along that concatenated axis.

Example: two channels are attached to the model, one (called SR) with 2 bins, the other (called CR) with 3 bins. Then model.channel_sub_idxs("SR") might return [0,1] or [3,4], depending on the order in which the channels have been concatenated when the model is prepared.

Note

This method is only available after the model has been prepared!

Parameters

channel_name (str) – the name of the channel

Return type

DeviceArrayBase

Returns

the array of indices of that channel

compile(compile_pred=False, compile_hess=False)[source]

JIT-compile the NLL function and the analytical NLL gradient

Note

This method is only available after the model has been prepared!

Calling this function is not strictly necessary - jitting and gradient tracing will happen anyway the first time the NLL, the gradient or the Hessian are called. This function merely allows to “get it done” all at once.

Parameters

  • compile_hess (bool) – Also compile Model.pred(), which can be useful if you expect to call it very often. However using Model.batch_pred() instead is probably better.

  • compile_hess – Also compile the Hessian matrix. This is typically quite slow, and takes at least as much time as computing a non-JITed Hessian once. Since the Hessian is typically only computed once (on the best-fit point), JITing is not done by default and only useful when you expect to call the Hessian very often (for instance, for studying coverage properties of intervals using toys).

concat_obs(nObs)[source]

Concatenate observed yields across all channels in the order expected by the prepared model

Note

This method is only available after the model has been prepared!

Parameters

nObs (Union[ndarray, DeviceArrayBase, Dict[str, ndarray], Dict[str, DeviceArrayBase]]) – either a 1D array with the observed yields in all bins and channels in the order expected by the prepared model, or a dictionary where the keys are the channel names and the values are 1D arrays with the observed yields in the corresponding channel

Return type

DeviceArrayBase

Returns

a 1D jnp.DeviceArray with the observed yields in all bins and channels in the order expected by the prepared model

correlate_nuisances(v1, v2, corr)[source]

Specify the correlation between a pair of nuisance parameters

Parameters

  • v1 (Union[Variable, Tuple[Variable, int]]) – either a Variable if the variable is a scalar, or a tuple (var, index) specifying the index of the component of var to be correlated if var is vector-valued

  • v2 (Union[Variable, Tuple[Variable, int]]) – as v1

  • corr (float) – the correlation coefficient, should lie in the interval ]-1, 1[

cov(values, nObs, nGlob=None)[source]

Compute the fit covariance matrix analytically

The order of the entries in the matrix follows the order of the variables in the model.

Note

This method is only available after the model has been prepared!

Parameters

  • values (DeviceArrayBase) – 1D array with the parameter values in the order expected by the prepared model, e.g. as returned by Model.values_from_dict()

  • nObs (Union[ndarray, DeviceArrayBase, Dict[str, ndarray], Dict[str, DeviceArrayBase]]) – either a 1D array with the observed yields in all bins and channels in the order expected by the prepared model, or a dictionary where the keys are the channel names and the values are 1D arrays with the observed yields in the corresponding channel

  • nGlob (Union[ndarray, DeviceArrayBase, None]) – a 1D array with the vector of global observables associated with the nuisance parameters. If not specified, it is generated from Model.default_glob().

Return type

DeviceArrayBase

Returns

the analytical covariance matrix (inverse Hessian) evaluated at values (as 2D jnp.DeviceArray)

default_glob()[source]

Return default vector of global observables

Note

This method is only available after the model has been prepared!

Return type

DeviceArrayBase

Returns

a 1D jnp.DeviceArray with the default vector of global observables corresponding to the nuisance parameters in the model (= all zeros)

enable_bblite()[source]

Enable the computation of bin-by-bin Monte-Carlo statistical uncertainties using the Barlow-Beeston lite method

When preparing the model, n Gaussian nuisance parameters will be added to the model, where n is the number of bins across all channels. The corresponding Variable is called bblite and has n components.

The (Gaussian) uncertainty in each bin is equal to sqrt(sumW2), where sumW2 is the sum of the sums of squared weight entries for that bin across all processes. If a process has no registered sumw2 array, it won’t contribute to the uncertainty (i.e. it has infinite statistics).

fit(nObs, nGlob=None, init=None, store_hess=True, freeze_params=None, method='trust-constr', minimizer_opts=None)[source]

Fit the model to the data

Minimize the NLL given the observed yields (nObs).

The default minimizer used is the trust-constr method in scipy, see the documentation for the overall minimizer options here and for the trust-constr method here.

Minimization starts from the default values of the parameters, and takes the bounds on the parameters into account.

This method can also be used e.g. for profiled NLL scans (in which case the POIs are fixed to the scanned values), or for nuisance parameter impacts (to obtain a new best-fit point when fixing a nuisance parameter to a specific value).

Specifying the profiled parameters is done using the freeze_params argument, which is a dictionary where the keys are Variable objects and the values can be either:

  • a number or a 1D array (with the same number of components as the variable in the key) holding the value to which the parameter should be fixed in the fit

  • a tuple (val, idxs), which allows to fix only specific components of vector-valued parameters (schematically, fix var[idxs]=val in the fit)

Note

This method is only available after the model has been prepared!

Parameters

  • nObs (Union[ndarray, DeviceArrayBase, Dict[str, ndarray], Dict[str, DeviceArrayBase]]) – either a 1D array with the observed yields in all bins and channels in the order expected by the prepared model, or a dictionary where the keys are the channel names and the values are 1D arrays with the observed yields in the corresponding channel

  • nGlob (Union[ndarray, DeviceArrayBase, None]) – a 1D array with the vector of global observables associated with the nuisance parameters. If not specified, it is generated from Model.default_glob().

  • init (Union[ndarray, DeviceArrayBase, None]) – specify initial point from which to minimize (if not specified, use default values of variables)

  • store_hess (bool) – compute the analytical covariance matrix (inverse Hessian) at the minimum and store it as the hess_inv attribute of the returned object.

  • freeze_params (Optional[Dict[Variable, Union[float, ndarray, DeviceArrayBase, Tuple[Union[ndarray, DeviceArrayBase], Union[List[int], ndarray, DeviceArrayBase]]]]]) – specification of the parameters to keep fixed in a profiled fit (see above for details)

  • method (str) – name of the scipy minimizer, see full list of accepted methods here

  • minimizer_opts (Optional[Dict[str, Any]]) – a dictionary of options passed to the minimizer, see the scipy documentation. In particular, the maxiter field sets the maximum number of iterations.

Return type

OptimizeResult

Returns

a scipy.optimize.OptimizeResult object (see here) containing in particular the fitted parameter values (x), the NLL value at the minimum (fun), and a success flag (success). See the trust-constr documentation linked above for a full list of attributes.

get_bblite_errs(values, sumw2)[source]

Compute the contribution to the predicted yields due to the fluctuated bin-by-bin statistics

Note

This method is only available after the model has been prepared!

Parameters

  • values (DeviceArrayBase) – a 1D jnp.DeviceArray with the parameter values

  • sumw2 (DeviceArrayBase) – a 1D jnp.DeviceArray with the sums of the sumw2 arrays in each bin

Return type

DeviceArrayBase

Returns

a 1D jnp.DeviceArray with the contribution of the BB-lite parameters to the yields in each bin

minos_bounds(var, best_fit, nObs, nGlob=None, sub_idx=0, level=0.5, init_up=None, init_down=None, freeze_params=None, minimizer_opts=None)[source]

Find the confidence interval on a parameter using the Minos method

Minos intervals are obtained by finding the points where the profiled likelihood increases by a given amount w.r.t. the minimum (best fit). Any parameter (or any component of a vector-valued parameter) can be profiled.

The returned up, down values are such that up >= x0[i] >= down, where x0[i] is the best-fit value of the chosen parameter.

This function is a convenience wrapper for calling Model.minos_direction(), see more details about some of the arguments there.

Bounds on the model’s parameters are taken into account.

Note

This method is only available after the model has been prepared!

Parameters

  • Variable – the Variable on which to find the bounds; for vector-valued parameters, specify the component using sub_idxs

  • best_fit (OptimizeResult) – the fit result as returned from Model.fit()

  • nObs (Union[ndarray, DeviceArrayBase, Dict[str, ndarray], Dict[str, DeviceArrayBase]]) – either a 1D array with the observed yields in all bins and channels in the order expected by the prepared model, or a dictionary where the keys are the channel names and the values are 1D arrays with the observed yields in the corresponding channel

  • nGlob (Union[ndarray, DeviceArrayBase, None]) – a 1D array with the vector of global observables associated with the nuisance parameters. If not specified, it is generated from Model.default_glob().

  • sub_idx (int) – the index of the component of the variable to consider if the variable is vector-valued

  • level (float) – increase in the profile NLL from the minimum

  • init_up (Union[ndarray, DeviceArrayBase, None]) – specify initial point from which to minimize for the up uncertainty

  • init_down (Union[ndarray, DeviceArrayBase, None]) – specify initial point from which to minimize for the down uncertainty

  • freeze_params (Optional[Dict[Variable, Union[float, ndarray, DeviceArrayBase, Tuple[Union[ndarray, DeviceArrayBase], Union[List[int], ndarray, DeviceArrayBase]]]]]) – specification of parameters to keep fixed in the fit (see docs of Model.fit() for details)

  • minimizer_opts (Optional[Dict[str, Any]]) – a dictionary of options passed to the mimimization algorithm

Return type

Tuple[DeviceArrayBase, DeviceArrayBase, OptimizeResult, OptimizeResult]

Returns

a tuple (up_bound, down_bound, up_result, down_result) where the first two entries are the upper and lower ends of the interval on the parameter and the latter two are the full scipy.optimize.OptimizeResult objects from the calls to Model.minos_direction() to find the upper and lower bound, respectively

minos_direction(direction, best_fit, nObs, nGlob=None, level=0.5, init=None, freeze_params=None, minimizer_opts=None)[source]

Find the level crossing of the profile NLL in a specific direction

Profiling the component i of x means finding the point prof(x_i,i) that minimizes NLL(x) while fixing x[i]=x_i. This functions finds the point x=prof(x_i,i) that satisfies NLL(prof(x_i,i)) = NLL0 + level, but is more general in that the direction in which the profiling is done is arbitrary.

This function can therefore be used to obtain the profiled, so-called Minos uncertainties on a parameter or on a linear combination of parameters.

In practice, this minimizes the function f(x) = -dot(direction, x) under the constraint that NLL(x) <= NLL0 + level, where NLL0 is the value of the NLL at the best fit. The minimization algorithm is scipy’s trust-constr, see Model.fit() for more details.

The initial point at which the minimization starts can be changed in case of convergence issues (by default, minimization starts from the best fit).

Bounds on the model’s parameters are taken into account at every step.

Note

This method is only available after the model has been prepared!

Parameters

  • direction (Union[ndarray, DeviceArrayBase]) – direction (in the form of a 1D array analogous to the model parameter array) in which to search the profile NLL crossing

  • best_fit (OptimizeResult) – the fit result as returned from Model.fit()

  • nObs (Union[ndarray, DeviceArrayBase, Dict[str, ndarray], Dict[str, DeviceArrayBase]]) – either a 1D array with the observed yields in all bins and channels in the order expected by the prepared model, or a dictionary where the keys are the channel names and the values are 1D arrays with the observed yields in the corresponding channel

  • nGlob (Union[ndarray, DeviceArrayBase, None]) – a 1D array with the vector of global observables associated with the nuisance parameters. If not specified, it is generated from Model.default_glob().

  • level (float) – increase in the profile NLL from the minimum

  • init (Union[ndarray, DeviceArrayBase, None]) – specify initial point from which to minimize

  • freeze_params (Optional[Dict[Variable, Union[float, ndarray, DeviceArrayBase, Tuple[Union[ndarray, DeviceArrayBase], Union[List[int], ndarray, DeviceArrayBase]]]]]) – specification of parameters to keep fixed in the fit (see docs of Model.fit() for details)

  • minimizer_opts (Optional[Dict[str, Any]]) – a dictionary of options passed to the mimimization algorithm

Return type

OptimizeResult

Returns

a scipy.optimize.OptimizeResult

nll(values, nObs, nGlob=None)[source]

Evaluate the negative log-likelihood (NLL) of the model

This is a convenience wrapper of the jitted method Model.nll_fast()

Note

This method is only available after the model has been prepared!

Parameters

  • values (Union[ndarray, DeviceArrayBase]) – 1D array with the parameter values in the order expected by the prepared model, e.g. as returned by Model.values_from_dict()

  • nObs (Union[ndarray, DeviceArrayBase, Dict[str, ndarray], Dict[str, DeviceArrayBase]]) – either a 1D array with the observed yields in all bins and channels in the order expected by the prepared model, or a dictionary where the keys are the channel names and the values are 1D arrays with the observed yields in the corresponding channel

  • nGlob (Union[ndarray, DeviceArrayBase, None]) – a 1D array with the vector of global observables associated with the nuisance parameters. If not specified, it is generated from Model.default_glob().

Return type

DeviceArrayBase

Returns

a scalar jnp.DeviceArray with the NLL

nll_crossings(direction, best_fit, nObs, nGlob=None, level=0.5, init=0.1, bracket=(0, 1), root_opts=None)[source]

Find the level crossing of the NLL in a specific direction

This function finds the solution s to NLL(x(s)) = NLL0 + level, where x(s) = x0 + s * direction, x0 is the best-fit point, and NLL0 = NLL(x0) is the value of the NLL at the minimum. The parameter s is constrained to the range [bracket[0], bracket[1]], and the algorithm starts with a=init.

The bracket and direction arguments should be chosen to make sure that NLL[x(bracket[1])] >= level.

Bounds on the model’s parameters are taken into account.

This function can be used to find statistical-only uncertainties from likelihood scans, e.g. to find the 1-sigma uncertainty on the i th parameter, set direction[j] to 1 if j=i and 0 otherwise, and level=0.5.

Note

This method is only available after the model has been prepared!

Parameters

  • direction (Union[ndarray, DeviceArrayBase]) – direction (in the form of a 1D array analogous to the model parameter array) in which to search the NLL crossing

  • best_fit (OptimizeResult) – the fit result as returned from Model.fit()

  • nObs (Union[ndarray, DeviceArrayBase, Dict[str, ndarray], Dict[str, DeviceArrayBase]]) – either a 1D array with the observed yields in all bins and channels in the order expected by the prepared model, or a dictionary where the keys are the channel names and the values are 1D arrays with the observed yields in the corresponding channel

  • nGlob (Union[ndarray, DeviceArrayBase, None]) – a 1D array with the vector of global observables associated with the nuisance parameters. If not specified, it is generated from Model.default_glob().

  • level (float) – increase in the NLL from the minimum

  • init (float) – initial position of the algorithm along the specified direction

  • bracket (Tuple[float, float]) – bounds on the range of values in which the crossing is searched

  • root_opts (Optional[Dict[str, Any]]) – a dictionary of options passed to the root-finding algorithm

Return type

RootResults

Returns

a scipy.optimize.RootResults whose attribute x is set to x = x0 + s * direction, with s the solution of the equation above

nll_fast(values, nObs, nGlob)[source]

Evaluate the negative log-likelihood (NLL) of the model

This method expects all the inputs to be properly ordered jnp.DeviceArray arrays. Use Model.nll() for a more convenient interface.

Note

This method is only available after the model has been prepared!

Parameters

  • values (DeviceArrayBase) – 1D jnp.DeviceArray with the parameter values in the order expected by the prepared model

  • nObs (DeviceArrayBase) – a 1D jnp.DeviceArray with the observed yields in all bins and channels in the order expected by the prepared model

  • nGlob (DeviceArrayBase) – a 1D jnp.DeviceArray array with the vector of global observables associated with the nuisance parameters

Return type

DeviceArrayBase

Returns

a scalar jnp.DeviceArray with the NLL

nll_grad(values, nObs, nGlob)[source]
nll_hess(values, nObs, nGlob)[source]
pred(values)[source]

Compute the summed yields of all processes given the parameter values, including the effect of the BB-lite parameters (if activated).

Note

This method is only available after the model has been prepared!

Parameters

values (DeviceArrayBase) – 1D array with the parameter values in the order expected by the prepared model, e.g. as returned by Model.values_from_dict()

Return type

DeviceArrayBase

Returns

a 1D jnp.DeviceArray containing the predicted yields across all channels and bins

prepare()[source]

Prepare the model for inference

This collects all the parameters, processes and channels that have been registered. A global parameter array is built, to be passed to the likelihood function as a single argument. Yield arrays for every process and channel are concatenated into a single large array, within a MergedProcess object accessible though the attribute Model.merged_process. Systematic uncertainties are concatenated across sources, processes and channels. The nuisance parameter covariance matrix is also built, accessible through the attribute Model.nuisance_cov_matrix.

sample_from_covariance(best_fit, nSamples, respect_bounds=True)[source]

Sample from the fit covariance matrix

Generate a batch of parameter values distributed as a multivariate normal whose mean is the best fit value and whose covariance matrix is the fit covariance matrix as estimated from the inverse Hessian of the NLL at the minimum.

Note

This method is only available after the model has been prepared!

Parameters

  • best_fit (OptimizeResult) – the fit result as returned from Model.fit(), it must contain the covariance matrix

  • nSamples (int) – number of samples to draw

  • respect_bounds (bool) – enforce the bounds on the parameters (i.e. truncate the generated values to respect the bounds)

Return type

DeviceArrayBase

Returns

a 2D jnp.DeviceArray where each row is a sampled parameter array (with the order of variables in the model)

split_obs(nObs)[source]

Split prediction array into channels

Split a full yields array, returned e.g. by model.pred(), into the different channels.

Note

This method is only available after the model has been prepared!

Parameters

nObs (Union[ndarray, DeviceArrayBase]) – array with the observed yields in all bins and channels (1D or 2D, in the case of batch predictions)

Return type

Dict[str, DeviceArrayBase]

Returns

a dictionary with the channel names as keys, and the yields in the channels (as 1D or 2D jnp.DeviceArray) as values

toy_glob(values, nToys=1)[source]

Generate random toys for the global observables associated with the model’s nuisance parameters

Note

This method is only available after the model has been prepared!

The toys are drawn from a multivariate normal distribution with means given by the values of the nuisance parameters in values and with covariance matrix fixed from the model’s Model.nuisance_cov_matrix attribute.

Parameters

  • values (DeviceArrayBase) – a 1D jnp.DeviceArray with the parameter values used to generate the toy

  • nToys (int) – the number of toys to generate

Return type

DeviceArrayBase

Returns

a 2D jnp.DeviceArray where each row contains a generated vector of global observables

toy_pred(values, nToys=1)[source]

Generate random toys for the predicted yields

Note

This method is only available after the model has been prepared!

The toys are drawn from Poisson distributions with means given by Model.pred() evaluated on values.

Parameters

  • values (DeviceArrayBase) – a 1D jnp.DeviceArray with the global parameter values used to generate the toy

  • nToys (int) – the number of toys to generate

Return type

DeviceArrayBase

Returns

a 2D jnp.DeviceArray where each row is a toy, i.e. the yields in all bins across all channels

values_from_dict(values=None)[source]

Build a global parameter vector

Build a 1D array with the values of all of the model’s parameters in the expected order.

By default, the initial values of the Variable objects are used, but these can be overridden using the values argument. This will not modify the “set” values of the model’s Variable objects.

Note

This method is only available after the model has been prepared!

Parameters

values (Optional[Dict[Variable, Union[float, List[float], ndarray, DeviceArrayBase]]]) – a dictionary where the keys are Variable objects and the values are the values to use for the corresponding variable

Return type

DeviceArrayBase

Returns

a 1D jnp.DeviceArray with all the parameter values in the order expected by the model

var_index(var)[source]

Return the index of a Variable within the compiled model

Note

This method is only available after the model has been prepared!

Parameters

var (Variable) – the Variable

Return type

int

Returns

the index of the variable

var_sub_index(var)[source]

Return the sub-indices of a Variable within the global parameter array of the compiled model

Note

This method is only available after the model has been prepared!

Parameters

var (Variable) – the Variable

Return type

DeviceArrayBase

Returns

the array of indices of the variable’s components within the global parameter array of the compiled model