qc_parser.py

"""
qc_parser.py

Parsers for ORCA input and output files.

This module provides functions to parse and validate ORCA calculations,
including atom coordinates and multiplicities.
"""

from pprint import pprint
import itertools as it
import numpy as np
from typing import Any, Callable
import os
import datetime
from abc import ABC, abstractmethod
from math import log10

### FILE PARSER FUNCTIONS ###
def read_coords(coords_str: str) -> list[dict[str, str|float]]:
    coords_list = coords_str.split()
    assert len(coords_list) % 4 == 0
    out = []
    new_coords: dict[str, str|float] = {}
    for val,key in zip(coords_list, it.cycle(["element", "x", "y", "z"])):
        if key == "element":
            new_coords[key] = val
        else:
            new_coords[key] = float(val)
        if key == "z":
            out.append(new_coords)
            new_coords = {}
    return out

def read_xyz_file(file: str) -> list[dict[str, str|float]]:
    with open(file, "r") as f:
        lines = f.readlines()
    if len(lines[0].split()) == 1:
        lines = lines[2:]
    return read_coords(" ".join(lines))

def read_input_file(file: str) -> dict[str, Any]:
    assert ".inp" in file
    filename = file.split("/")[-1]
    relpath = file[:file.find(filename)]
    asterisk_encountered = False
    coords: list[str] = []
    out: dict[str, Any] = {}
    out["keywords"] = []
    out["path"] = relpath
    out["filename"] = filename.replace(".inp", "")
    with open(file, "r") as f:
        for line in f:
            line = line.strip()
            if len(line) == 0:
                continue
            elif line[0] == "!":
                out["keywords"] += line[1:].lower().split()
            elif line[0] == "*" and not asterisk_encountered:
                line = line[1:].split()
                out["charge"] = int(line[1])
                out["multiplicity"] = int(line[2])
                asterisk_encountered = True
    out["method"] = "unknown"
    for method in ["rhf", "uhf", "hf", "rohf", "rks", "uks", "ks"]:
        if True in [True for x in out["keywords"] if method == x.lower()]:
            out["method"] = method
    
    return out

