Skip to content

Commit

Permalink
Implemented tt OOs interface for ILC/CLIC and homogoneized the implem…
Browse files Browse the repository at this point in the history 
…entation in general.
  • Loading branch information
@arossia94
arossia94 committed Dec 20, 2024
1 parent 0cef568 commit 26d070f
Showing 1 changed file with 252 additions and 12 deletions.
264 changes: 252 additions & 12 deletions external_chi2/optimal_observables/interface_oos.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,6 @@
current_file_path = pathlib.Path(__file__).resolve().parent



class OptimalWWFCC161:
def __init__(self, coefficients, rgemat=None):
oo_wc_basis = ["OpD", "OpWB", "OWWW", "Opl1", "Ope", "O3pl1"]
Expand Down Expand Up @@ -125,6 +124,7 @@ def compute_chi2(self, coefficient_values):

return chi2_value


class OptimalWWCEPC161:
def __init__(self, coefficients, rgemat=None):
oo_wc_basis = ["OpD", "OpWB", "OWWW", "Opl1", "Ope", "O3pl1"]
Expand Down Expand Up @@ -244,6 +244,7 @@ def compute_chi2(self, coefficient_values):

return chi2_value


class OptimalWWILC250:
def __init__(self, coefficients, rgemat=None):
oo_wc_basis = ["OpD", "OpWB", "OWWW", "Opl1", "Ope", "O3pl1"]
Expand All @@ -266,7 +267,7 @@ def __init__(self, coefficients, rgemat=None):

incovs_reordered = []
for path in self.datasets.values():
invcov = np.loadtxt(current_file_path / path.format(collider=collider))
invcov = np.loadtxt(current_file_path / path)
temp = jnp.einsum("ij, jk, kl", self.project.T, invcov, self.project)
incovs_reordered.append(temp)
self.incov_tot = jnp.sum(jnp.array(incovs_reordered), axis=0)
Expand Down Expand Up @@ -311,7 +312,7 @@ def __init__(self, coefficients, rgemat=None):

incovs_reordered = []
for path in self.datasets.values():
invcov = np.loadtxt(current_file_path / path.format(collider=collider))
invcov = np.loadtxt(current_file_path / path)
temp = jnp.einsum("ij, jk, kl", self.project.T, invcov, self.project)
incovs_reordered.append(temp)
self.incov_tot = jnp.sum(jnp.array(incovs_reordered), axis=0)
Expand Down Expand Up @@ -356,7 +357,7 @@ def __init__(self, coefficients, rgemat=None):

incovs_reordered = []
for path in self.datasets.values():
invcov = np.loadtxt(current_file_path / path.format(collider=collider))
invcov = np.loadtxt(current_file_path / path)
temp = jnp.einsum("ij, jk, kl", self.project.T, invcov, self.project)
incovs_reordered.append(temp)
self.incov_tot = jnp.sum(jnp.array(incovs_reordered), axis=0)
Expand Down Expand Up @@ -397,7 +398,7 @@ def __init__(self, coefficients, rgemat=None):

incovs_reordered = []
for path in self.datasets.values():
invcov = np.loadtxt(current_file_path / path.format(collider=collider))
invcov = np.loadtxt(current_file_path / path)
temp = jnp.einsum("ij, jk, kl", self.project.T, invcov, self.project)
incovs_reordered.append(temp)
self.incov_tot = jnp.sum(jnp.array(incovs_reordered), axis=0)
Expand Down Expand Up @@ -438,7 +439,7 @@ def __init__(self, coefficients, rgemat=None):

incovs_reordered = []
for path in self.datasets.values():
invcov = np.loadtxt(current_file_path / path.format(collider=collider))
invcov = np.loadtxt(current_file_path / path)
temp = jnp.einsum("ij, jk, kl", self.project.T, invcov, self.project)
incovs_reordered.append(temp)
self.incov_tot = jnp.sum(jnp.array(incovs_reordered), axis=0)
Expand Down Expand Up @@ -479,7 +480,7 @@ def __init__(self, coefficients, rgemat=None):

