pyerrors.input.hadrons

  1import os
  2from collections import Counter
  3import h5py
  4from pathlib import Path
  5import numpy as np
  6from ..obs import Obs, CObs
  7from ..correlators import Corr
  8from ..dirac import epsilon_tensor_rank4
  9from .misc import fit_t0
 10
 11
 12def _get_files(path, filestem, idl):
 13    ls = os.listdir(path)
 14
 15    # Clean up file list
 16    files = list(filter(lambda x: x.startswith(filestem + "."), ls))
 17
 18    if not files:
 19        raise Exception('No files starting with', filestem, 'in folder', path)
 20
 21    def get_cnfg_number(n):
 22        return int(n.replace(".h5", "")[len(filestem) + 1:])  # From python 3.9 onward the safer 'removesuffix' method can be used.
 23
 24    # Sort according to configuration number
 25    files.sort(key=get_cnfg_number)
 26
 27    cnfg_numbers = []
 28    filtered_files = []
 29    for line in files:
 30        no = get_cnfg_number(line)
 31        if idl:
 32            if no in list(idl):
 33                filtered_files.append(line)
 34                cnfg_numbers.append(no)
 35        else:
 36            filtered_files.append(line)
 37            cnfg_numbers.append(no)
 38
 39    if idl:
 40        if Counter(list(idl)) != Counter(cnfg_numbers):
 41            raise Exception("Not all configurations specified in idl found, configurations " + str(list(Counter(list(idl)) - Counter(cnfg_numbers))) + " are missing.")
 42
 43    # Check that configurations are evenly spaced
 44    dc = np.unique(np.diff(cnfg_numbers))
 45    if np.any(dc < 0):
 46        raise Exception("Unsorted files")
 47    if len(dc) == 1:
 48        idx = range(cnfg_numbers[0], cnfg_numbers[-1] + dc[0], dc[0])
 49    elif idl:
 50        idx = idl
 51    else:
 52        raise Exception("Configurations are not evenly spaced. Provide an idl if you want to proceed with this set of configurations.")
 53
 54    return filtered_files, idx
 55
 56
 57def read_hd5(filestem, ens_id, group, attrs=None, idl=None, part="real"):
 58    r'''Read hadrons hdf5 file and extract entry based on attributes.
 59
 60    Parameters
 61    -----------------
 62    filestem : str
 63        Full namestem of the files to read, including the full path.
 64    ens_id : str
 65        name of the ensemble, required for internal bookkeeping
 66    group : str
 67        label of the group to be extracted.
 68    attrs : dict or int
 69        Dictionary containing the attributes. For example
 70        ```python
 71        attrs = {"gamma_snk": "Gamma5",
 72                 "gamma_src": "Gamma5"}
 73         ```
 74        Alternatively an integer can be specified to identify the sub group.
 75        This is discouraged as the order in the file is not guaranteed.
 76    idl : range
 77        If specified only configurations in the given range are read in.
 78    part: str
 79        string specifying whether to extract the real part ('real'),
 80        the imaginary part ('imag') or a complex correlator ('complex').
 81        Default 'real'.
 82
 83    Returns
 84    -------
 85    corr : Corr
 86        Correlator of the source sink combination in question.
 87    '''
 88
 89    path_obj = Path(filestem)
 90    path = path_obj.parent.as_posix()
 91    filestem = path_obj.stem
 92
 93    files, idx = _get_files(path, filestem, idl)
 94
 95    if isinstance(attrs, dict):
 96        h5file = h5py.File(path + '/' + files[0], "r")
 97        entry = None
 98        for key in h5file[group].keys():
 99            if attrs.items() <= {k: v[0].decode() for k, v in h5file[group][key].attrs.items()}.items():