def read_output_file(file: str) -> dict[str, Any]:
    assert ".out" in file
    filename = file.split("/")[-1]
    relpath = file[:file.find(filename)]
    reading_coords = False
    coords: list[str] = []
    out: dict[str, Any] = {}
    out["coordinates"] = []
    out["path"] = relpath
    out["filename"] = filename.replace(".out", "")
    ### SET DEFAULTS ###
    out["no_error"] = False
    out["opt_converged"] = False
    ### READ OUTPUT FILE ###
    with open(file, "r") as f:
        for line in f:
            if "HURRAY" in line:
                out["opt_converged"] = True
            elif "ORCA TERMINATED NORMALLY" in line:
                out["no_error"] = True
            # This line overwrites the scf-energy several times in geometry optimizations
            # Still, the last one is the correct one
            elif "FINAL SINGLE POINT ENERGY" in line:
                out["scf_energy"] = float(line.split()[-1])
            elif "Total Enthalpy" in line:
                out["enthalpy"] = float(line.split()[-2])
            elif "VIBRATIONAL FREQUENCIES" in line:
                frequencies = []
                for i in range(5):
                    line = next(f)
                while len(line.split()) > 0 and line.split()[2] == "cm**-1":
                    frequencies.append(float(line.split()[1]))
                    line = next(f)
                out["frequencies"] = frequencies
                if any(x < 0 for x in frequencies):
                    out["negative_frequencies"] = True
                else:
                    out["negative_frequencies"] = False
            # Reading coordinates
            elif "CARTESIAN COORDINATES (ANGSTROEM)" in line:
                line = next(f)
                reading_coords = True
            elif line.strip() == "" and reading_coords:
                reading_coords = False
                out["coordinates"].append(read_coords(" ".join(coords)))
                coords = []
            elif reading_coords:
                coords.append(line)
            # Reading basis set information
            elif "Number of basis functions" in line:
                out["n_basis"] = int(line.split()[-1])
            elif "BASIS SET IN INPUT FORMAT" in line:
                out["basis"] = {}
                element = ""
                line = next(f)
                line = next(f)
                while "------------------" not in line:
                    if "NewGTO" in line:
                        primitive_dict = {"S": 0, "P": 0, "D": 0, "F": 0, "G": 0, "H": 0, "I": 0}
                        element = line.split()[-1]
                        out["basis"][element] = {}
                    elif "end;" in line:
                        element = ""
                    elif element != "":
                        angular_momentum = line.split()[0]
                        if not out["basis"][element].get(angular_momentum, False):
                            out["basis"][element][angular_momentum] = {}
                        primitives = int(line.split()[1])
                        primitive_dict[angular_momentum] += 1
                        out["basis"][element][angular_momentum][primitive_dict[angular_momentum]] = {
                            "exponents" : [],
                            "coefficients" : []
                        }
                        for i in range(primitives):
                            line = next(f)
                            spl = line.split()
                            out["basis"][element][angular_momentum][primitive_dict[angular_momentum]]["exponents"].append(float(spl[1]))
                            out["basis"][element][angular_momentum][primitive_dict[angular_momentum]]["coefficients"].append(float(spl[2]))
                    line = next(f)
            elif "Number of Electrons" in line:
                out["n_electrons"] = int(line.split()[-1])
            elif "MOLECULAR ORBITALS" in line:
                out["mos"] = {}
                line = next(f)
                atoms: list[int] = []
                orbitals: list[str] = []
                for spin in ("alpha", "beta"):
                    line = next(f)
                    blocks = []
                    energies = []
                    colno = 0
                    electrons = 0.0
                    while colno < out["n_basis"]:
                        line = next(f)
                        energies += [float(x) for x in line.split()]
                        line = next(f)
                        electrons += sum([float(x) for x in line.split()])
                        line = next(f)
                        block = []
                        for i in range(out["n_basis"]):
                            line = next(f)
                            l = line.split()
                            atom = "".join([x for x in l[0] if x.isdigit()])
                            if len(atoms) < out["n_basis"]:
                                atoms.append(int(atom))
                            orbital = l[1]
                            if len(orbitals) < out["n_basis"]:
                                orbitals.append(orbital)
                            width = len(l) - 2
                            block.append([float(x) for x in l[2:]])
                        blocks.append(np.array(block))
                        colno += width
                        line = next(f)
                    out["mos"][spin] = {}
                    out["mos"][spin]["energies"] = energies
                    out["mos"][spin]["coefficients"] = np.concatenate(blocks, axis=1).T
                    if electrons == out["n_electrons"]:
                        out["mos"]["beta"] = {}
                        out["mos"]["beta"]["energies"] = energies
                        out["mos"]["beta"]["coefficients"] = np.concatenate(blocks, axis=1).T
                        break
                out["mos"]["atoms"] = atoms
                out["mos"]["orbitals"] = orbitals
            # Density is read twice for geometry optimizations;
            # the optimized one is kept, so no problem here
            elif "DENSITY" == line.strip():
                line = next(f)
                blocks = []
                colno = 0
                while colno < out["n_basis"]:
                    line = next(f)
                    block = []
                    for i in range(out["n_basis"]):
                        line = next(f)
                        width = len(line.split()) - 1
                        block.append([float(x) for x in line.split()[1:]])
                    blocks.append(np.array(block))
                    colno += width
                out["density"] = np.concatenate(blocks, axis=1)

    
    sum_formula: dict[str, int] = {}
    if "coordinates" not in out:
        raise ValueError(f"No coordinates found in input file {file}.")
    for element in [coord["element"] for coord in out["coordinates"][0]]:
        if element in sum_formula:
            sum_formula[element] += 1
        else:
            sum_formula[element] = 1
    out["sum_formula"] = "".join([f"{key}{value}" for key,value in sorted(sum_formula.items(), key=lambda x: x[0])])

    return out

def read_qc_file(file: str) -> dict[str, Any]:
    """
    Reads ORCA calculation files and returns a dictionary containing important properties.
    Provide an input or output file, the other one is read automatically.

    Args:
        file (str): The path to the quantum chemistry file.

    Returns:
        dict[str, Any]: A dictionary containing the read properties.

    Raises:
        AssertionError: If the file does not have a ".inp" or ".out" extension.

    """
    assert ".inp" in file or ".out" in file
    filename = file.split("/")[-1]
    relpath = file[:file.find(filename)]
    if relpath == "":
        relpath = "./"
    out: dict[str, Any] = {}
    if ".inp" in file:
        infile = read_input_file(file)
        if filename.replace(".inp", ".out") in os.listdir(relpath):
            outfile = read_output_file(relpath + filename.replace(".inp", ".out"))
        else:
            outfile = {}
    else:
        if filename.replace(".out", ".inp") in os.listdir(relpath):
            infile = read_input_file(relpath + filename.replace(".out", ".inp"))
        else:
            infile = {}
        outfile = read_output_file(file)
    
    if infile == {} or outfile == {}:
        raise FileNotFoundError(f"Incomplete calculation for {file}. Input or output is missing.")

    return {**infile, **outfile}

