diff --git a/src/cabinetry/fit/__init__.py b/src/cabinetry/fit/__init__.py index 3283c2b0..d6e0536c 100644 --- a/src/cabinetry/fit/__init__.py +++ b/src/cabinetry/fit/__init__.py @@ -8,7 +8,6 @@ import pyhf import scipy.optimize import scipy.stats -import re from cabinetry import model_utils from cabinetry.fit.results_containers import ( @@ -87,10 +86,7 @@ def _fit_model_pyhf( bestfit = pyhf.tensorlib.to_numpy(result[:, 0]) uncertainty = pyhf.tensorlib.to_numpy(result[:, 1]) - labels = model.config.par_names() - _mod_dict = dict(model.config.modifiers) - _clean_labels = [re.sub(r"\[.*\]", "", label) for label in labels] - types = [_mod_dict[n] if n in _mod_dict else None for n in _clean_labels] + labels, types = model_utils._labels_modifiers(model) corr_mat = pyhf.tensorlib.to_numpy(corr_mat) best_twice_nll = float(best_twice_nll) # convert 0-dim np.ndarray to float @@ -150,7 +146,7 @@ def _fit_model_custom( fix_pars = fix_pars or model.config.suggested_fixed() par_bounds = par_bounds or model.config.suggested_bounds() - labels = model.config.par_names() + labels, types = model_utils._labels_modifiers(model) # set initial step size to 0 for fixed parameters # this will cause the associated parameter uncertainties to be 0 post-fit @@ -195,6 +191,7 @@ def twice_nll_func(pars: np.ndarray) -> Any: bestfit, uncertainty, labels, + types, corr_mat, best_twice_nll, minos_uncertainty=minos_results, diff --git a/src/cabinetry/fit/results_containers.py b/src/cabinetry/fit/results_containers.py index 2c8a3857..8f04af17 100644 --- a/src/cabinetry/fit/results_containers.py +++ b/src/cabinetry/fit/results_containers.py @@ -1,6 +1,6 @@ """Provides containers for inference results.""" -from typing import Dict, List, NamedTuple, Tuple +from typing import Dict, List, NamedTuple, Optional, Tuple import numpy as np @@ -24,7 +24,7 @@ class FitResults(NamedTuple): bestfit: np.ndarray uncertainty: np.ndarray labels: List[str] - types: List[str] + types: List[Optional[str]] corr_mat: np.ndarray best_twice_nll: float goodness_of_fit: float = -1 diff --git a/src/cabinetry/model_utils.py b/src/cabinetry/model_utils.py index 827f4e8a..37ab38b0 100644 --- a/src/cabinetry/model_utils.py +++ b/src/cabinetry/model_utils.py @@ -2,6 +2,7 @@ import json import logging +import re from typing import Any, Dict, List, NamedTuple, Optional, Tuple, Union import awkward as ak @@ -516,6 +517,17 @@ def _filter_channels( return filtered_channels +def _labels_modifiers( + model: pyhf.pdf.Model, +) -> Tuple[List[str], List[Optional[str]]]: + """ """ + labels = model.config.par_names() + _mod_dict = dict(model.config.modifiers) + _clean_labels = [re.sub(r"\[.*\]", "", label) for label in labels] + types = [_mod_dict[n] if n in _mod_dict else None for n in _clean_labels] + return labels, types + + def match_fit_results(model: pyhf.pdf.Model, fit_results: FitResults) -> FitResults: """Matches results from a fit to a model by adding or removing parameters as needed. @@ -543,7 +555,7 @@ def match_fit_results(model: pyhf.pdf.Model, fit_results: FitResults) -> FitResu bestfit = asimov_parameters(model) # Asimov parameter values for target model uncertainty = prefit_uncertainties(model) # pre-fit uncertainties for target model - labels = model.config.par_names() # labels for target model + labels, types = _labels_modifiers(model) # indices of parameters in current fit results, or None if they are missing indices_for_corr: List[Optional[int]] = [None] * len(labels) @@ -577,6 +589,7 @@ def match_fit_results(model: pyhf.pdf.Model, fit_results: FitResults) -> FitResu np.asarray(bestfit), np.asarray(uncertainty), labels, + types, corr_mat, fit_results.best_twice_nll, fit_results.goodness_of_fit, diff --git a/src/cabinetry/visualize/__init__.py b/src/cabinetry/visualize/__init__.py index a339d20f..6387282e 100644 --- a/src/cabinetry/visualize/__init__.py +++ b/src/cabinetry/visualize/__init__.py @@ -1,9 +1,9 @@ """High-level entry point for visualizing fit models and inference results.""" +import fnmatch import glob import logging import pathlib -import fnmatch from typing import Any, Dict, List, Optional, Tuple, Union import matplotlib as mpl @@ -442,7 +442,7 @@ def pulls( *, figure_folder: Union[str, pathlib.Path] = "figures", exclude: Optional[Union[str, List[str], Tuple[str, ...]]] = None, - exclude_by_type: Optional[Union[str, List[str]]] = ["staterror"], + exclude_by_type: Optional[List[str]] = None, close_figure: bool = True, save_figure: bool = True, ) -> mpl.figure.Figure: @@ -489,6 +489,8 @@ def pulls( ) # exclude by type + if exclude_by_type is None: + exclude_by_type = ["staterror"] exclude_set.update( [ label diff --git a/src/cabinetry/visualize/plot_result.py b/src/cabinetry/visualize/plot_result.py index 8f3708a4..755bcfe0 100644 --- a/src/cabinetry/visualize/plot_result.py +++ b/src/cabinetry/visualize/plot_result.py @@ -91,7 +91,7 @@ def pulls( matplotlib.figure.Figure: the pull figure """ num_pars = len(bestfit) - numeric = np.zeros(num_pars, dtype=bool) if numeric is None else numeric + numeric = np.zeros(num_pars, dtype=bool) if numeric is None else np.asarray(numeric) y_positions = np.arange(num_pars)[::-1] fig, ax = plt.subplots(figsize=(6, 1 + num_pars / 4), dpi=100) diff --git a/tests/cli/test_cli.py b/tests/cli/test_cli.py index afb5acd4..74a4c9db 100644 --- a/tests/cli/test_cli.py +++ b/tests/cli/test_cli.py @@ -95,7 +95,12 @@ def test_workspace(mock_validate, mock_build, cli_helpers, tmp_path): @mock.patch( "cabinetry.fit.fit", return_value=fit.FitResults( - np.asarray([1.0]), np.asarray([0.1]), ["label"], np.asarray([[1.0]]), 1.0 + np.asarray([1.0]), + np.asarray([0.1]), + ["label"], + ["type"], + np.asarray([[1.0]]), + 1.0, ), autospec=True, ) @@ -109,8 +114,9 @@ def test_fit(mock_utils, mock_fit, mock_pulls, mock_corrmat, tmp_path): bestfit = np.asarray([1.0]) uncertainty = np.asarray([0.1]) labels = ["label"] + types = ["type"] corr_mat = np.asarray([[1.0]]) - fit_results = fit.FitResults(bestfit, uncertainty, labels, corr_mat, 1.0) + fit_results = fit.FitResults(bestfit, uncertainty, labels, types, corr_mat, 1.0) workspace_path = str(tmp_path / "workspace.json") @@ -204,7 +210,12 @@ def test_fit(mock_utils, mock_fit, mock_pulls, mock_corrmat, tmp_path): @mock.patch( "cabinetry.fit.fit", return_value=fit.FitResults( - np.asarray([1.0]), np.asarray([0.1]), ["label"], np.asarray([[1.0]]), 1.0 + np.asarray([1.0]), + np.asarray([0.1]), + ["label"], + ["type"], + np.asarray([[1.0]]), + 1.0, ), autospec=True, ) @@ -218,8 +229,9 @@ def test_ranking(mock_utils, mock_fit, mock_rank, mock_vis, tmp_path): bestfit = np.asarray([1.0]) uncertainty = np.asarray([0.1]) labels = ["label"] + types = ["type"] corr_mat = np.asarray([[1.0]]) - fit_results = fit.FitResults(bestfit, uncertainty, labels, corr_mat, 1.0) + fit_results = fit.FitResults(bestfit, uncertainty, labels, types, corr_mat, 1.0) workspace_path = str(tmp_path / "workspace.json") @@ -467,7 +479,12 @@ def test_significance(mock_utils, mock_sig, tmp_path): @mock.patch( "cabinetry.fit.fit", return_value=fit.FitResults( - np.asarray([1.0]), np.asarray([0.1]), ["label"], np.asarray([[1.0]]), 1.0 + np.asarray([1.0]), + np.asarray([0.1]), + ["label"], + ["type"], + np.asarray([[1.0]]), + 1.0, ), autospec=True, ) @@ -508,7 +525,12 @@ def test_data_mc( config_path = str(tmp_path / "config.yml") cli_helpers.write_config(config_path, config) fit_results = fit.FitResults( - np.asarray([1.0]), np.asarray([0.1]), ["label"], np.asarray([[1.0]]), 1.0 + np.asarray([1.0]), + np.asarray([0.1]), + ["label"], + ["type"], + np.asarray([[1.0]]), + 1.0, ) result = runner.invoke( diff --git a/tests/fit/test_fit.py b/tests/fit/test_fit.py index d2c901b1..63d15e57 100644 --- a/tests/fit/test_fit.py +++ b/tests/fit/test_fit.py @@ -16,7 +16,8 @@ def test_print_results(caplog): bestfit = np.asarray([1.0, 2.0]) uncertainty = np.asarray([0.1, 0.3]) labels = ["param_A", "param_B"] - fit_results = fit.FitResults(bestfit, uncertainty, labels, np.empty(0), 0.0) + types = ["normsys", "shapesys"] + fit_results = fit.FitResults(bestfit, uncertainty, labels, types, np.empty(0), 0.0) fit.print_results(fit_results) assert "param_A = 1.0000 +/- 0.1000" in [rec.message for rec in caplog.records] @@ -143,13 +144,13 @@ def test__fit_model_custom(mock_minos, example_spec, example_spec_multibin): @mock.patch( "cabinetry.fit._fit_model_custom", return_value=fit.FitResults( - np.asarray([1.2]), np.asarray([0.2]), ["par"], np.empty(0), 2.0 + np.asarray([1.2]), np.asarray([0.2]), ["par"], ["normsys"], np.empty(0), 2.0 ), ) @mock.patch( "cabinetry.fit._fit_model_pyhf", return_value=fit.FitResults( - np.asarray([1.1]), np.asarray([0.2]), ["par"], np.empty(0), 2.0 + np.asarray([1.1]), np.asarray([0.2]), ["par"], ["normsys"], np.empty(0), 2.0 ), ) def test__fit_model(mock_pyhf, mock_custom, example_spec): @@ -295,7 +296,7 @@ def test__goodness_of_fit( @mock.patch( "cabinetry.fit._fit_model", return_value=fit.FitResults( - np.asarray([1.0]), np.asarray([0.1]), ["par"], np.empty(0), 2.0 + np.asarray([1.0]), np.asarray([0.1]), ["par"], ["normsys"], np.empty(0), 2.0 ), ) def test_fit(mock_fit, mock_print, mock_gof): @@ -382,33 +383,78 @@ def test_fit(mock_fit, mock_print, mock_gof): "cabinetry.fit._fit_model", side_effect=[ fit.FitResults( - np.asarray([0.9, 1.3]), np.asarray([0.1, 0.1]), ["a", "b"], np.empty(0), 0.0 + np.asarray([0.9, 1.3]), + np.asarray([0.1, 0.1]), + ["a", "b"], + ["normsys", "normsys"], + np.empty(0), + 0.0, ), fit.FitResults( - np.asarray([0.9, 0.7]), np.asarray([0.1, 0.1]), ["a", "b"], np.empty(0), 0.0 + np.asarray([0.9, 0.7]), + np.asarray([0.1, 0.1]), + ["a", "b"], + ["normsys", "normsys"], + np.empty(0), + 0.0, ), fit.FitResults( - np.asarray([0.9, 1.2]), np.asarray([0.1, 0.1]), ["a", "b"], np.empty(0), 0.0 + np.asarray([0.9, 1.2]), + np.asarray([0.1, 0.1]), + ["a", "b"], + ["normsys", "normsys"], + np.empty(0), + 0.0, ), fit.FitResults( - np.asarray([0.9, 0.8]), np.asarray([0.1, 0.1]), ["a", "b"], np.empty(0), 0.0 + np.asarray([0.9, 0.8]), + np.asarray([0.1, 0.1]), + ["a", "b"], + ["normsys", "normsys"], + np.empty(0), + 0.0, ), # for second ranking call with fixed parameter fit.FitResults( - np.asarray([0.9, 1.2]), np.asarray([0.1, 0.1]), ["a", "b"], np.empty(0), 0.0 + np.asarray([0.9, 1.2]), + np.asarray([0.1, 0.1]), + ["a", "b"], + ["normsys", "normsys"], + np.empty(0), + 0.0, ), fit.FitResults( - np.asarray([0.9, 0.8]), np.asarray([0.1, 0.1]), ["a", "b"], np.empty(0), 0.0 + np.asarray([0.9, 0.8]), + np.asarray([0.1, 0.1]), + ["a", "b"], + ["normsys", "normsys"], + np.empty(0), + 0.0, ), # for third ranking call without reference results fit.FitResults( - np.asarray([0.9, 1.0]), np.asarray([0.3, 0.3]), ["a", "b"], np.empty(0), 0.0 + np.asarray([0.9, 1.0]), + np.asarray([0.3, 0.3]), + ["a", "b"], + ["normsys", "normsys"], + np.empty(0), + 0.0, ), fit.FitResults( - np.asarray([0.9, 1.3]), np.asarray([0.1, 0.1]), ["a", "b"], np.empty(0), 0.0 + np.asarray([0.9, 1.3]), + np.asarray([0.1, 0.1]), + ["a", "b"], + ["normsys", "normsys"], + np.empty(0), + 0.0, ), fit.FitResults( - np.asarray([0.9, 0.7]), np.asarray([0.1, 0.1]), ["a", "b"], np.empty(0), 0.0 + np.asarray([0.9, 0.7]), + np.asarray([0.1, 0.1]), + ["a", "b"], + ["normsys", "normsys"], + np.empty(0), + 0.0, ), ], ) @@ -417,7 +463,8 @@ def test_ranking(mock_fit, example_spec): bestfit = np.asarray([0.9, 1.0]) uncertainty = np.asarray([0.02, 0.1]) labels = ["staterror", "mu"] - fit_results = fit.FitResults(bestfit, uncertainty, labels, np.empty(0), 0.0) + types = ["staterror", "normfactor"] + fit_results = fit.FitResults(bestfit, uncertainty, labels, types, np.empty(0), 0.0) model, data = model_utils.model_and_data(example_spec) ranking_results = fit.ranking(model, data, fit_results=fit_results) @@ -500,16 +547,16 @@ def test_ranking(mock_fit, example_spec): "cabinetry.fit._fit_model", side_effect=[ fit.FitResults( - np.asarray([0.9, 1.3]), np.asarray([0.1, 0.1]), [], np.empty(0), 8.0 + np.asarray([0.9, 1.3]), np.asarray([0.1, 0.1]), [], [], np.empty(0), 8.0 ) ] # nominal fit + [ - fit.FitResults(np.empty(0), np.empty(0), [], np.empty(0), abs(i) + 8) + fit.FitResults(np.empty(0), np.empty(0), [], [], np.empty(0), abs(i) + 8) for i in np.linspace(-5, 5, 11) ] # fits in scan + [ fit.FitResults( - np.asarray([0.9, 1.3]), np.asarray([0.1, 0.1]), [], np.empty(0), 2.0 + np.asarray([0.9, 1.3]), np.asarray([0.1, 0.1]), [], [], np.empty(0), 2.0 ) ] * 6, # fits for custom parameter range diff --git a/tests/fit/test_fit_results_containers.py b/tests/fit/test_fit_results_containers.py index 0739663b..c9746849 100644 --- a/tests/fit/test_fit_results_containers.py +++ b/tests/fit/test_fit_results_containers.py @@ -9,9 +9,12 @@ def test_FitResults(): bestfit = np.asarray([1.0]) uncertainty = np.asarray([0.1]) labels = ["par_a"] + types = [None] corr_mat = np.asarray([[1.0]]) best_twice_nll = 2.0 - fit_results = fit.FitResults(bestfit, uncertainty, labels, corr_mat, best_twice_nll) + fit_results = fit.FitResults( + bestfit, uncertainty, labels, types, corr_mat, best_twice_nll + ) assert np.allclose(fit_results.bestfit, bestfit) assert np.allclose(fit_results.uncertainty, uncertainty) assert fit_results.labels == labels @@ -30,7 +33,13 @@ def test_RankingResults(): postfit_up = np.asarray([0.2]) postfit_down = np.asarray([-0.2]) ranking_results = fit.RankingResults( - bestfit, uncertainty, labels, prefit_up, prefit_down, postfit_up, postfit_down + bestfit, + uncertainty, + labels, + prefit_up, + prefit_down, + postfit_up, + postfit_down, ) assert np.allclose(ranking_results.bestfit, bestfit) assert np.allclose(ranking_results.uncertainty, uncertainty) @@ -100,7 +109,8 @@ def test_print_results(caplog): bestfit = np.asarray([1.0, 2.0]) uncertainty = np.asarray([0.1, 0.3]) labels = ["param_A", "param_B"] - fit_results = fit.FitResults(bestfit, uncertainty, labels, np.empty(0), 0.0) + types = [None, None] + fit_results = fit.FitResults(bestfit, uncertainty, labels, types, np.empty(0), 0.0) fit.print_results(fit_results) assert "param_A = 1.0000 +/- 0.1000" in [rec.message for rec in caplog.records] diff --git a/tests/test_model_utils.py b/tests/test_model_utils.py index ca5e40bb..9624f2f2 100644 --- a/tests/test_model_utils.py +++ b/tests/test_model_utils.py @@ -265,6 +265,7 @@ def test_prediction( np.asarray([1.01, 1.1]), np.asarray([0.03, 0.1]), ["staterror_Signal-Region[0]", "Signal strength"], + ["staterror", "normfactor"], np.asarray([[1.0, 0.2], [0.2, 1.0]]), 0.0, ) @@ -298,6 +299,7 @@ def test_prediction( np.asarray([1.01, 1.1]), np.asarray([0.03, 0.1]), ["a", "b"], + ["staterror", "normfactor"], np.asarray([[1.0, 0.2], [0.2, 1.0]]), 0.0, ) @@ -394,6 +396,7 @@ def test_match_fit_results(mock_pars, mock_uncs): np.asarray([1.0, 2.0, 3.0]), np.asarray([0.1, 0.2, 0.3]), ["par_a", "par_b", "par_c"], + [None, None, None], np.asarray([[1.0, 0.2, 0.5], [0.2, 1.0, 0.1], [0.5, 0.1, 1.0]]), 5.0, 0.1, diff --git a/tests/visualize/reference/pulls.png b/tests/visualize/reference/pulls.png index 8fad0631..9759e843 100644 Binary files a/tests/visualize/reference/pulls.png and b/tests/visualize/reference/pulls.png differ diff --git a/tests/visualize/test_visualize.py b/tests/visualize/test_visualize.py index 303eb54c..714fef87 100644 --- a/tests/visualize/test_visualize.py +++ b/tests/visualize/test_visualize.py @@ -360,10 +360,11 @@ def test_correlation_matrix(mock_draw): corr_mat = np.asarray([[1.0, 0.2, 0.1], [0.2, 1.0, 0.1], [0.1, 0.1, 1.0]]) corr_mat_pruned = np.asarray([[1.0, 0.2], [0.2, 1.0]]) labels = ["a", "b", "c"] + types = [None, None, None] labels_pruned = ["a", "b"] folder_path = "tmp" figure_path = pathlib.Path(folder_path) / "correlation_matrix.pdf" - fit_results = fit.FitResults(np.empty(0), np.empty(0), labels, corr_mat, 1.0) + fit_results = fit.FitResults(np.empty(0), np.empty(0), labels, types, corr_mat, 1.0) # pruning with threshold fig = visualize.correlation_matrix( @@ -382,7 +383,7 @@ def test_correlation_matrix(mock_draw): # close figure, do not save corr_mat_fixed = np.asarray([[1.0, 0.2, 0.0], [0.2, 1.0, 0.0], [0.0, 0.0, 0.0]]) fit_results_fixed = fit.FitResults( - np.empty(0), np.empty(0), labels, corr_mat_fixed, 1.0 + np.empty(0), np.empty(0), labels, types, corr_mat_fixed, 1.0 ) _ = visualize.correlation_matrix( fit_results_fixed, @@ -407,9 +408,10 @@ def test_pulls(mock_draw): bestfit = np.asarray([0.8, 1.0, 1.05, 1.1]) uncertainty = np.asarray([0.9, 1.0, 0.03, 0.7]) labels = ["a", "b", "staterror_region[0]", "c"] + types = [None, None, "staterror", None] exclude = ["a"] folder_path = "tmp" - fit_results = fit.FitResults(bestfit, uncertainty, labels, np.empty(0), 1.0) + fit_results = fit.FitResults(bestfit, uncertainty, labels, types, np.empty(0), 1.0) filtered_bestfit = np.asarray([1.0, 1.1]) filtered_uncertainty = np.asarray([1.0, 0.7]) @@ -429,6 +431,7 @@ def test_pulls(mock_draw): for i in range(len(filtered_labels)) ] ) + assert np.allclose(mock_draw.call_args[1].pop("numeric"), np.array([False, False])) assert mock_draw.call_args[1] == {"figure_path": figure_path, "close_figure": True} # filtering single parameter instead of list @@ -462,6 +465,7 @@ def test_pulls(mock_draw): for i in range(len(labels_expected)) ] ) + assert np.allclose(mock_draw.call_args[1].pop("numeric"), np.array([False, False])) assert mock_draw.call_args[1] == {"figure_path": None, "close_figure": False}