100                if entry is None:
101                    entry = key
102                else:
103                    raise ValueError("More than one fitting entry found. More constraint on attributes needed.")
104        h5file.close()
105        if entry is None:
106            raise ValueError(f"Entry with attributes {attrs} not found.")
107    elif isinstance(attrs, int):
108        entry = group + f"_{attrs}"
109    else:
110        raise TypeError("Invalid type for 'attrs'. Needs to be dict or int.")
111
112    corr_data = []
113    infos = []
114    for hd5_file in files:
115        h5file = h5py.File(path + '/' + hd5_file, "r")
116        if not group + '/' + entry in h5file:
117            raise Exception("Entry '" + entry + "' not contained in the files.")
118        raw_data = h5file[group + '/' + entry + '/corr']
119        real_data = raw_data[:].view("complex")
120        corr_data.append(real_data)
121        if not infos:
122            for k, i in h5file[group + '/' + entry].attrs.items():
123                infos.append(k + ': ' + i[0].decode())
124        h5file.close()
125    corr_data = np.array(corr_data)
126
127    if part == "complex":
128        l_obs = []
129        for c in corr_data.T:
130            l_obs.append(CObs(Obs([c.real], [ens_id], idl=[idx]),
131                              Obs([c.imag], [ens_id], idl=[idx])))
132    else:
133        corr_data = getattr(corr_data, part)
134        l_obs = []
135        for c in corr_data.T:
136            l_obs.append(Obs([c], [ens_id], idl=[idx]))
137
138    corr = Corr(l_obs)
139    corr.tag = r", ".join(infos)
140    return corr
141
142
143def read_meson_hd5(path, filestem, ens_id, meson='meson_0', idl=None, gammas=None):
144    r'''Read hadrons meson hdf5 file and extract the meson labeled 'meson'
145
146    Parameters
147    -----------------
148    path : str
149        path to the files to read
150    filestem : str
151        namestem of the files to read
152    ens_id : str
153        name of the ensemble, required for internal bookkeeping
154    meson : str
155        label of the meson to be extracted, standard value meson_0 which
156        corresponds to the pseudoscalar pseudoscalar two-point function.
157    gammas : tuple of strings
158        Instrad of a meson label one can also provide a tuple of two strings
159        indicating the gamma matrices at sink and source (gamma_snk, gamma_src).
160        ("Gamma5", "Gamma5") corresponds to the pseudoscalar pseudoscalar
161        two-point function. The gammas argument dominateds over meson.
162    idl : range
163        If specified only configurations in the given range are read in.
164
165    Returns
166    -------
167    corr : Corr
168        Correlator of the source sink combination in question.
169    '''
170    if gammas is None:
171        attrs = int(meson.rsplit('_', 1)[-1])
172    else:
173        if len(gammas) != 2:
174            raise ValueError("'gammas' needs to have exactly two entries")
175        attrs = {"gamma_snk": gammas[0],
176                 "gamma_src": gammas[1]}
177    return read_hd5(filestem=path + "/" + filestem, ens_id=ens_id,
178                    group=meson.rsplit('_', 1)[0], attrs=attrs, idl=idl,
179                    part="real")
180
181
182def _extract_real_arrays(path, files, tree, keys):
183    corr_data = {}
184    for key in keys:
185        corr_data[key] = []
186    for hd5_file in files:
187        h5file = h5py.File(path + '/' + hd5_file, "r")
188        for key in keys:
189            if not tree + '/' + key in h5file:
190                raise Exception("Entry '" + key + "' not contained in the files.")
191            raw_data = h5file[tree + '/' + key + '/data']
192            real_data = raw_data[:].astype(np.double)
193            corr_data[key].append(real_data)
194        h5file.close()
195    for key in keys:
196        corr_data[key] = np.array(corr_data[key])
197    return corr_data
198
199
200def extract_t0_hd5(path, filestem, ens_id, obs='Clover energy density', fit_range=5, idl=None, **kwargs):
201    r'''Read hadrons FlowObservables hdf5 file and extract t0
202
203    Parameters
204    -----------------
205    path : str
206        path to the files to read
207    filestem : str
208        namestem of the files to read
209    ens_id : str
210        name of the ensemble, required for internal bookkeeping
211    obs : str
212        label of the observable from which t0 should be extracted.
213        Options: 'Clover energy density' and 'Plaquette energy density'
214    fit_range : int
215        Number of data points left and right of the zero
216        crossing to be included in the linear fit. (Default: 5)
217    idl : range
218        If specified only configurations in the given range are read in.
219    plot_fit : bool
220        If true, the fit for the extraction of t0 is shown together with the data.
221    '''
222
223    files, idx = _get_files(path, filestem, idl)
224    tree = "FlowObservables"
225
226    h5file = h5py.File(path + '/' + files[0], "r")
227    obs_key = None
228    for key in h5file[tree].keys():
229        if obs == h5file[tree][key].attrs["description"][0].decode():
230            obs_key = key
231            break
232    h5file.close()
233    if obs_key is None:
234        raise Exception(f"Observable {obs} not found.")
235
236    corr_data = _extract_real_arrays(path, files, tree, ["FlowObservables_0", obs_key])
237
238    if not np.allclose(corr_data["FlowObservables_0"][0], corr_data["FlowObservables_0"][:]):
239        raise Exception("Not all flow times were equal.")
240
241    t2E_dict = {}
242    for t2, dat in zip(corr_data["FlowObservables_0"][0], corr_data[obs_key].T):
243        t2E_dict[t2] = Obs([dat], [ens_id], idl=[idx]) - 0.3
244
245    return fit_t0(t2E_dict, fit_range, plot_fit=kwargs.get('plot_fit'))
246
247
248def read_DistillationContraction_hd5(path, ens_id, diagrams=["direct"], idl=None):
249    """Read hadrons DistillationContraction hdf5 files in given directory structure
250
251    Parameters
252    -----------------
253    path : str
254        path to the directories to read
255    ens_id : str
256        name of the ensemble, required for internal bookkeeping
257    diagrams : list
258        List of strings of the diagrams to extract, e.g. ["direct", "box", "cross"].
259    idl : range
260        If specified only configurations in the given range are read in.
261
262    Returns
263    -------
264    result : dict
265        extracted DistillationContration data
266    """
267
268    res_dict = {}
269
270    directories, idx = _get_files(path, "data", idl)
271
272    explore_path = Path(path + "/" + directories[0])
273
274    for explore_file in explore_path.iterdir():
275        if explore_file.is_file():
276            stem = explore_file.with_suffix("").with_suffix("").as_posix().split("/")[-1]
277        else:
278            continue
279
280        file_list = []
281        for dir in directories:
282            tmp_path = Path(path + "/" + dir)
283            file_list.append((tmp_path / stem).as_posix() + tmp_path.suffix + ".h5")
284
285        corr_data = {}
286
287        for diagram in diagrams:
288            corr_data[diagram] = []
289
290        try:
291            for n_file, (hd5_file, n_traj) in enumerate(zip(file_list, list(idx))):
292                h5file = h5py.File(hd5_file)
293
294                if n_file == 0:
295                    if h5file["DistillationContraction/Metadata"].attrs.get("TimeSources")[0].decode() != "0...":
296                        raise Exception("Routine is only implemented for files containing inversions on all timeslices.")
297
298                    Nt = h5file["DistillationContraction/Metadata"].attrs.get("Nt")[0]
299
300                    identifier = []
301                    for in_file in range(len(h5file["DistillationContraction/Metadata/DmfInputFiles"].attrs.keys()) - 1):
302                        encoded_info = h5file["DistillationContraction/Metadata/DmfInputFiles"].attrs.get("DmfInputFiles_" + str(in_file))
303                        full_info = encoded_info[0].decode().split("/")[-1].replace(".h5", "").split("_")
304                        my_tuple = (full_info[0], full_info[1][1:], full_info[2], full_info[3])
305                        identifier.append(my_tuple)
306                    identifier = tuple(identifier)
307                    # "DistillationContraction/Metadata/DmfSuffix" contains info about different quarks, irrelevant in the SU(3) case.
308
309                for diagram in diagrams:
310
311                    if diagram == "triangle" and "Identity" not in str(identifier):
312                        part = "im"
313                    else:
314                        part = "re"
315
316                    real_data = np.zeros(Nt)
317                    for x0 in range(Nt):
318                        raw_data = h5file["DistillationContraction/Correlators/" + diagram + "/" + str(x0)][:][part].astype(np.double)
319                        real_data += np.roll(raw_data, -x0)
320                    real_data /= Nt
321
322                    corr_data[diagram].append(real_data)
323                h5file.close()
324
325            res_dict[str(identifier)] = {}
326
327            for diagram in diagrams:
328
329                tmp_data = np.array(corr_data[diagram])
330
331                l_obs = []
332                for c in tmp_data.T:
333                    l_obs.append(Obs([c], [ens_id], idl=[idx]))
334
335                corr = Corr(l_obs)
336                corr.tag = str(identifier)
337
338                res_dict[str(identifier)][diagram] = corr
339        except FileNotFoundError:
340            print("Skip", stem)
341
342    return res_dict
343
344
345class Npr_matrix(np.ndarray):
346
347    def __new__(cls, input_array, mom_in=None, mom_out=None):
348        obj = np.asarray(input_array).view(cls)
349        obj.mom_in = mom_in
350        obj.mom_out = mom_out
351        return obj
352
353    @property
354    def g5H(self):
355        """Gamma_5 hermitean conjugate
356
357        Uses the fact that the propagator is gamma5 hermitean, so just the
358        in and out momenta of the propagator are exchanged.
359        """
360        return Npr_matrix(self,
361                          mom_in=self.mom_out,
362                          mom_out=self.mom_in)
363
364    def _propagate_mom(self, other, name):
365        s_mom = getattr(self, name, None)
366        o_mom = getattr(other, name, None)
367        if s_mom is not None and o_mom is not None:
368            if not np.allclose(s_mom, o_mom):
369                raise Exception(name + ' does not match.')
370        return o_mom if o_mom is not None else s_mom
371
372    def __matmul__(self, other):
373        return self.__new__(Npr_matrix,
374                            super().__matmul__(other),
375                            self._propagate_mom(other, 'mom_in'),
376                            self._propagate_mom(other, 'mom_out'))
377
378    def __array_finalize__(self, obj):
379        if obj is None:
380            return
381        self.mom_in = getattr(obj, 'mom_in', None)
382        self.mom_out = getattr(obj, 'mom_out', None)
383
384
385def read_ExternalLeg_hd5(path, filestem, ens_id, idl=None):
386    """Read hadrons ExternalLeg hdf5 file and output an array of CObs
387
388    Parameters
389    ----------
390    path : str
391        path to the files to read
392    filestem : str
393        namestem of the files to read
394    ens_id : str
395        name of the ensemble, required for internal bookkeeping
396    idl : range
397        If specified only configurations in the given range are read in.
398
399    Returns
400    -------
401    result : Npr_matrix
402        read Cobs-matrix
403    """
404
405    files, idx = _get_files(path, filestem, idl)
406
407    mom = None
408
409    corr_data = []
410    for hd5_file in files:
411        file = h5py.File(path + '/' + hd5_file, "r")
412        raw_data = file['ExternalLeg/corr'][0][0].view('complex')
413        corr_data.append(raw_data)
414        if mom is None:
415            mom = np.array(str(file['ExternalLeg/info'].attrs['pIn'])[3:-2].strip().split(), dtype=float)
416        file.close()
417    corr_data = np.array(corr_data)
418
419    rolled_array = np.rollaxis(corr_data, 0, 5)
420
421    matrix = np.empty((rolled_array.shape[:-1]), dtype=object)
422    for si, sj, ci, cj in np.ndindex(rolled_array.shape[:-1]):
423        real = Obs([rolled_array[si, sj, ci, cj].real], [ens_id], idl=[idx])
424        imag = Obs([rolled_array[si, sj, ci, cj].imag], [ens_id], idl=[idx])
425        matrix[si, sj, ci, cj] = CObs(real, imag)
426
427    return Npr_matrix(matrix, mom_in=mom)
428
429
430def read_Bilinear_hd5(path, filestem, ens_id, idl=None):
431    """Read hadrons Bilinear hdf5 file and output an array of CObs
432
433    Parameters
434    ----------
435    path : str
436        path to the files to read
437    filestem : str
438        namestem of the files to read
439    ens_id : str
440        name of the ensemble, required for internal bookkeeping
441    idl : range
442        If specified only configurations in the given range are read in.
443
444    Returns
445    -------
446    result_dict: dict[Npr_matrix]
447        extracted Bilinears
448    """
449
450    files, idx = _get_files(path, filestem, idl)
451
452    mom_in = None
453    mom_out = None
454
455    corr_data = {}
456    for hd5_file in files:
457        file = h5py.File(path + '/' + hd5_file, "r")
458        for i in range(16):
459            name = file['Bilinear/Bilinear_' + str(i) + '/info'].attrs['gamma'][0].decode('UTF-8')
460            if name not in corr_data:
461                corr_data[name] = []
462            raw_data = file['Bilinear/Bilinear_' + str(i) + '/corr'][0][0].view('complex')
463            corr_data[name].append(raw_data)
464            if mom_in is None:
465                mom_in = np.array(str(file['Bilinear/Bilinear_' + str(i) + '/info'].attrs['pIn'])[3:-2].strip().split(), dtype=float)
466            if mom_out is None:
467                mom_out = np.array(str(file['Bilinear/Bilinear_' + str(i) + '/info'].attrs['pOut'])[3:-2].strip().split(), dtype=float)
468
469        file.close()
470
471    result_dict = {}
472
473    for key, data in corr_data.items():
474        local_data = np.array(data)
475
476        rolled_array = np.rollaxis(local_data, 0, 5)
477
478        matrix = np.empty((rolled_array.shape[:-1]), dtype=object)
479        for si, sj, ci, cj in np.ndindex(rolled_array.shape[:-1]):
480            real = Obs([rolled_array[si, sj, ci, cj].real], [ens_id], idl=[idx])
481            imag = Obs([rolled_array[si, sj, ci, cj].imag], [ens_id], idl=[idx])
482            matrix[si, sj, ci, cj] = CObs(real, imag)
483
484        result_dict[key] = Npr_matrix(matrix, mom_in=mom_in, mom_out=mom_out)
485
486    return result_dict
487
488
489def read_Fourquark_hd5(path, filestem, ens_id, idl=None, vertices=["VA", "AV"]):
490    """Read hadrons FourquarkFullyConnected hdf5 file and output an array of CObs
491
492    Parameters
493    ----------
494    path : str
495        path to the files to read
496    filestem : str
497        namestem of the files to read
498    ens_id : str
499        name of the ensemble, required for internal bookkeeping
500    idl : range
501        If specified only configurations in the given range are read in.
502    vertices : list
503        Vertex functions to be extracted.
504
505    Returns
506    -------
507    result_dict : dict
508        extracted fourquark matrizes
509    """
510
511    files, idx = _get_files(path, filestem, idl)
512
513    mom_in = None
514    mom_out = None
515
516    vertex_names = []
517    for vertex in vertices:
518        vertex_names += _get_lorentz_names(vertex)
519
520    corr_data = {}
521
522    tree = 'FourQuarkFullyConnected/FourQuarkFullyConnected_'
523
524    for hd5_file in files:
525        file = h5py.File(path + '/' + hd5_file, "r")
526
527        for i in range(32):
528            name = (file[tree + str(i) + '/info'].attrs['gammaA'][0].decode('UTF-8'), file[tree + str(i) + '/info'].attrs['gammaB'][0].decode('UTF-8'))
529            if name in vertex_names:
530                if name not in corr_data:
531                    corr_data[name] = []
532                raw_data = file[tree + str(i) + '/corr'][0][0].view('complex')
533                corr_data[name].append(raw_data)
534                if mom_in is None:
535                    mom_in = np.array(str(file[tree + str(i) + '/info'].attrs['pIn'])[3:-2].strip().split(), dtype=float)
536                if mom_out is None:
537                    mom_out = np.array(str(file[tree + str(i) + '/info'].attrs['pOut'])[3:-2].strip().split(), dtype=float)
538
539        file.close()
540
541    intermediate_dict = {}
542
543    for vertex in vertices:
544        lorentz_names = _get_lorentz_names(vertex)
545        for v_name in lorentz_names:
546            if v_name in [('SigmaXY', 'SigmaZT'),
547                          ('SigmaXT', 'SigmaYZ'),
548                          ('SigmaYZ', 'SigmaXT'),
549                          ('SigmaZT', 'SigmaXY')]:
550                sign = -1
551            else:
552                sign = 1
553            if vertex not in intermediate_dict:
554                intermediate_dict[vertex] = sign * np.array(corr_data[v_name])
555            else:
556                intermediate_dict[vertex] += sign * np.array(corr_data[v_name])
557
558    result_dict = {}
559
560    for key, data in intermediate_dict.items():
561
562        rolled_array = np.moveaxis(data, 0, 8)
563
564        matrix = np.empty((rolled_array.shape[:-1]), dtype=object)
565        for index in np.ndindex(rolled_array.shape[:-1]):
566            real = Obs([rolled_array[index].real], [ens_id], idl=[idx])
567            imag = Obs([rolled_array[index].imag], [ens_id], idl=[idx])
568            matrix[index] = CObs(real, imag)
569
570        result_dict[key] = Npr_matrix(matrix, mom_in=mom_in, mom_out=mom_out)
571
572    return result_dict
573
574
575def _get_lorentz_names(name):
576    lorentz_index = ['X', 'Y', 'Z', 'T']
577
578    res = []
579
580    if name == "TT":
581        for i in range(4):
582            for j in range(i + 1, 4):
583                res.append(("Sigma" + lorentz_index[i] + lorentz_index[j], "Sigma" + lorentz_index[i] + lorentz_index[j]))
584        return res
585
586    if name == "TTtilde":
587        for i in range(4):
588            for j in range(i + 1, 4):
589                for k in range(4):
590                    for o in range(k + 1, 4):
591                        fac = epsilon_tensor_rank4(i, j, k, o)
592                        if not np.isclose(fac, 0.0):
593                            res.append(("Sigma" + lorentz_index[i] + lorentz_index[j], "Sigma" + lorentz_index[k] + lorentz_index[o]))
594        return res
595
596    assert len(name) == 2
597
598    if 'S' in name or 'P' in name:
599        if not set(name) <= set(['S', 'P']):
600            raise Exception("'" + name + "' is not a Lorentz scalar")
601
602        g_names = {'S': 'Identity',
603                   'P': 'Gamma5'}
604
605        res.append((g_names[name[0]], g_names[name[1]]))
606
607    else:
608        if not set(name) <= set(['V', 'A']):
609            raise Exception("'" + name + "' is not a Lorentz scalar")
610
611        for ind in lorentz_index:
612            res.append(('Gamma' + ind + (name[0] == 'A') * 'Gamma5',
613                        'Gamma' + ind + (name[1] == 'A') * 'Gamma5'))
614
615    return res