### ABSTRACT CLASS: Logger ###
class Logger(ABC):
    """
    Abstract base class for logging the results of a quantum chemistry calculation.

    This class provides a common interface for logging the results of a QC calculation
    to different output formats.

    Attributes:
        title (str): A title for the log.
        false_only (bool): If True, only log incorrect results.

    Methods:
        log(correct: bool, file: str, property: str, found: Any, expected: Any) -> None:
            Logs a result.

        header() -> None:
            Writes a header to the log.

        log_plain(message: str) -> None:
            Writes a plain message to the log.

        check_calculations(path: str, calculation_refs: dict, tolerance: float = 0.001) -> None:
            Checks the results of a directory of calculations against a reference dictionary.
        
        check_number(file: str, property: str, found: Any, expected: Any) -> None:
            Logs a numeric property check.

        check_dict(file: str, property: str, found: Dict, expected: Dict) -> None:
            Logs a dictionary property check.
    
    """
    title: str

    def __init__(self, title: str, false_only: bool = False) -> None:
        self.title = title
        self.false_only = false_only

    @abstractmethod
    def log(self, correct: bool, file: str, property: str, found: Any, expected: Any) -> None:
        pass

    @abstractmethod
    def header(self) -> None:
        pass

    @abstractmethod
    def log_plain(self, message: str) -> None:
        pass

    def check_calculations(self, path: str, calculation_refs: dict, tolerance: float = 0.001) -> None:
        PREC = 5
        self.header()
        checked = []
        if path[-1] != "/":
            path += "/"
        for file in [file for file in os.listdir(path) if ".inp" in file]:
            stud = read_qc_file(path + file)
            file = file.replace(".inp", "").replace(".out", "")
            if stud["sum_formula"] in calculation_refs.keys():
                ref = calculation_refs[stud["sum_formula"]]
            else:
                self.log(False, file, "sum_formula", stud["sum_formula"], "sum_formula not in reference")
                continue
            for key in ref.keys():
                if stud.get(key, None) is None:
                    self.log(False, file, key, "Key not found", "")
                    self.log_plain(f"    --> Incomplete output?")
                    break
                if not hasattr(ref[key], "__iter__"):
                    if isinstance(ref[key], (int, float)):
                        if abs(ref[key] - stud[key]) < tolerance:
                            correct = True
                        else:
                            correct = False
                        self.log(correct, file, key, round(stud[key], PREC), round(ref[key], PREC))
                        continue
                    elif ref[key] == stud[key]:
                        correct = True
                    else:
                        correct = False
                    self.log(correct, file, key, stud[key], ref[key])
                    continue
                else:
                    for must_have in ref[key]:
                        if must_have.lower() in [x.lower() for x in stud[key]]:
                            correct = True
                        else:
                            correct = False
                        self.log(correct, file, key, " ".join(stud[key]), must_have)
                        continue
            if any("opt" in x.lower() for x in stud["keywords"]):
                self.log(stud["opt_converged"], file, "opt_converged", f"{stud['opt_converged']}", "True")
            checked.append(stud["sum_formula"])
        for sum_formula in calculation_refs.keys():
            if sum_formula not in checked:
                self.log(False, "missing calculation", sum_formula, "", "")
    
    def check_calculations_new(self, path: str, calculation_refs: list, tolerance: float = 0.001) -> None:
        PREC = 5
        self.header()
        checked = []
        if path[-1] != "/":
            path += "/"
        for file in [file for file in os.listdir(path) if ".inp" in file]:
            try:
                stud = read_qc_file(path + file)
            except FileNotFoundError:
                self.log_plain(f"Calculation {file} seems incomplete. Skipping...")
                continue
            file = file.replace(".inp", "").replace(".out", "")
            found = False
            for ref in calculation_refs:
                if ref["identifier"](stud):
                    found = True
                    break
                    # print("found!")
                    # print(stud)
            if not found:
                self.log(False, file, "", "", "Unidentified calculation found.")
                continue
            self.log_plain(ref["description"])
            for key in ref.keys():
                if key == "identifier" or key == "description":
                    continue
                if stud.get(key, None) is None:
                    self.log(False, file, key, "Key not found", "")
                    self.log_plain(f"    --> Incomplete output?")
                    break
                if not hasattr(ref[key], "__iter__"):
                    if isinstance(ref[key], (int, float)):
                        if abs(ref[key] - stud[key]) < tolerance:
                            correct = True
                        else:
                            correct = False
                        self.log(correct, file, key, round(stud[key], PREC), round(ref[key], PREC))
                        continue
                    elif ref[key] == stud[key]:
                        correct = True
                    else:
                        correct = False
                    self.log(correct, file, key, stud[key], ref[key])
                    continue
                else:
                    for must_have in ref[key]:
                        if must_have.lower() in [x.lower() for x in stud[key]]:
                            correct = True
                        else:
                            correct = False
                        self.log(correct, file, key, " ".join(stud[key]), must_have)
                        continue
            if any("opt" in x.lower() for x in stud["keywords"]):
                self.log(stud["opt_converged"], file, "opt_converged", f"{stud['opt_converged']}", "True")
            checked.append(stud["sum_formula"])

    def check_raw_number(self, student: float|int, reference: float|int, tolerance: float, property: str, task: str) -> bool:
        if abs(student - reference) < tolerance:
            correct = True
        else:
            correct = False
        self.log(correct, task, property, round(student, int(log10(1/tolerance))), round(reference, int(log10(1/tolerance))))
        return correct

    def check_number(self, student: float|int, references: list[dict[str,Any]], tolerance: float, property: str, task: str) -> bool:
        """
        Compare a single float/int number from a student with a list of reference values.
        The list of reference values should contain dictionaries with keys 'ref', 'correct' and optionally 'reason'.
        If the student number matches one of the references within the given tolerance, the 'correct' key in the reference is used to determine if the student answer is correct.
        If the student number does not match any of the references, the answer is considered incorrect.
        :param student: The number from the student.
        :param references: A list of reference values in dictionaries.
        :param tolerance: The tolerance within which student and reference values are considered the same.
        :param property: The name of the property that is being checked.
        :param task: The name of the task that is being checked.
        """
        for reference in references:
            if abs(student - reference["ref"]) < tolerance:
                self.log(reference["correct"], task, property, round(student, int(log10(1/tolerance))), round(references[0]["ref"], int(log10(1/tolerance))))
                if not reference["correct"]:
                    self.log_plain(f"    {reference['reason']}")
                return reference["correct"]
        self.log(False, task, property, round(student, int(log10(1/tolerance)+2)), round(references[0]["ref"], int(log10(1/tolerance)+2)))
        self.log_plain("    No reference found")
        return False
        
    def check_dict(self, student: dict, references: dict[str, dict], tolerance: float, property: str, task: str) -> None:
        for key, val in student.items():
            if key not in references.keys():
                self.log(False, task, key, "Key not in reference", "")
                self.log_plain(f"    Value is {val}")
                continue
            self.check_number(val, references[key], tolerance, key, task)
        for key in references.keys():
            if key not in student.keys():
                self.log(False, task, key, "Key not found", "")
    
    def check_list(self, student: list, references: list, tolerance: float, property: str, task: str) -> None:
        check_results = []
        for val, ref in zip(student, references):
            # self.check_raw_number(val, ref, tolerance, property, task)
            if abs(val - ref) < tolerance:
                correct = True
            else:
                correct = False
            check_results.append([correct, "", "", val, ref])
        all_correct = False if False in [x[0] for x in check_results] else True
        self.log(all_correct, task, property, "", "")
        for check_result in check_results:
            self.log(*check_result)
            # self.log(correct, task, property, round(val, int(log10(1/tolerance))), round(ref, int(log10(1/tolerance))))


