diff --git a/src/pyhf/optimize/mixins.py b/src/pyhf/optimize/mixins.py index eb718ea40b..0504ae1d18 100644 --- a/src/pyhf/optimize/mixins.py +++ b/src/pyhf/optimize/mixins.py @@ -78,6 +78,7 @@ def _internal_postprocess( return_uncertainties=False, uncertainties=None, hess_inv=None, + calc_correlations=False, fixed_vals=None, init_pars=None, ): @@ -122,31 +123,59 @@ def _internal_postprocess( if return_uncertainties: fitted_pars = tensorlib.stack([fitted_pars, uncertainties], axis=1) - correlations = getattr(fitresult, 'corr', None) - if correlations is not None: + cov = getattr(fitresult, 'hess_inv', None) + if cov is None and hess_inv is not None: + cov = hess_inv + + # we also need to edit the covariance matrix to zero-out uncertainties! + if fixed_vals is not None: + if using_minuit: + fixed_bools = fitresult.minuit.fixed + else: + fixed_bools = [False] * len(init_pars) + # Convert fixed_bools to a numpy array and reshape to make it a column vector + fixed_mask = tensorlib.reshape( + tensorlib.astensor(fixed_bools, dtype="bool"), (-1, 1) + ) + # Create 2D masks for rows and columns + row_mask = fixed_mask + col_mask = tensorlib.transpose(fixed_mask) + + # Use logical OR to combine the masks + final_mask = row_mask | col_mask + + # Use np.where to set elements of the covariance matrix to 0 where the mask is True + cov = tensorlib.where( + final_mask, tensorlib.astensor(0.0), tensorlib.astensor(cov) + ) + + correlations_from_fit = getattr(fitresult, 'corr', None) + if correlations_from_fit is None and calc_correlations: + correlations_from_fit = cov / tensorlib.outer(uncertainties, uncertainties) + correlations_from_fit = tensorlib.where( + tensorlib.isfinite(correlations_from_fit), + correlations_from_fit, + tensorlib.astensor(0.0), + ) + + if correlations_from_fit is not None: _zeros = tensorlib.zeros(num_fixed_pars) # possibly a more elegant way to do this stitched_columns = [ stitch_pars(tensorlib.astensor(column), stitch_with=_zeros) - for column in zip(*correlations) + for column in zip(*correlations_from_fit) ] stitched_rows = [ stitch_pars(tensorlib.astensor(row), stitch_with=_zeros) for row in zip(*stitched_columns) ] - correlations = tensorlib.stack(stitched_rows, axis=1) - - if hess_inv is not None: - if hasattr(fitresult, 'hess_inv'): - if fitresult.hess_inv is None: - fitresult.hess_inv = hess_inv - else: - fitresult.hess_inv = hess_inv + correlations_from_fit = tensorlib.stack(stitched_rows, axis=1) fitresult.x = fitted_pars fitresult.fun = tensorlib.astensor(fitresult.fun) fitresult.unc = uncertainties - fitresult.corr = correlations + fitresult.hess_inv = cov + fitresult.corr = correlations_from_fit return fitresult @@ -235,9 +264,11 @@ def minimize( all_pars = stitch_pars(tensorlib.astensor(result.x)) hess_inv = tensorlib.fisher_cov(pdf, all_pars, data) uncertainties = tensorlib.sqrt(tensorlib.diagonal(hess_inv)) + calc_correlations = True else: hess_inv = None uncertainties = None + calc_correlations = False # uncerts are set to 0 in here for fixed pars result = self._internal_postprocess( @@ -247,6 +278,7 @@ def minimize( return_uncertainties=return_uncertainties, uncertainties=uncertainties, hess_inv=hess_inv, + calc_correlations=calc_correlations, fixed_vals=fixed_vals, init_pars=init_pars, ) diff --git a/tests/test_optim.py b/tests/test_optim.py index c580816703..bddb688b58 100644 --- a/tests/test_optim.py +++ b/tests/test_optim.py @@ -421,13 +421,44 @@ def test_optim_uncerts_autodiff(backend, source, spec, mu): return_uncertainties=True, ) assert result.shape == (2, 2) - # TODO: add proper numerical test for autodiff uncerts + # TODO: add proper numerical test for autodiff uncerts (does not match minuit at all) # assert pytest.approx([0.26418431, 0.0]) == pyhf.tensorlib.tolist(result[:, 1]) @pytest.mark.parametrize('mu', [1.0], ids=['mu=1']) @pytest.mark.only_numpy_minuit -def test_optim_correlations(backend, source, spec, mu): +def test_optim_correlations_minuit(backend, source, spec, mu): + pdf = pyhf.Model(spec) + data = source['bindata']['data'] + pdf.config.auxdata + + init_pars = pdf.config.suggested_init() + par_bounds = pdf.config.suggested_bounds() + + optim = pyhf.optimizer + + result = optim.minimize(pyhf.infer.mle.twice_nll, data, pdf, init_pars, par_bounds) + assert pyhf.tensorlib.tolist(result) + + result, correlations = optim.minimize( + pyhf.infer.mle.twice_nll, + data, + pdf, + init_pars, + par_bounds, + [(pdf.config.poi_index, mu)], + return_correlations=True, + ) + assert result.shape == (2,) + assert correlations.shape == (2, 2) + assert pyhf.tensorlib.tolist(result) + assert pyhf.tensorlib.tolist(correlations) + + assert np.allclose([[1.0, 0.0], [0.0, 0.0]], pyhf.tensorlib.tolist(correlations)) + + +@pytest.mark.parametrize('mu', [1.0], ids=['mu=1']) +@pytest.mark.skip_numpy +def test_optim_correlations_autodiff(backend, source, spec, mu): pdf = pyhf.Model(spec) data = source['bindata']['data'] + pdf.config.auxdata