From 04b0c3babc214b1f0a87285d8bfcc0e85e952acc Mon Sep 17 00:00:00 2001 From: Lars Henkelmann Date: Mon, 16 Aug 2021 18:29:38 +0100 Subject: [PATCH] feat: implement return_fitted_pars kwarg on test statistics --- src/pyhf/infer/test_statistics.py | 111 ++++++++++++++++++++++++++---- 1 file changed, 99 insertions(+), 12 deletions(-) diff --git a/src/pyhf/infer/test_statistics.py b/src/pyhf/infer/test_statistics.py index 6ac4417bf7..06944eaef5 100644 --- a/src/pyhf/infer/test_statistics.py +++ b/src/pyhf/infer/test_statistics.py @@ -13,7 +13,9 @@ def __dir__(): return __all__ -def _qmu_like(mu, data, pdf, init_pars, par_bounds, fixed_params): +def _qmu_like( + mu, data, pdf, init_pars, par_bounds, fixed_params, return_fitted_pars=False +): """ Clipped version of _tmu_like where the returned test statistic is 0 if muhat > 0 else tmu_like_stat. @@ -22,12 +24,14 @@ def _qmu_like(mu, data, pdf, init_pars, par_bounds, fixed_params): qmu_tilde. Otherwise this is qmu (no tilde). """ tensorlib, optimizer = get_backend() - tmu_like_stat, (_, muhatbhat) = _tmu_like( + tmu_like_stat, (mubhathat, muhatbhat) = _tmu_like( mu, data, pdf, init_pars, par_bounds, fixed_params, return_fitted_pars=True ) qmu_like_stat = tensorlib.where( muhatbhat[pdf.config.poi_index] > mu, tensorlib.astensor(0.0), tmu_like_stat ) + if return_fitted_pars: + return qmu_like_stat, (mubhathat, muhatbhat) return qmu_like_stat @@ -56,7 +60,7 @@ def _tmu_like( return tmu_like_stat -def qmu(mu, data, pdf, init_pars, par_bounds, fixed_params): +def qmu(mu, data, pdf, init_pars, par_bounds, fixed_params, return_fitted_pars=False): r""" The test statistic, :math:`q_{\mu}`, for establishing an upper limit on the strength parameter, :math:`\mu`, as defiend in @@ -94,6 +98,10 @@ def qmu(mu, data, pdf, init_pars, par_bounds, fixed_params): >>> pyhf.infer.test_statistics.qmu(test_mu, data, model, init_pars, par_bounds, fixed_params) array(3.9549891) + or, also return the best-fit parameter arrays: + >>> pyhf.infer.test_statistics.qmu(test_mu, data, model, init_pars, par_bounds, fixed_params, return_fitted_pars = True) + (array(3.9549891), (array([1. , 0.97224597, 0.87553894]), array([-0.06679525, 1.00555369, 0.96930896]))) + Args: mu (Number or Tensor): The signal strength parameter data (Tensor): The data to be considered @@ -104,9 +112,14 @@ def qmu(mu, data, pdf, init_pars, par_bounds, fixed_params): The shape should be ``(n, 2)`` for ``n`` model parameters. fixed_params (:obj:`list` of :obj:`bool`): The flag to set a parameter constant to its starting value during minimization. + return_fitted_pars (:obj:`bool`): Return the best-fit parameter tensors + the fixed-POI and unconstrained fits have converged on + (i.e. :math:`\mu, \hat{\hat{\theta}}` and :math:`\hat{\mu}, \hat{\theta}`) Returns: Float: The calculated test statistic, :math:`q_{\mu}` + Tuple of two Tensors: (only if ``return_fitted_pars`` is ``True``) the best-fit parameter tensors, + :math:`\mu, \hat{\hat{\theta}}` and :math:`\hat{\mu}, \hat{\theta}` """ if pdf.config.poi_index is None: raise UnspecifiedPOI( @@ -117,10 +130,20 @@ def qmu(mu, data, pdf, init_pars, par_bounds, fixed_params): 'qmu test statistic used for fit configuration with POI bounded at zero.\n' + 'Use the qmu_tilde test statistic (pyhf.infer.test_statistics.qmu_tilde) instead.' ) - return _qmu_like(mu, data, pdf, init_pars, par_bounds, fixed_params) + return _qmu_like( + mu, + data, + pdf, + init_pars, + par_bounds, + fixed_params, + return_fitted_pars=return_fitted_pars, + ) -def qmu_tilde(mu, data, pdf, init_pars, par_bounds, fixed_params): +def qmu_tilde( + mu, data, pdf, init_pars, par_bounds, fixed_params, return_fitted_pars=False +): r""" The "alternative" test statistic, :math:`\tilde{q}_{\mu}`, for establishing an upper limit on the strength parameter, :math:`\mu`, for models with @@ -163,6 +186,10 @@ def qmu_tilde(mu, data, pdf, init_pars, par_bounds, fixed_params): >>> pyhf.infer.test_statistics.qmu_tilde(test_mu, data, model, init_pars, par_bounds, fixed_params) array(3.93824492) + or, also return the best-fit parameter arrays: + >>> pyhf.infer.test_statistics.qmu_tilde(test_mu, data, model, init_pars, par_bounds, fixed_params, return_fitted_pars = True) + (array(3.93824492), (array([1. , 0.97224597, 0.87553894]), array([0. , 1.0030512 , 0.96266961]))) + Args: mu (Number or Tensor): The signal strength parameter data (:obj:`tensor`): The data to be considered @@ -173,9 +200,14 @@ def qmu_tilde(mu, data, pdf, init_pars, par_bounds, fixed_params): The shape should be ``(n, 2)`` for ``n`` model parameters. fixed_params (:obj:`list` of :obj:`bool`): The flag to set a parameter constant to its starting value during minimization. + return_fitted_pars (:obj:`bool`): Return the best-fit parameter tensors + the fixed-POI and unconstrained fits have converged on + (i.e. :math:`\mu, \hat{\hat{\theta}}` and :math:`\hat{\mu}, \hat{\theta}`) Returns: Float: The calculated test statistic, :math:`\tilde{q}_{\mu}` + Tuple of two Tensors: (only if ``return_fitted_pars`` is ``True``) the best-fit parameter tensors, + :math:`\mu, \hat{\hat{\theta}}` and :math:`\hat{\mu}, \hat{\theta}` """ if pdf.config.poi_index is None: raise UnspecifiedPOI( @@ -186,10 +218,18 @@ def qmu_tilde(mu, data, pdf, init_pars, par_bounds, fixed_params): 'qmu_tilde test statistic used for fit configuration with POI not bounded at zero.\n' + 'Use the qmu test statistic (pyhf.infer.test_statistics.qmu) instead.' ) - return _qmu_like(mu, data, pdf, init_pars, par_bounds, fixed_params) + return _qmu_like( + mu, + data, + pdf, + init_pars, + par_bounds, + fixed_params, + return_fitted_pars=return_fitted_pars, + ) -def tmu(mu, data, pdf, init_pars, par_bounds, fixed_params): +def tmu(mu, data, pdf, init_pars, par_bounds, fixed_params, return_fitted_pars=False): r""" The test statistic, :math:`t_{\mu}`, for establishing a two-sided interval on the strength parameter, :math:`\mu`, as defiend in Equation (8) @@ -221,6 +261,10 @@ def tmu(mu, data, pdf, init_pars, par_bounds, fixed_params): >>> pyhf.infer.test_statistics.tmu(test_mu, data, model, init_pars, par_bounds, fixed_params) array(3.9549891) + or, also return the best-fit parameter arrays: + >>> pyhf.infer.test_statistics.tmu(test_mu, data, model, init_pars, par_bounds, fixed_params, return_fitted_pars = True) + (array(3.9549891), (array([1. , 0.97224597, 0.87553894]), array([-0.06679525, 1.00555369, 0.96930896]))) + Args: mu (Number or Tensor): The signal strength parameter data (Tensor): The data to be considered @@ -231,9 +275,14 @@ def tmu(mu, data, pdf, init_pars, par_bounds, fixed_params): The shape should be ``(n, 2)`` for ``n`` model parameters. fixed_params (:obj:`list` of :obj:`bool`): The flag to set a parameter constant to its starting value during minimization. + return_fitted_pars (:obj:`bool`): Return the best-fit parameter tensors + the fixed-POI and unconstrained fits have converged on + (i.e. :math:`\mu, \hat{\hat{\theta}}` and :math:`\hat{\mu}, \hat{\theta}`) Returns: Float: The calculated test statistic, :math:`t_{\mu}` + Tuple of two Tensors: (only if ``return_fitted_pars`` is ``True``) the best-fit parameter tensors, + :math:`\mu, \hat{\hat{\theta}}` and :math:`\hat{\mu}, \hat{\theta}` """ if pdf.config.poi_index is None: raise UnspecifiedPOI( @@ -244,10 +293,20 @@ def tmu(mu, data, pdf, init_pars, par_bounds, fixed_params): 'tmu test statistic used for fit configuration with POI bounded at zero.\n' + 'Use the tmu_tilde test statistic (pyhf.infer.test_statistics.tmu_tilde) instead.' ) - return _tmu_like(mu, data, pdf, init_pars, par_bounds, fixed_params) + return _tmu_like( + mu, + data, + pdf, + init_pars, + par_bounds, + fixed_params, + return_fitted_pars=return_fitted_pars, + ) -def tmu_tilde(mu, data, pdf, init_pars, par_bounds, fixed_params): +def tmu_tilde( + mu, data, pdf, init_pars, par_bounds, fixed_params, return_fitted_pars=False +): r""" The test statistic, :math:`\tilde{t}_{\mu}`, for establishing a two-sided interval on the strength parameter, :math:`\mu`, for models with @@ -284,6 +343,10 @@ def tmu_tilde(mu, data, pdf, init_pars, par_bounds, fixed_params): >>> pyhf.infer.test_statistics.tmu_tilde(test_mu, data, model, init_pars, par_bounds, fixed_params) array(3.93824492) + or, also return the best-fit parameter arrays: + >>> pyhf.infer.test_statistics.tmu_tilde(test_mu, data, model, init_pars, par_bounds, fixed_params, return_fitted_pars = True) + (array(3.93824492), (array([1. , 0.97224597, 0.87553894]), array([0. , 1.0030512 , 0.96266961]))) + Args: mu (Number or Tensor): The signal strength parameter data (:obj:`tensor`): The data to be considered @@ -294,9 +357,14 @@ def tmu_tilde(mu, data, pdf, init_pars, par_bounds, fixed_params): The shape should be ``(n, 2)`` for ``n`` model parameters. fixed_params (:obj:`list` of :obj:`bool`): The flag to set a parameter constant to its starting value during minimization. + return_fitted_pars (:obj:`bool`): Return the best-fit parameter tensors + the fixed-POI and unconstrained fits have converged on + (i.e. :math:`\mu, \hat{\hat{\theta}}` and :math:`\hat{\mu}, \hat{\theta}`) Returns: Float: The calculated test statistic, :math:`\tilde{t}_{\mu}` + Tuple of two Tensors: (only if ``return_fitted_pars`` is ``True``) the best-fit parameter tensors, + :math:`\mu, \hat{\hat{\theta}}` and :math:`\hat{\mu}, \hat{\theta}` """ if pdf.config.poi_index is None: raise UnspecifiedPOI( @@ -307,10 +375,18 @@ def tmu_tilde(mu, data, pdf, init_pars, par_bounds, fixed_params): 'tmu_tilde test statistic used for fit configuration with POI not bounded at zero.\n' + 'Use the tmu test statistic (pyhf.infer.test_statistics.tmu) instead.' ) - return _tmu_like(mu, data, pdf, init_pars, par_bounds, fixed_params) + return _tmu_like( + mu, + data, + pdf, + init_pars, + par_bounds, + fixed_params, + return_fitted_pars=return_fitted_pars, + ) -def q0(mu, data, pdf, init_pars, par_bounds, fixed_params): +def q0(mu, data, pdf, init_pars, par_bounds, fixed_params, return_fitted_pars=False): r""" The test statistic, :math:`q_{0}`, for discovery of a positive signal as defined in Equation (12) in :xref:`arXiv:1007.1727`, for :math:`\mu=0`. @@ -340,6 +416,10 @@ def q0(mu, data, pdf, init_pars, par_bounds, fixed_params): >>> pyhf.infer.test_statistics.q0(test_mu, data, model, init_pars, par_bounds, fixed_params) array(2.98339447) + or, also return the best-fit parameter arrays: + >>> pyhf.infer.test_statistics.q0(test_mu, data, model, init_pars, par_bounds, fixed_params, return_fitted_pars = True) + (array(2.98339447), (array([0. , 1.03050845, 1.12128752]), array([0.95260667, 0.99635345, 1.02140172]))) + Args: mu (Number or Tensor): The signal strength parameter (must be set to zero) data (Tensor): The data to be considered @@ -350,9 +430,14 @@ def q0(mu, data, pdf, init_pars, par_bounds, fixed_params): The shape should be ``(n, 2)`` for ``n`` model parameters. fixed_params (:obj:`list` of :obj:`bool`): The flag to set a parameter constant to its starting value during minimization. + return_fitted_pars (:obj:`bool`): Return the best-fit parameter tensors + the fixed-POI and unconstrained fits have converged on + (i.e. :math:`\mu, \hat{\hat{\theta}}` and :math:`\hat{\mu}, \hat{\theta}`) Returns: Float: The calculated test statistic, :math:`q_{0}` + Tuple of two Tensors: (only if ``return_fitted_pars`` is ``True``) the best-fit parameter tensors, + :math:`\mu, \hat{\hat{\theta}}` and :math:`\hat{\mu}, \hat{\theta}` """ if pdf.config.poi_index is None: @@ -367,10 +452,12 @@ def q0(mu, data, pdf, init_pars, par_bounds, fixed_params): tensorlib, optimizer = get_backend() - tmu_like_stat, (_, muhatbhat) = _tmu_like( + tmu_like_stat, (mubhathat, muhatbhat) = _tmu_like( mu, data, pdf, init_pars, par_bounds, fixed_params, return_fitted_pars=True ) q0_stat = tensorlib.where( muhatbhat[pdf.config.poi_index] < 0, tensorlib.astensor(0.0), tmu_like_stat ) + if return_fitted_pars: + return q0_stat, (mubhathat, muhatbhat) return q0_stat