fix: satisfy precommit

Author: andrzejnovak

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,

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

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,

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

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)

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(

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