From 8234ac037b4367f012c94359022bcd9c9a20b5b2 Mon Sep 17 00:00:00 2001 From: Fabian Joswig Date: Wed, 24 Nov 2021 17:15:30 +0000 Subject: [PATCH] feat: hadrons input routines now accept a range of configs as optional argument --- pyerrors/input/hadrons.py | 37 ++++++++++++++++++++++++------------- 1 file changed, 24 insertions(+), 13 deletions(-) diff --git a/pyerrors/input/hadrons.py b/pyerrors/input/hadrons.py index b006d81b..45c14479 100644 --- a/pyerrors/input/hadrons.py +++ b/pyerrors/input/hadrons.py @@ -1,6 +1,3 @@ -#!/usr/bin/env python -# coding: utf-8 - import os import h5py import numpy as np @@ -8,7 +5,7 @@ from ..obs import Obs, CObs from ..correlators import Corr -def _get_files(path, filestem): +def _get_files(path, filestem, idl): ls = os.listdir(path) # Clean up file list @@ -23,11 +20,19 @@ def _get_files(path, filestem): # Sort according to configuration number files.sort(key=get_cnfg_number) - # Check that configurations are evenly spaced cnfg_numbers = [] + filtered_files = [] for line in files: - cnfg_numbers.append(get_cnfg_number(line)) + no = get_cnfg_number(line) + if idl: + if no in list(idl): + filtered_files.append(line) + cnfg_numbers.append(no) + else: + filtered_files.append(line) + cnfg_numbers.append(no) + # Check that configurations are evenly spaced dc = np.unique(np.diff(cnfg_numbers)) if np.any(dc < 0): raise Exception("Unsorted files") @@ -36,10 +41,10 @@ def _get_files(path, filestem): else: raise Exception('Configurations are not evenly spaced.') - return files, idx + return filtered_files, idx -def read_meson_hd5(path, filestem, ens_id, meson='meson_0', tree='meson'): +def read_meson_hd5(path, filestem, ens_id, meson='meson_0', tree='meson', idl=None): """Read hadrons meson hdf5 file and extract the meson labeled 'meson' Parameters @@ -57,9 +62,11 @@ def read_meson_hd5(path, filestem, ens_id, meson='meson_0', tree='meson'): Label of the upmost directory in the hdf5 file, default 'meson' for outputs of the Meson module. Can be altered to read input from other modules with similar structures. + idl : range + If specified only conifgurations in the given range are read in. """ - files, idx = _get_files(path, filestem) + files, idx = _get_files(path, filestem, idl) corr_data = [] infos = [] @@ -123,7 +130,7 @@ class Npr_matrix(np.ndarray): self.mom_out = getattr(obj, 'mom_out', None) -def read_ExternalLeg_hd5(path, filestem, ens_id, order='F'): +def read_ExternalLeg_hd5(path, filestem, ens_id, order='F', idl=None): """Read hadrons ExternalLeg hdf5 file and output an array of CObs Parameters @@ -134,9 +141,11 @@ def read_ExternalLeg_hd5(path, filestem, ens_id, order='F'): order -- order in which the array is to be reshaped, 'F' for the first index changing fastest (9 4x4 matrices) default. 'C' for the last index changing fastest (16 3x3 matrices), + idl : range + If specified only conifgurations in the given range are read in. """ - files, idx = _get_files(path, filestem) + files, idx = _get_files(path, filestem, idl) mom = None @@ -161,7 +170,7 @@ def read_ExternalLeg_hd5(path, filestem, ens_id, order='F'): return Npr_matrix(matrix.swapaxes(1, 2).reshape((12, 12), order=order), mom_in=mom) -def read_Bilinear_hd5(path, filestem, ens_id, order='F'): +def read_Bilinear_hd5(path, filestem, ens_id, order='F', idl=None): """Read hadrons Bilinear hdf5 file and output an array of CObs Parameters @@ -172,9 +181,11 @@ def read_Bilinear_hd5(path, filestem, ens_id, order='F'): order -- order in which the array is to be reshaped, 'F' for the first index changing fastest (9 4x4 matrices) default. 'C' for the last index changing fastest (16 3x3 matrices), + idl : range + If specified only conifgurations in the given range are read in. """ - files, idx = _get_files(path, filestem) + files, idx = _get_files(path, filestem, idl) mom_in = None mom_out = None