Merge branch 'develop' into documentation

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