incovs_reordered = []
for path in self.datasets.values():
invcov = np.loadtxt(current_file_path / path.format(collider=collider))
invcov = np.loadtxt(current_file_path / path)
temp = jnp.einsum("ij, jk, kl", self.project.T, invcov, self.project)
incovs_reordered.append(temp)
self.incov_tot = jnp.sum(jnp.array(incovs_reordered), axis=0)
Expand All @@ -502,7 +503,7 @@ def compute_chi2(self, coefficient_values):
return chi2_value


class OptimalttFCC:
class OptimalttFCC365:
def __init__(self, coefficients, rgemat=None):
oo_tt_wc_basis = ["OpQM", "Opt", "OtW", "OtZ"]

Expand All @@ -511,7 +512,7 @@ def __init__(self, coefficients, rgemat=None):
if op in coefficients.name:
self.project[i, np.argwhere(coefficients.name == op)[0, 0]] = 1

collider="FCCee"
collider = "FCCee"
self.datasets = {"{collider}_tt_365": "invcov_{collider}_tt_365GeV.dat"}

incovs_reordered = []
Expand Down Expand Up @@ -539,7 +540,7 @@ def compute_chi2(self, coefficient_values):
return chi2_value


class OptimalttCEPC:
class OptimalttCEPC365:
def __init__(self, coefficients, rgemat=None):
oo_tt_wc_basis = ["OpQM", "Opt", "OtW", "OtZ"]

Expand All @@ -548,7 +549,7 @@ def __init__(self, coefficients, rgemat=None):
if op in coefficients.name:
self.project[i, np.argwhere(coefficients.name == op)[0, 0]] = 1

collider="CEPC"
collider = "CEPC"
self.datasets = {"{collider}_tt_365": "invcov_{collider}_tt_365GeV.dat"}

incovs_reordered = []
Expand All @@ -573,4 +574,243 @@ def compute_chi2(self, coefficient_values):
"i, ij, j", coefficient_values, self.incov_tot, coefficient_values
)

return chi2_value
return chi2_value


class OptimalttCLIC380:
def __init__(self, coefficients, rgemat=None):
oo_wc_basis = ["OpQM", "Opt", "OtW", "OtZ"]

self.project = np.zeros((len(oo_wc_basis), coefficients.size))
for i, op in enumerate(oo_wc_basis):
if op in coefficients.name:
self.project[i, np.argwhere(coefficients.name == op)[0, 0]] = 1

self.datasets = {
"CLIC_tt_wbwb_380_neg80_0": "invcov_CLIC_tt_wbwb_380_neg80_0.dat",
"CLIC_tt_wbwb_380_pos80_0": "invcov_CLIC_tt_wbwb_380_pos80_0.dat",
}

incovs_reordered = []
for path in self.datasets.values():
invcov = np.loadtxt(current_file_path / path)
temp = jnp.einsum("ij, jk, kl", self.project.T, invcov, self.project)
incovs_reordered.append(temp)
self.incov_tot = jnp.sum(jnp.array(incovs_reordered), axis=0)

self.rgemat = rgemat

if self.rgemat is not None:
# multiply the RGE matrix as well
self.incov_tot = jnp.einsum(
"ij, jk, kl", self.rgemat.T, self.incov_tot, self.rgemat
)

self.n_dat = len(oo_wc_basis)

def compute_chi2(self, coefficient_values):
chi2_value = jnp.einsum(
"i, ij, j", coefficient_values, self.incov_tot, coefficient_values
)
return chi2_value


class OptimalttCLIC1500:
def __init__(self, coefficients, rgemat=None):
oo_wc_basis = ["OpQM", "Opt", "OtW", "OtZ"]

self.project = np.zeros((len(oo_wc_basis), coefficients.size))
for i, op in enumerate(oo_wc_basis):
if op in coefficients.name:
self.project[i, np.argwhere(coefficients.name == op)[0, 0]] = 1

