diff --git a/mdevaluate/__init__.py b/mdevaluate/__init__.py index dd41f92..ef59f43 100644 --- a/mdevaluate/__init__.py +++ b/mdevaluate/__init__.py @@ -73,3 +73,12 @@ def open(directory='', topology='*.tpr', trajectory='*.xtc', cached=False, except reader.NojumpError: reader.generate_nojump_matrixes(coords) return coords + +def open_energy(file, energies=None): + """Reads an gromacs energy file and output the data in a pandas DataFrame. + Args: + file: Filename of the energy file + energies (opt.): Specify energies to extract from the energy file + """ + df = reader.energy_reader(file, energies=energies) + return df diff --git a/mdevaluate/reader.py b/mdevaluate/reader.py index 0cea171..e07971a 100755 --- a/mdevaluate/reader.py +++ b/mdevaluate/reader.py @@ -181,6 +181,44 @@ def correct_nojump_matrixes_for_whole(trajectory): for d in range(3): reader.nojump_matrixes[d][0] = cor[:, d] save_nojump_matrixes(reader) + + +def energy_reader(file, energies=None): + read_energy(file, energies=None): + """Reads an gromacs energy file and output the data in a pandas DataFrame. + Args: + file: Filename of the energy file + energies (opt.): Specify energies to extract from the energy file + """ + if energies is None: + energies = np.arange(1, 100).astype('str') + directory = file.rsplit("/", 1)[0] + ps = subprocess.Popen(("echo", *energies), stdout=subprocess.PIPE) + try: + subprocess.run(("gmx", "energy", "-f", file, "-o", f"{directory}/tmp.xvg", "-quiet"), stdin=ps.stdout) + except FileNotFoundError: + print("No GROMACS found!") + ps.wait() + labels = [] + is_legend = False + with open(f"{directory}/tmp.xvg") as f: + for i, line in enumerate(f): + if line.split(" ")[0] == "@": + if re.search("s\d+", line.split()[1]): + is_legend = True + labels.append(line.split('"')[1]) + elif is_legend: + header = i + break + + data = np.genfromtxt(f"{directory}/tmp.xvg", skip_header=header) + + df = pd.DataFrame({"Time":data[:,0]}) + for i, label in enumerate(labels): + tmp_df = pd.DataFrame({label:data[:,i+1]}) + df = pd.concat([df, tmp_df], axis=1) + subprocess.run(("rm", f"{directory}/tmp.xvg")) + return df class BaseReader: