diff --git a/pyxtal/interface/latte_calculator.py b/pyxtal/interface/latte_calculator.py new file mode 100644 index 00000000..3596f37c --- /dev/null +++ b/pyxtal/interface/latte_calculator.py @@ -0,0 +1,316 @@ +from ase.io import read, write +from ase.optimize import FIRE +from ase.constraints import UnitCellFilter +from ase.calculators.calculator import Calculator, all_changes + +from collections import Counter +from pathlib import Path +import numpy as np +import os +import subprocess +import re + +# ===================== 1. LATTE Calculator ===================== + +class LATTECalculator(Calculator): + """ + Minimal ASE calculator that: + 1) writes bl/inputblock.dat + 2) runs LATTE + 3) reads energy, stress, forces from + latte.log + lastsystem.cfg (final frame only) + + Requirements in workdir: + ./LATTE_DOUBLE + TBparam/control.in + MDcontroller + bl/ (auto-created) + """ + + implemented_properties = ["energy", "forces", "stress"] + + def __init__(self, exe="./LATTE_DOUBLE", workdir=".", timeout=3600, **kwargs): + super().__init__(**kwargs) + self.workdir = os.path.abspath(workdir) + self.exe = os.path.abspath(os.path.join(self.workdir, exe)) + self.timeout = timeout + self.logfile = os.path.join(self.workdir, "latte.log") + self.cfgfile = os.path.join(self.workdir, "lastsystem.cfg") + + def calculation_required(self, atoms, properties): + return True + + def calculate(self, atoms=None, properties=("energy",), system_changes=all_changes): + if atoms is None: + atoms = self.atoms + else: + self.atoms = atoms.copy() + + # clear previous results + self.results = {} + + # write input for *current* geometry + self._write_inputblock(self.atoms) + self._set_single_point_mode() # RELAX=0 in control.in + + # run LATTE and parse outputs + self._run_latte() + + E = self._read_energy_from_log() + F = self._read_forces_from_cfg() + S = self._read_stress_from_log() + + self.results["energy"] = E + self.results["forces"] = F + if S is not None: + self.results["stress"] = S + + # Write inputblock.dat + def _write_inputblock(self, atoms): + bl = os.path.join(self.workdir, "bl") + os.makedirs(bl, exist_ok=True) + fpath = os.path.join(bl, "inputblock.dat") + + cell = atoms.cell.array + pos = atoms.get_positions() + syms = atoms.get_chemical_symbols() + + with open(fpath, "w") as f: + f.write(f"{len(atoms)}\n") + for v in cell: + f.write(f"{v[0]:.10f} {v[1]:.10f} {v[2]:.10f}\n") + for s, r in zip(syms, pos): + f.write(f"{s} {r[0]:.10f} {r[1]:.10f} {r[2]:.10f}\n") + + # Force single-point mode in control.in + def _set_single_point_mode(self): + cpath = os.path.join(self.workdir, "TBparam", "control.in") + if not os.path.exists(cpath): + print("WARNING: No control.in found.") + return + + lines = open(cpath).read().splitlines() + new = [] + updated_relax = False + + for line in lines: + if line.strip().startswith("RELAX"): + updated_relax = True + new.append("RELAX= 0 RELAXTYPE= SD MAXITER= 10 RLXFTOL= 0.1\n") + else: + new.append(line + "\n") + + if not updated_relax: + new.append("\nRELAX= 0 RELAXTYPE= SD MAXITER= 10 RLXFTOL= 0.1\n") + + open(cpath, "w").write("".join(new)) + + # Run LATTE + def _run_latte(self): + with open(self.logfile, "w") as f: + subprocess.run( + [self.exe], + cwd=self.workdir, + stdout=f, stderr=subprocess.STDOUT, + timeout=self.timeout, check=True + ) + + # Read total energy from latte.log + def _read_energy_from_log(self): + lines = open(self.logfile).read().splitlines() + E = None + pattern = re.compile(r"FREE ENERGY\s*=\s*([-0-9.Ee]+)") + for l in lines: + m = pattern.search(l) + if m: + E = float(m.group(1)) + if E is None: + raise RuntimeError("Energy not found in log!") + return E + + # Read stress tensor from latte.log (GPa → eV/A³) + def _read_stress_from_log(self): + lines = open(self.logfile).read().splitlines() + S = None + for i, l in enumerate(lines): + if "Stress tensor (GPa)" in l: + m = [] + for j in range(1, 4): + parts = lines[i+j].replace("#","").split() + m.append([float(x) for x in parts[:3]]) + S = np.array(m) + break + + if S is None: + return None + + # convert GPa → eV/ų and flip sign (LATTE → ASE convention) + conv = 0.00625 + return -(S*conv) + + # Read the LAST FRAME from lastsystem.cfg + def _read_forces_from_cfg(self): + if not os.path.exists(self.cfgfile): + raise RuntimeError("lastsystem.cfg not found") + + lines = open(self.cfgfile).read().splitlines() + + # number of atoms + natoms = None + for l in lines: + if "Number of particles" in l: + natoms = int(l.split("=")[1]) + break + + if natoms is None: + raise RuntimeError("Could not find 'Number of particles' in lastsystem.cfg") + + # H0 matrix + H = np.zeros((3, 3)) + for l in lines: + if "H0(" in l: + tmp = l.replace("=", " ").replace("A", "").split() + ij = tmp[0][3:-1] + i, j = ij.split(",") + H[int(i) - 1, int(j) - 1] = float(tmp[1]) + + # collect all numeric atom-lines (≥7 floats) + all_atoms = [] + for l in lines: + p = l.split() + if len(p) < 7: + continue + try: + nums = [float(x) for x in p[:7]] + all_atoms.append(nums) + except ValueError: + continue + + if len(all_atoms) < natoms: + raise RuntimeError("Not enough atom lines found in lastsystem.cfg") + + # last natoms rows = final forces + block = all_atoms[-natoms:] + + frac = np.array([b[:3] for b in block]) + forces = np.array([b[3:6] for b in block]) + + coords_cart = frac @ H # not used but kept for completeness + + return forces + +# ===================== 2. Reorder CIF & fix formula ===================== + +def reorder_cif_and_fix_formula(infile, outfile, desired_order=("O", "H", "C")): + """ + Read CIF, reorder atoms by desired_order, write new CIF, + and fix _chemical_formula_sum line. Returns reordered Atoms. + """ + atoms = read(infile) + + # Reorder atoms + symbols = np.array(atoms.get_chemical_symbols()) + indices = [] + for el in desired_order: + idx = np.where(symbols == el)[0] + indices.extend(idx.tolist()) + atoms_reordered = atoms[indices] + + # Write reordered CIF + write(outfile, atoms_reordered, format="cif") + + # Count elements + counts = Counter(atoms_reordered.get_chemical_symbols()) + parts = [] + for el in desired_order: + if el in counts: + n = counts[el] + parts.append(f"{el}{n}" if n != 1 else el) + custom_formula = " ".join(parts) + + # Patch _chemical_formula_sum + text = Path(outfile).read_text() + lines = text.splitlines() + + new_lines = [] + formula_found = False + for line in lines: + if line.lower().startswith("_chemical_formula_sum"): + new_lines.append(f"_chemical_formula_sum '{custom_formula}'") + formula_found = True + else: + new_lines.append(line) + + if not formula_found: + inserted = [] + inserted_flag = False + for line in new_lines: + inserted.append(line) + if (not inserted_flag) and line.lower().startswith("data_"): + inserted.append(f"_chemical_formula_sum '{custom_formula}'") + inserted_flag = True + new_lines = inserted + + Path(outfile).write_text("\n".join(new_lines)) + + print("Wrote reordered CIF to:", outfile) + print("Formula set to:", custom_formula) + return atoms_reordered + +# ===================== 3. Stress logger ===================== + +def log_stress(atoms, filename="stress_log.txt"): + """ + Append latest stress tensor from atoms.calc.results["stress"] + to filename without forcing a new LATTE evaluation. + """ + calc = atoms.calc + S = calc.results.get("stress", None) + if S is None: + # fall back to ASE call if needed (this will trigger a calc) + S = atoms.get_stress(voigt=False) + + S = np.array(S) + + first_time = not os.path.exists(filename) + with open(filename, "a") as f: + if first_time: + f.write("# xx yy zz yz xz xy\n") + if S.shape == (3, 3): + xx, yy, zz = S[0, 0], S[1, 1], S[2, 2] + yz, xz, xy = S[1, 2], S[0, 2], S[0, 1] + f.write(f"{xx: .6f} {yy: .6f} {zz: .6f} {yz: .6f} {xz: .6f} {xy: .6f}\n") + else: + # assume already 6-component + f.write(" ".join(f"{x: .6f}" for x in S) + "\n") + +# ===================== 4. Main workflow ===================== + +if __name__ == "__main__": + # --- user settings --- + infile = "bddc_exp.cif" # Original CIF + rearranged_cif = "rearranged_bddc.cif" # Reordered output CIF + latte_workdir = "." # Folder with LATTE_DOUBLE, TBparam, MDcontroller + + # 1) Reorder CIF + fix formula, get reordered Atoms + desired_order = ("O", "H", "C") #(you can see the order through running with original one for 1 step and observe from "lastsystem.cfg") + atoms = reorder_cif_and_fix_formula(infile, rearranged_cif, desired_order) + + # 2) Attach LATTE calculator + calc = LATTECalculator(exe="./LATTE_DOUBLE", workdir=latte_workdir) + atoms.set_calculator(calc) + + # 3) Cell + atomic relaxation with FIRE + ucf = UnitCellFilter(atoms) + opt = FIRE(ucf, dt=0.2, maxmove=0.2, logfile="latte_opt.log") + + # Attach stress logger (uses closure over 'atoms') + opt.attach(lambda: log_stress(atoms, filename="stress_log.txt"), interval=10) + + opt.run(fmax=0.1) + + # 4) Save relaxed structure + write("Relaxed_structure.cif", atoms) + print("Relaxed structure written to Relaxed_structure.cif") +