self.datasets = {
"CLIC_tt_wbwb_1500_neg80_0": "invcov_CLIC_tt_wbwb_1500_neg80_0.dat",
"CLIC_tt_wbwb_1500_pos80_0": "invcov_CLIC_tt_wbwb_1500_pos80_0.dat",
}

incovs_reordered = []
for path in self.datasets.values():
invcov = np.loadtxt(current_file_path / path)
temp = jnp.einsum("ij, jk, kl", self.project.T, invcov, self.project)
incovs_reordered.append(temp)
self.incov_tot = jnp.sum(jnp.array(incovs_reordered), axis=0)

self.rgemat = rgemat

if self.rgemat is not None:
# multiply the RGE matrix as well
self.incov_tot = jnp.einsum(
"ij, jk, kl", self.rgemat.T, self.incov_tot, self.rgemat
)

self.n_dat = len(oo_wc_basis)

def compute_chi2(self, coefficient_values):
chi2_value = jnp.einsum(
"i, ij, j", coefficient_values, self.incov_tot, coefficient_values
)

return chi2_value


class OptimalttCLIC3000:
def __init__(self, coefficients, rgemat=None):
oo_wc_basis = ["OpQM", "Opt", "OtW", "OtZ"]

self.project = np.zeros((len(oo_wc_basis), coefficients.size))
for i, op in enumerate(oo_wc_basis):
if op in coefficients.name:
self.project[i, np.argwhere(coefficients.name == op)[0, 0]] = 1

self.datasets = {
"CLIC_tt_wbwb_3000_neg80_0": "invcov_CLIC_tt_wbwb_3000_neg80_0.dat",
"CLIC_tt_wbwb_3000_pos80_0": "invcov_CLIC_tt_wbwb_3000_pos80_0.dat",
}

incovs_reordered = []
for path in self.datasets.values():
invcov = np.loadtxt(current_file_path / path)
temp = jnp.einsum("ij, jk, kl", self.project.T, invcov, self.project)
incovs_reordered.append(temp)
self.incov_tot = jnp.sum(jnp.array(incovs_reordered), axis=0)

self.rgemat = rgemat

if self.rgemat is not None:
# multiply the RGE matrix as well
self.incov_tot = jnp.einsum(
"ij, jk, kl", self.rgemat.T, self.incov_tot, self.rgemat
)

self.n_dat = len(oo_wc_basis)

def compute_chi2(self, coefficient_values):
chi2_value = jnp.einsum(
"i, ij, j", coefficient_values, self.incov_tot, coefficient_values
)

return chi2_value


class OptimalttILC350:
def __init__(self, coefficients, rgemat=None):
oo_wc_basis = ["OpQM", "Opt", "OtW", "OtZ"]

self.project = np.zeros((len(oo_wc_basis), coefficients.size))
for i, op in enumerate(oo_wc_basis):
if op in coefficients.name:
self.project[i, np.argwhere(coefficients.name == op)[0, 0]] = 1

self.datasets = {
"ILC_tt_wbwb_350_neg80_neg30": "invcov_ILC_tt_wbwb_350_neg80_neg30.dat",
"ILC_tt_wbwb_350_pos80_neg30": "invcov_ILC_tt_wbwb_350_pos80_neg30.dat",
"ILC_tt_wbwb_350_neg80_pos30": "invcov_ILC_tt_wbwb_350_neg80_pos30.dat",
"ILC_tt_wbwb_350_pos80_pos30": "invcov_ILC_tt_wbwb_350_pos80_pos30.dat",
}

incovs_reordered = []
for path in self.datasets.values():
invcov = np.loadtxt(current_file_path / path)
temp = jnp.einsum("ij, jk, kl", self.project.T, invcov, self.project)
incovs_reordered.append(temp)
self.incov_tot = jnp.sum(jnp.array(incovs_reordered), axis=0)

self.rgemat = rgemat

