use benedict to make loops easier

This commit is contained in:
Justus Kuhlmann 2023-10-09 16:06:13 +00:00
parent eb58920e34
commit fa3987479b

View file

@ -4,6 +4,8 @@ import re
import numpy as np # Thinly-wrapped numpy import numpy as np # Thinly-wrapped numpy
from ..obs import Obs from ..obs import Obs
from .utils import sort_names, check_idl from .utils import sort_names, check_idl
from benedict import benedict
from copy import deepcopy
def read_sfcf(path, prefix, name, quarks='.*', corr_type="bi", noffset=0, wf=0, wf2=0, version="1.0c", cfg_separator="n", silent=False, **kwargs): def read_sfcf(path, prefix, name, quarks='.*', corr_type="bi", noffset=0, wf=0, wf2=0, version="1.0c", cfg_separator="n", silent=False, **kwargs):
@ -65,7 +67,9 @@ def read_sfcf(path, prefix, name, quarks='.*', corr_type="bi", noffset=0, wf=0,
list of Observables with length T, observable per timeslice. list of Observables with length T, observable per timeslice.
bb-type correlators have length 1. bb-type correlators have length 1.
""" """
ret = read_sfcf_multi(path, prefix, [name], quarks_list=[quarks], corr_type_list=[corr_type], noffset_list=[noffset], wf_list=[wf], wf2_list=[wf2], version=version, cfg_separator=cfg_separator, silent=silent, nice_output=True, **kwargs) ret = read_sfcf_multi(path, prefix, [name], quarks_list=[quarks], corr_type_list=[corr_type],
noffset_list=[noffset], wf_list=[wf], wf2_list=[wf2], version=version,
cfg_separator=cfg_separator, silent=silent, nice_output=True, **kwargs)
return ret return ret
@ -195,32 +199,11 @@ def read_sfcf_multi(path, prefix, name_list, quarks_list=['.*'], corr_type_list=
idl = [] idl = []
# setup dect structures # setup dict structures
intern = {} intern = _setup_intern_dict(name_list, corr_type_list, quarks_list, noffset_list, wf_list, wf2_list)
for name, corr_type in zip(name_list, corr_type_list):
intern[name] = {}
b2b, single = _extract_corr_type(corr_type)
intern[name]["b2b"] = b2b
intern[name]["single"] = single
intern[name]["spec"] = {}
for quarks in quarks_list:
intern[name]["spec"][quarks] = {}
for off in noffset_list:
intern[name]["spec"][quarks][str(off)] = {}
for w in wf_list:
intern[name]["spec"][quarks][str(off)][str(w)] = {}
for w2 in wf2_list:
intern[name]["spec"][quarks][str(off)][str(w)][str(w2)] = {}
internal_ret_dict = {} ret_key_paths, internal_ret_dict = _setup_ret_dict(name_list, quarks_list, noffset_list, wf_list, wf2_list)
for name in name_list: dict_template = deepcopy(internal_ret_dict)
internal_ret_dict[name] = {}
for quarks in quarks_list:
internal_ret_dict[name][quarks] = {}
for off in noffset_list:
internal_ret_dict[name][quarks][str(off)] = {}
for w in wf_list:
internal_ret_dict[name][quarks][str(off)][str(w)] = {}
if not appended: if not appended:
for i, item in enumerate(ls): for i, item in enumerate(ls):
@ -249,20 +232,21 @@ def read_sfcf_multi(path, prefix, name_list, quarks_list=['.*'], corr_type_list=
if i == 0: if i == 0:
if version != "0.0" and compact: if version != "0.0" and compact:
file = path + '/' + item + '/' + sub_ls[0] file = path + '/' + item + '/' + sub_ls[0]
for k in ret_key_paths:
for name in name_list: ori_keys = k.split("/")
name = ori_keys[0]
quarks = ori_keys[1]
off = ori_keys[2]
w = ori_keys[3]
w2 = ori_keys[4]
if version == "0.0" or not compact: if version == "0.0" or not compact:
file = path + '/' + item + '/' + sub_ls[0] + '/' + name file = path + '/' + item + '/' + sub_ls[0] + '/' + name
for quarks in quarks_list:
for off in noffset_list:
for w in wf_list:
for w2 in wf2_list:
# here, we want to find the place within the file, # here, we want to find the place within the file,
# where the correlator we need is stored. # where the correlator we need is stored.
# to do so, the pattern needed is put together # to do so, the pattern needed is put together
# from the input values # from the input values
intern[name]["spec"][quarks][str(off)][str(w)][str(w2)]["pattern"] = _make_pattern(version, name, off, w, w2, intern[name]['b2b'], quarks) intern[name]["spec"][quarks][str(off)][str(w)][str(w2)]["pattern"] = _make_pattern(version, name, off, w, w2, intern[name]['b2b'], quarks)
start_read, T = _find_correlator(file, version, intern[name]["spec"][quarks][str(off)][str(w)][str(w2)]["pattern"], intern[name]['b2b'], silent=silent) start_read, T = _find_correlator(file, version, intern[name]["spec"][quarks][off][w][w2]["pattern"], intern[name]['b2b'], silent=silent)
intern[name]["spec"][quarks][str(off)][str(w)][str(w2)]["start"] = start_read intern[name]["spec"][quarks][str(off)][str(w)][str(w2)]["start"] = start_read
intern[name]["T"] = T intern[name]["T"] = T
# preparing the datastructure # preparing the datastructure
@ -270,39 +254,31 @@ def read_sfcf_multi(path, prefix, name_list, quarks_list=['.*'], corr_type_list=
deltas = [] deltas = []
for j in range(intern[name]["T"]): for j in range(intern[name]["T"]):
deltas.append([]) deltas.append([])
internal_ret_dict[name][quarks][str(off)][str(w)][str(w2)] = deltas internal_ret_dict[k] = deltas
if compact: if compact:
rep_deltas = _read_compact_rep(path, item, sub_ls, intern, im) rep_deltas = _read_compact_rep(path, item, sub_ls, intern, ret_key_paths, im, dict_template)
for name in name_list: for k in ret_key_paths:
for quarks in quarks_list: name = k.split("/")[0]
for off in noffset_list:
for w in wf_list:
for w2 in wf2_list:
for t in range(intern[name]["T"]): for t in range(intern[name]["T"]):
internal_ret_dict[name][quarks][str(off)][str(w)][str(w2)][t].append(rep_deltas[name][quarks][str(off)][str(w)][str(w2)][t]) internal_ret_dict[k][t].append(rep_deltas[k][t])
else: else:
for name in name_list: for name in internal_ret_dict:
rep_data = [] rep_data = []
for cnfg, subitem in enumerate(sub_ls): for cnfg, subitem in enumerate(sub_ls):
for quarks in quarks_list: for k in _get_deep_paths(internal_ret_dict.subset(name)):
for off in noffset_list:
for w in wf_list:
for w2 in wf2_list:
cfg_path = path + '/' + item + '/' + subitem cfg_path = path + '/' + item + '/' + subitem
file_data = _read_o_file(cfg_path, name, intern, version, im) file_data = _read_o_file(cfg_path, name, intern, version, im, dict_template)
rep_data.append(file_data) rep_data.append(file_data)
for quarks in quarks_list: for k in _get_deep_paths(internal_ret_dict.subset(name)[name]):
for off in noffset_list:
for w in wf_list:
for w2 in wf2_list:
for t in range(intern[name]["T"]): for t in range(intern[name]["T"]):
internal_ret_dict[name][quarks][str(off)][str(w)][str(w2)][t].append([]) internal_ret_dict[name][k][t].append([])
for cfg in range(no_cfg): for cfg in range(no_cfg):
internal_ret_dict[name][quarks][str(off)][str(w)][str(w2)][t][i].append(rep_data[cfg][quarks][str(off)][str(w)][str(w2)][t]) internal_ret_dict[name][k][t][i].append(rep_data[cfg][k][t])
else: else:
for name in name_list:
for name in internal_ret_dict:
if "files" in kwargs: if "files" in kwargs:
ls = kwargs.get("files") ls = kwargs.get("files")
else: else:
@ -310,11 +286,14 @@ def read_sfcf_multi(path, prefix, name_list, quarks_list=['.*'], corr_type_list=
for exc in name_ls: for exc in name_ls:
if not fnmatch.fnmatch(exc, prefix + '*.' + name): if not fnmatch.fnmatch(exc, prefix + '*.' + name):
name_ls = list(set(name_ls) - set([exc])) name_ls = list(set(name_ls) - set([exc]))
name_ls = sort_names(name_ls) name_ls = sort_names(name_ls)
for quarks in quarks_list: for k in _get_deep_paths(internal_ret_dict.subset(name)[name]):
for off in noffset_list: ori_keys = k.split("/")
for w in wf_list: quarks = ori_keys[0]
for w2 in wf2_list: off = ori_keys[1]
w = ori_keys[2]
w2 = ori_keys[3]
pattern = _make_pattern(version, name, off, w, w2, intern[name]['b2b'], quarks) pattern = _make_pattern(version, name, off, w, w2, intern[name]['b2b'], quarks)
deltas = [] deltas = []
for rep, file in enumerate(name_ls): for rep, file in enumerate(name_ls):
@ -327,7 +306,7 @@ def read_sfcf_multi(path, prefix, name_list, quarks_list=['.*'], corr_type_list=
deltas.append([]) deltas.append([])
for t in range(intern[name]['T']): for t in range(intern[name]['T']):
deltas[t].append(rep_data[t]) deltas[t].append(rep_data[t])
internal_ret_dict[name][quarks][str(off)][str(w)][str(w2)] = deltas internal_ret_dict[name][k] = deltas
if name == name_list[0]: if name == name_list[0]:
idl.append(rep_idl) idl.append(rep_idl)
@ -344,37 +323,68 @@ def read_sfcf_multi(path, prefix, name_list, quarks_list=['.*'], corr_type_list=
if not silent: if not silent:
print("Done") print("Done")
result_dict = {} result_dict = deepcopy(dict_template)
for name in name_list: for k in ret_key_paths:
result_dict[name] = {} name = k.split("/")[0]
for quarks in quarks_list:
result_dict[name][quarks] = {}
for off in noffset_list:
result_dict[name][quarks][str(off)] = {}
for w in wf_list:
result_dict[name][quarks][str(off)][str(w)] = {}
for w2 in wf2_list:
result = [] result = []
for t in range(intern[name]["T"]): for t in range(intern[name]["T"]):
result.append(Obs(internal_ret_dict[name][quarks][str(off)][str(w)][str(w2)][t], new_names, idl=idl)) result.append(Obs(internal_ret_dict[k][t], new_names, idl=idl))
result_dict[name][quarks][str(off)][str(w)][str(w2)] = result print(result)
result_dict[k] = result
if nice_output: if nice_output:
result_dict = _reduce_dict(result_dict) result_dict = _reduce_dict(result_dict)
return result_dict return result_dict
def _read_o_file(cfg_path, name, intern, version, im): def _setup_ret_dict(name_list, quarks_list, noffset_list, wf_list, wf2_list):
internal_ret_dict = {}
for name in name_list:
internal_ret_dict[name] = {}
for quarks in quarks_list:
internal_ret_dict[name][quarks] = {}
for off in noffset_list:
internal_ret_dict[name][quarks][str(off)] = {}
for w in wf_list:
internal_ret_dict[name][quarks][str(off)][str(w)] = {}
for w2 in wf2_list:
internal_ret_dict[name][quarks][str(off)][str(w)][str(w2)] = []
internal_ret_dict = benedict(internal_ret_dict, keypath_separator="/")
ret_key_paths = _get_deep_paths(internal_ret_dict)
return ret_key_paths, internal_ret_dict
def _setup_intern_dict(name_list, corr_type_list, quarks_list, noffset_list, wf_list, wf2_list):
intern = {}
for name, corr_type in zip(name_list, corr_type_list):
intern[name] = {}
b2b, single = _extract_corr_type(corr_type)
intern[name]["b2b"] = b2b
intern[name]["single"] = single
intern[name]["spec"] = {}
for quarks in quarks_list:
intern[name]["spec"][quarks] = {}
for off in noffset_list:
intern[name]["spec"][quarks][str(off)] = {}
for w in wf_list:
intern[name]["spec"][quarks][str(off)][str(w)] = {}
for w2 in wf2_list:
intern[name]["spec"][quarks][str(off)][str(w)][str(w2)] = {}
return intern
def _read_o_file(cfg_path, name, intern, version, im, template):
file = cfg_path + '/' + name file = cfg_path + '/' + name
return_vals = {} return_vals = deepcopy(template)
return_vals = return_vals.subset(name)[name]
print(return_vals)
with open(file) as fp: with open(file) as fp:
lines = fp.readlines() lines = fp.readlines()
for quarks in intern[name]["spec"].keys(): for k in _get_deep_paths(return_vals):
return_vals[quarks] = {} ori_keys = k.split("/")
for off in intern[name]["spec"][quarks].keys(): quarks = ori_keys[0]
return_vals[quarks][off] = {} off = ori_keys[1]
for w in intern[name]["spec"][quarks][off].keys(): w = ori_keys[2]
return_vals[quarks][off][w] = {} w2 = ori_keys[3]
for w2 in intern[name]["spec"][quarks][off][w].keys():
T = intern[name]["T"] T = intern[name]["T"]
start_read = intern[name]["spec"][quarks][off][w][w2]["start"] start_read = intern[name]["spec"][quarks][off][w][w2]["start"]
deltas = [] deltas = []
@ -468,19 +478,17 @@ def _find_correlator(file_name, version, pattern, b2b, silent=False):
return start_read, T return start_read, T
def _read_compact_file(rep_path, cfg_file, intern, im): def _read_compact_file(rep_path, cfg_file, intern, im, template):
return_vals = {} return_vals = deepcopy(template)
with open(rep_path + cfg_file) as fp: with open(rep_path + cfg_file) as fp:
lines = fp.readlines() lines = fp.readlines()
for name in intern.keys(): for key in _get_deep_paths(return_vals):
return_vals[name] = {} ori_keys = key.split("/")
for quarks in intern[name]["spec"].keys(): name = ori_keys[0]
return_vals[name][quarks] = {} quarks = ori_keys[1]
for off in intern[name]["spec"][quarks].keys(): off = ori_keys[2]
return_vals[name][quarks][off] = {} w = ori_keys[3]
for w in intern[name]["spec"][quarks][off].keys(): w2 = ori_keys[4]
return_vals[name][quarks][off][w] = {}
for w2 in intern[name]["spec"][quarks][off][w].keys():
T = intern[name]["T"] T = intern[name]["T"]
start_read = intern[name]["spec"][quarks][off][w][w2]["start"] start_read = intern[name]["spec"][quarks][off][w][w2]["start"]
# check, if the correlator is in fact # check, if the correlator is in fact
@ -496,39 +504,29 @@ def _read_compact_file(rep_path, cfg_file, intern, im):
for k in range(6, T + 6): for k in range(6, T + 6):
floats = list(map(float, corr_lines[k].split())) floats = list(map(float, corr_lines[k].split()))
t_vals.append(floats[-2:][im]) t_vals.append(floats[-2:][im])
return_vals[name][quarks][off][w][w2] = t_vals return_vals[key] = t_vals
return return_vals return return_vals
def _read_compact_rep(path, rep, sub_ls, intern, im_list): def _read_compact_rep(path, rep, sub_ls, intern, keys, im_list, template):
rep_path = path + '/' + rep + '/' rep_path = path + '/' + rep + '/'
no_cfg = len(sub_ls) no_cfg = len(sub_ls)
print(template)
return_vals = {} return_vals = deepcopy(template)
for name in intern.keys(): for k in keys:
return_vals[name] = {} name = k.split("/")[0]
for quarks in intern[name]["spec"].keys():
return_vals[name][quarks] = {}
for off in intern[name]["spec"][quarks].keys():
return_vals[name][quarks][off] = {}
for w in intern[name]["spec"][quarks][off].keys():
return_vals[name][quarks][off][w] = {}
for w2 in intern[name]["spec"][quarks][off][w].keys():
deltas = [] deltas = []
for t in range(intern[name]["T"]): for t in range(intern[name]["T"]):
deltas.append(np.zeros(no_cfg)) deltas.append(np.zeros(no_cfg))
return_vals[name][quarks][off][w][w2] = deltas return_vals[k] = deltas
for cfg in range(no_cfg): for cfg in range(no_cfg):
cfg_file = sub_ls[cfg] cfg_file = sub_ls[cfg]
cfg_data = _read_compact_file(rep_path, cfg_file, intern, im_list) cfg_data = _read_compact_file(rep_path, cfg_file, intern, im_list, template)
for name in intern.keys(): for k in keys:
for quarks in intern[name]["spec"].keys(): name = k.split("/")[0]
for off in intern[name]["spec"][quarks].keys():
for w in intern[name]["spec"][quarks][off].keys():
for w2 in intern[name]["spec"][quarks][off][w].keys():
for t in range(intern[name]["T"]): for t in range(intern[name]["T"]):
return_vals[name][quarks][off][w][w2][t][cfg] = cfg_data[name][quarks][off][w][w2][t] return_vals[k][t][cfg] = cfg_data[k][t]
print(template)
return return_vals return return_vals
@ -641,3 +639,13 @@ def _reduce_dict(d):
if len(list(d.keys())) == 1: if len(list(d.keys())) == 1:
ret = d[list(d.keys())[0]] ret = d[list(d.keys())[0]]
return ret return ret
def _get_deep_paths(f, sep='/'):
max_count = max([k.count(sep) for k in f.keypaths()])
deep_paths = []
for k in f.keypaths():
if k.count(sep) == max_count:
deep_paths.append(k)
return deep_paths