gambit_plotting_tools/example_2D_profile_like_hdf5_multiple.py at main · GambitBSM/gambit_plotting_tools · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
from copy import deepcopy
import numpy as np
import matplotlib.pyplot as plt

import gambit_plotting_tools.gambit_plot_utils as plot_utils
import gambit_plotting_tools.gambit_plot_settings as gambit_plot_settings
from gambit_plotting_tools.annotate import add_header


#
# Read files
#

hdf5_file_and_group_names = [
    ("./example_data/results_run1.hdf5", "data"),
    ("./example_data/results_run2.hdf5", "data"),
]

# Create a list of tuples of the form (shorthand key, (full dataset name, dataset type))
datasets = [
    ("LogLike", ("LogLike", float)),
    ("mu",      ("#NormalDist_parameters @NormalDist::primary_parameters::mu", float)),
    ("sigma",   ("#NormalDist_parameters @NormalDist::primary_parameters::sigma", float)),
]

# Now create our main data dictionary by reading the hdf5 files
data = plot_utils.read_hdf5_datasets(hdf5_file_and_group_names, datasets, filter_invalid_points=True)


#
# Add any derived datasets we might be interested in
#

data["sigma2"] = data["sigma"]**2
data["log_mu"] = np.log(data["mu"])


#
# Make 2D profile likelihood plots
#

confidence_levels = [0.683, 0.954]
likelihood_ratio_contour_values = plot_utils.get_2D_likelihood_ratio_levels(confidence_levels)

# Make a plot for every combination of (x_key, y_key, z_key)

x_keys = [
    "mu",
    "log_mu",
]

y_keys = [
    "sigma",
    "sigma2"
]

z_keys = [
    "LogLike",
]

# Set some bounds manually?
dataset_bounds = {
    "mu": [15, 30],
    "sigma": [0, 5],
}

# Specify some pretty plot labels?
plot_labels = {
    "mu": r"$\mu$ (unit)",
    "sigma": r"$\sigma$ (unit)",
    "sigma2": r"$\sigma^2$ (unit)",
    "log_mu": r"$\log{\mu}$ (unit)",
}

# Number of bins used for profiling
xy_bins = (100, 100)

for z_key in z_keys:
    for x_key in x_keys:
        for y_key in y_keys:

            # If a pretty plot label is not given, just use the key
            x_label = plot_labels.get(x_key, x_key)
            y_label = plot_labels.get(y_key, y_key)
            z_label = plot_labels.get(z_key, z_key)
            labels = (x_label, y_label, z_label)

            # If variable bounds are not specified, use the full range from the data
            x_bounds = dataset_bounds.get(x_key, [np.min(data[x_key]), np.max(data[x_key])])
            y_bounds = dataset_bounds.get(y_key, [np.min(data[y_key]), np.max(data[y_key])])
            xy_bounds = (x_bounds, y_bounds)

            # Copy default GAMBIT plot settings (and make changes if necessary)
            plot_settings = deepcopy(gambit_plot_settings.plot_settings)

            # Create 2D profile likelihood figure
            fig, ax, cbar_ax = plot_utils.plot_2D_profile(
                data[x_key],
                data[y_key],
                data[z_key],
                labels,
                xy_bins,
                xy_bounds=xy_bounds,
                z_is_loglike=True,
                plot_likelihood_ratio=True,
                contour_levels=likelihood_ratio_contour_values,
                missing_value_color=plot_settings["colormap"](0.0),
                add_max_likelihood_marker = True,
                plot_settings=plot_settings,
            )

            # Add header
            header_text = r"$1\sigma$ and $2\sigma$ CL regions."
            if plt.rcParams.get("text.usetex"):
                header_text += r" \textsf{GAMBIT} 2.5"
            else:
                header_text += r" GAMBIT 2.5"
            add_header(header_text, ax=ax)

            # Save to file
            output_path = f"./plots/2D_profile__{x_key}__{y_key}__{z_key}.pdf"
            plot_utils.create_folders_if_not_exist(output_path)
            plt.savefig(output_path)
            plt.close()
            print(f"Wrote file: {output_path}")