if self.rgemat is not None:
# multiply the RGE matrix as well
self.incov_tot = jnp.einsum(
"ij, jk, kl", self.rgemat.T, self.incov_tot, self.rgemat
)

self.n_dat = len(oo_wc_basis)

def compute_chi2(self, coefficient_values):
chi2_value = jnp.einsum(
"i, ij, j", coefficient_values, self.incov_tot, coefficient_values
)

return chi2_value


class OptimalttILC500:
def __init__(self, coefficients, rgemat=None):
oo_wc_basis = ["OpQM", "Opt", "OtW", "OtZ"]

self.project = np.zeros((len(oo_wc_basis), coefficients.size))
for i, op in enumerate(oo_wc_basis):
if op in coefficients.name:
self.project[i, np.argwhere(coefficients.name == op)[0, 0]] = 1

self.datasets = {
"ILC_tt_wbwb_500_neg80_neg30": "invcov_ILC_tt_wbwb_500_neg80_neg30.dat",
"ILC_tt_wbwb_500_pos80_neg30": "invcov_ILC_tt_wbwb_500_pos80_neg30.dat",
"ILC_tt_wbwb_500_neg80_pos30": "invcov_ILC_tt_wbwb_500_neg80_pos30.dat",
"ILC_tt_wbwb_500_pos80_pos30": "invcov_ILC_tt_wbwb_500_pos80_pos30.dat",
}

incovs_reordered = []
for path in self.datasets.values():
invcov = np.loadtxt(current_file_path / path)
temp = jnp.einsum("ij, jk, kl", self.project.T, invcov, self.project)
incovs_reordered.append(temp)
self.incov_tot = jnp.sum(jnp.array(incovs_reordered), axis=0)

self.rgemat = rgemat

if self.rgemat is not None:
# multiply the RGE matrix as well
self.incov_tot = jnp.einsum(
"ij, jk, kl", self.rgemat.T, self.incov_tot, self.rgemat
)

self.n_dat = len(oo_wc_basis)

def compute_chi2(self, coefficient_values):
chi2_value = jnp.einsum(
"i, ij, j", coefficient_values, self.incov_tot, coefficient_values
)

return chi2_value


class OptimalttILC1000:
def __init__(self, coefficients, rgemat=None):
oo_wc_basis = ["OpQM", "Opt", "OtW", "OtZ"]

self.project = np.zeros((len(oo_wc_basis), coefficients.size))
for i, op in enumerate(oo_wc_basis):
if op in coefficients.name:
self.project[i, np.argwhere(coefficients.name == op)[0, 0]] = 1

self.datasets = {
"ILC_tt_wbwb_1000_neg80_neg20": "invcov_ILC_tt_wbwb_1000_neg80_neg20.dat",
"ILC_tt_wbwb_1000_pos80_neg20": "invcov_ILC_tt_wbwb_1000_pos80_neg20.dat",
"ILC_tt_wbwb_1000_neg80_pos20": "invcov_ILC_tt_wbwb_1000_neg80_pos20.dat",
"ILC_tt_wbwb_1000_pos80_pos20": "invcov_ILC_tt_wbwb_1000_pos80_pos20.dat",
}

incovs_reordered = []
for path in self.datasets.values():
invcov = np.loadtxt(current_file_path / path)
temp = jnp.einsum("ij, jk, kl", self.project.T, invcov, self.project)
incovs_reordered.append(temp)
self.incov_tot = jnp.sum(jnp.array(incovs_reordered), axis=0)

self.rgemat = rgemat

if self.rgemat is not None:
# multiply the RGE matrix as well
self.incov_tot = jnp.einsum(
"ij, jk, kl", self.rgemat.T, self.incov_tot, self.rgemat
)

self.n_dat = len(oo_wc_basis)

def compute_chi2(self, coefficient_values):
chi2_value = jnp.einsum(
"i, ij, j", coefficient_values, self.incov_tot, coefficient_values
)

return chi2_value

0 comments on commit 26d070f

Please sign in to comment.