Skip to content

Commit

Permalink
fix: satisfy precommit
Browse files Browse the repository at this point in the history
  • Loading branch information
andrzejnovak committed Jun 9, 2022
1 parent 9ed8358 commit 6a3ab96
Show file tree
Hide file tree
Showing 11 changed files with 139 additions and 41 deletions.
9 changes: 3 additions & 6 deletions src/cabinetry/fit/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down
4 changes: 2 additions & 2 deletions src/cabinetry/fit/results_containers.py
Original file line number Diff line number Diff line change
@@ -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

Expand All @@ -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
Expand Down
15 changes: 14 additions & 1 deletion src/cabinetry/model_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

import json
import logging
import re
from typing import Any, Dict, List, NamedTuple, Optional, Tuple, Union

import awkward as ak
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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,
Expand Down
6 changes: 4 additions & 2 deletions src/cabinetry/visualize/__init__.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -489,6 +489,8 @@ def pulls(
)

# exclude by type
if exclude_by_type is None:
exclude_by_type = ["staterror"]
exclude_set.update(
[
label
Expand Down
2 changes: 1 addition & 1 deletion src/cabinetry/visualize/plot_result.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
34 changes: 28 additions & 6 deletions tests/cli/test_cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)
Expand All @@ -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")

Expand Down Expand Up @@ -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,
)
Expand All @@ -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")

Expand Down Expand Up @@ -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,
)
Expand Down Expand Up @@ -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(
Expand Down
81 changes: 64 additions & 17 deletions tests/fit/test_fit.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down Expand Up @@ -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):
Expand Down Expand Up @@ -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):
Expand Down Expand Up @@ -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,
),
],
)
Expand All @@ -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)

Expand Down Expand Up @@ -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
Expand Down
Loading

0 comments on commit 6a3ab96

Please sign in to comment.