corrlib/backlogger/input/sfcf.py
2024-06-19 12:55:42 +00:00

195 lines
6.4 KiB
Python

import pyerrors as pe
import datalad.api as dl
def read_param(path, project, file_in_project):
"""
Read the parameters from the sfcf file.
Parameters
----------
file : str
The path to the sfcf file.
Returns
-------
dict
The parameters in the sfcf file.
"""
file = path + "/projects/" + project + '/' + file_in_project
dl.get(file, dataset=path)
with open(file, 'r') as f:
lines = f.readlines()
params = {}
params['wf_offsets'] = []
params['wf_basis'] = []
params['wf_coeff'] = []
params['qr'] = {}
params['mrr'] = []
params['crr'] = []
params['qs'] = {}
params['mrs'] = []
params['crs'] = []
for line in lines:
if line.startswith('#'):
continue
if line.startswith('\n'):
continue
if line.startswith('wf_offsets'):
num_wf_offsets = line.split()[1]
for i in range(int(num_wf_offsets)):
params['wf_offsets'].append([float(x) for x in lines[lines.index(line) + i + 1].split("#")[0].split()])
if line.startswith('wf_basis'):
num_wf_basis = line.split()[1]
for i in range(int(num_wf_basis)):
params['wf_basis'].append([float(x) for x in lines[lines.index(line) + i + 1].split("#")[0].split()])
if line.startswith('wf_coeff'):
num_wf_coeff = line.split()[1]
for i in range(int(num_wf_coeff)):
params['wf_coeff'].append([float(x) for x in lines[lines.index(line) + i + 1].split("#")[0].split()])
if line.startswith('qr'):
num_qr = line.split()[1]
for i in range(int(num_qr)):
dat = lines[lines.index(line) + i + 1].split("#")[0].strip().split()[:-1]
params['qr'][dat[0]] = {}
params['qr'][dat[0]]['mass'] = float(dat[1])
params['qr'][dat[0]]['thetas'] = [float(x) for x in dat[2:5]]
if line.startswith('mrr'):
num_mrr = line.split()[1]
for i in range(int(num_mrr)):
params['mrr'].append(lines[lines.index(line) + i + 1].split("#")[0].strip())
if line.startswith('crr'):
num_crr = line.split()[1]
for i in range(int(num_crr)):
params['crr'].append(lines[lines.index(line) + i + 1].split("#")[0].strip())
if line.startswith('qs'):
num_qs = line.split()[1]
for i in range(int(num_qs)):
dat = lines[lines.index(line) + i + 1].split("#")[0].strip().split()[:-1]
params['qs'][dat[0]] = {}
params['qs'][dat[0]]['mass'] = float(dat[1])
params['qs'][dat[0]]['thetas'] = [float(x) for x in dat[2:5]]
if line.startswith('mrs'):
num_mrs = line.split()[1]
for i in range(int(num_mrs)):
params['mrs'].append(lines[lines.index(line) + i + 1].split("#")[0].strip())
if line.startswith('crs'):
num_crs = line.split()[1]
for i in range(int(num_crs)):
params['mrs'].append(lines[lines.index(line) + i + 1].split("#")[0].strip())
# catch standard cases
if params['wf_offsets'] == []:
params['wf_offsets'] = [[0, 0, 0]]
if params['wf_coeff'] == []:
params['wf_coeff'] = [[0, -1]]
return params
def _map_params(params, spec_list):
"""
Map the extracted parameters to the extracted data.
"""
new_specs = {}
# quarks
quarks = spec_list[0].split(" ")
new_specs['quarks'] = (params['qr'][quarks[0]], params['qr'][quarks[1]])
# offset
new_specs['offset'] = (params['wf_offsets'][int(spec_list[1])])
# wf1
contribs = []
for i, coeff in enumerate(params['wf_coeff'][int(spec_list[2])]):
if not coeff == 0:
contrib = (coeff, params['wf_basis'][i])
contribs.append(contrib)
new_specs['wf1'] = contribs
if len(spec_list) == 4:
# wf2
contribs = []
for i, coeff in enumerate(params['wf_coeff'][int(spec_list[3])]):
if not coeff == 0:
contrib = (coeff, params['wf_basis'][i])
contribs.append(contrib)
new_specs['wf2'] = contribs
return new_specs
def read_data(path, project, dir_in_project, prefix, param, version='1.0c', cfg_seperator='n', sep='/', **kwargs):
"""
Extract the data from the sfcf file.
Returns
-------
dict
The data from the sfcf file.
"""
names = kwargs.get('names', None)
corr_types = {
'f_A': 'bi',
'f_P': 'bi',
'g_A': 'bi',
'g_P': 'bi',
'f_1': 'bb',
'k_1': 'bb',
}
directory = path + "/projects/" + project + '/' + dir_in_project
dl.get(directory, dataset=path)
corr_type_list = []
for corr_name in param['crr']:
if corr_name not in corr_types:
raise ValueError('Correlator type not known.')
corr_type_list.append(corr_types[corr_name])
if not param['crr'] == []:
if names is not None:
data_crr = pe.input.sfcf.read_sfcf_multi(directory, prefix, param['crr'], param['mrr'], corr_type_list, range(len(param['wf_offsets'])),
range(len(param['wf_basis'])), range(len(param['wf_basis'])), version, cfg_seperator, keyed_out=True, names=names)
else:
data_crr = pe.input.sfcf.read_sfcf_multi(directory, prefix, param['crr'], param['mrr'], corr_type_list, range(len(param['wf_offsets'])),
range(len(param['wf_basis'])), range(len(param['wf_basis'])), version, cfg_seperator, keyed_out=True)
if not param['crs'] == []:
data_crs = pe.input.sfcf.read_sfcf_multi(directory, param['crs'])
data = {}
if not param['crr'] == []:
for key in data_crr.keys():
data[key] = data_crr[key]
if not param['crs'] == []:
for key in data_crs.keys():
data[key] = data_crs[key]
# sort data by correlator
sorted_data = {}
for key in data.keys():
key_parts = key.split(sep)
corr = key_parts[0]
if corr_types[corr] == 'bi':
specs = _map_params(param, key_parts[1:-1])
else:
specs = _map_params(param, key_parts[1:])
if corr not in sorted_data:
sorted_data[corr] = {}
sorted_data[corr][sep.join(key_parts[1:])] = data[key]
return sorted_data, specs