import os
import scipy.optimize as opt
from scipy.stats import ncx2
import numpy as np
from ..core import load_data as ld
from ..core import chisquared as chi2
from collections.abc import Iterable
[docs]class CHI2_PROFILES(object):
def __init__(self, config):
self.config = config
self.rep = int(config["rep"])
# Load data
self.loaded_datasets = ld.load_datasets(config, self.rep)
# Tuple of coefficient labels and values
self.coefficients = ld.aggregate_coefficients(
config, self.loaded_datasets, self.rep
)
for k in config["coefficients"]:
if k not in self.coefficients.labels:
raise NotImplementedError(
k + " not in fitted coefficients, please comment out in setup file."
)
# Get indice locations for quadratic corrections
if self.config["HOlambda"] == "HO":
self.config["HOindex1"] = []
self.config["HOindex2"] = []
for coeff in self.loaded_datasets.HOcorrectionsDICT:
idx1 = np.where(self.coefficients.labels == coeff.split("*")[0])[0][0]
idx2 = np.where(self.coefficients.labels == coeff.split("*")[1])[0][0]
self.config["HOindex1"].append(idx1)
self.config["HOindex2"].append(idx2)
self.config["HOindex1"] = np.array(self.config["HOindex1"])
self.config["HOindex2"] = np.array(self.config["HOindex2"])
# For closure test we produce pseudodata where theory is SM (i.e. known)
# and add noise to simulate the experimental data
self.smzero = np.zeros(len(self.coefficients.values))
# If one wants to do a CT to non-SM here is
# where you would change the initial coeffs
# smzero[5] = 10
if config["CT0L"]:
self.setup_CT()
[docs] def setup_CT(self):
"""Set up the closure tests"""
level = self.config["ctlevel"]
print(f"Running a Level %d closure test to SM" % level)
ct0l = ld.generate_closure(
self.config,
self.loaded_datasets,
self.coefficients.values,
self.smzero,
self.coefficients.labels,
level,
)
self.loaded_datasets = ct0l
[docs] def chi2_func(self, params):
"""Wrap the chi2 in a function for scipy optimiser. Pass noise and
data info as args. Log the chi2 value and values of the coefficients."""
loaded_data = self.loaded_datasets
coefficients = self.coefficients
chi2_info = chi2.compute_total_chi2(
self.config, loaded_data, params, coefficients.labels
)
current_chi2 = chi2_info[0]
self.npts = chi2_info[1]
return current_chi2
[docs] def set_constraints(self):
"""Sets parameter constraints"""
for k in self.config["coefficients"].keys():
idx = np.where(self.coefficients.labels == k)[0][0]
if self.config["coefficients"][k]["fixed"] is False:
continue
elif self.config["coefficients"][k]["fixed"] is True:
self.coefficients.values[:][idx] = self.config["coefficients"][k][
"value"
]
else:
prop_const = self.config["coefficients"][k]["value"]
if isinstance(prop_const, Iterable):
self.coefficients.values[:][idx] = 0.0
for i, const in enumerate(prop_const):
idx_fixed = np.where(
self.coefficients.labels
== self.config["coefficients"][k]["fixed"][i]
)[0][0]
self.coefficients.values[:][idx] += (
const * self.coefficients.values[:][idx_fixed]
)
else:
idx_fixed = np.where(
self.coefficients.labels
== self.config["coefficients"][k]["fixed"]
)[0][0]
self.coefficients.values[:][idx] = (
prop_const * self.coefficients.values[:][idx_fixed]
)
[docs] def get_chi2_profiles(self):
""" Scans the chi2 space to find suitable coefficient bounds """
completed_coeffs = [
os.listdir(self.config["results_path"])[i][:-4]
for i in range(len(os.listdir(self.config["results_path"])))
]
npar = len(self.coefficients.values)
for i in range(npar):
self.coefficients.values[i] = 0.0
# Initialize coefficient range and chi2
nx = 5
ny = nx
minx_val = 0
maxy_val = 0
# Change values here if you want search to extend outside -100 < c < 100
miny_val = -50
maxx_val = 50
X = np.linspace(minx_val, maxx_val, nx)
Y = np.linspace(miny_val, maxy_val, ny)
# Initialize large chi2
chi2_max = 1000
chi2_min = 1000
eps = 0.001
# Bounds are specified by chi2 = 2
# Change value below if you want something lower or higher
target_chi2 = 2 + eps
print("=" * 80)
print("Starting chi2 profile analysis")
print("=" * 80)
for k in self.config["coefficients"].keys():
if k in completed_coeffs:
continue
if k == "OpD" and "OpDC" in completed_coeffs:
continue
print("Scanning %s..." % k)
idx = np.where(self.coefficients.labels == k)[0][0]
print("Scanning upper limit...")
while chi2_max > target_chi2:
for x in range(X.size):
self.coefficients.values[idx] = X[x]
self.set_constraints()
chi2 = self.chi2_func(self.coefficients.values)
if chi2 / self.npts < (target_chi2 - eps):
minx_val = X[x]
if chi2 / self.npts > (target_chi2 - eps):
maxx_val = X[x]
chi2_max = chi2 / self.npts
break
if minx_val != 0 or maxx_val != 50:
X = np.linspace(minx_val, maxx_val, nx)
if minx_val == maxx_val:
print("No max limit value found for c < %d" % maxx_val)
break
if chi2_max < 1000:
print("Found max value:")
print("Maximum = %0.8f\t Chi2/Ndat = %0.3f" % (maxx_val, chi2_max))
print("Scanning lower limit...")
# Reset Values in case of constraints
for i in range(npar):
self.coefficients.values[i] = 0.0
while chi2_min > target_chi2:
for y in range(Y.size):
self.coefficients.values[idx] = Y[Y.size - 1 - y]
self.set_constraints()
chi2 = self.chi2_func(self.coefficients.values)
if chi2 / self.npts < (target_chi2 - eps):
maxy_val = Y[Y.size - 1 - y]
if chi2 / self.npts > (target_chi2 - eps):
miny_val = Y[Y.size - 1 - y]
chi2_min = chi2 / self.npts
break
if maxy_val > -50 or miny_val != 0:
Y = np.linspace(miny_val, maxy_val, ny)
if miny_val == maxy_val:
print("No min limit found for c > %d" % miny_val)
break
if chi2_min < 1000:
print("Found min value")
print("Minimum = %0.8f\t Chi2/Ndat = %0.3f" % (miny_val, chi2_min))
print("Scan complete!")
print("Writing file...")
# Reset values
for i in range(npar):
self.coefficients.values[i] = 0.0
temp = np.linspace(miny_val, maxx_val, 20)
if k == "OpD":
name = "/%sC.txt" % k
else:
name = "/%s.txt" % k
f = open(self.config["results_path"] + name, "w")
for t in temp:
self.coefficients.values[idx] = t
self.set_constraints()
chi2 = self.chi2_func(self.coefficients.values) / self.npts
f.write("%0.8f\t%0.3f\n" % (t, chi2))
# Reset values
for i in range(npar):
self.coefficients.values[i] = 0.0
miny_val = -50
maxy_val = 0
minx_val = 0
maxx_val = 50
X = np.linspace(minx_val, maxx_val, nx)
Y = np.linspace(miny_val, maxy_val, ny)
chi2_max = 1000
chi2_min = 1000
print("Finished!")
print("=" * 80)