pyerrors.input.hadrons

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