typehints for hadron read functions

This commit is contained in:
Justus Kuhlmann 2025-11-03 16:43:44 +00:00
commit 53192d31a9

View file

@ -142,7 +142,7 @@ def read_hd5(filestem: str, ens_id: str, group: str, attrs: Optional[Union[dict,
return corr
def read_meson_hd5(path, filestem, ens_id, meson='meson_0', idl=None, gammas=None):
def read_meson_hd5(path: str, filestem: str, ens_id: str, meson: str='meson_0', idl: Optional[range]=None, gammas: Optional[tuple[str, ...]]=None):
r'''Read hadrons meson hdf5 file and extract the meson labeled 'meson'
Parameters
@ -156,7 +156,7 @@ def read_meson_hd5(path, filestem, ens_id, meson='meson_0', idl=None, gammas=Non
meson : str
label of the meson to be extracted, standard value meson_0 which
corresponds to the pseudoscalar pseudoscalar two-point function.
gammas : tuple of strings
gammas : tuple[str]
Instrad of a meson label one can also provide a tuple of two strings
indicating the gamma matrices at sink and source (gamma_snk, gamma_src).
("Gamma5", "Gamma5") corresponds to the pseudoscalar pseudoscalar
@ -181,7 +181,7 @@ def read_meson_hd5(path, filestem, ens_id, meson='meson_0', idl=None, gammas=Non
part="real")
def _extract_real_arrays(path, files, tree, keys):
def _extract_real_arrays(path: str, files: list[str], tree: str, keys: list[str]):
corr_data = {}
for key in keys:
corr_data[key] = []
@ -199,7 +199,7 @@ def _extract_real_arrays(path, files, tree, keys):
return corr_data
def extract_t0_hd5(path, filestem, ens_id, obs='Clover energy density', fit_range=5, idl=None, **kwargs):
def extract_t0_hd5(path: str, filestem: str, ens_id: str, obs='Clover energy density', fit_range: int=5, idl: Optional[range]=None, **kwargs):
r'''Read hadrons FlowObservables hdf5 file and extract t0
Parameters
@ -247,7 +247,7 @@ def extract_t0_hd5(path, filestem, ens_id, obs='Clover energy density', fit_rang
return fit_t0(t2E_dict, fit_range, plot_fit=kwargs.get('plot_fit'))
def read_DistillationContraction_hd5(path, ens_id, diagrams=["direct"], idl=None):
def read_DistillationContraction_hd5(path: str, ens_id: str, diagrams: list[str]=["direct"], idl: Optional[range]=None):
"""Read hadrons DistillationContraction hdf5 files in given directory structure
Parameters
@ -256,7 +256,7 @@ def read_DistillationContraction_hd5(path, ens_id, diagrams=["direct"], idl=None
path to the directories to read
ens_id : str
name of the ensemble, required for internal bookkeeping
diagrams : list
diagrams : list[str]
List of strings of the diagrams to extract, e.g. ["direct", "box", "cross"].
idl : range
If specified only configurations in the given range are read in.
@ -384,7 +384,7 @@ class Npr_matrix(np.ndarray):
self.mom_out = getattr(obj, 'mom_out', None)
def read_ExternalLeg_hd5(path, filestem, ens_id, idl=None):
def read_ExternalLeg_hd5(path: str, filestem: str, ens_id: str, idl: Optional[range]=None):
"""Read hadrons ExternalLeg hdf5 file and output an array of CObs
Parameters
@ -429,7 +429,7 @@ def read_ExternalLeg_hd5(path, filestem, ens_id, idl=None):
return Npr_matrix(matrix, mom_in=mom)
def read_Bilinear_hd5(path, filestem, ens_id, idl=None):
def read_Bilinear_hd5(path: str, filestem: str, ens_id: str, idl: Optional[range]=None):
"""Read hadrons Bilinear hdf5 file and output an array of CObs
Parameters
@ -488,7 +488,7 @@ def read_Bilinear_hd5(path, filestem, ens_id, idl=None):
return result_dict
def read_Fourquark_hd5(path, filestem, ens_id, idl=None, vertices=["VA", "AV"]):
def read_Fourquark_hd5(path: str, filestem: str, ens_id: str, idl: Optional[range]=None, vertices: list[str]=["VA", "AV"]):
"""Read hadrons FourquarkFullyConnected hdf5 file and output an array of CObs
Parameters
@ -574,7 +574,7 @@ def read_Fourquark_hd5(path, filestem, ens_id, idl=None, vertices=["VA", "AV"]):
return result_dict
def _get_lorentz_names(name):
def _get_lorentz_names(name: str):
lorentz_index = ['X', 'Y', 'Z', 'T']
res = []