baseline_optimize.py

import numpy as np
from pymatgen.core import Structure
from chgnet.model.model import CHGNet
from chgnet.model import StructOptimizer
import datetime
from io import StringIO 
import sys
from operator import add
from pathlib import Path
from tempfile import TemporaryDirectory
import pickle
import argparse
from pandas import DataFrame

from ase.io import read
from ase.calculators.calculator import *
from ase.geometry.cell import *
from ase.optimize.sciopt import OptimizerConvergenceError

from analyze_res import save_stats

def relax(atoms, optimizer_class, fmax, steps, relax_cell=False):

    class Capturing(list):
        def __enter__(self):
            self._stdout = sys.stdout
            sys.stdout = self._stringio = StringIO()
            return self
        def __exit__(self, *args):
            self.extend(self._stringio.getvalue().splitlines())
            del self._stringio    # free up some memory
            sys.stdout = self._stdout

    relaxer = StructOptimizer(optimizer_class=optimizer_class, use_device='cpu')
    step_list = []
    e_step_list = []
    time_list = []
    energy_list = []
    fmax_list = []
    traj = {}

    start_time = None

    with Capturing() as output:

        try:
            start_t = datetime.datetime.now()
            prediction = relaxer.relax(atoms.copy(), fmax=fmax, steps=steps, relax_cell=relax_cell, verbose=True, save_path="tmp.pkl")
            temp_atoms = prediction['final_structure']
            traj = {
                    "energy": prediction['trajectory'].energies,
                    "forces": prediction['trajectory'].forces,
                    "stresses": prediction['trajectory'].stresses,
                    "magmoms": prediction['trajectory'].magmoms,
                    "atom_positions": prediction['trajectory'].atom_positions,
                    "cell": prediction['trajectory'].cells,
                    "atomic_number": prediction['trajectory'].atoms.get_atomic_numbers(),
                    }
        
            with TemporaryDirectory() as temp_dir:
                temp_atoms.to_file(f"{temp_dir}/relaxed.cif")
                new_atoms = read(f"{temp_dir}/relaxed.cif")

        except (OptimizerConvergenceError, RuntimeError, IndexError) as e:
            print(f"Optimizer {optimizer_class} for {atoms}, error: {e}", sys.stderr)
            opt_t = (datetime.datetime.now() - start_t).total_seconds()
            return failed_res_list(opt_t), None, traj
    
    for j in range(1, len(output)):
        
        try:

            separ_str = output[j].replace("  ", " ").replace("[ ", "[").replace(" ]", "]").split()

            if start_time == None:
                start_time = datetime.datetime.strptime(separ_str[2], "%H:%M:%S")

            energy_list.append(float(separ_str[3])/len(atoms))
            fmax_list.append(float(separ_str[4]))
            time_list.append((datetime.datetime.strptime(separ_str[2], "%H:%M:%S") - start_time).total_seconds())

            step_list.append(j - 1)
            if "[" in separ_str[1]:
                e_step_list.append(int(separ_str[1].split("[")[1].split("]")[0]))
            else:
                e_step_list.append(step_list[-1])

        except ValueError as e:
            print(f"Value error in the output list: {e}", sys.stderr)
            continue

    return {"Step": step_list, "E_step": e_step_list, "Energy": energy_list, "Time": time_list, "Fmax": fmax_list}, new_atoms, traj


def failed_res_list(secs):

    return {f"Step": [0, 1], "E_step": [0, 1], "Energy": [0, 0], "Time": [0, secs], "Fmax": [1, 1]}


def check_atoms(atoms, a_len):

    if atoms == None:
        return False
    
    # if structure lost some atoms
    if len(atoms.get_atomic_numbers()) < a_len:
        return False
    
    return True


