Skip to content

Commit

Permalink
feat: implement return_fitted_pars kwarg on test statistics
Browse files Browse the repository at this point in the history
  • Loading branch information
@lhenkelm
lhenkelm committed Aug 16, 2021
1 parent 6ba3aa9 commit 04b0c3b
Showing 1 changed file with 99 additions and 12 deletions.
111 changes: 99 additions & 12 deletions src/pyhf/infer/test_statistics.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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


Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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(
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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(
Expand All @@ -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)
Expand Down Expand Up @@ -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
Expand All @@ -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(
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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(
Expand All @@ -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`.
Expand Down Expand Up @@ -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
Expand All @@ -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:
Expand All @@ -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

0 comments on commit 04b0c3b

Please sign in to comment.