A script to collect and plot the results of a forcebalance water optimisation.

water_analysis.py

"""
Analyses the optimize out file from a forcebalance run.

This is quite specific for the results of a water optimzation to multiple liquid property targets.
This will produce a csv file with all of the results per iteration for the properties and target parameters.
Graphs will also be produced for each iteration showing the csv data.
"""
import matplotlib.pyplot as plt
from typing import List, Dict, Tuple
import pandas as pd
import numpy as np
from collections import defaultdict
import os
import re
import click


def parse_block(lines: List[str]) -> Tuple[str, Dict[str, np.array]]:
    """
    Parse the block of lines into a dictionary of results.

    Args:
        lines: The list of lines which should be parsed.

    Returns:
        A dictionary of the properties extracted from the results.
    """
    data = defaultdict(list)
    prop_name = lines[0].strip()
    ansi_escape = re.compile(r'\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])')
    prop_name = ansi_escape.sub("", prop_name)
    prop_name = prop_name.split()[1:-1]
    prop_name = " ".join(prop_name)
    for line in lines[3:]:
        ls = line.split()
        data["Temperature K"].append(float(ls[0]))
        data["Pressure"].append(float(ls[1]))
        data["Pressure Unit"].append(ls[2])
        data["Reference"].append(float(ls[3]))
        data["Result"].append(float(ls[4]))
        data["Stdev"].append(float(ls[6]))

    return prop_name, data


def plot_block(prop_name: str, file_name: str, data: Dict[str, np.array], model_name: str):
    """

    :param prop_name:
    :param data:
    :return:
    """
    # only grab the standard pressure
    for i, p in enumerate(data["Pressure"]):
        if p != 1:
            break
    plt.plot(data["Temperature K"][:i], data["Reference"][:i], label="EXP")
    plt.errorbar(data["Temperature K"][:i], data["Result"][:i], yerr=data["Stdev"][:i], label=model_name)
    plt.title(prop_name)
    plt.xlabel("Temperature K")
    plt.ylabel(prop_name)
    plt.legend()
    plt.savefig(file_name)
    plt.close()


def read_parameters(lines: List[str]) -> Dict[str, float]:
    """
    Read the parameters from a parameter block.
    :param lines:
    :return:
    """
    return {}


def read_objective():
    pass


@click.command()
@click.argument("file_path", type=click.Path(exists=True))
@click.argument("model_name", type=click.STRING)
def main(file_path: str , model_name: str):
    """
    The main function to read the file and make the csv and plots.
    """
    available_properties = ["Liquid Density (kg m^-3)", "Liquid Enthalpy of Vaporization (kJ mol^-1)", "Liquid Thermal Expansion Coefficient (10^-4 K^-1) ",
                            "Liquid Isothermal Compressibility (10^-6 bar^-1)", "Liquid Isobaric Heat Capacity (cal mol^-1 K^-1)",
                            "Liquid Dielectric Constant "]

    # open the file
    with open(file_path) as f:
        lines = f.readlines()
        start_line = 0
        reading = False
        iteration = -1
        dataframes = []
        for i, line in enumerate(lines):
            # sometimes the object rises and so it is re-evaluated at the last point.
            if "Iteration" in line or "Objective function rises!" in line:
                iteration += 1
                # reset
                dataframes = []
                print("reading data for iteration", line.split()[3])
                os.makedirs(f"iteration_{iteration}", exist_ok=True)
                
            for prop in available_properties:
                if prop in line:
                    start_line = i
                    reading = True
            if "--------------------------------------------------------------------------------------------" in line and reading:
                property_name, data = parse_block(lines[start_line:i - 1])
                df = pd.DataFrame(data)
                df["Property"] = [property_name for _ in df.index]
                print(df)
                dataframes.append(df)
                plot_block(property_name, os.path.join(f"iteration_{iteration}", f"{property_name}.pdf"), data, model_name)
                print(property_name, data)
                start_line = 0
                reading = False
                if len(dataframes) == len(available_properties):
                    # write out the csv
                    master = pd.concat(dataframes)
                    master.to_csv(os.path.join(f"iteration_{iteration}", "properties.csv"))


if __name__ == "__main__":
    main()