Merge pull request #133 from jkuhl-uni/feat/read_ms5_xsf

first version of skript to read xsf data

pyerrors/input/openQCD.py

@@ -8,6 +8,8 @@ import matplotlib.pyplot as plt
 from matplotlib import gridspec
 from ..obs import Obs
 from ..fits import fit_lin
+from ..obs import CObs
+from ..correlators import Corr
 
 
 def read_rwms(path, prefix, version='2.0', names=None, **kwargs):
@@ -985,3 +987,147 @@ def read_qtop_sector(path, prefix, c, target=0, **kwargs):
     qtop = read_qtop(path, prefix, c, **kwargs)
 
     return qtop_projection(qtop, target=target)
+
+
+def read_ms5_xsf(path, prefix, qc, corr, sep="r", **kwargs):
+    """
+    Read data from files in the specified directory with the specified prefix and quark combination extension, and return a `Corr` object containing the data.
+
+    Parameters
+    ----------
+    path : str
+        The directory to search for the files in.
+    prefix : str
+        The prefix to match the files against.
+    qc : str
+        The quark combination extension to match the files against.
+    corr : str
+        The correlator to extract data for.
+    sep : str, optional
+        The separator to use when parsing the replika names.
+    **kwargs
+        Additional keyword arguments. The following keyword arguments are recognized:
+
+        - names (List[str]): A list of names to use for the replicas.
+
+    Returns
+    -------
+    Corr
+        A complex valued `Corr` object containing the data read from the files. In case of boudary to bulk correlators.
+    or
+    CObs
+        A complex valued `CObs` object containing the data read from the files. In case of boudary to boundary correlators.
+
+
+    Raises
+    ------
+    FileNotFoundError
+        If no files matching the specified prefix and quark combination extension are found in the specified directory.
+    IOError
+        If there is an error reading a file.
+    struct.error
+        If there is an error unpacking binary data.
+    """
+
+    found = []
+    files = []
+    names = []
+    for (dirpath, dirnames, filenames) in os.walk(path + "/"):
+        found.extend(filenames)
+        break
+
+    for f in found:
+        if fnmatch.fnmatch(f, prefix + "*.ms5_xsf_" + qc + ".dat"):
+            files.append(f)
+            if not sep == "":
+                names.append(prefix + "|r" + f.split(".")[0].split(sep)[1])
+            else:
+                names.append(prefix)
+    files = sorted(files)
+
+    if "names" in kwargs:
+        names = kwargs.get("names")
+    else:
+        names = sorted(names)
+
+    cnfgs = []
+    realsamples = []
+    imagsamples = []
+    repnum = 0
+    for file in files:
+        with open(path + "/" + file, "rb") as fp:
+
+            t = fp.read(8)
+            kappa = struct.unpack('d', t)[0]
+            t = fp.read(8)
+            csw = struct.unpack('d', t)[0]
+            t = fp.read(8)
+            dF = struct.unpack('d', t)[0]
+            t = fp.read(8)
+            zF = struct.unpack('d', t)[0]
+
+            t = fp.read(4)
+            tmax = struct.unpack('i', t)[0]
+            t = fp.read(4)
+            bnd = struct.unpack('i', t)[0]
+
+            placesBI = ["gS", "gP",
+                        "gA", "gV",
+                        "gVt", "lA",
+                        "lV", "lVt",
+                        "lT", "lTt"]
+            placesBB = ["g1", "l1"]
+
+            # the chunks have the following structure:
+            # confignumber, 10x timedependent complex correlators as doubles, 2x timeindependent complex correlators as doubles
+
+            chunksize = 4 + (8 * 2 * tmax * 10) + (8 * 2 * 2)
+            packstr = '=i' + ('d' * 2 * tmax * 10) + ('d' * 2 * 2)
+            cnfgs.append([])
+            realsamples.append([])
+            imagsamples.append([])
+            for t in range(tmax):
+                realsamples[repnum].append([])
+                imagsamples[repnum].append([])
+
+            while True:
+                cnfgt = fp.read(chunksize)
+                if not cnfgt:
+                    break
+                asascii = struct.unpack(packstr, cnfgt)
+                cnfg = asascii[0]
+                cnfgs[repnum].append(cnfg)
+
+                if corr not in placesBB:
+                    tmpcorr = asascii[1 + 2 * tmax * placesBI.index(corr):1 + 2 * tmax * placesBI.index(corr) + 2 * tmax]
+                else:
+                    tmpcorr = asascii[1 + 2 * tmax * len(placesBI) + 2 * placesBB.index(corr):1 + 2 * tmax * len(placesBI) + 2 * placesBB.index(corr) + 2]
+
+                corrres = [[], []]
+                for i in range(len(tmpcorr)):
+                    corrres[i % 2].append(tmpcorr[i])
+                for t in range(int(len(tmpcorr) / 2)):
+                    realsamples[repnum][t].append(corrres[0][t])
+                for t in range(int(len(tmpcorr) / 2)):
+                    imagsamples[repnum][t].append(corrres[1][t])
+        repnum += 1
+
+    s = "Read correlator " + corr + " from " + str(repnum) + " replika with " + str(len(realsamples[0][t]))
+    for rep in range(1, repnum):
+        s += ", " + str(len(realsamples[rep][t]))
+    s += " samples"
+    print(s)
+    print("Asserted run parameters:\n T:", tmax, "kappa:", kappa, "csw:", csw, "dF:", dF, "zF:", zF, "bnd:", bnd)
+
+    # we have the data now... but we need to re format the whole thing and put it into Corr objects.
+
+    compObs = []
+
+    for t in range(int(len(tmpcorr) / 2)):
+        compObs.append(CObs(Obs([realsamples[rep][t] for rep in range(repnum)], names=names, idl=cnfgs),
+                            Obs([imagsamples[rep][t] for rep in range(repnum)], names=names, idl=cnfgs)))
+
+    if len(compObs) == 1:
+        return compObs[0]
+    else:
+        return Corr(compObs)