From d91baf0a3895246570dd0513c9693fba500c93af Mon Sep 17 00:00:00 2001 From: Andrzej Date: Thu, 9 Jun 2022 12:35:44 +0200 Subject: [PATCH] fix: satisfy precommit --- src/cabinetry/fit/__init__.py | 9 +-- src/cabinetry/model_utils.py | 15 ++++- src/cabinetry/visualize/__init__.py | 6 +- src/cabinetry/visualize/plot_result.py | 2 +- tests/cli/test_cli.py | 34 +++++++++-- tests/fit/test_fit.py | 79 ++++++++++++++++++++------ tests/test_model_utils.py | 3 + 7 files changed, 116 insertions(+), 32 deletions(-) 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/model_utils.py b/src/cabinetry/model_utils.py index 827f4e8a..eb913923 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[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..130f1eae 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[Union[str, 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..8b0b511c 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,6 +463,7 @@ def test_ranking(mock_fit, example_spec): bestfit = np.asarray([0.9, 1.0]) uncertainty = np.asarray([0.02, 0.1]) labels = ["staterror", "mu"] + labels = ["staterror", "normfactor"] fit_results = fit.FitResults(bestfit, uncertainty, labels, 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/test_model_utils.py b/tests/test_model_utils.py index ca5e40bb..23afddda 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"], + ["normfactor", "normfactor", "normfactor"], np.asarray([[1.0, 0.2, 0.5], [0.2, 1.0, 0.1], [0.5, 0.1, 1.0]]), 5.0, 0.1,