### CLASS: BilloLogger ###
class BilloLogger(Logger):
    """
    Logger that prints the results of a QC calculation to the console.

    This class implements the Logger interface and provides a simple way to log
    the results of a QC calculation to the console. It prints each result on a
    separate line and includes information about whether the result is correct or
    not.

    Attributes:
        title (str): A title for the log.
        false_only (bool): If True, only log incorrect results.

    Methods:
        log(correct: bool, file: str, property: str, found: Any, expected: Any) -> None:
            Logs a result.

        header() -> None:
            Writes a header to the log.

        log_plain(message: str) -> None:
            Writes a plain message to the log.
    """
    standard_line = "{checkbox} {file:<20} {property:>20}: {found:>15} {expected:>15}"
    total_checks = 0
    total_correct = 0

    def log(self, correct: bool, file: str, property: str, found: Any, expected: Any) -> None:
        self.total_checks += 1
        if correct:
            self.total_correct += 1
        if self.false_only and correct:
            return
        print(self.standard_line.format(checkbox = '[\u2713]' if correct else '[\u2717]',
                                        file=file,
                                        property=property,
                                        found=found,
                                        expected=expected))
    
    def header(self) -> None:
        print(self.title)
        print(datetime.datetime.now())
        print(self.standard_line.format(checkbox = "[ ]",
                                        file="File",
                                        property="Property",
                                        found="Found",
                                        expected="Expected"))
    
    def log_plain(self, message: str) -> None:
        print(message)
    
    def footer(self) -> None:
        print(f"Total checks: {self.total_checks}\nTotal correct: {self.total_correct}\n\n")