def chgnet_optimization_logs(opti, 
                             starting_structures_dir="starting_structures", 
                             final_structures_dir="final_structures", 
                             baselines_dir="baselines"):

    chgnet = CHGNet.load()

    cif_files = [f for f in os.listdir(starting_structures_dir) if os.path.isfile(f"{starting_structures_dir}/{f}") and f.endswith(".cif")]
    cif_files.sort()
    
    steps = 1000

    Path(baselines_dir).mkdir(parents=True, exist_ok=True)

    res_list = {"Step": [], "E_step": [], "Energy": [], "Time": [], "Fmax": []}
    failures_n = 0
    total_time = 0
    relax_cell = False

    Path(f"{final_structures_dir}").mkdir(parents=True, exist_ok=True)
    
    for cif in cif_files:
        
        atoms = Structure.from_file(f"{starting_structures_dir}/{cif}")

        if opti == "FIRE+BFGSLineSearch":

            start_time = datetime.datetime.now()
            res_list_1, new_atoms, traj = relax(atoms, "FIRE", 0.1, 250, relax_cell=relax_cell)

            if not check_atoms(new_atoms, len(atoms.atomic_numbers)):
                opt_t = (datetime.datetime.now() - start_time).total_seconds()
                res_list_1 = failed_res_list(opt_t)

            else:

                res_list_2, new_atoms, traj_2 = relax(new_atoms, "BFGSLineSearch", 0.05, 750, relax_cell=relax_cell)

                # first check if the final structure has not lost any atoms
                if not check_atoms(new_atoms, len(atoms.atomic_numbers)):
                    opt_t = (datetime.datetime.now() - start_time).total_seconds()
                    res_list_1 = failed_res_list(opt_t)

                else:
                    fire_steps = len(res_list_1["Step"])

                    res_list_1["Step"].extend(map(add, res_list_2["Step"], [fire_steps] * len(res_list_2["Step"])))
                    res_list_1["E_step"].extend(map(add, res_list_2["E_step"], [fire_steps] * len(res_list_2["E_step"])))
                    res_list_1["Energy"].extend(res_list_2["Energy"])
                    res_list_1["Time"].extend(map(add, res_list_2["Time"], [res_list_1["Time"][-1]] * len(res_list_2["Time"])))
                    res_list_1["Fmax"].extend(res_list_2["Fmax"])
                
                    try:
                        for key in traj:
                            traj[key] = np.concatenate((traj[key], traj_2[key]))

                    except ValueError as e:
                        print(f"Value error in the output list: {e}")
                        for key in traj:
                            print(f"For Key: {key}, FIRE len: {len(traj[key])}, BFGS len: {len(traj_2[key])}")
                            print("")
                        quit()

            total_time += (datetime.datetime.now() - start_time).total_seconds()

        else:

            start_time = datetime.datetime.now()
            res_list_1, new_atoms, traj = relax(atoms, opti, 0.05, steps, relax_cell=relax_cell)
            total_time += (datetime.datetime.now() - start_time).total_seconds()

            if not check_atoms(new_atoms, len(atoms.atomic_numbers)):
                opt_t = (datetime.datetime.now() - start_time).total_seconds()
                res_list_1 = failed_res_list(opt_t)

        if len(res_list_1["Step"]) == 2 and res_list_1["Fmax"][1] == 1:
            failures_n += 1

        else:
            new_atoms.write(f"{final_structures_dir}/{cif}")
            pickle.dump(traj, open(f"{final_structures_dir}/{cif.replace('.cif', '.pkl')}", "wb"))

        for k in res_list:
            res_list[k].extend(res_list_1[k])

        df = DataFrame(res_list)
        df.to_csv(f"{baselines_dir}/{opti}_opt_traj.csv")

    save_stats(df, baselines_dir)
        

if __name__ == "__main__":

    parser = argparse.ArgumentParser(description='The optimization of a test set with a baseline')
    parser.add_argument("--method", type=str, default="BFGSLineSearch",
                        help="specific method to optimize, can be BFGS, BFGSLineSearch, MDMin, SciPyFminCG, FIRE, or FIRE+BFGSLineSearch")
    parser.add_argument("--input_dir", type=str, default="starting_structures/SrTiO3x8",
                        help="input dir with starting structures")
    parser.add_argument("--output_dir", type=str, default="baseline_output",
                        help="The output dir to save the optimized structures and the energy history")

    args = parser.parse_args()

    opti = args.method
    starting_structures_dir = args.input_dir
    final_structures_dir = str(Path(args.output_dir) / f"final_structures")
    baselines_dir = args.output_dir

    chgnet_optimization_logs(opti, starting_structures_dir=starting_structures_dir, 
                                        final_structures_dir=final_structures_dir, 
                                        baselines_dir=baselines_dir)